diff --git a/pysus/management/sync.py b/pysus/management/sync.py index 945bb9da..fa84db25 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 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() diff --git a/pysus/tests/management/test_sync.py b/pysus/tests/management/test_sync.py index 128e77c7..c72a7bc3 100644 --- a/pysus/tests/management/test_sync.py +++ b/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): 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