fix(saude): never reuse or publish a partial CKAN download

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
 
 

pysus/management/sync.py

@@ -604,19 +604,39 @@ 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):
             try:
                 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

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(