From e1bd441c2fb00d141286c7bf0f5f843552a7eb3c Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Lu=C3=A3=20Bida=20Vacaro?= Date: Thu, 1 Oct 2026 15:47:13 -0300 Subject: [PATCH 1/4] test(ftp): cover stale-session liveness and reset paths From f6ac6e5f24668d63b6eb42e956d49ae43d7a6ba7 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Lu=C3=A3=20Bida=20Vacaro?= Date: Thu, 1 Oct 2026 17:10:40 -0300 Subject: [PATCH 2/4] fix(ftp): resume interrupted downloads with REST Large DATASUS transfers are throttled and dropped mid-stream; restarting them from zero on every retry meant they could never finish. Resume from the partial file via REST (retrbinary rest=), keep FTP partials between retries, and stop abandoning the worker thread so a stalled transfer is joined before the session is reset. --- pysus/management/sync.py | 74 ++++++++++++++++++++--------- pysus/tests/management/test_sync.py | 58 ++++++++++++++++++++++ 2 files changed, 109 insertions(+), 23 deletions(-) diff --git a/pysus/management/sync.py b/pysus/management/sync.py index 945bb9da..5841a343 100644 --- a/pysus/management/sync.py +++ b/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 != "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() diff --git a/pysus/tests/management/test_sync.py b/pysus/tests/management/test_sync.py index 128e77c7..718b2cd9 100644 --- a/pysus/tests/management/test_sync.py +++ b/pysus/tests/management/test_sync.py @@ -218,6 +218,64 @@ 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() + + class TestLiveFTP: @pytest.mark.asyncio async def test_connects_when_no_session(self): From d0b595c77f2374d9ac80e8c6f106f9c9631486f3 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Lu=C3=A3=20Bida=20Vacaro?= Date: Thu, 1 Oct 2026 19:05:15 -0300 Subject: [PATCH 3/4] fix(ftp): keep FTP partials across retries (origin is uppercase) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The resume guard compared the client origin to lowercase "ftp", but the FTP client reports "FTP", so every retry deleted the partial and the download restarted from zero — defeats REST resume for large files. --- pysus/management/sync.py | 2 +- pysus/tests/management/test_sync.py | 47 +++++++++++++++++++++++++++++ 2 files changed, 48 insertions(+), 1 deletion(-) diff --git a/pysus/management/sync.py b/pysus/management/sync.py index 5841a343..fa84db25 100644 --- a/pysus/management/sync.py +++ b/pysus/management/sync.py @@ -615,7 +615,7 @@ async def _download_raw_with_retry( # 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 != "ftp": + if origin is None or origin.upper() != "FTP": self._cleanup_local(output) wait_time = 2**attempt + (attempt * 2) error( diff --git a/pysus/tests/management/test_sync.py b/pysus/tests/management/test_sync.py index 718b2cd9..c72a7bc3 100644 --- a/pysus/tests/management/test_sync.py +++ b/pysus/tests/management/test_sync.py @@ -276,6 +276,53 @@ def retrbinary(self, cmd, callback, rest=None): 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): From 61d11c299c2860c3b5924a98f5b021893c1f7c34 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Lu=C3=A3=20Bida=20Vacaro?= Date: Fri, 2 Oct 2026 12:57:56 -0300 Subject: [PATCH 4/4] test(ftp): accept rest kwarg and size() in resume-path mocks _retr_with_resume now calls retrbinary(..., rest=...) and reads ftp.size(); update the two mocks that predate it. --- pysus/tests/management/test_sync_internal.py | 2 +- pysus/tests/management/test_sync_more.py | 3 ++- 2 files changed, 3 insertions(+), 2 deletions(-) diff --git a/pysus/tests/management/test_sync_internal.py b/pysus/tests/management/test_sync_internal.py index 2a913c65..7f31bcf8 100644 --- a/pysus/tests/management/test_sync_internal.py +++ b/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 diff --git a/pysus/tests/management/test_sync_more.py b/pysus/tests/management/test_sync_more.py index ed4f78ce..70df5ad9 100644 --- a/pysus/tests/management/test_sync_more.py +++ b/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