From 6784e1dd800542fec47b71179e4589d4d6d0da92 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Lu=C3=A3=20Bida=20Vacaro?= Date: Fri, 28 Aug 2026 18:17:08 -0300 Subject: [PATCH] fix: fix readthedocs build --- poetry.lock | 15 +- pyproject.toml | 1 + pysus/cli/management.py | 100 ++++++--- pysus/data/dbf_reader.py | 48 ++-- pysus/management/records.py | 92 +++++++- pysus/management/sync.py | 220 +++++++++++++++++-- pysus/tests/cli/test_management.py | 93 ++------ pysus/tests/management/test_records.py | 113 ++++++++++ pysus/tests/management/test_sync_internal.py | 77 ++++++- 9 files changed, 606 insertions(+), 153 deletions(-) diff --git a/poetry.lock b/poetry.lock index 922c58e2..e6c6898f 100644 --- a/poetry.lock +++ b/poetry.lock @@ -748,7 +748,7 @@ files = [ {file = "colorama-0.4.6-py2.py3-none-any.whl", hash = "sha256:4f1d9991f5acc0ca119f9d443620b77f9d6b33703e51011c16baf57afb285fc6"}, {file = "colorama-0.4.6.tar.gz", hash = "sha256:08695f5cb7ed6e0531a20572697297273c47b8cae5a63ffc6d6ed5c201be6e44"}, ] -markers = {main = "sys_platform == \"win32\" or platform_system == \"Windows\"", dev = "platform_system == \"Windows\" or sys_platform == \"win32\"", docs = "sys_platform == \"win32\""} +markers = {main = "platform_system == \"Windows\" or sys_platform == \"win32\"", dev = "platform_system == \"Windows\" or sys_platform == \"win32\"", docs = "sys_platform == \"win32\""} [[package]] name = "comm" @@ -1635,9 +1635,10 @@ files = [ name = "gitdb" version = "4.0.12" description = "Git Object Database" -optional = false +optional = true python-versions = ">=3.7" groups = ["main"] +markers = "extra == \"web\"" files = [ {file = "gitdb-4.0.12-py3-none-any.whl", hash = "sha256:67073e15955400952c6565cc3e707c554a4eea2e428946f7a4c162fab9bd9bcf"}, {file = "gitdb-4.0.12.tar.gz", hash = "sha256:5ef71f855d191a3326fcfbc0d5da835f26b13fbcba60c32c21091c349ffdb571"}, @@ -1650,9 +1651,10 @@ smmap = ">=3.0.1,<6" name = "gitpython" version = "3.1.58" description = "GitPython is a Python library used to interact with Git repositories" -optional = false +optional = true python-versions = ">=3.7" groups = ["main"] +markers = "extra == \"web\"" files = [ {file = "gitpython-3.1.58-py3-none-any.whl", hash = "sha256:d331e722577f0fd7fc1f857419b3ecc07af66282b933d2a4d95f84a042fdd50f"}, {file = "gitpython-3.1.58.tar.gz", hash = "sha256:621416df10ef3fd0e19fabf9172ddeed0fa704d353d04f194eec56a625a95b22"}, @@ -4160,7 +4162,7 @@ version = "6.0.3" description = "YAML parser and emitter for Python" optional = false python-versions = ">=3.8" -groups = ["dev", "docs"] +groups = ["main", "dev", "docs"] files = [ {file = "PyYAML-6.0.3-cp38-cp38-macosx_10_13_x86_64.whl", hash = "sha256:c2514fceb77bc5e7a2f7adfaa1feb2fb311607c9cb518dbc378688ec73d8292f"}, {file = "PyYAML-6.0.3-cp38-cp38-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:9c57bb8c96f6d1808c030b1687b9b5fb476abaa47f0db9c0101f5e9f394e97f4"}, @@ -4802,9 +4804,10 @@ files = [ name = "smmap" version = "5.0.3" description = "A pure Python implementation of a sliding window memory map manager" -optional = false +optional = true python-versions = ">=3.7" groups = ["main"] +markers = "extra == \"web\"" files = [ {file = "smmap-5.0.3-py3-none-any.whl", hash = "sha256:c106e05d5a61449cf6ba9a1e650227ecfb141590d2a98412103ff35d89fc7b2f"}, {file = "smmap-5.0.3.tar.gz", hash = "sha256:4d9debb8b99007ae47165abc08670bd74cb74b5227dda7f643eccc4e9eb5642c"}, @@ -5799,4 +5802,4 @@ web = ["streamlit"] [metadata] lock-version = "2.1" python-versions = ">=3.11,<3.14" -content-hash = "0f4729b9d1190e478e60c86ffde6950adbcf0758074466ed27c00a874195eecb" +content-hash = "e977e2830aeca92c8938c314aa38478261d87d1ab172aa480b87389ab20add4b" diff --git a/pyproject.toml b/pyproject.toml index 166b513e..88cc02fb 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -37,6 +37,7 @@ aioftp = "^0.21.4" dbfread = "2.0.7" bigtree = "^0.12.2" pyreaddbc = ">=2.0.4" +PyYAML = "^6.0.1" dotenv = "^0.9.9" boto3 = "^1.42.89" typer = "^0.24.1" diff --git a/pysus/cli/management.py b/pysus/cli/management.py index 5f7e2d9c..b4eb1f19 100644 --- a/pysus/cli/management.py +++ b/pysus/cli/management.py @@ -114,18 +114,16 @@ def check( help="Path to a file with ACCESS_KEY/SECRET_KEY/DADOSGOV_TOKEN", ), ): - """Check every source against the S3 databases. - - By default this is a dry run: it only reports which files would need - to be updated/uploaded (``needs_update``) and which are already at - the most updated version (``skipped``), without touching S3. Pass - ``--apply`` to actually download, convert, upload and catalog the - outdated files. Use ``--json`` to stream machine-readable results. - - A run is resumable: each completed file is appended to a journal - (``--resume``, or derived from ``--reupload-before``), and a paused - run can be resumed with the same command to skip already-processed - files. + """Check a database against its FTP origin to see if it needs updating. + + By default this is a dry run: it classifies every mirrored file as + ``missing`` (not on S3), ``outdated`` (the FTP file is more recent or + has a different size) or ``current``, and prints a per-database table + with a "needs update" / "up to date" verdict — without touching S3. + + Pass ``--apply`` to actually download, convert, upload and catalog the + outstanding files (the full sync pipeline). Pass ``--json`` to stream + machine-readable results. """ from pysus.api.client import _run_sync from pysus.management.sync import SyncEngine @@ -137,9 +135,67 @@ def check( dadosgov_token=env.get("DADOSGOV_TOKEN"), ) + datasets = [d.upper() for d in name] if name else None + + def _flush() -> None: + import sys + + sys.stdout.flush() + sys.stderr.flush() + + def _print_check_report(checks) -> None: + total_missing = sum(len(c.missing) for c in checks.values()) + total_outdated = sum(len(c.outdated) for c in checks.values()) + total_current = sum(len(c.current) for c in checks.values()) + + typer.echo( + f"{'DATABASE':<30}{'MISSING':>9}{'OUTDATED':>10}" + f"{'CURRENT':>9} STATUS" + ) + typer.echo("-" * 78) + for ds in sorted(checks): + c = checks[ds] + verdict = "needs update" if c.needs_update else "up to date" + typer.echo( + f"{ds:<30}{len(c.missing):>9}{len(c.outdated):>10}" + f"{len(c.current):>9} {verdict}" + ) + typer.echo("-" * 78) + typer.echo( + f"{'TOTAL':<30}{total_missing:>9}{total_outdated:>10}" + f"{total_current:>9}" + ) + if total_missing or total_outdated: + typer.echo( + f"\n{total_missing + total_outdated} of " + f"{total_missing + total_outdated + total_current} file(s) " + "need updating — re-run with --apply to mirror them." + ) + + def _print_check_json(checks) -> None: + for ds in sorted(checks): + typer.echo(json.dumps({"dataset": ds, **checks[ds].summary()})) + + async def _run_check() -> None: + await engine.__aenter__(lock=False) + try: + checks = await engine.check(datasets=datasets) + finally: + await engine.__aexit__(None, None, None) + _flush() + if json_out: + _print_check_json(checks) + else: + _print_check_report(checks) + _flush() + + if not apply: + _run_sync(_run_check()) + return + journal = _journal_path(resume, reupload_before) resume_keys = set() - if apply and journal is not None and journal.exists(): + if journal is not None and journal.exists(): resume_keys = load_journal_keys(journal) counts: dict[str, int] = {} @@ -172,25 +228,19 @@ def on_outcome(outcome) -> None: if total % 500 == 0: typer.echo(f"progress: {counts}", err=True) - def _flush() -> None: - import sys - - sys.stdout.flush() - sys.stderr.flush() - - async def _run(): + async def _run() -> dict[str, int]: async with engine: report = await engine.run( - datasets=[d.upper() for d in name] if name else None, + datasets=datasets, force=force, reupload_before=_parse_date(reupload_before), - dry_run=not apply, + dry_run=False, workers=workers, ftp_connections=ftp_connections, - checkpoint_every=checkpoint_every if apply else None, + checkpoint_every=checkpoint_every, on_outcome=on_outcome, resume=resume_keys or None, - journal=journal if apply else None, + journal=journal, ) summary = report.summary() _flush() @@ -211,5 +261,5 @@ async def _run(): return summary summary = _run_sync(_run()) - if summary["failed"]: + if summary and summary["failed"]: raise typer.Exit(code=1) diff --git a/pysus/data/dbf_reader.py b/pysus/data/dbf_reader.py index 4730f48e..832bd7a8 100644 --- a/pysus/data/dbf_reader.py +++ b/pysus/data/dbf_reader.py @@ -256,31 +256,29 @@ def stream_dbf_fast( dtype = schema.build_dtype() with open(path, "rb") as fh: - fh.seek(schema.header_len) - raw = fh.read(n * rl) - - for start in range(0, n, chunk_size): - end = min(start + chunk_size, n) - chunk_n = end - start - chunk_raw = raw[start * rl : end * rl] - - records: np.ndarray = np.frombuffer( - chunk_raw, dtype=dtype, count=chunk_n - ) - records = records[records["_deleted"] != b"*"] # skip deleted rows - chunk_n = len(records) - - data = {} - for fld in schema.fields: - col: np.ndarray = records[fld.name] - decoded = np.empty(chunk_n, dtype=object) - for i in range(chunk_n): - val = col[i] - b = val if isinstance(val, bytes) else val.tobytes() - decoded[i] = _decode(b) - data[fld.name] = decoded - - yield pd.DataFrame(data) + for start in range(0, n, chunk_size): + end = min(start + chunk_size, n) + chunk_n = end - start + fh.seek(schema.header_len + start * rl) + chunk_raw = fh.read(chunk_n * rl) + + records: np.ndarray = np.frombuffer( + chunk_raw, dtype=dtype, count=chunk_n + ) + records = records[records["_deleted"] != b"*"] # skip deleted rows + chunk_n = len(records) + + data = {} + for fld in schema.fields: + col: np.ndarray = records[fld.name] + decoded = np.empty(chunk_n, dtype=object) + for i in range(chunk_n): + val = col[i] + b = val if isinstance(val, bytes) else val.tobytes() + decoded[i] = _decode(b) + data[fld.name] = decoded + + yield pd.DataFrame(data) def _find_field(schema: DBFSchema, name: str) -> DBFField: diff --git a/pysus/management/records.py b/pysus/management/records.py index 2a78fda9..6dc20165 100644 --- a/pysus/management/records.py +++ b/pysus/management/records.py @@ -335,6 +335,87 @@ def to_dict(self) -> dict[str, Any]: } +#: Freshness classification for a mirrored file vs its origin. +FRESH_MISSING = "missing" # absent from the S3 mirror +FRESH_OUTDATED = "outdated" # present but the origin file is more recent +FRESH_CURRENT = "current" # present and up to date + + +def freshness_status(comparison: FileComparison) -> tuple[str, str]: + """Classify a logical file against its mirrored S3 (ducklake) copy. + + A file that exists on an origin (FTP/DadosGov/Saude) but has no S3 + artifact is ``missing``. When an S3 artifact exists, the mirror is + ``outdated`` if any origin record is *newer* (its ``modified`` origin + date is later than the recorded ``source_modified``) **or** has a + *different* ``size`` than the recorded ``source_size``. Otherwise it + is ``current``. + + Returns a ``(status, reason)`` tuple where *status* is one of + :data:`FRESH_MISSING`, :data:`FRESH_OUTDATED` or :data:`FRESH_CURRENT`. + """ + s3 = comparison._pick("ducklake") + if s3 is None: + return FRESH_MISSING, "no mirror artifact in the S3 catalog" + + reasons: list[str] = [] + for record in comparison.records: + if record.origin == "ducklake": + continue + # 1) origin modification date ("origin date") newer than mirrored. + if ( + record.modified is not None + and s3.source_modified is not None + and record.modified > s3.source_modified + ): + reasons.append( + f"{record.origin} modified {record.modified:%Y-%m-%d} is " + f"newer than the mirrored {s3.source_modified:%Y-%m-%d}" + ) + # 2) origin size differs from the mirrored source size. + if record.size and s3.source_size and record.size != s3.source_size: + reasons.append( + f"{record.origin} size {record.size} differs from the " + f"mirrored {s3.source_size}" + ) + + if reasons: + return FRESH_OUTDATED, "; ".join(reasons) + return FRESH_CURRENT, "mirror is up to date" + + +@dataclass +class DatabaseCheck: + """Aggregated freshness check for one database.""" + + dataset: str + missing: list[str] = field(default_factory=list) + outdated: list[str] = field(default_factory=list) + current: list[str] = field(default_factory=list) + + @property + def needs_update(self) -> bool: + """True when any file is missing or outdated in this database.""" + return bool(self.missing or self.outdated) + + def add(self, status: str, label: str, reason: str = "") -> None: + bucket = { + FRESH_MISSING: self.missing, + FRESH_OUTDATED: self.outdated, + FRESH_CURRENT: self.current, + }.get(status) + if bucket is not None: + bucket.append(label if not reason else f"{label} — {reason}") + + def summary(self) -> dict[str, int]: + return { + "missing": len(self.missing), + "outdated": len(self.outdated), + "current": len(self.current), + "needs_update": self.needs_update, + } + + @dataclass class SnapshotDiff: """Difference between two snapshots of the same origin.""" @@ -432,10 +513,13 @@ def write_journal_line(path: Path, outcome: SyncOutcome) -> None: def load_journal_keys(path: Path) -> set[IdentityKey]: - """Return identity keys already processed in a prior run. + """Return identity keys already transferred in a prior run. - Both ``uploaded`` and ``failed`` entries are included so a resumed run - skips files that were already transferred or could not be downloaded. + Only ``uploaded`` entries are included so a resumed run skips files + that were transferred and never re-downloads them. ``failed`` entries + are deliberately excluded: transient failures (e.g. a throttled FTP + server or a dropped connection) must be retried on the next run, not + dropped permanently. """ keys: set[IdentityKey] = set() if not path.exists(): @@ -448,7 +532,7 @@ def load_journal_keys(path: Path) -> set[IdentityKey]: data = json.loads(line) except json.JSONDecodeError: continue - if data.get("status") not in ("uploaded", "failed"): + if data.get("status") != "uploaded": continue try: keys.add( diff --git a/pysus/management/sync.py b/pysus/management/sync.py index 8ef0a8aa..0d73d7e9 100644 --- a/pysus/management/sync.py +++ b/pysus/management/sync.py @@ -19,9 +19,10 @@ from __future__ import annotations import asyncio +import shutil from collections.abc import Callable from datetime import datetime -from logging import error, warning +from logging import error, info, warning from pathlib import Path from typing import TYPE_CHECKING, Any from uuid import uuid4 @@ -38,12 +39,14 @@ from .inventory import Inventory from .records import ( DOWNLOAD_PRIORITY, + DatabaseCheck, FileComparison, FileRecord, IdentityKey, SyncOutcome, SyncReport, compose_s3_key, + freshness_status, write_journal_line, ) @@ -74,6 +77,64 @@ #: from stalling the pipeline. Set to 0 to disable the limit. _MAX_FILE_SIZE: int = 500 * 1024 * 1024 # 500 MB +#: Minimum per-file in-flight "weight" (bytes). Files count for at least +#: this much towards the concurrency budget, so a flood of tiny files can +#: never spawn an unbounded number of simultaneous conversions. +_SMALL_FILE_BYTES: int = 64 * 1024 * 1024 # 64 MB + +#: Total in-flight conversion/download weight budget (bytes). Concurrency +#: shrinks as queued files grow: dozens of small files run in parallel, +#: while a handful of multi-GB DATASUS archives are processed almost one +#: at a time. Keeps peak RAM/disk usage bounded on low-resource hosts. +_CONCURRENCY_BUDGET: int = 2500 * 1024 * 1024 # 2.5 GB + + +class _WeightGate: + """Bound how many pipeline slots are in use by byte weight. + + ``acquire(weight)`` reserves a slot only while enough *budget* is + free; ``release(weight)`` returns it. Small files (weight capped at + ``_SMALL_FILE_BYTES``) fit many at once, large files consume most or + all of the budget and therefore serialize the pipeline. + """ + + def __init__(self, budget: int = _CONCURRENCY_BUDGET) -> None: + self._budget = budget + self._used = 0 + self._cond = asyncio.Condition() + + @staticmethod + def weight_of(size: object, budget: int = _CONCURRENCY_BUDGET) -> int: + """Return the gate weight for a file of ``size`` bytes. + + Weights are capped at the budget so an oversized file (e.g. a + multi-GB vaccine CSV well above 2.5 GB) can still acquire alone + instead of deadlocking in ``adjust``. + """ + if isinstance(size, (int, float)) and size > 0: + return min(max(int(size), _SMALL_FILE_BYTES), budget) + return min(_SMALL_FILE_BYTES, budget) + + async def acquire(self, weight: int) -> int: + async with self._cond: + while self._used + weight > self._budget: + await self._cond.wait() + self._used += weight + return weight + + async def release(self, weight: int) -> None: + async with self._cond: + self._used -= weight + self._cond.notify_all() + + async def adjust(self, old_weight: int, new_weight: int) -> int: + """Swap ``old_weight`` for ``new_weight`` in the current usage.""" + if new_weight == old_weight: + return new_weight + await self.release(old_weight) + await self.acquire(new_weight) + return new_weight + class SyncEngine: """Orchestrates inventory → compare → download → parquet → catalog.""" @@ -120,7 +181,7 @@ def writer(self) -> CatalogWriter: # ------------------------------------------------------------------ # lifecycle # ------------------------------------------------------------------ - async def __aenter__(self) -> SyncEngine: + async def __aenter__(self, lock: bool = True) -> SyncEngine: if self.pysus is None: from pysus.api.client import PySUS @@ -134,7 +195,8 @@ async def __aenter__(self) -> SyncEngine: access_key=self.access_key, secret_key=self.secret_key, ) - self._acquire_sync_lock() + if lock: + self._acquire_sync_lock() return self def _acquire_sync_lock(self) -> None: @@ -421,8 +483,7 @@ async def _download_raw_with_retry( last_error: Exception | None = None for attempt in range(max_retries): try: - await self._download_once(file, output, ftp_client) - return output + return await self._download_once(file, output, ftp_client) except _RETRYABLE as exc: last_error = exc self._cleanup_local(output) @@ -443,8 +504,14 @@ async def _download_once( file: BaseRemoteFile, output: Path, ftp_client: Any | None = None, - ) -> None: - """Perform one raw download to *output*.""" + ) -> Path: + """Perform one raw download to *output*. + + Returns the local path that actually received the bytes. Some + origins (e.g. CKAN/Saude resources) ignore the requested filename + and write to their own derived name in ``output.parent``, so the + caller must use the returned path for conversion and cleanup. + """ from anyio import to_thread if ( @@ -480,7 +547,7 @@ def _retr(): try: with anyio.fail_after(_DOWNLOAD_TIMEOUT): await to_thread.run_sync(_retr) - return + return output except Exception: # noqa try: ftp.quit() @@ -501,8 +568,8 @@ def _direct_retr(): with anyio.fail_after(_DOWNLOAD_TIMEOUT): await to_thread.run_sync(_direct_retr) - return - await file._download(output=output) + return output + return await file._download(output=output) @staticmethod def _cleanup_local(path: Path) -> None: @@ -511,6 +578,34 @@ def _cleanup_local(path: Path) -> None: except OSError: pass + @staticmethod + def _cleanup_stale_tmp() -> None: + """Remove leftover files in the tmp directory from prior runs. + + Parquets, DBCs, DBFs, CSVs and ``*.tmp_extract`` extraction + directories left behind by crashed (e.g. OOM-killed) processes are + deleted before a new run starts. + """ + tmp = Path(CACHEPATH) / "management" / "tmp" + if not tmp.is_dir(): + return + removed = 0 + for p in tmp.iterdir(): + if p.is_dir() and p.name.endswith(".tmp_extract"): + try: + shutil.rmtree(p) + removed += 1 + except OSError: + pass + elif p.is_file(): + try: + p.unlink() + removed += 1 + except OSError: + pass + if removed: + info(f"Cleaned {removed} stale files from {tmp}") + @staticmethod def _is_current( file: BaseRemoteFile, @@ -550,8 +645,8 @@ async def run( save_snapshots: bool = True, checkpoint_every: int | None = None, on_outcome: Callable[[SyncOutcome], None] | None = None, - workers: int = 4, - ftp_connections: int = 4, + workers: int = 16, + ftp_connections: int = 8, origins: tuple[str, ...] | None = None, reupload_before: datetime | None = None, resume: set[IdentityKey] | None = None, @@ -592,6 +687,9 @@ async def run( report = SyncReport(dataset=",".join(datasets) if datasets else None) active_origins = origins or ("ducklake", "ftp", "dadosgov", "saude") + if not dry_run: + self._cleanup_stale_tmp() + def emit(outcome: SyncOutcome) -> None: """Record an outcome: report + callback + resume journal.""" report.outcomes.append(outcome) @@ -704,9 +802,23 @@ async def collect_with_retry(origin: str, datasets=None, **kwargs): continue parallel.append((comparison, record)) + # Process smallest files first: tiny artifacts are cheap to + # convert and upload, so they drain fast with full parallelism, + # and the multi-GB DATASUS archives are left for last when the + # weight gate naturally serializes them. + parallel.sort( + key=lambda cf: ( + cf[1].file.size + if isinstance(cf[1].file.size, (int, float)) + else float("inf") + ) + ) + ftp_items = [(c, r) for c, r in parallel if r.origin == "ftp"] gov_items = [(c, r) for c, r in parallel if r.origin != "ftp"] + gate = _WeightGate() + ftp_pool: list[Any] = [] if ftp_items: from pysus.api.ftp.client import FTP @@ -723,23 +835,36 @@ async def ftp_downloader( client: Any, items: list[tuple[FileComparison, FileRecord]] ) -> None: for comparison, record in items: + weight = gate.weight_of(record.file.size) + try: + await gate.acquire(weight) + except Exception: # noqa: BLE001 — budget is advisory + weight = 0 try: raw = await self._download_raw_with_retry( record.file, ftp_client=client ) - await raw_queue.put((comparison, record, raw, None)) + if weight: + try: + weight = await gate.adjust( + weight, gate.weight_of(raw.stat().st_size) + ) + except Exception: # noqa: BLE001 — keep old weight + pass + await raw_queue.put((comparison, record, raw, None, weight)) except Exception as exc: # noqa - await raw_queue.put((comparison, record, None, str(exc))) - await raw_queue.put(None) + if weight: + await gate.release(weight) + await raw_queue.put((comparison, record, None, str(exc), 0)) async def raw_processor() -> None: while True: entry = await raw_queue.get() + if entry is None: + return + comparison, record, raw, err, weight = entry raw_path = None try: - if entry is None: - return - comparison, record, raw, err = entry raw_path = raw if err is not None: await write_queue.put((comparison, record, None, err)) @@ -749,9 +874,11 @@ async def raw_processor() -> None: ) await write_queue.put((comparison, record, payload, None)) except Exception as exc: # noqa - comparison, record, _, _ = entry + comparison, record, _, _, _ = entry await write_queue.put((comparison, record, None, str(exc))) finally: + if weight: + await gate.release(weight) if raw_path: self._cleanup_local(raw_path) raw_queue.task_done() @@ -759,10 +886,22 @@ async def raw_processor() -> None: async def gov_worker() -> None: while gov_items: comparison, record = gov_items.pop() + weight = gate.weight_of(record.file.size) + try: + await gate.acquire(weight) + except Exception: # noqa: BLE001 — budget is advisory + weight = 0 raw_path = None try: raw = await self._download_raw_with_retry(record.file) raw_path = raw + if weight: + try: + weight = await gate.adjust( + weight, gate.weight_of(raw.stat().st_size) + ) + except Exception: # noqa: BLE001 — keep old weight + pass payload = await self._convert_and_upload( record.file, raw, callback=callback ) @@ -770,6 +909,8 @@ async def gov_worker() -> None: except Exception as exc: # noqa await write_queue.put((comparison, record, None, str(exc))) finally: + if weight: + await gate.release(weight) if raw_path: self._cleanup_local(raw_path) @@ -907,6 +1048,47 @@ async def catalog_writer() -> None: return report + async def check( + self, + datasets: list[str] | None = None, + ) -> dict[str, DatabaseCheck]: + """Check which mirrored files need updating, per database. + + Lists the origin(s) and the DuckLake catalog and classifies every + logical file via :func:`~pysus.management.records.freshness_status`: + + * ``missing`` — present on the origin but absent from the S3 catalog; + * ``outdated`` — mirrored, but the origin file is more recent + (newer ``origin date``) or has a different size; + * ``current`` — mirrored and up to date. + + No downloads, uploads or catalog writes are performed. Returns a + mapping of ``{dataset: DatabaseCheck}``. A database ``needs_update`` + when it has any missing or outdated files. + """ + if isinstance(datasets, str): + datasets = [datasets.upper()] + + records: list[FileRecord] = await self.inventory.collect( + "ducklake", datasets + ) + records += await self.inventory.collect("ftp", datasets) + if self.dadosgov_token: + records += await self.inventory.collect( + "dadosgov", datasets, dadosgov_token=self.dadosgov_token + ) + + comparisons = self.comparator.compare(records) + + checks: dict[str, DatabaseCheck] = {} + for comparison in comparisons: + ds = checks.setdefault( + comparison.key.dataset, DatabaseCheck(comparison.key.dataset) + ) + status, reason = freshness_status(comparison) + ds.add(status, self._label(comparison), reason) + return checks + async def _convert_and_upload( self, file: BaseRemoteFile, diff --git a/pysus/tests/cli/test_management.py b/pysus/tests/cli/test_management.py index 07ab7c01..60c80041 100644 --- a/pysus/tests/cli/test_management.py +++ b/pysus/tests/cli/test_management.py @@ -13,6 +13,7 @@ def _make_engine(summary: dict): engine.__aenter__ = AsyncMock(return_value=engine) engine.__aexit__ = AsyncMock(return_value=None) engine.run = AsyncMock(return_value=MagicMock(summary=lambda: summary)) + engine.check = AsyncMock(return_value={}) return engine @@ -33,11 +34,9 @@ def test_check_all_databases(self, mock_env): with patch("pysus.management.sync.SyncEngine", return_value=engine): result = runner.invoke(app, ["management", "check"]) assert result.exit_code == 0 - assert "needs_update: 1" in result.output - assert "skipped: 2" in result.output - _, kwargs = engine.run.call_args - assert kwargs["datasets"] is None - assert kwargs["dry_run"] is True + assert "DATABASE" in result.output + assert "TOTAL" in result.output + engine.check.assert_awaited_once_with(datasets=None) @patch("pysus.cli.management._load_env") def test_check_specific_databases(self, mock_env): @@ -57,8 +56,9 @@ def test_check_specific_databases(self, mock_env): app, ["management", "check", "SINAN", "SINASC", "SIM"] ) assert result.exit_code == 0 - _, kwargs = engine.run.call_args - assert kwargs["datasets"] == ["SINAN", "SINASC", "SIM"] + engine.check.assert_awaited_once_with( + datasets=["SINAN", "SINASC", "SIM"] + ) @patch("pysus.cli.management._load_env") def test_check_apply_disables_dry_run(self, mock_env): @@ -83,64 +83,24 @@ def test_check_apply_disables_dry_run(self, mock_env): assert kwargs["checkpoint_every"] == 500 @patch("pysus.cli.management._load_env") - def test_check_json_streams_outcomes(self, mock_env): + def test_check_json_streams_freshness(self, mock_env): mock_env.return_value = {"ACCESS_KEY": "ak", "SECRET_KEY": "sk"} - engine = MagicMock() - engine.__aenter__ = AsyncMock(return_value=engine) - engine.__aexit__ = AsyncMock(return_value=None) - engine.run = AsyncMock( - return_value=MagicMock( - summary=lambda: { - "total": 1, - "needs_update": 1, - "uploaded": 0, - "skipped": 0, - "failed": 0, - "needs_token": 0, - } - ) - ) - - from pysus.management.records import IdentityKey, SyncOutcome - - key = IdentityKey( - dataset="SINAN", - group="DENG", - year=2025, - month=None, - state=None, - stem="dengbr25", - ) + from pysus.management.records import DatabaseCheck - async def fake_run(**kwargs): - kwargs["on_outcome"]( - SyncOutcome( - key=key, - origin="ftp", - status="needs_update", - detail="SINAN/DENG/2025/-/dengbr25 (ftp)", - ) - ) - return MagicMock( - summary=lambda: { - "total": 1, - "needs_update": 1, - "uploaded": 0, - "skipped": 0, - "failed": 0, - "needs_token": 0, - } - ) + chk = DatabaseCheck(dataset="SINAN") + chk.add("missing", "SINAN/DENG/-/DENGBR20") + engine = _make_engine({}) + engine.check = AsyncMock(return_value={"SINAN": chk}) - engine.run = fake_run with patch("pysus.management.sync.SyncEngine", return_value=engine): result = runner.invoke(app, ["management", "check", "--json"]) assert result.exit_code == 0 - assert '"status": "needs_update"' in result.output assert '"dataset": "SINAN"' in result.output + assert '"missing": 1' in result.output + assert '"needs_update": true' in result.output @patch("pysus.cli.management._load_env") - def test_check_exit_code_on_failure(self, mock_env): + def test_check_apply_exit_code_on_failure(self, mock_env): mock_env.return_value = {"ACCESS_KEY": "ak", "SECRET_KEY": "sk"} engine = _make_engine( { @@ -153,24 +113,15 @@ def test_check_exit_code_on_failure(self, mock_env): } ) with patch("pysus.management.sync.SyncEngine", return_value=engine): - result = runner.invoke(app, ["management", "check", "SINAN"]) + result = runner.invoke( + app, ["management", "check", "SINAN", "--apply"] + ) assert result.exit_code == 1 @patch("pysus.cli.management._load_env") def test_check_reupload_before(self, mock_env): - from datetime import datetime - mock_env.return_value = {"ACCESS_KEY": "ak", "SECRET_KEY": "sk"} - engine = _make_engine( - { - "total": 1, - "needs_update": 1, - "uploaded": 0, - "skipped": 0, - "failed": 0, - "needs_token": 0, - } - ) + engine = _make_engine({}) with patch("pysus.management.sync.SyncEngine", return_value=engine): result = runner.invoke( app, @@ -183,9 +134,7 @@ def test_check_reupload_before(self, mock_env): ], ) assert result.exit_code == 0 - _, kwargs = engine.run.call_args - assert kwargs["reupload_before"] == datetime(2026, 7, 6) - assert kwargs["dry_run"] is True + engine.check.assert_awaited_once_with(datasets=["SINAN"]) @patch("pysus.cli.management._load_env") def test_check_reupload_before_with_apply(self, mock_env): diff --git a/pysus/tests/management/test_records.py b/pysus/tests/management/test_records.py index 9df4a356..fd93fd67 100644 --- a/pysus/tests/management/test_records.py +++ b/pysus/tests/management/test_records.py @@ -4,6 +4,10 @@ from pysus.management.records import ( DOWNLOAD_PRIORITY, + FRESH_CURRENT, + FRESH_MISSING, + FRESH_OUTDATED, + DatabaseCheck, FileComparison, FileRecord, IdentityKey, @@ -14,6 +18,7 @@ canonical_group, compose_s3_key, format_of, + freshness_status, parquet_key, stem_of, ) @@ -368,3 +373,111 @@ def test_summary(self): assert s["skipped"] == 1 assert s["failed"] == 1 assert s["needs_token"] == 1 + + +class TestFreshnessStatus: + @staticmethod + def _rec(origin, modified, size, **kw): + data = { + "origin": origin, + "dataset": "SINAN", + "name": "DENGBR20.dbc", + "path": "x", + "modified": modified, + "size": size, + } + data.update(kw) + return FileRecord(**data) + + @staticmethod + def _cmp(*records): + s3 = next((r for r in records if r.origin == "ducklake"), None) + if s3 is None: + s3 = records[0] + return FileComparison(key=s3.identity_key(), records=list(records)) + + def test_missing_when_no_s3_artifact(self): + cmp = self._cmp( + self._rec("ftp", datetime(2026, 1, 1), 100), + ) + status, _ = freshness_status(cmp) + assert status == FRESH_MISSING + + def test_current_when_same_date_and_size(self): + s3 = self._rec( + "ducklake", + datetime(2026, 1, 2), + 10, + source_modified=datetime(2026, 1, 1), + source_size=100, + ) + cmp = self._cmp(s3, self._rec("ftp", datetime(2026, 1, 1), 100)) + status, reason = freshness_status(cmp) + assert status == FRESH_CURRENT + + def test_outdated_when_origin_date_newer(self): + s3 = self._rec( + "ducklake", + datetime(2026, 1, 2), + 10, + source_modified=datetime(2026, 1, 1), + source_size=100, + ) + cmp = self._cmp(s3, self._rec("ftp", datetime(2026, 3, 1), 100)) + status, reason = freshness_status(cmp) + assert status == FRESH_OUTDATED + assert "newer" in reason + + def test_outdated_when_size_differs(self): + s3 = self._rec( + "ducklake", + datetime(2026, 1, 2), + 10, + source_modified=datetime(2026, 1, 1), + source_size=100, + ) + cmp = self._cmp(s3, self._rec("ftp", datetime(2026, 1, 1), 999)) + status, reason = freshness_status(cmp) + assert status == FRESH_OUTDATED + assert "size" in reason + + def test_outdated_when_source_date_missing_on_s3(self): + s3 = self._rec( + "ducklake", + datetime(2026, 1, 2), + 10, + source_modified=None, + source_size=100, + ) + cmp = self._cmp(s3, self._rec("ftp", datetime(2026, 1, 1), 100)) + status, _ = freshness_status(cmp) + assert status == FRESH_CURRENT + + +class TestDatabaseCheck: + def test_tracks_buckets(self): + chk = DatabaseCheck(dataset="SINAN") + chk.add(FRESH_MISSING, "a") + chk.add(FRESH_OUTDATED, "b", "reason") + chk.add(FRESH_CURRENT, "c") + assert chk.needs_update is True + assert chk.missing == ["a"] + assert chk.outdated == ["b — reason"] + assert chk.current == ["c"] + + def test_uptodate_when_no_missing_or_outdated(self): + chk = DatabaseCheck(dataset="SINAN") + chk.add(FRESH_CURRENT, "a") + assert chk.needs_update is False + + def test_summary(self): + chk = DatabaseCheck(dataset="SINAN") + chk.add(FRESH_MISSING, "a") + chk.add(FRESH_OUTDATED, "b") + chk.add(FRESH_CURRENT, "c") + assert chk.summary() == { + "missing": 1, + "outdated": 1, + "current": 1, + "needs_update": True, + } diff --git a/pysus/tests/management/test_sync_internal.py b/pysus/tests/management/test_sync_internal.py index 039fd023..88c12436 100644 --- a/pysus/tests/management/test_sync_internal.py +++ b/pysus/tests/management/test_sync_internal.py @@ -633,7 +633,8 @@ def test_journal_roundtrip(self, tmp_path): keys = load_journal_keys(journal) assert keys == {key} - def test_load_journal_includes_failed(self, tmp_path): + def test_load_journal_excludes_failed(self, tmp_path): + """A failed entry must be retried on the next run, not skipped.""" from pysus.management.records import ( IdentityKey, SyncOutcome, @@ -655,7 +656,7 @@ def test_load_journal_includes_failed(self, tmp_path): SyncOutcome(key=key, origin="ftp", status="failed"), ) keys = load_journal_keys(journal) - assert keys == {key} + assert keys == set() def test_load_journal_missing_file(self, tmp_path): from pysus.management.records import load_journal_keys @@ -726,3 +727,75 @@ async def test_resume_skips_uploaded_key(self, engine, monkeypatch): assert report.summary()["uploaded"] == 0 assert report.summary()["skipped"] == 1 engine._convert_and_upload.assert_not_awaited() + + +class TestCheck: + @pytest.mark.asyncio + async def test_check_classifies_missing_outdated_current(self): + from datetime import datetime + + engine = SyncEngine(access_key="ak", secret_key="sk") + engine.dadosgov_token = None + + s3_rec = _record( + "ducklake", + "DENGBR20.parquet", + year=2020, + modified=datetime(2026, 1, 2), + source_modified=datetime(2026, 1, 1), + source_size=100, + ) + current = _record( + "ftp", + "DENGBR20.dbc", + year=2020, + modified=datetime(2026, 1, 1), + size=100, + ) + # a mirrored file whose FTP origin was updated afterwards + s3_chik = _record( + "ducklake", + "CHIKBR22.parquet", + year=2022, + modified=datetime(2026, 1, 2), + source_modified=datetime(2026, 1, 1), + source_size=50, + ) + outdated = _record( + "ftp", + "CHIKBR22.dbc", + year=2022, + modified=datetime(2026, 5, 1), + size=50, + ) + missing = _record( + "ftp", + "DENGBR00.dbc", + year=2000, + modified=datetime(2026, 1, 1), + size=30, + ) + + # ducklake + ftp records returned by inventory.collect + records = { + "ducklake": [s3_rec, s3_chik], + "ftp": [current, outdated, missing], + } + mock_inv = MagicMock() + mock_inv.collect = AsyncMock( + side_effect=lambda origin, datasets=None, **kw: records.get( + origin, [] + ) + ) + + with patch.object(engine, "_require_pysus", return_value=MagicMock()): + with patch( + "pysus.management.sync.Inventory", return_value=mock_inv + ): + checks = await engine.check(datasets=["SINAN"]) + + sinan = checks["SINAN"] + assert len(sinan.missing) == 1 + assert len(sinan.outdated) == 1 + assert len(sinan.current) == 1 + assert sinan.needs_update is True