Fix/ftp stale session

pysus/management/sync.py

@@ -211,6 +211,47 @@ async def run(self, coro) -> Any:
         )
 
 
+def _retr_with_resume(
+    ftp: Any,
+    remote_path: str,
+    output: Path,
+    watch: Any,
+) -> int:
+    """Download *remote_path* into *output*, resuming a partial file.
+
+    DATASUS throttles and drops large transfers mid-stream. ``REST`` (the
+    ``rest`` argument of ``retrbinary``) lets a retry continue from the
+    bytes already written instead of re-downloading the whole file. If the
+    server refuses ``REST`` the partial is discarded so the next attempt
+    restarts cleanly.
+    """
+    total = ftp.size(remote_path) or 0
+    offset = output.stat().st_size if output.exists() else 0
+    if total and offset >= total:
+        return total
+    received = 0
+
+    def _write(chunk: bytes) -> int:
+        nonlocal received
+        written = f.write(chunk)
+        received += len(chunk)
+        watch.poke()
+        return written
+
+    try:
+        with open(output, "ab" if offset else "wb") as f:
+            ftp.retrbinary(f"RETR {remote_path}", _write, rest=offset or None)
+    except Exception:  # noqa: BLE001
+        if offset and received == 0:
+            # REST was refused; drop the partial so the retry restarts.
+            try:
+                output.unlink(missing_ok=True)
+            except OSError:
+                pass
+        raise
+    return total
+
+
 class SyncEngine:
     """Orchestrates inventory → compare → download → parquet → catalog."""
 
@@ -571,7 +612,11 @@ async def _download_raw_with_retry(
                 return await self._download_once(file, output, ftp_client)
             except _RETRYABLE as exc:
                 last_error = exc
-                self._cleanup_local(output)
+                # Keep FTP partials so the retry can resume with REST; other
+                # origins are re-downloaded from scratch.
+                origin = getattr(getattr(file, "client", None), "name", None)
+                if origin is None or origin.upper() != "FTP":
+                    self._cleanup_local(output)
                 wait_time = 2**attempt + (attempt * 2)
                 error(
                     f"Download attempt {attempt + 1}/{max_retries} failed "
@@ -656,21 +701,12 @@ async def _download_once(
             watch = _StallWatch(file.basename)
 
             def _retr() -> int:
-                total = ftp.size(remote_path) or 0
-                with open(output, "wb") as f:
-
-                    def _write(chunk: bytes) -> int:
-                        written = f.write(chunk)
-                        watch.poke()
-                        return written
-
-                    ftp.retrbinary(f"RETR {remote_path}", _write)
-                return total
+                return _retr_with_resume(ftp, remote_path, output, watch)
 
             try:
                 with anyio.fail_after(_DOWNLOAD_TIMEOUT):
                     await watch.run(
-                        to_thread.run_sync(_retr, abandon_on_cancel=True)
+                        to_thread.run_sync(_retr, abandon_on_cancel=False)
                     )
                 return output
             except Exception:  # noqa
@@ -690,20 +726,12 @@ def _write(chunk: bytes) -> int:
             watch = _StallWatch(file.basename)
 
             def _direct_retr() -> None:
-                with open(output, "wb") as f:
-
-                    def _write(chunk: bytes) -> int:
-                        written = f.write(chunk)
-                        watch.poke()
-                        return written
-
-                    ftp.retrbinary(f"RETR {remote_path}", _write)
+                _retr_with_resume(ftp, remote_path, output, watch)
 
+            fetch = to_thread.run_sync(_direct_retr, abandon_on_cancel=False)
             try:
                 with anyio.fail_after(_DOWNLOAD_TIMEOUT):
-                    await watch.run(
-                        to_thread.run_sync(_direct_retr, abandon_on_cancel=True)
-                    )
+                    await watch.run(fetch)
             except Exception:  # noqa
                 try:
                     ftp.close()

pysus/tests/management/test_sync.py

@@ -218,6 +218,111 @@ async def connect(self):
         self._ftp = self._next
 
 
+def test_retr_with_resume_uses_rest_offset(tmp_path):
+    from pysus.management.sync import _retr_with_resume, _StallWatch
+
+    out = tmp_path / "f.bin"
+    out.write_bytes(b"x" * 1000)
+
+    class _FTP:
+        def __init__(self):
+            self.rests = []
+
+        def size(self, path):
+            return 5000
+
+        def retrbinary(self, cmd, callback, rest=None):
+            self.rests.append(rest)
+            callback(b"y" * (5000 - rest))
+
+    ftp = _FTP()
+    total = _retr_with_resume(ftp, "/f", out, _StallWatch("f", timeout=60))
+    assert total == 5000
+    assert ftp.rests == [1000]
+    assert out.stat().st_size == 5000
+
+
+def test_retr_with_resume_skips_when_already_complete(tmp_path):
+    from pysus.management.sync import _retr_with_resume, _StallWatch
+
+    out = tmp_path / "f.bin"
+    out.write_bytes(b"x" * 5000)
+
+    class _FTP:
+        def size(self, path):
+            return 5000
+
+        def retrbinary(self, *args, **kwargs):
+            raise AssertionError("must not transfer a complete file")
+
+    assert _retr_with_resume(_FTP(), "/f", out, _StallWatch("f")) == 5000
+
+
+def test_retr_with_resume_drops_partial_when_rest_refused(tmp_path):
+    from pysus.management.sync import _retr_with_resume, _StallWatch
+
+    out = tmp_path / "f.bin"
+    out.write_bytes(b"x" * 1000)
+
+    class _FTP:
+        def size(self, path):
+            return 5000
+
+        def retrbinary(self, cmd, callback, rest=None):
+            raise OSError("550 REST not understood")
+
+    with pytest.raises(OSError):
+        _retr_with_resume(_FTP(), "/f", out, _StallWatch("f", timeout=60))
+    assert not out.exists()
+
+
+@pytest.mark.asyncio
+async def test_download_retry_keeps_ftp_partial(tmp_path, monkeypatch):
+    from pysus.management import sync as sync_mod
+
+    monkeypatch.setattr(sync_mod, "CACHEPATH", tmp_path)
+    engine = SyncEngine()
+    file = MagicMock()
+    file.basename = "x.dbc"
+    file.client.name = "FTP"
+    seen = []
+
+    async def fake_once(file, output, ftp_client=None):
+        seen.append(output.exists())
+        if len(seen) == 1:
+            output.write_bytes(b"1234")
+            raise ConnectionResetError("boom")
+        return output
+
+    monkeypatch.setattr(engine, "_download_once", fake_once)
+    result = await engine._download_raw_with_retry(file, max_retries=2)
+    assert seen == [False, True]
+    assert result.exists()
+
+
+@pytest.mark.asyncio
+async def test_download_retry_clears_non_ftp_partial(tmp_path, monkeypatch):
+    from pysus.management import sync as sync_mod
+
+    monkeypatch.setattr(sync_mod, "CACHEPATH", tmp_path)
+    engine = SyncEngine()
+    file = MagicMock()
+    file.basename = "x.csv.zip"
+    file.client.name = "DADOSGOV"
+    seen = []
+
+    async def fake_once(file, output, ftp_client=None):
+        seen.append(output.exists())
+        if len(seen) == 1:
+            output.write_bytes(b"1234")
+            raise ConnectionResetError("boom")
+        return output
+
+    monkeypatch.setattr(engine, "_download_once", fake_once)
+    await engine._download_raw_with_retry(file, max_retries=2)
+    assert seen == [False, False]
+
+
 class TestLiveFTP:
     @pytest.mark.asyncio
     async def test_connects_when_no_session(self):

pysus/tests/management/test_sync_internal.py

@@ -90,7 +90,7 @@ async def test_download_once_ftp_pooled(self, engine, tmp_path):
         client = MagicMock()
         client.ftp = ftp
 
-        def _retrbinary(cmd, cb):
+        def _retrbinary(cmd, cb, rest=None):
             cb(b"hello")
 
         ftp.retrbinary = _retrbinary

pysus/tests/management/test_sync_more.py

@@ -238,8 +238,9 @@ async def test_rejects_oversized_file(self, engine, tmp_path):
     @pytest.mark.asyncio
     async def test_direct_ftp_retr(self, engine, tmp_path):
         ftp = MagicMock()
+        ftp.size.return_value = 5
 
-        def _retrbinary(cmd, cb):
+        def _retrbinary(cmd, cb, rest=None):
             cb(b"hello")
 
         ftp.retrbinary = _retrbinary