From 58dfe440e9fd19234c94e600d6a54cf73c32cd53 Mon Sep 17 00:00:00 2001 From: devgtv Date: Tue, 29 Sep 2026 15:44:28 -0300 Subject: [PATCH] fix(saude): never reuse or publish a partial CKAN download download_resource() short-circuits on dest_path.exists() but streams directly into that path, so a mid-stream drop leaves a truncated file that the next call happily returns, and the engine then hashes, converts and uploads it to S3 as the official artifact. CKAN/Saude also ignore the requested filename and write their own derived name inside output.parent, so the retry's cleanup of output deleted a path that never existed. All Saude workers share management/tmp as dest_dir, so two packages publishing a same-named resource also resolved to the same dest_path and interleaved their writes. download_resource() now streams into a sibling ...part file and os.replace()s it on success, unlinking it on any BaseException, so the rename is atomic. _download_raw_with_retry() gives each attempt its own scratch directory and drops the whole directory on retry, which is the only reliable cleanup when the origin chose the filename itself. Rebased onto main: the FTP origin keeps its attempt directory, because _retr_with_resume reads the existing bytes off it to resume with REST -- wiping it between attempts would restart every RETR from byte 0. The per-attempt directory therefore applies to the origins that need it, and a test pins both behaviours. --- pysus/api/saude/download.py | 41 +++++-- pysus/management/sync.py | 57 ++++++++- pysus/tests/api/saude/test_download.py | 140 +++++++++++++++++++++++ pysus/tests/management/test_sync_more.py | 121 ++++++++++++++++++++ 4 files changed, 344 insertions(+), 15 deletions(-) diff --git a/pysus/api/saude/download.py b/pysus/api/saude/download.py index ef74bb18..9ddcba5c 100644 --- a/pysus/api/saude/download.py +++ b/pysus/api/saude/download.py @@ -10,10 +10,12 @@ from __future__ import annotations +import os import re from collections.abc import Callable from pathlib import Path from urllib.parse import urlparse +from uuid import uuid4 from .errors import ResourceNotFound from .resources import CKANPackage, Resource @@ -112,16 +114,35 @@ async def download_resource( if dest_path.exists() and not overwrite: return dest_path - async with client.stream("GET", resource.url) as response: - response.raise_for_status() - total = int(response.headers.get("Content-Length", 0)) - downloaded = 0 - with dest_path.open("wb") as fh: - async for chunk in response.aiter_bytes(chunk_size=_DOWNLOAD_CHUNK): - fh.write(chunk) - downloaded += len(chunk) - if progress: - progress(downloaded, total) + # Stream into a sibling temp file and rename on success. Writing + # straight to dest_path would leave a truncated file behind when the + # connection drops mid-transfer, and the `dest_path.exists()` reuse + # above would then hand that partial file to the next attempt. Two + # packages publishing a resource with the same name also resolve to the + # same dest_path, so the rename makes the final write atomic instead of + # interleaving concurrent writers. + tmp_path = dest_path.with_name(f".{dest_path.name}.{uuid4().hex[:8]}.part") + try: + async with client.stream("GET", resource.url) as response: + response.raise_for_status() + total = int(response.headers.get("Content-Length", 0)) + downloaded = 0 + with tmp_path.open("wb") as fh: + async for chunk in response.aiter_bytes( + chunk_size=_DOWNLOAD_CHUNK + ): + fh.write(chunk) + downloaded += len(chunk) + if progress: + progress(downloaded, total) + os.replace(tmp_path, dest_path) + except BaseException: + # Never leave a partial file that the reuse check could pick up. + try: + tmp_path.unlink(missing_ok=True) + except OSError: + pass + raise return dest_path diff --git a/pysus/management/sync.py b/pysus/management/sync.py index fa84db25..581dff46 100644 --- a/pysus/management/sync.py +++ b/pysus/management/sync.py @@ -604,7 +604,17 @@ async def _download_raw_with_retry( """ raw_dir = Path(CACHEPATH) / "management" / "tmp" raw_dir.mkdir(parents=True, exist_ok=True) - output = raw_dir / f"{uuid4().hex[:8]}-{file.basename}" + # Each attempt gets its own directory. CKAN/Saude origins ignore + # the requested filename and derive their own from the resource + # name inside ``output.parent``, so with a shared directory a + # failed attempt could not be cleaned up reliably: the retry + # deleted ``output`` (which never existed) and left the real + # partial file behind for the next attempt to pick up. A + # per-attempt directory also stops two packages that publish a + # resource with the same name from writing the same file. + attempt_dir = raw_dir / uuid4().hex[:8] + attempt_dir.mkdir(parents=True, exist_ok=True) + output = attempt_dir / file.basename last_error: Exception | None = None for attempt in range(max_retries): @@ -612,11 +622,21 @@ async def _download_raw_with_retry( return await self._download_once(file, output, ftp_client) except _RETRYABLE as exc: last_error = exc - # Keep FTP partials so the retry can resume with REST; other - # origins are re-downloaded from scratch. + # Keep FTP partials so the retry can resume with REST: the + # attempt directory is the resume point, so it stays as-is. origin = getattr(getattr(file, "client", None), "name", None) if origin is None or origin.upper() != "FTP": - self._cleanup_local(output) + # Other origins are re-downloaded from scratch, and + # CKAN/Saude name the file they write themselves, so + # only dropping the whole attempt directory cleans up + # reliably. + self._cleanup_attempt(attempt_dir) + # The last attempt leaves nothing behind, so do not create + # a scratch directory that will never be used or removed. + if attempt < max_retries - 1: + attempt_dir = raw_dir / uuid4().hex[:8] + attempt_dir.mkdir(parents=True, exist_ok=True) + output = attempt_dir / file.basename wait_time = 2**attempt + (attempt * 2) error( f"Download attempt {attempt + 1}/{max_retries} failed " @@ -758,6 +778,30 @@ def _cleanup_local(path: Path) -> None: except OSError: pass + @staticmethod + def _is_attempt_dir(name: str) -> bool: + """Whether *name* is a per-attempt download scratch directory. + + They are 8 hex characters, which no file or ``.tmp_extract`` + directory in this tmp area uses, so leftovers from a killed run + can be swept up between runs. + """ + return len(name) == 8 and all(c in "0123456789abcdef" for c in name) + + @staticmethod + def _cleanup_attempt(attempt_dir: Path) -> None: + """Remove a download attempt's scratch directory. + + CKAN/Saude origins name the file they write themselves, so the + caller cannot know which path received the bytes; removing the + whole per-attempt directory is the only reliable way to ensure no + partial file survives into the next attempt. + """ + try: + shutil.rmtree(attempt_dir, ignore_errors=True) + except OSError: + pass + @staticmethod def _cleanup_stale_tmp() -> None: """Remove leftover files in the tmp directory from prior runs. @@ -771,7 +815,10 @@ def _cleanup_stale_tmp() -> None: return removed = 0 for p in tmp.iterdir(): - if p.is_dir() and p.name.endswith(".tmp_extract"): + if p.is_dir() and ( + p.name.endswith(".tmp_extract") + or SyncEngine._is_attempt_dir(p.name) + ): try: shutil.rmtree(p) removed += 1 diff --git a/pysus/tests/api/saude/test_download.py b/pysus/tests/api/saude/test_download.py index 9e5eb7df..f011ff71 100644 --- a/pysus/tests/api/saude/test_download.py +++ b/pysus/tests/api/saude/test_download.py @@ -335,6 +335,146 @@ async def test_select_by_fmt( assert path.exists() +class _FlakyTransport(httpx.AsyncBaseTransport): + """Yields ``good_bytes`` and then fails, once, like a dropped socket. + + The failure happens *after* real bytes have been handed to the caller, + which is what leaves a truncated file behind when the destination is + written in place. + """ + + def __init__(self, head: bytes, tail: bytes, fail_once: bool = True): + self.head = head + self.tail = tail + self.fail_once = fail_once + self.calls = 0 + + async def handle_async_request(self, request: httpx.Request): + self.calls += 1 + should_fail = self.fail_once and self.calls == 1 + return httpx.Response( + 200, + content=_flaky_stream( + self.head, self.tail, fail=should_fail, request=request + ), + headers={"Content-Length": str(len(self.head) + len(self.tail))}, + ) + + +async def _flaky_stream( + head: bytes, tail: bytes, fail: bool, request: httpx.Request +): + yield head + if fail: + raise httpx.ReadError("connection reset", request=request) + yield tail + + +class TestDownloadResourcePartialTransfer: + """A dropped connection must not leave a file the retry can reuse. + + ``download_resource`` short-circuits on ``dest_path.exists()``. If a + failed transfer left a truncated file there, the next attempt would + return it as a complete download and the caller would convert and + upload the partial bytes. + """ + + @pytest.mark.asyncio + async def test_failure_mid_stream_leaves_no_file( + self, saude_dataset_page_props, tmp_path + ): + from pysus.api.saude.resources import CKANPackage + + package = CKANPackage.model_validate(saude_dataset_page_props) + csv_resource = next(r for r in package.resources if r.format == "CSV") + + head, tail = b"col_a,col_b\n1,2\n", b"3,4\n" + transport = _FlakyTransport(head, tail) + + async with httpx.AsyncClient(transport=transport) as client: + with pytest.raises(httpx.ReadError): + await download_resource( + client, + package, + resource_id=csv_resource.id, + dest_dir=tmp_path, + ) + + # Nothing at the derived name, and no leftover .part scratch file. + assert list(tmp_path.iterdir()) == [] + + @pytest.mark.asyncio + async def test_retry_after_failure_downloads_full_content( + self, saude_dataset_page_props, tmp_path + ): + from pysus.api.saude.resources import CKANPackage + + package = CKANPackage.model_validate(saude_dataset_page_props) + csv_resource = next(r for r in package.resources if r.format == "CSV") + head, tail = b"col_a,col_b\n1,2\n", b"3,4\n" + transport = _FlakyTransport(head, tail) + + async with httpx.AsyncClient(transport=transport) as client: + with pytest.raises(httpx.ReadError): + await download_resource( + client, + package, + resource_id=csv_resource.id, + dest_dir=tmp_path, + ) + # Second attempt must not short-circuit on a partial file. + path = await download_resource( + client, + package, + resource_id=csv_resource.id, + dest_dir=tmp_path, + ) + + assert path.read_bytes() == head + tail + + @pytest.mark.asyncio + async def test_http_error_leaves_no_file( + self, saude_dataset_page_props, tmp_path + ): + from pysus.api.saude.resources import CKANPackage + + package = CKANPackage.model_validate(saude_dataset_page_props) + csv_resource = next(r for r in package.resources if r.format == "CSV") + + def handler(request: httpx.Request) -> httpx.Response: + return httpx.Response(404) + + async with httpx.AsyncClient( + transport=httpx.MockTransport(handler) + ) as client: + with pytest.raises(httpx.HTTPStatusError): + await download_resource( + client, + package, + resource_id=csv_resource.id, + dest_dir=tmp_path, + ) + + assert list(tmp_path.iterdir()) == [] + + @pytest.mark.asyncio + async def test_complete_download_leaves_no_scratch_file( + self, saude_dataset_page_props, tmp_path, mocked_saude + ): + from pysus.api.saude.resources import CKANPackage + + package = CKANPackage.model_validate(saude_dataset_page_props) + csv_resource = next(r for r in package.resources if r.format == "CSV") + async with httpx.AsyncClient(transport=mocked_saude) as client: + await download_resource( + client, + package, + resource_id=csv_resource.id, + dest_dir=tmp_path, + ) + assert not list(tmp_path.glob(".*.part")) + + class TestDownloadDatasetEdgeCases: @pytest.mark.asyncio async def test_no_dest_dir( diff --git a/pysus/tests/management/test_sync_more.py b/pysus/tests/management/test_sync_more.py index 70df5ad9..5cbbaa30 100644 --- a/pysus/tests/management/test_sync_more.py +++ b/pysus/tests/management/test_sync_more.py @@ -9,6 +9,7 @@ import asyncio import os from datetime import datetime +from pathlib import Path from unittest.mock import AsyncMock, MagicMock, PropertyMock, patch import pytest @@ -65,6 +66,126 @@ def _adapter(close_raises=False): return a +class TestDownloadRawWithRetry: + """A failed attempt must not leave bytes behind for the next one. + + CKAN/Saude origins ignore the requested filename and write to a name + they derive from the resource inside ``output.parent``. Cleaning up + only ``output`` therefore deleted a path that never existed and left + the real partial file in place. + """ + + @pytest.mark.asyncio + async def test_failed_attempt_does_not_leak_partial_file( + self, engine, tmp_path, monkeypatch + ): + monkeypatch.setattr("pysus.management.sync.CACHEPATH", tmp_path) + + attempts: list[Path] = [] + + async def _download_once(file, output, ftp_client=None): + # Emulate an origin that picks its own filename and fails + # after writing a partial file. + actual = output.parent / "derived-name.csv" + actual.write_bytes(b"partial") + attempts.append(actual) + raise OSError("boom") + + engine._download_once = AsyncMock(side_effect=_download_once) + + with pytest.raises(RuntimeError, match="after 2 attempts"): + await engine._download_raw_with_retry( + _remote_file(client_name="saude"), max_retries=2 + ) + + # Each attempt got a fresh directory, and none of them survived. + assert len(attempts) == 2 + assert attempts[0] != attempts[1] + for p in attempts: + assert not p.exists() + tmp_dir = tmp_path / "management" / "tmp" + assert not list(tmp_dir.iterdir()) + + @pytest.mark.asyncio + async def test_successful_retry_returns_fresh_path( + self, engine, tmp_path, monkeypatch + ): + monkeypatch.setattr("pysus.management.sync.CACHEPATH", tmp_path) + monkeypatch.setattr("pysus.management.sync.asyncio.sleep", AsyncMock()) + seen: list[Path] = [] + + async def _download_once(file, output, ftp_client=None): + seen.append(output) + if len(seen) == 1: + (output.parent / "derived-name.csv").write_bytes(b"partial") + raise OSError("boom") + output.parent.mkdir(parents=True, exist_ok=True) + output.write_bytes(b"complete") + return output + + engine._download_once = AsyncMock(side_effect=_download_once) + + result = await engine._download_raw_with_retry( + _remote_file(client_name="saude"), max_retries=2 + ) + + assert result.read_bytes() == b"complete" + assert seen[0] != seen[1] + + @pytest.mark.asyncio + async def test_ftp_attempt_dir_survives_so_the_retry_resumes( + self, engine, tmp_path, monkeypatch + ): + """FTP writes to ``output`` itself, so the dir is the resume point. + + Wiping it between attempts would restart every RETR from byte 0 and + defeat the REST resume the pooled FTP path depends on. + """ + monkeypatch.setattr("pysus.management.sync.CACHEPATH", tmp_path) + monkeypatch.setattr("pysus.management.sync.asyncio.sleep", AsyncMock()) + seen: list[Path] = [] + + async def _download_once(file, output, ftp_client=None): + seen.append(output) + output.write_bytes(b"partial") + raise ConnectionResetError("drop") + + engine._download_once = AsyncMock(side_effect=_download_once) + + with pytest.raises(RuntimeError, match="after 2 attempts"): + await engine._download_raw_with_retry( + _remote_file(client_name="ftp"), max_retries=2 + ) + + assert len(seen) == 2 + assert seen[0] == seen[1] + assert seen[0].exists() + + def test_is_attempt_dir_recognises_only_hex_dirs(self): + assert SyncEngine._is_attempt_dir("0a1b2c3d") is True + assert SyncEngine._is_attempt_dir("12345678") is True + assert SyncEngine._is_attempt_dir("X.tmp_extract") is False + assert SyncEngine._is_attempt_dir("short") is False + assert SyncEngine._is_attempt_dir("zzzzzzzz") is False + + @pytest.mark.asyncio + async def test_stale_attempt_dirs_are_swept(self, tmp_path, monkeypatch): + monkeypatch.setattr("pysus.management.sync.CACHEPATH", tmp_path) + tmp = tmp_path / "management" / "tmp" + stale = tmp / "deadbeef" + stale.mkdir(parents=True) + (stale / "leftover.csv").write_bytes(b"x") + extract = tmp / "keepme.tmp_extract" + extract.mkdir(parents=True) + (extract / "part.bin").write_bytes(b"y") + + SyncEngine._cleanup_stale_tmp() + + assert not stale.exists() + assert not extract.exists() + assert not list(tmp.iterdir()) + + class TestWeightGate: def test_weight_of_caps_at_budget(self): from pysus.management.sync import _WeightGate