diff --git a/worker/lyra_worker/pipeline.py b/worker/lyra_worker/pipeline.py index fdc3dc8..eb966bb 100644 --- a/worker/lyra_worker/pipeline.py +++ b/worker/lyra_worker/pipeline.py @@ -1,14 +1,11 @@ -import time +from typing import Sequence import psycopg -# (state, currentStage) pairs the job moves through, in order, after claim. -_STAGES = [ - ("matched", "rank"), - ("downloading", "download"), - ("tagging", "tag"), - ("imported", "import"), -] +from lyra_worker.adapters.base import SourceAdapter +from lyra_worker.quality import quality_class +from lyra_worker.ranker import rank_candidates +from lyra_worker.types import Candidate, MBTarget def _set_state(conn: psycopg.Connection, job_id: str, state: str, stage: str) -> None: @@ -20,17 +17,144 @@ def _set_state(conn: psycopg.Connection, job_id: str, state: str, stage: str) -> conn.commit() -def run_pipeline(conn: psycopg.Connection, job_id: str, stage_delay: float = 0.0) -> None: - """Fake acquisition: walk the claimed job through to 'imported'.""" - for state, stage in _STAGES: - if stage_delay: - time.sleep(stage_delay) - _set_state(conn, job_id, state, stage) +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 "Request" SET status = \'completed\' ' - 'WHERE id = (SELECT "requestId" FROM "Job" WHERE id = %s)', + '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 _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: + with conn.cursor() as cur: + cur.execute( + 'UPDATE "Candidate" SET chosen = true WHERE "jobId" = %s AND "sourceRef" = %s', + (job_id, source_ref), + ) + 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 NOTHING', + (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], + min_confidence: float = 0.7, + dest_root: str = "/music", +) -> None: + """Real staged acquisition using source-agnostic adapters. Fakes in this plan.""" + # 1. intake + _set_state(conn, job_id, "matching", "intake") + target = _load_target(conn, job_id) + if _already_in_library(conn, target): + _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 + _set_state(conn, job_id, "matching", "match") + candidates: list[Candidate] = [] + for adapter in adapters: + candidates.extend(adapter.search(target)) + + # 3. rank + _set_state(conn, job_id, "matched", "rank") + ranked = rank_candidates(target, candidates, min_confidence) + _persist_candidates(conn, job_id, ranked) + if not ranked: + _fail(conn, job_id, "no candidate above confidence threshold") + return + + # 4. download (fall-through) + _set_state(conn, job_id, "downloading", "download") + by_source = {a.name: a for a in adapters} + dest = f"{dest_root}/{target.artist}/{target.album}" + winner = None + result = None + for candidate in ranked: + adapter = by_source.get(candidate.source) + if adapter is None: + continue + result = adapter.download(candidate, dest, lambda _pct: None) + if result.ok: + winner = candidate + _mark_chosen(conn, job_id, candidate.source_ref) + break + if winner is None or result is None or not result.ok: + _fail(conn, job_id, "all downloads failed") + return + + # 5. tag (integrity check; real tagging is Plan 4) + _set_state(conn, job_id, "tagging", "tag") + if result.track_count != winner.track_count: + _fail(conn, job_id, "incomplete download") + return + + # 6. import + _set_state(conn, job_id, "imported", "import") + _import(conn, job_id, target, winner, result.path or dest) diff --git a/worker/tests/test_pipeline.py b/worker/tests/test_pipeline.py index 99b9162..7252d3d 100644 --- a/worker/tests/test_pipeline.py +++ b/worker/tests/test_pipeline.py @@ -1,20 +1,101 @@ +from lyra_worker.adapters.fakes import FakeQobuz, FakeSoulseek, FakeYouTube, FailingAdapter from lyra_worker.claim import claim_next from lyra_worker.pipeline import run_pipeline from tests.conftest import insert_request -def test_pipeline_drives_job_to_imported(conn): - job_id = insert_request(conn) - claim_next(conn) - - run_pipeline(conn, job_id, stage_delay=0.0) - +def _job_state(conn, job_id): with conn.cursor() as cur: cur.execute('SELECT state, "currentStage" FROM "Job" WHERE id = %s', (job_id,)) - state, stage = cur.fetchone() - cur.execute('SELECT status FROM "Request" WHERE id = (SELECT "requestId" FROM "Job" WHERE id = %s)', (job_id,)) - req_status = cur.fetchone()[0] + return cur.fetchone() - assert state == "imported" - assert stage == "import" - assert req_status == "completed" + +def _request_status(conn, job_id): + with conn.cursor() as cur: + cur.execute( + 'SELECT status FROM "Request" WHERE id = (SELECT "requestId" FROM "Job" WHERE id = %s)', + (job_id,), + ) + return cur.fetchone()[0] + + +def test_success_picks_best_source_and_imports(conn): + job_id = insert_request(conn, artist="Radiohead", album="In Rainbows") + claim_next(conn) + run_pipeline(conn, job_id, [FakeYouTube(), FakeSoulseek(), FakeQobuz()], dest_root="/tmp/lib") + + assert _job_state(conn, job_id) == ("imported", "import") + assert _request_status(conn, job_id) == "completed" + with conn.cursor() as cur: + cur.execute('SELECT source FROM "Candidate" WHERE "jobId" = %s AND chosen = true', (job_id,)) + assert cur.fetchone()[0] == "qobuz" # highest quality won + cur.execute('SELECT source, format FROM "LibraryItem" WHERE artist = %s', ("Radiohead",)) + row = cur.fetchone() + assert row[0] == "qobuz" + + +def test_falls_through_when_best_download_fails(conn): + job_id = insert_request(conn, artist="Radiohead", album="In Rainbows") + claim_next(conn) + # FailingAdapter is source 'qobuz' (tier 0) whose download fails; Soulseek should win. + run_pipeline(conn, job_id, [FailingAdapter(), FakeSoulseek()], dest_root="/tmp/lib") + + assert _job_state(conn, job_id) == ("imported", "import") + with conn.cursor() as cur: + cur.execute('SELECT source FROM "Candidate" WHERE "jobId" = %s AND chosen = true', (job_id,)) + assert cur.fetchone()[0] == "soulseek" + + +def test_no_match_goes_to_needs_attention(conn): + job_id = insert_request(conn, artist="Radiohead", album="In Rainbows") + claim_next(conn) + run_pipeline(conn, job_id, [FakeQobuz(matches=False)], dest_root="/tmp/lib") + + assert _job_state(conn, job_id)[0] == "needs_attention" + assert _request_status(conn, job_id) == "needs_attention" + + +def test_all_downloads_fail_goes_to_needs_attention(conn): + job_id = insert_request(conn, artist="Radiohead", album="In Rainbows") + claim_next(conn) + run_pipeline(conn, job_id, [FailingAdapter()], dest_root="/tmp/lib") + + assert _job_state(conn, job_id)[0] == "needs_attention" + + +def test_incomplete_download_goes_to_needs_attention(conn): + from lyra_worker.adapters.fakes import FakeQobuz + from lyra_worker.types import DownloadResult + + class TruncatingQobuz(FakeQobuz): + def download(self, candidate, dest, on_progress): + on_progress(1.0) + # reports fewer tracks than the candidate promised -> integrity check fails + return DownloadResult(ok=True, path=dest, track_count=candidate.track_count - 1) + + job_id = insert_request(conn, artist="Radiohead", album="In Rainbows") + claim_next(conn) + run_pipeline(conn, job_id, [TruncatingQobuz()], dest_root="/tmp/lib") + + assert _job_state(conn, job_id)[0] == "needs_attention" + assert _request_status(conn, job_id) == "needs_attention" + + +def test_dedupe_skips_already_in_library(conn): + # First acquisition + job1 = insert_request(conn, artist="Radiohead", album="In Rainbows") + claim_next(conn) + run_pipeline(conn, job1, [FakeQobuz()], dest_root="/tmp/lib") + # Second request for the same album + job2 = insert_request(conn, artist="Radiohead", album="In Rainbows") + claim_next(conn) + run_pipeline(conn, job2, [FakeQobuz()], dest_root="/tmp/lib") + + assert _job_state(conn, job2) == ("imported", "import") + with conn.cursor() as cur: + # dedupe path creates no candidates for the second job + cur.execute('SELECT count(*) FROM "Candidate" WHERE "jobId" = %s', (job2,)) + assert cur.fetchone()[0] == 0 + cur.execute('SELECT count(*) FROM "LibraryItem" WHERE artist = %s AND album = %s', + ("Radiohead", "In Rainbows")) + assert cur.fetchone()[0] == 1 # not duplicated