import os import shutil import threading from dataclasses import replace from typing import Sequence import psycopg from lyra_worker.adapters.base import SourceAdapter from lyra_worker.confidence import score_confidence from lyra_worker.library import album_dir, import_album, staging_dir from lyra_worker.quality import quality_class from lyra_worker.ranker import rank_candidates from lyra_worker.types import Candidate, MBTarget _AUDIO_EXT = {".flac", ".mp3", ".m4a", ".opus", ".ogg", ".aac", ".wav"} def _count_staged_audio(staging: str) -> int: """Count audio files anywhere under the staging dir (streamrip nests them in a subfolder).""" n = 0 for _root, _dirs, files in os.walk(staging): for f in files: if os.path.splitext(f)[1].lower() in _AUDIO_EXT: n += 1 return n def _poll_download_progress(job_id: str, staging: str, expected: int, stop: "threading.Event") -> None: """Until `stop`, periodically write Job.downloadProgress = files-in-staging / expected (capped 0.99). Uses its own short-lived connection (psycopg conns aren't shareable across threads). A missing DSN or expected<=0 just no-ops.""" dsn = os.environ.get("DATABASE_URL") if not dsn or expected <= 0: return conn = psycopg.connect(dsn) try: while not stop.is_set(): frac = min(_count_staged_audio(staging) / expected, 0.99) with conn.cursor() as cur: cur.execute('UPDATE "Job" SET "downloadProgress" = %s WHERE id = %s', (frac, job_id)) conn.commit() stop.wait(1.5) except Exception as e: # a progress poller must never affect the job print(f"pipeline: download progress poller error: {e}", flush=True) finally: conn.close() def _set_state(conn: psycopg.Connection, job_id: str, state: str, stage: str) -> None: with conn.cursor() as cur: cur.execute( 'UPDATE "Job" SET state = %s, "currentStage" = %s, "updatedAt" = now() WHERE id = %s', (state, stage, job_id), ) conn.commit() def _set_download_progress(conn: psycopg.Connection, job_id: str, frac: float) -> None: with conn.cursor() as cur: cur.execute('UPDATE "Job" SET "downloadProgress" = %s WHERE id = %s', (frac, job_id)) conn.commit() def _request_id(conn: psycopg.Connection, job_id: str) -> str: with conn.cursor() as cur: cur.execute('SELECT "requestId" FROM "Job" WHERE id = %s', (job_id,)) return cur.fetchone()[0] def _fail(conn: psycopg.Connection, job_id: str, reason: str) -> None: with conn.cursor() as cur: cur.execute( 'UPDATE "Job" SET state = \'needs_attention\', error = %s, "updatedAt" = now() WHERE id = %s', (reason, job_id), ) cur.execute( 'UPDATE "Request" SET status = \'needs_attention\' WHERE id = (SELECT "requestId" FROM "Job" WHERE id = %s)', (job_id,), ) conn.commit() def _load_target(conn: psycopg.Connection, job_id: str) -> MBTarget: with conn.cursor() as cur: cur.execute( 'SELECT artist, album FROM "Request" WHERE id = (SELECT "requestId" FROM "Job" WHERE id = %s)', (job_id,), ) artist, album = cur.fetchone() return MBTarget(artist=artist, album=album) def _already_in_library(conn: psycopg.Connection, target: MBTarget) -> bool: with conn.cursor() as cur: cur.execute( 'SELECT 1 FROM "LibraryItem" WHERE artist = %s AND album = %s LIMIT 1', (target.artist, target.album), ) return cur.fetchone() is not None def _is_upgrade_job(conn: psycopg.Connection, job_id: str) -> bool: with conn.cursor() as cur: cur.execute( 'SELECT "monitoredReleaseId" IS NOT NULL FROM "Request" ' 'WHERE id = (SELECT "requestId" FROM "Job" WHERE id = %s)', (job_id,), ) row = cur.fetchone() return bool(row and row[0]) def _already_in_library_at_cutoff(conn: psycopg.Connection, target: MBTarget, cutoff: int) -> bool: with conn.cursor() as cur: cur.execute( 'SELECT 1 FROM "LibraryItem" WHERE artist = %s AND album = %s ' 'AND "qualityClass" >= %s LIMIT 1', (target.artist, target.album, cutoff), ) return cur.fetchone() is not None def _library_quality(conn: psycopg.Connection, target: MBTarget) -> int | None: with conn.cursor() as cur: cur.execute( 'SELECT "qualityClass" FROM "LibraryItem" WHERE artist = %s AND album = %s', (target.artist, target.album), ) row = cur.fetchone() return row[0] if row else None def _persist_candidates(conn: psycopg.Connection, job_id: str, candidates: list[Candidate]) -> None: with conn.cursor() as cur: for c in candidates: cur.execute( 'INSERT INTO "Candidate" (id, "jobId", source, format, "qualityClass", ' '"trackCount", confidence, "sourceRef", chosen, "createdAt") ' "VALUES (gen_random_uuid()::text, %s, %s, %s, %s, %s, %s, %s, false, now())", (job_id, c.source, c.quality.fmt, quality_class(c.quality), c.track_count, c.confidence, c.source_ref), ) conn.commit() def _mark_chosen(conn: psycopg.Connection, job_id: str, source_ref: str) -> None: """Mark exactly one candidate (the one now being downloaded) as chosen, clearing any prior choice. Called at each download attempt so `chosen` reflects the source in flight.""" with conn.cursor() as cur: cur.execute( 'UPDATE "Candidate" SET chosen = ("sourceRef" = %s) WHERE "jobId" = %s', (source_ref, job_id), ) conn.commit() def _import(conn: psycopg.Connection, job_id: str, target: MBTarget, winner: Candidate, path: str) -> None: request_id = _request_id(conn, job_id) with conn.cursor() as cur: cur.execute( 'INSERT INTO "LibraryItem" (id, "requestId", artist, album, path, source, ' 'format, "qualityClass", "importedAt") ' "VALUES (gen_random_uuid()::text, %s, %s, %s, %s, %s, %s, %s, now()) " 'ON CONFLICT (artist, album) DO UPDATE SET ' ' "requestId" = EXCLUDED."requestId", path = EXCLUDED.path, source = EXCLUDED.source, ' ' format = EXCLUDED.format, "qualityClass" = EXCLUDED."qualityClass", ' ' "importedAt" = now() ' 'WHERE EXCLUDED."qualityClass" > "LibraryItem"."qualityClass"', (request_id, target.artist, target.album, path, winner.source, winner.quality.fmt, quality_class(winner.quality)), ) cur.execute( "UPDATE \"Request\" SET status = 'completed' WHERE id = %s", (request_id,) ) conn.commit() def run_pipeline( conn: psycopg.Connection, job_id: str, adapters: Sequence[SourceAdapter], resolver=None, tagger=None, min_confidence: float = 0.7, dest_root: str = "/music", staging_root: str | None = None, upgrade_cutoff: int | None = None, ) -> None: """Real staged acquisition using source-agnostic adapters. Fakes in this plan.""" names = [a.name for a in adapters] if len(names) != len(set(names)): raise ValueError(f"adapter names must be unique, got {names}") # 1. intake _set_state(conn, job_id, "matching", "intake") target = _load_target(conn, job_id) if resolver is not None: resolved = resolver.resolve(target.artist, target.album) if resolved is not None: target = resolved # dedupe: a monitor-driven upgrade job proceeds unless the existing copy already meets # the cutoff; a plain request short-circuits on any existing copy (slice-1 behavior). if upgrade_cutoff is not None and _is_upgrade_job(conn, job_id): _dedupe_hit = _already_in_library_at_cutoff(conn, target, upgrade_cutoff) else: _dedupe_hit = _already_in_library(conn, target) if _dedupe_hit: _set_state(conn, job_id, "imported", "import") with conn.cursor() as cur: cur.execute( "UPDATE \"Request\" SET status = 'completed' WHERE id = %s", (_request_id(conn, job_id),), ) conn.commit() return # 2. match — search all adapters; score + persist EVERY candidate found _set_state(conn, job_id, "matching", "match") found: list[Candidate] = [] for adapter in adapters: try: results = adapter.search(target) except Exception as e: # a down source contributes no candidates, never crashes the job print(f"pipeline: adapter {adapter.name} search failed: {e}", flush=True) continue for c in results: found.append(replace(c, source_tier=adapter.tier, confidence=score_confidence(target, c))) _persist_candidates(conn, job_id, found) # 3. rank _set_state(conn, job_id, "matched", "rank") ranked = rank_candidates(target, found, min_confidence) if not ranked: _fail(conn, job_id, "no candidate above confidence threshold") return # 4. download (fall-through) into an isolated per-job staging dir _set_state(conn, job_id, "downloading", "download") by_source = {a.name: a for a in adapters} staging = staging_dir(staging_root or f"{dest_root}/.staging", job_id) winner = None result = None _set_download_progress(conn, job_id, 0.0) # reset for this run expected = target.track_count or (ranked[0].track_count if ranked else 0) _stop = threading.Event() _poller = threading.Thread( target=_poll_download_progress, args=(job_id, staging, expected, _stop), daemon=True ) _poller.start() try: try: for candidate in ranked: adapter = by_source.get(candidate.source) if adapter is None: continue _mark_chosen(conn, job_id, candidate.source_ref) # reflect the source now in flight shutil.rmtree(staging, ignore_errors=True) # clean slate per attempt result = adapter.download(candidate, staging, lambda _pct: None) if result.ok: winner = candidate break finally: _stop.set() _poller.join(timeout=3) if winner is None or result is None or not result.ok: _fail(conn, job_id, "all downloads failed") return _set_download_progress(conn, job_id, 1.0) # download done — UI shows 100% into Finishing # 5. verify completeness on the staging copy, then promote + tag _set_state(conn, job_id, "tagging", "tag") expected = target.track_count if target.track_count is not None else winner.track_count if result.track_count < expected: _fail(conn, job_id, "incomplete download") return final = album_dir(dest_root, target) existing_q = _library_quality(conn, target) new_q = quality_class(winner.quality) _set_state(conn, job_id, "imported", "import") if existing_q is None or new_q > existing_q: import_album(staging, final) # clean audio(+cover) move, atomically replacing any prior copy if tagger is not None: try: tagger.tag_album(final, target) except Exception as e: # a tagging failure must not discard a good download print(f"pipeline: tagging failed for job {job_id}: {e}", flush=True) _import(conn, job_id, target, winner, final) else: # an existing copy is already at >= this quality — keep it untouched, just complete the request with conn.cursor() as cur: cur.execute( "UPDATE \"Request\" SET status = 'completed' WHERE id = %s", (_request_id(conn, job_id),), ) conn.commit() finally: shutil.rmtree(staging, ignore_errors=True) # a partial/failed download never persists