From 29604db9cd982a7f0d20f7bcb48d0d769124d0dc Mon Sep 17 00:00:00 2001 From: Jonathan Date: Sat, 11 Jul 2026 15:13:31 +0200 Subject: [PATCH] feat: download to staging + atomically promote complete album into the library Rewrite run_pipeline's download->import section to download into a per-job staging dir (cleared before each candidate), verify completeness there, promote only a complete album into the library via import_album, tag the final dir, then DB-import. The staging dir is always removed in a try/finally, so a partial/failed download never touches or corrupts the library, and a retry can no longer leave orphan files behind. Also point test_upgrade_e2e.py's run_pipeline call at dest_root="/tmp/lib" (matching every other pipeline test) since import_album now unconditionally creates the final dir, which requires a writable dest_root. --- worker/lyra_worker/pipeline.py | 66 +++++++++++++------------ worker/tests/test_staging_pipeline.py | 70 +++++++++++++++++++++++++++ worker/tests/test_tag_pipeline.py | 6 ++- worker/tests/test_upgrade_e2e.py | 2 +- 4 files changed, 111 insertions(+), 33 deletions(-) create mode 100644 worker/tests/test_staging_pipeline.py diff --git a/worker/lyra_worker/pipeline.py b/worker/lyra_worker/pipeline.py index 3aad183..25b6ef1 100644 --- a/worker/lyra_worker/pipeline.py +++ b/worker/lyra_worker/pipeline.py @@ -1,3 +1,4 @@ +import shutil from dataclasses import replace from typing import Sequence @@ -5,7 +6,7 @@ import psycopg from lyra_worker.adapters.base import SourceAdapter from lyra_worker.confidence import score_confidence -from lyra_worker.library import album_dir +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 @@ -181,40 +182,45 @@ def run_pipeline( _fail(conn, job_id, "no candidate above confidence threshold") return - # 4. download (fall-through) + # 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} - dest = album_dir(dest_root, target) + staging = staging_dir(dest_root, job_id) 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 + try: + for candidate in ranked: + adapter = by_source.get(candidate.source) + if adapter is None: + continue + shutil.rmtree(staging, ignore_errors=True) # clean slate per attempt + result = adapter.download(candidate, staging, 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 against the canonical track count when known) - _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 + # 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 - library_path = result.path or dest + final = album_dir(dest_root, target) + import_album(staging, final) # clean audio(+cover) move into the library, replacing any prior copy - if tagger is not None: - try: - tagger.tag_album(library_path, 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) + 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) - # 6. import - _set_state(conn, job_id, "imported", "import") - _import(conn, job_id, target, winner, library_path) + # 6. import (DB) + _set_state(conn, job_id, "imported", "import") + _import(conn, job_id, target, winner, final) + finally: + shutil.rmtree(staging, ignore_errors=True) # a partial/failed download never persists diff --git a/worker/tests/test_staging_pipeline.py b/worker/tests/test_staging_pipeline.py new file mode 100644 index 0000000..089f95e --- /dev/null +++ b/worker/tests/test_staging_pipeline.py @@ -0,0 +1,70 @@ +import os + +from lyra_worker.claim import claim_next +from lyra_worker.pipeline import run_pipeline +from lyra_worker.types import Candidate, DownloadResult, Quality +from tests.conftest import insert_request + +_Q = Quality(fmt="FLAC", lossless=True, bit_depth=24, sample_rate=96000) + + +class FlakyAdapter: + """Writes a partial (1-track, oddly-named) set on the first download, the full 3-track set + after. The first-attempt file is named UNLIKE the real tracks — mirroring streamrip's + inconsistent per-attempt naming, so an un-isolated download would leave a visible orphan.""" + + name = "qobuz" + tier = 0 + + def __init__(self): + self.calls = 0 + + def health(self): + return True + + def search(self, target): + return [Candidate(source="qobuz", source_ref="r", matched_artist=target.artist, + matched_album=target.album, quality=_Q, track_count=3, source_tier=0)] + + def download(self, candidate, dest, on_progress): + self.calls += 1 + os.makedirs(dest, exist_ok=True) + if self.calls == 1: + open(os.path.join(dest, "orphan-a.flac"), "wb").close() # partial, distinct name + return DownloadResult(ok=True, path=dest, track_count=1) + for i in range(3): + open(os.path.join(dest, f"{i} track.flac"), "wb").close() + return DownloadResult(ok=True, path=dest, track_count=3) + + +def test_partial_download_leaves_library_untouched(conn, tmp_path): + adapter = FlakyAdapter() # first call is partial + job_id = insert_request(conn, artist="Radiohead", album="In Rainbows") + claim_next(conn) + run_pipeline(conn, job_id, [adapter], dest_root=str(tmp_path)) + + with conn.cursor() as cur: + cur.execute('SELECT state FROM "Job" WHERE id = %s', (job_id,)) + assert cur.fetchone()[0] == "needs_attention" # completeness check rejected the partial + assert not (tmp_path / "Radiohead" / "In Rainbows").exists() # library never touched + staging_root = tmp_path / ".staging" + assert not staging_root.exists() or list(staging_root.iterdir()) == [] # staging cleaned + + +def test_retry_imports_clean_with_no_orphans(conn, tmp_path): + adapter = FlakyAdapter() + # attempt 1: partial -> needs_attention (no library write) + job1 = insert_request(conn, artist="Radiohead", album="In Rainbows") + claim_next(conn) + run_pipeline(conn, job1, [adapter], dest_root=str(tmp_path)) + # attempt 2: full -> imported + job2 = insert_request(conn, artist="Radiohead", album="In Rainbows") + claim_next(conn) + run_pipeline(conn, job2, [adapter], dest_root=str(tmp_path)) + + with conn.cursor() as cur: + cur.execute('SELECT state FROM "Job" WHERE id = %s', (job2,)) + assert cur.fetchone()[0] == "imported" + final = tmp_path / "Radiohead" / "In Rainbows" + # exactly the 3 real tracks — the attempt-1 "orphan-a.flac" is NOT present + assert sorted(os.listdir(final)) == ["0 track.flac", "1 track.flac", "2 track.flac"] diff --git a/worker/tests/test_tag_pipeline.py b/worker/tests/test_tag_pipeline.py index 14ecad0..dade2be 100644 --- a/worker/tests/test_tag_pipeline.py +++ b/worker/tests/test_tag_pipeline.py @@ -45,5 +45,7 @@ def test_download_dir_receives_files_and_tagger_runs(conn, tmp_path): with conn.cursor() as cur: cur.execute('SELECT state FROM "Job" WHERE id = %s', (job_id,)) assert cur.fetchone()[0] == "imported" - assert tagger.seen == written["dir"] # tagger ran on the exact directory the download wrote to - assert os.path.isdir(tagger.seen) + final = os.path.join(str(tmp_path), "John Mayer", "Continuum") + assert tagger.seen == final # tagger runs on the final library dir + assert os.path.isfile(os.path.join(final, "track.opus")) # the downloaded file was promoted into it + assert not os.path.exists(written["dir"]) # the staging download dir was cleaned up diff --git a/worker/tests/test_upgrade_e2e.py b/worker/tests/test_upgrade_e2e.py index b20f832..6e6f17d 100644 --- a/worker/tests/test_upgrade_e2e.py +++ b/worker/tests/test_upgrade_e2e.py @@ -87,7 +87,7 @@ def test_below_cutoff_copy_is_upgraded_end_to_end(conn): assert job_id is not None # upgrade_cutoff=2 so intake does NOT short-circuit on the existing quality-1 copy. - run_pipeline(conn, job_id, [QobuzAdapter(FakeQobuzClient())], upgrade_cutoff=2) + run_pipeline(conn, job_id, [QobuzAdapter(FakeQobuzClient())], upgrade_cutoff=2, dest_root="/tmp/lib") assert _job_state(conn, job_id) == "imported" path, quality, source = _library_row(conn, "John Mayer", "Continuum")