diff --git a/backend/app/services/external_fetch.py b/backend/app/services/external_fetch.py index 4707858..9a2b711 100644 --- a/backend/app/services/external_fetch.py +++ b/backend/app/services/external_fetch.py @@ -241,11 +241,19 @@ def _fetch_gdrive(url: str, dest_dir: Path, *, read_timeout: float, def _fetch_mega(url: str, dest_dir: Path, *, read_timeout: float, total_timeout: float, should_stop: Callable[[], bool]) -> FetchResult: - before = set(dest_dir.iterdir()) if dest_dir.exists() else set() + # A folder link keeps its subfolders, so files are found at any depth. + # Dotfiles are megatools' `.megatmp.*` staging names. + def _complete_files() -> set[Path]: + return { + p for p in dest_dir.rglob("*") + if p.is_file() and not p.name.startswith(".") + } if dest_dir.exists() else set() + + before = _complete_files() # megatools is a subprocess: its timeout IS a total wall-clock cap (the read # timeout has no analogue here), so the total budget applies directly. _run_mega_get(url, str(dest_dir), timeout=total_timeout) - new = [p for p in dest_dir.iterdir() if p not in before and p.is_file()] + new = sorted(_complete_files() - before) if not new: return FetchResult(error="mega-get wrote no new file") return FetchResult(files=new, bytes=sum(p.stat().st_size for p in new)) diff --git a/backend/app/tasks/external.py b/backend/app/tasks/external.py index a99dacc..43ff98f 100644 --- a/backend/app/tasks/external.py +++ b/backend/app/tasks/external.py @@ -33,6 +33,7 @@ from ..services.external_fetch import fetch_external from ..services.importer import Importer from ..services.thumbnailer import Thumbnailer from ._sync_engine import sync_session_factory as _sync_session_factory +from .maintenance import is_temp_staging_file log = logging.getLogger(__name__) @@ -189,10 +190,9 @@ def fetch_external_link(self, link_id: int, _serialize_waits: int = 0) -> dict: / str(post.external_post_id) / str(link_id) ) try: - # Clear what an earlier, interrupted attempt left behind, so this fetch - # starts from the link's imported files only. - with SessionLocal() as session: - _discard_unrecorded(session, post_dir) + # Clear the staging files an interrupted attempt left behind. Complete + # files it left stay, and are imported once this fetch succeeds. + _discard_partials(post_dir) result = fetch_external(host, url, post_dir) with SessionLocal() as session: link = session.get(ExternalLink, link_id) @@ -203,7 +203,7 @@ def fetch_external_link(self, link_id: int, _serialize_waits: int = 0) -> dict: ) _record_failure(link, attempts, result.error or "fetch failed") session.commit() - _discard_unrecorded(session, post_dir) + _discard_partials(post_dir) return {"link_id": link_id, "error": result.error} log.info( @@ -215,7 +215,8 @@ def fetch_external_link(self, link_id: int, _serialize_waits: int = 0) -> dict: post = session.get(Post, link.post_id) artist = session.get(Artist, post.artist_id) source = session.get(Source, post.source_id) if post.source_id else None - image_ids = _route_files(session, result.files, post, platform, artist, source) + files = _unimported_files(session, post_dir, first=result.files) + image_ids = _route_files(session, files, post, platform, artist, source) link.status = "downloaded" link.last_error = None link.completed_at = datetime.now(UTC) @@ -245,7 +246,7 @@ def fetch_external_link(self, link_id: int, _serialize_waits: int = 0) -> dict: if link is not None and link.status == "downloading": _record_failure(link, attempts, f"{type(exc).__name__}: {exc}") session.commit() - _discard_unrecorded(session, post_dir) + _discard_partials(post_dir) raise finally: if lock is not None: @@ -334,25 +335,37 @@ def _route_files(session, files, post, platform, artist, source) -> list[int]: return image_ids -def _discard_unrecorded(session, link_dir: Path) -> int: - """Delete every file in a link's staging dir that no ImageRecord points at, - then the dirs left empty. The dir belongs to one link, so anything in it - without a record is a partial or abandoned download from an earlier - attempt (megatools' `.megatmp.*`, a killed transfer, a rejected file). - Images imported in place stay. Returns the number of files removed.""" +def _unimported_files(session, link_dir: Path, *, first=()) -> list[Path]: + """Every complete file under a link's dir, subfolders included, that no + ImageRecord points at: what this fetch wrote plus anything an earlier + attempt finished writing but never imported. `first` (the fetch's own + files) leads the list. Staging files and sidecars are left out.""" if not link_dir.is_dir(): - return 0 - files = [p for p in link_dir.rglob("*") if p.is_file()] - if not files: - return 0 + return [] + found = { + p for p in link_dir.rglob("*") + if p.is_file() and not is_temp_staging_file(p) and p.suffix.lower() != ".json" + } + ordered = [p for p in first if p in found] + sorted(found.difference(first)) recorded = set(session.execute( select(ImageRecord.path).where( - ImageRecord.path.in_([str(p) for p in files]) + ImageRecord.path.in_([str(p) for p in ordered]) ) - ).scalars()) + ).scalars()) if ordered else set() + return [p for p in ordered if str(p) not in recorded] + + +def _discard_partials(link_dir: Path) -> int: + """Delete the staging files under a link's dir (megatools' `.megatmp.*`, an + HTTP fetch's `.part`), then the dirs left empty. Every fetcher writes to a + staging name and renames on completion, so a staging file is a transfer + that never finished. Complete files stay: they may be the only copy of art + whose upstream has since gone. Returns the number of files removed.""" + if not link_dir.is_dir(): + return 0 removed = 0 - for p in files: - if str(p) in recorded: + for p in link_dir.rglob("*"): + if not (p.is_file() and is_temp_staging_file(p)): continue try: p.unlink(missing_ok=True) @@ -363,7 +376,7 @@ def _discard_unrecorded(session, link_dir: Path) -> int: try: d.rmdir() except OSError: - pass # not empty: it holds an imported file + pass # not empty: it holds a complete file try: link_dir.rmdir() except OSError: @@ -381,7 +394,8 @@ def sweep_external_links() -> dict: First it recovers links whose worker died mid-fetch: a link 'downloading' for longer than the task's hard limit has no worker left, so it counts as a failed attempt and goes back to the retry path (dead-lettered at the - threshold like any other failure). The next attempt clears its leftovers.""" + threshold like any other failure). The next attempt clears its staging + files and imports the complete ones.""" SessionLocal = _sync_session_factory() with SessionLocal() as session: now = datetime.now(UTC) diff --git a/tests/test_external_fetch.py b/tests/test_external_fetch.py index c456f75..6dd0cab 100644 --- a/tests/test_external_fetch.py +++ b/tests/test_external_fetch.py @@ -108,6 +108,27 @@ def test_mega_collects_new_files(tmp_path, monkeypatch): assert res.files[0].name == "film.mp4" +def test_mega_folder_link_collects_files_in_subfolders(tmp_path, monkeypatch): + """A folder link keeps the folder's tree; every file in it is collected, + files already there before the fetch and staging files are not.""" + (tmp_path / "old.png").write_bytes(b"earlier") + + def fake_mega(url, out_dir, *, timeout): + nested = Path(out_dir) / "images" / "2020" + nested.mkdir(parents=True) + (nested / "kara.png").write_bytes(b"png") + (Path(out_dir) / "videos").mkdir() + (Path(out_dir) / "videos" / "loop.webm").write_bytes(b"webm") + (nested / ".megatmp.AbCd").write_bytes(b"half") + + monkeypatch.setattr(ef, "_run_mega_get", fake_mega) + res = fetch_external("mega", "https://mega.nz/folder/x#k", tmp_path) + assert res.ok + assert sorted(p.relative_to(tmp_path).as_posix() for p in res.files) == [ + "images/2020/kara.png", "videos/loop.webm", + ] + + def test_unsupported_host(tmp_path): res = fetch_external("nope", "https://x", tmp_path) assert not res.ok diff --git a/tests/test_external_worker.py b/tests/test_external_worker.py index dc661e8..bd9500b 100644 --- a/tests/test_external_worker.py +++ b/tests/test_external_worker.py @@ -377,6 +377,74 @@ def test_a_rejected_download_is_removed(db_sync, tmp_path, monkeypatch): assert db_sync.get(ExternalLink, link_id).status == "downloaded" +def _other_jpeg_bytes(): + # A second structured image, unlike _structured_jpeg_bytes, so the two + # don't near-duplicate each other at import. + im = Image.new("RGB", (256, 256)) + px = im.load() + for y in range(256): + for x in range(256): + px[x, y] = ((x * x) % 256, (y * 3) % 256, (x ^ y) % 256) + buf = io.BytesIO() + im.save(buf, "JPEG") + return buf.getvalue() + + +def test_a_failed_fetch_keeps_complete_files(db_sync, tmp_path, monkeypatch): + """A failed fetch removes staging files only. A complete file may be the + only copy left once the host has dropped the upload.""" + _, link = _seed(db_sync, host="mega") + link_id = link.id + monkeypatch.setattr(ext, "IMAGES_ROOT", tmp_path) + monkeypatch.setattr(ext, "_redis", lambda: _FakeRedis()) + + def fake_fetch(host, url, dest_dir, **kwargs): + (dest_dir / "images").mkdir(parents=True, exist_ok=True) + (dest_dir / "images" / "done.jpg").write_bytes(_structured_jpeg_bytes()) + (dest_dir / "images" / ".megatmp.AbCd").write_bytes(b"half") + return FetchResult(error="API call 'f' failed: ENOENT") + + monkeypatch.setattr(ext, "fetch_external", fake_fetch) + ext.fetch_external_link(link_id) + + link_dir = _link_dir(tmp_path, "mega", link_id) + assert (link_dir / "images" / "done.jpg").exists() + assert not (link_dir / "images" / ".megatmp.AbCd").exists() + + +def test_a_successful_fetch_imports_what_an_earlier_attempt_finished( + db_sync, tmp_path, monkeypatch, +): + """Complete files an earlier attempt wrote, in subfolders too, are imported + with the new fetch's files instead of being deleted or left behind.""" + _, link = _seed(db_sync, host="mega") + link_id = link.id + monkeypatch.setattr(ext, "IMAGES_ROOT", tmp_path) + monkeypatch.setattr(ext, "_redis", lambda: _FakeRedis()) + link_dir = _link_dir(tmp_path, "mega", link_id) + (link_dir / "pack").mkdir(parents=True) + earlier = link_dir / "pack" / "earlier.jpg" + earlier.write_bytes(_structured_jpeg_bytes()) + + def fake_fetch(host, url, dest_dir, **kwargs): + f = dest_dir / "new.jpg" + f.write_bytes(_other_jpeg_bytes()) + return FetchResult(files=[f], bytes=f.stat().st_size) + + monkeypatch.setattr(ext, "fetch_external", fake_fetch) + import backend.app.tasks.ml as ml_mod + import backend.app.tasks.thumbnail as thumb_mod + monkeypatch.setattr(ml_mod.embed_image, "delay", lambda i: None) + monkeypatch.setattr(thumb_mod.generate_thumbnail, "delay", lambda i: None) + + out = ext.fetch_external_link(link_id) + + assert out["images"] == 2 + db_sync.expire_all() + paths = set(db_sync.execute(select(ImageRecord.path)).scalars()) + assert {str(earlier), str(link_dir / "new.jpg")} <= paths + + def test_sweep_recovers_a_link_whose_worker_died(db_sync, monkeypatch): """A link 'downloading' past the task's hard limit (or with no claim time) counts as a failed attempt and is retried; a fresh claim is left alone."""