From 4ef4a85711bd20ffdc15fde89f276efbd9be4d13 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Lu=C3=A3=20Bida=20Vacaro?= Date: Wed, 12 Aug 2026 20:30:04 -0300 Subject: [PATCH 1/5] feat(management): cross-client sync pipeline with hierarchical S3 layout Add the pysus.management workflow package: inventory, comparison, catalog persistence and sync engine across FTP, DadosGov and DuckLake/S3. - records: origin-agnostic FileRecord, format-independent stems and hierarchical S3 keys (compose_s3_key) with download priority policy - inventory: collect listings from all clients + JSON snapshots and diffs - compare: logical identity grouping (dataset, year, stem) with csv/json/xml triplet dedup and content fingerprints - catalog: parameterized upserts into the per-dataset catalogs - sync: SyncEngine (inventory -> compare -> download -> parquet -> upload -> catalog) with checkpointing and freshness-based re-uploads - normalize/relayout: bucket migration to the hierarchical layout keeping alias objects at former keys for backwards compatibility - ducklake: download_http follows alias markers; adapters upload only dirty catalogs and never silently fall back to stale local copies - cli: streamlit web share option for Google Colab Tests: 723 passing, including key composition, alias following and comparison reporting suites. --- pysus/api/ducklake/catalog/adapters.py | 12 +- pysus/api/ducklake/functional.py | 30 + pysus/cli/__init__.py | 64 +- pysus/management/__init__.py | 68 +++ pysus/management/catalog.py | 367 ++++++++++++ pysus/management/client.py | 335 +---------- pysus/management/compare.py | 203 +++++++ pysus/management/inventory.py | 257 ++++++++ pysus/management/normalize.py | 623 ++++++++++++++++++++ pysus/management/records.py | 396 +++++++++++++ pysus/management/report.py | 111 ++++ pysus/management/scripts/__init__.py | 0 pysus/management/scripts/compare_clients.py | 141 +++++ pysus/management/scripts/relayout_bucket.py | 180 ++++++ pysus/management/scripts/sync_clients.py | 90 +++ pysus/management/sync.py | 538 +++++++++++++++++ pysus/tests/api/ducklake/test_client.py | 3 + pysus/tests/api/ducklake/test_functional.py | 78 ++- pysus/tests/management/__init__.py | 0 pysus/tests/management/test_compare.py | 176 ++++++ pysus/tests/management/test_records.py | 209 +++++++ pysus/tests/management/test_report.py | 77 +++ pysus/tests/management/test_sync.py | 149 +++++ 23 files changed, 3785 insertions(+), 322 deletions(-) create mode 100644 pysus/management/catalog.py create mode 100644 pysus/management/compare.py create mode 100644 pysus/management/inventory.py create mode 100644 pysus/management/normalize.py create mode 100644 pysus/management/records.py create mode 100644 pysus/management/report.py create mode 100644 pysus/management/scripts/__init__.py create mode 100644 pysus/management/scripts/compare_clients.py create mode 100644 pysus/management/scripts/relayout_bucket.py create mode 100644 pysus/management/scripts/sync_clients.py create mode 100644 pysus/management/sync.py create mode 100644 pysus/tests/management/__init__.py create mode 100644 pysus/tests/management/test_compare.py create mode 100644 pysus/tests/management/test_records.py create mode 100644 pysus/tests/management/test_report.py create mode 100644 pysus/tests/management/test_sync.py diff --git a/pysus/api/ducklake/catalog/adapters.py b/pysus/api/ducklake/catalog/adapters.py index 62edc56e..f21abf0f 100644 --- a/pysus/api/ducklake/catalog/adapters.py +++ b/pysus/api/ducklake/catalog/adapters.py @@ -36,10 +36,11 @@ def __init__( **data, ) -> None: self._engine = engine - self._session_factory = None + self._session_factory: sessionmaker[Session] | None = None self.cache_dir.mkdir(parents=True, exist_ok=True) self.credentials = credentials self.update_on_close = update_on_close + self._local_dirty = False @property def remote_url(self) -> str: @@ -79,6 +80,12 @@ async def connect( force=True, callback=callback, ) + self._local_dirty = False + self._engine = await to_thread.run_sync(self.setup_engine) + self._session_factory = sessionmaker(bind=self._engine) + return + + if self._local_dirty: self._engine = await to_thread.run_sync(self.setup_engine) self._session_factory = sessionmaker(bind=self._engine) return @@ -229,8 +236,9 @@ async def _upload_catalog(self) -> None: ) async def close(self, update: bool = False) -> None: - if update: + if update and self._local_dirty: await self._upload_catalog() + self._local_dirty = False if self._engine: await to_thread.run_sync(self._engine.dispose) diff --git a/pysus/api/ducklake/functional.py b/pysus/api/ducklake/functional.py index 4666522d..ef7eb020 100644 --- a/pysus/api/ducklake/functional.py +++ b/pysus/api/ducklake/functional.py @@ -8,6 +8,35 @@ from botocore.config import Config from pysus.api import types +ALIAS_META_HEADER = "x-amz-meta-pysus-alias" +MAX_ALIAS_HOPS = 5 + + +def alias_marker(target_key: str) -> str: + """Return the alias marker content pointing to *target_key*.""" + import json + + return json.dumps({"pysus-alias": target_key}) + + +def _url_for_key(key: str) -> str: + key = str(key).replace("\\", "/") + return f"https://{types.S3_ENDPOINT}/{types.S3_BUCKET}/{key}" + + +async def _resolve_alias(url: str, client: httpx.AsyncClient) -> str: + """Follow ``pysus-alias`` marker objects up to ``MAX_ALIAS_HOPS``.""" + for _ in range(MAX_ALIAS_HOPS): + head = await client.head(url) + if head.status_code == 404: + return url + head.raise_for_status() + target = head.headers.get(ALIAS_META_HEADER) + if not target or not isinstance(target, str): + return url + url = _url_for_key(target) + raise RuntimeError(f"Too many alias hops resolving {url}") + async def download_http( remote_path: str, @@ -38,6 +67,7 @@ async def download_http( limits=limits, timeout=timeout, ) as client: + url = await _resolve_alias(url, client) async with client.stream("GET", url) as r: r.raise_for_status() total = int(r.headers.get("Content-Length", 0)) diff --git a/pysus/cli/__init__.py b/pysus/cli/__init__.py index 29135914..e6b514d9 100644 --- a/pysus/cli/__init__.py +++ b/pysus/cli/__init__.py @@ -1,11 +1,20 @@ +import os +import sys +import webbrowser + import typer from pysus import __version__ app = typer.Typer(help="PySUS CLI") +def _is_colab() -> bool: + return "COLAB_RELEASE_TAG" in os.environ + + @app.command() def version(): + """Print the installed PySUS version.""" print(__version__) @@ -17,24 +26,24 @@ def web( "--port", help="Port to bind the server to", ), + share: bool = typer.Option( # noqa: B008 + False, + "--share", + help="When running in Google Colab, print the proxied URL", + ), ): - """Launch the local Streamlit visual interface.""" + """Launch the Streamlit web interface.""" try: - import streamlit.web.bootstrap as bootstrap # noqa - from streamlit.runtime.scriptrunner import get_script_run_ctx # noqa - except ImportError: + from streamlit.web import cli as stcli + except ImportError as exc: raise ImportError( "The HTTP UI requires extra dependencies. " "Install them with: pip install pysus[web]" - ) - import os - import sys - import webbrowser - - app_path = os.path.join(os.path.dirname(__file__), "..", "web", "app.py") - app_path = os.path.abspath(app_path) + ) from exc - from streamlit.web import cli as stcli + app_path = os.path.abspath( + os.path.join(os.path.dirname(__file__), "..", "web", "app.py") + ) sys.argv = [ "streamlit", @@ -45,12 +54,37 @@ def web( "--server.headless", "true", "--server.address", - "localhost", + "0.0.0.0", ] - webbrowser.open(f"http://localhost:{port}") + if _is_colab() and share: + import threading + + def _run_streamlit(): + stcli.main() + + t = threading.Thread(target=_run_streamlit, daemon=True) + t.start() + + import time + + time.sleep(3) + + try: + from google.colab.output import eval_js - stcli.main() + url = eval_js(f"google.colab.kernel.proxyPort({port})") + print(f"\nPySUS web interface running at:\n{url}\n") + except ImportError: + print( + "\nGoogle Colab detected but google.colab is not available.\n" + f"Open http://localhost:{port} manually.\n" + ) + t.join() + else: + if not _is_colab(): + webbrowser.open(f"http://localhost:{port}") + stcli.main() if __name__ == "__main__": diff --git a/pysus/management/__init__.py b/pysus/management/__init__.py index e69de29b..fe1e71a9 100644 --- a/pysus/management/__init__.py +++ b/pysus/management/__init__.py @@ -0,0 +1,68 @@ +"""PySUS management: cross-client file tracking, comparison and sync. + +The management package implements the workflow that keeps the S3 bucket +and its DuckLake catalogs in sync with the FTP and DadosGov clients: + +* :mod:`records` — origin-agnostic file records, identity keys and + sync reports; +* :mod:`inventory` — collect listings from every client + snapshots; +* :mod:`compare` — cross-client identity grouping and content + fingerprints; +* :mod:`catalog` — parameterized metadata upserts into DuckLake; +* :mod:`sync` — the end-to-end pipeline (inventory → compare → + download → parquet → upload → catalog); +* :mod:`normalize` — S3 key canonicalization utilities; +* :mod:`client` — ``CatalogManager`` facade for single-file uploads. +""" + +from .catalog import CatalogWriter, sha256_of # noqa +from .client import CatalogManager # noqa +from .compare import Comparator, content_fingerprint # noqa +from .inventory import Inventory # noqa +from .normalize import BucketNormalizer # noqa +from .records import ( # noqa + DOWNLOAD_PRIORITY, + KEY_MISSING, + NATIONAL_STATE, + FileComparison, + FileRecord, + IdentityKey, + SnapshotDiff, + SyncOutcome, + SyncReport, + base_stem, + canonical_dataset, + canonical_group, + compose_s3_key, + format_of, + parquet_key, + stem_of, +) +from .sync import SyncEngine # noqa + +__all__ = [ + "BucketNormalizer", + "CatalogManager", + "CatalogWriter", + "Comparator", + "DOWNLOAD_PRIORITY", + "FileComparison", + "FileRecord", + "IdentityKey", + "Inventory", + "KEY_MISSING", + "NATIONAL_STATE", + "SnapshotDiff", + "SyncEngine", + "SyncOutcome", + "SyncReport", + "base_stem", + "canonical_dataset", + "canonical_group", + "compose_s3_key", + "content_fingerprint", + "format_of", + "parquet_key", + "sha256_of", + "stem_of", +] diff --git a/pysus/management/catalog.py b/pysus/management/catalog.py new file mode 100644 index 00000000..9ef98e98 --- /dev/null +++ b/pysus/management/catalog.py @@ -0,0 +1,367 @@ +"""Catalog persistence: write file metadata into the DuckLake duckdb catalogs. + +The writer uses parameterized upserts against the DuckDB engines managed by +the DuckLake adapters. It is idempotent (keyed on the S3 ``path``), keeps +the legacy columns in sync, and adds a few management columns +(``origin``, ``format``, ``sha256``) when missing, so the catalog stays the +single source of truth for every file tracked by the workflow. +""" + +from __future__ import annotations + +import hashlib +from datetime import datetime +from pathlib import Path +from typing import TYPE_CHECKING, Any + +from pysus.api.errors import CatalogError + +if TYPE_CHECKING: # pragma: no cover + from pysus.api.ducklake.client import DuckLake + +_ARROW_TO_SQL = { + "int64": "BIGINT", + "int32": "INTEGER", + "int16": "INTEGER", + "int8": "INTEGER", + "uint64": "BIGINT", + "uint32": "BIGINT", + "uint16": "INTEGER", + "uint8": "INTEGER", + "double": "DOUBLE", + "float": "FLOAT", + "bool": "BOOLEAN", + "timestamp[us]": "TIMESTAMP", + "timestamp[ns]": "TIMESTAMP", + "date32[day]": "DATE", + "string": "VARCHAR", + "large_string": "VARCHAR", + "binary": "BLOB", +} + +_FILES_BASE_COLUMNS = ( + "dataset_id", + "group_id", + "path", + "size", + "rows", + "modified", + "origin_modified", + "origin_size", + "origin_path", + "year", + "month", + "state", +) + + +def sha256_of(path: Path) -> str: + """Compute the sha256 digest of a local file.""" + digest = hashlib.sha256() + with open(path, "rb") as f: + while chunk := f.read(1024 * 1024): + digest.update(chunk) + return digest.hexdigest() + + +class CatalogWriter: + """Upsert dataset/group/file/column metadata into the DuckLake catalogs.""" + + def __init__(self, ducklake: DuckLake): + self.ducklake = ducklake + + # ------------------------------------------------------------------ + # low-level plumbing + # ------------------------------------------------------------------ + @property + def _catalog_engine(self): + engine = self.ducklake._catalog_adap._engine + if engine is None: + raise CatalogError("DuckLake catalog engine is not initialized") + return engine + + @property + def _columns_engine(self): + engine = self.ducklake._columns_adap._engine + if engine is None: + raise CatalogError("DuckLake columns engine is not initialized") + return engine + + def _has_column(self, cursor, table: str, column: str) -> bool: + cursor.execute( + "SELECT 1 FROM information_schema.columns " + "WHERE table_schema = 'pysus' AND table_name = ? " + "AND column_name = ?", + (table, column), + ) + return cursor.fetchone() is not None + + def _ensure_column( + self, cursor, table: str, column: str, definition: str + ) -> None: + if self._has_column(cursor, table, column): + return + cursor.execute( + f'ALTER TABLE pysus.{table} ADD COLUMN "{column}" {definition}' + ) + + def _ensure_management_columns(self, catalog_cursor) -> None: + self._ensure_column(catalog_cursor, "files", "origin", "VARCHAR") + self._ensure_column(catalog_cursor, "files", "format", "VARCHAR") + + # ------------------------------------------------------------------ + # datasets & groups + # ------------------------------------------------------------------ + def ensure_dataset( + self, + cursor, + name: str, + long_name: str, + description: str | None = None, + ) -> int: + """Return the dataset id, creating the row if needed.""" + name = name.strip().lower() + cursor.execute( + "SELECT id, long_name, description FROM pysus.datasets " + "WHERE name = ?", + (name,), + ) + row = cursor.fetchone() + if row: + dataset_id = row[0] + if row[1] != long_name or row[2] != description: + cursor.execute( + "UPDATE pysus.datasets SET long_name = ?, " + "description = ? WHERE id = ?", + (long_name, description, dataset_id), + ) + return int(dataset_id) + + cursor.execute("SELECT MAX(id) FROM pysus.datasets") + max_row = cursor.fetchone() + dataset_id = (max_row[0] or 0) + 1 + cursor.execute( + "INSERT INTO pysus.datasets (id, name, long_name, description) " + "VALUES (?, ?, ?, ?)", + (dataset_id, name, long_name, description), + ) + return int(dataset_id) + + def ensure_group( + self, + cursor, + dataset_id: int, + name: str | None, + long_name: str | None = None, + description: str | None = None, + ) -> int | None: + """Return the group id for (dataset, name), creating it if needed.""" + if not name: + return None + + name = name.strip().upper() + cursor.execute( + "SELECT id, long_name, description FROM pysus.dataset_groups " + "WHERE dataset_id = ? AND name = ?", + (dataset_id, name), + ) + row = cursor.fetchone() + if row: + group_id = row[0] + if row[1] != long_name or row[2] != description: + cursor.execute( + "UPDATE pysus.dataset_groups SET long_name = ?, " + "description = ? WHERE id = ?", + (long_name, description, group_id), + ) + return int(group_id) + + cursor.execute("SELECT MAX(id) FROM pysus.dataset_groups") + max_row = cursor.fetchone() + group_id = (max_row[0] or 0) + 1 + cursor.execute( + "INSERT INTO pysus.dataset_groups (id, dataset_id, name, " + "long_name, description) VALUES (?, ?, ?, ?, ?)", + (group_id, dataset_id, name, long_name, description), + ) + return int(group_id) + + # ------------------------------------------------------------------ + # files + # ------------------------------------------------------------------ + def get_file(self, cursor, path: str) -> tuple[int, datetime | None] | None: + """Return ``(id, origin_modified)`` for the S3 *path*, if present.""" + cursor.execute( + "SELECT id, origin_modified FROM pysus.files WHERE path = ?", + (path,), + ) + row = cursor.fetchone() + if not row: + return None + return int(row[0]), row[1] + + def delete_file(self, cursor, file_id: int) -> None: + cursor.execute( + "DELETE FROM pysus.file_columns WHERE file_id = ?", (file_id,) + ) + cursor.execute("DELETE FROM pysus.files WHERE id = ?", (file_id,)) + + def upsert_file( + self, + cursor, + *, + dataset_id: int, + group_id: int | None, + path: str, + size: int, + rows: int, + modified: datetime | None, + origin_modified: datetime | None, + origin_size: int, + origin_path: str, + year: int | None, + month: int | None, + state: str | None, + origin: str | None = None, + format: str | None = None, + sha256: str | None = None, + file_type: str | None = None, + ) -> tuple[int, bool]: + """Insert or update the file row keyed on S3 *path*. + + Returns ``(file_id, created)``. + """ + existing = self.get_file(cursor, path) + if existing: + file_id, _ = existing + sets = [ + "size = ?", + "rows = ?", + "modified = ?", + "origin_modified = ?", + "origin_size = ?", + "origin_path = ?", + "year = ?", + "month = ?", + "state = ?", + ] + update_values: list[Any] = [ + size, + rows, + modified or datetime.now(), + origin_modified, + origin_size, + origin_path, + year, + month, + state, + ] + if origin is not None: + sets.append("origin = ?") + update_values.append(origin) + if format is not None: + sets.append("format = ?") + update_values.append(format) + if sha256 is not None: + sets.append("sha256 = ?") + update_values.append(sha256) + if file_type is not None: + sets.append("type = ?") + update_values.append(file_type) + update_values.append(file_id) + cursor.execute( + f"UPDATE pysus.files SET {', '.join(sets)} WHERE id = ?", + update_values, + ) + return file_id, False + + cursor.execute("SELECT MAX(id) FROM pysus.files") + max_row = cursor.fetchone() + file_id = (max_row[0] or 0) + 1 + columns = ["id", *_FILES_BASE_COLUMNS] + values: list[Any] = [ + file_id, + dataset_id, + group_id, + path, + size, + rows, + modified or datetime.now(), + origin_modified, + origin_size, + origin_path, + year, + month, + state, + ] + if origin is not None: + columns.append("origin") + values.append(origin) + if format is not None: + columns.append("format") + values.append(format) + if sha256 is not None: + columns.append("sha256") + values.append(sha256) + if file_type is not None: + columns.append("type") + values.append(file_type) + placeholders = ", ".join("?" for _ in columns) + cursor.execute( + f"INSERT INTO pysus.files ({', '.join(columns)}) " + f"VALUES ({placeholders})", + values, + ) + return file_id, True + + # ------------------------------------------------------------------ + # columns + # ------------------------------------------------------------------ + def link_columns( + self, + dataset_cursor, + columns_cursor, + file_id: int, + schema, + dataset_id: int, + ) -> None: + """Get-or-create column definitions for *schema* and link them. + + Column definitions live in the columns catalog; the + ``file_columns`` links live next to the files in the per-dataset + catalog. + """ + column_ids: list[int] = [] + for col_name in schema.names: + arrow_type = str(schema.field(col_name).type) + sql_type = _ARROW_TO_SQL.get(arrow_type, "VARCHAR") + + columns_cursor.execute( + "SELECT id FROM pysus.dataset_columns " + "WHERE dataset_id = ? AND name = ?", + (dataset_id, col_name), + ) + existing = columns_cursor.fetchone() + if existing: + column_ids.append(existing[0]) + continue + + columns_cursor.execute("SELECT MAX(id) FROM pysus.dataset_columns") + max_row = columns_cursor.fetchone() + new_id = (max_row[0] or 0) + 1 + columns_cursor.execute( + "INSERT INTO pysus.dataset_columns (id, dataset_id, name, " + "type, nullable) VALUES (?, ?, ?, ?, true)", + (new_id, dataset_id, col_name, sql_type), + ) + column_ids.append(new_id) + + dataset_cursor.execute( + "DELETE FROM pysus.file_columns WHERE file_id = ?", (file_id,) + ) + for column_id in column_ids: + dataset_cursor.execute( + "INSERT INTO pysus.file_columns (file_id, column_id) " + "VALUES (?, ?)", + (file_id, column_id), + ) diff --git a/pysus/management/client.py b/pysus/management/client.py index cc41c578..f5e9822d 100644 --- a/pysus/management/client.py +++ b/pysus/management/client.py @@ -1,331 +1,58 @@ -import asyncio +"""Catalog management entry point. + +``CatalogManager`` is a thin facade over +:class:`~pysus.management.sync.SyncEngine`: it downloads a remote file, +converts it to parquet, uploads it to S3 and upserts the metadata into the +DuckLake catalog. The full cross-client workflow (inventory → compare → +sync) lives in :mod:`pysus.management.sync`. +""" + +from __future__ import annotations + import os from collections.abc import Callable -from logging import error -from pathlib import Path -from pysus.api.client import PySUS from pysus.api.dadosgov.models import File as APIFile -from pysus.api.ducklake.functional import upload_s3 -from pysus.api.extensions import Parquet from pysus.api.ftp.models import File as FTPFile -from pysus.api.models import BaseRemoteFile + +from .sync import SyncEngine class CatalogManager: + """Upload + catalog single files into the PySUS S3 bucket.""" + def __init__( self, access_key: str | None = None, secret_key: str | None = None, dadosgov_token: str | None = None, ): - self.pysus = PySUS() - self.access_key = access_key or os.getenv("ACCESS_KEY") - self.secret_key = secret_key or os.getenv("SECRET_KEY") - self.dadosgov_token = dadosgov_token or os.getenv("DADOSGOV_TOKEN") - - if not self.access_key or not self.secret_key: + self.engine = SyncEngine( + access_key=access_key or os.getenv("ACCESS_KEY"), + secret_key=secret_key or os.getenv("SECRET_KEY"), + dadosgov_token=dadosgov_token or os.getenv("DADOSGOV_TOKEN"), + ) + if not self.engine.access_key or not self.engine.secret_key: raise ValueError("s3 credentials are needed") - async def __aenter__(self): - await self.pysus.__aenter__() - ducklake = await self.pysus.get_ducklake() - await ducklake.login( - access_key=self.access_key, secret_key=self.secret_key - ) + async def __aenter__(self) -> CatalogManager: + await self.engine.__aenter__() return self - async def __aexit__(self, exc_type, exc_val, exc_tb): - try: - if not exc_type: - ducklake = self.pysus._ducklake - if ducklake: - await ducklake.close(update_catalog=True) - finally: - await self.pysus.__aexit__(exc_type, exc_val, exc_tb) + async def __aexit__(self, exc_type, exc_val, exc_tb) -> None: + await self.engine.__aexit__(exc_type, exc_val, exc_tb) async def upload( self, file: FTPFile | APIFile, callback: Callable[[int, int], None] | None = None, - ) -> None: - if not self.pysus._ducklake: - raise ConnectionError("DuckLake is not connected") - - remote_file_path = Path(file.path) - s3_key = ( - f"public/data/{file.client.name.lower()}" - f"/{file.dataset.name.lower()}" - f"/{remote_file_path.with_suffix('.parquet').name}" - ) - - dataset_id = None - group_id = None - - catalog_engine = self.pysus._ducklake._catalog_adap._engine - columns_engine = self.pysus._ducklake._columns_adap._engine - - if not catalog_engine or not columns_engine: - raise ConnectionError( - "DuckLake database engines are not initialized" - ) - - catalog_conn = catalog_engine.raw_connection() - columns_conn = columns_engine.raw_connection() - - with catalog_conn, columns_conn: - catalog_cursor = catalog_conn.cursor() - columns_cursor = columns_conn.cursor() - - try: - dataset_name = file.dataset.name.lower() - is_ftp = file.client.name.lower() == "ftp" - - catalog_cursor.execute( - "SELECT id FROM pysus.datasets WHERE name = ?", - (dataset_name,), - ) - row = catalog_cursor.fetchone() - - if row: - dataset_id = row[0] - origin_val = "'FTP'" if is_ftp else "'API'" - catalog_cursor.execute( - f"UPDATE pysus.datasets SET origin = {origin_val} " - f"WHERE id = {dataset_id}" - ) - else: - catalog_cursor.execute("SELECT MAX(id) FROM pysus.datasets") - max_id_row = catalog_cursor.fetchone() - max_id = max_id_row[0] if max_id_row else None - dataset_id = (max_id or 0) + 1 - origin_val = "'FTP'" if is_ftp else "'API'" - catalog_cursor.execute( - f"INSERT INTO pysus.datasets (id, name, long_name, " - f"origin) VALUES ({dataset_id}, '{dataset_name}', " - f"'{file.dataset.long_name}', {origin_val})" - ) - - if file.group: - group_name = file.group.name - catalog_cursor.execute( - "SELECT id FROM pysus.dataset_groups " - "WHERE name = ? AND dataset_id = ?", - (group_name, dataset_id), - ) - row = catalog_cursor.fetchone() - if row: - group_id = row[0] - else: - catalog_cursor.execute( - "SELECT MAX(id) FROM pysus.dataset_groups" - ) - max_id_row = catalog_cursor.fetchone() - max_id = max_id_row[0] if max_id_row else None - group_id = (max_id or 0) + 1 - long_name = file.dataset.group_definitions.get( - group_name.upper(), group_name - ) - catalog_cursor.execute( - f"INSERT INTO pysus.dataset_groups (id, " - f"dataset_id, name, long_name) VALUES ({group_id}," - f" {dataset_id}, '{group_name}', '{long_name}')" - ) - - group_val = "NULL" if group_id is None else str(group_id) - - catalog_cursor.execute( - "SELECT id, group_id FROM pysus.files WHERE path = ?", - (s3_key,), - ) - row = catalog_cursor.fetchone() - - if row: - file_id, db_group_id = row - group_mismatch = db_group_id != group_id - should_upload = self._should_upload_raw( - catalog_cursor, - file_id, - file, - ) - - if not should_upload and not group_mismatch: - return - - columns_cursor.execute( - "DELETE FROM pysus.file_columns WHERE file_id = ?", - (file_id,), - ) - catalog_cursor.execute( - "DELETE FROM pysus.files WHERE id = ?", - (file_id,), - ) - else: - catalog_cursor.execute("SELECT MAX(id) FROM pysus.files") - max_id_row = catalog_cursor.fetchone() - max_id = max_id_row[0] if max_id_row else None - file_id = (max_id or 0) + 1 - - parquet_ext = await self._download_with_retry(file, callback) - await self._upload_to_s3(parquet_ext.path, s3_key) - - year_val = "NULL" if file.year is None else str(file.year) - month_val = "NULL" if file.month is None else str(file.month) - state_val = "NULL" if file.state is None else f"'{file.state}'" - - catalog_cursor.execute( - f"INSERT INTO pysus.files (id, dataset_id, group_id, " - f"path, size, rows, modified, origin_modified, " - f"origin_path, year, month, state) VALUES ({file_id}, " - f"{dataset_id}, {group_val}, '{s3_key}', " - f"{parquet_ext.size}, {parquet_ext.rows}, " - f"CURRENT_TIMESTAMP, '{file.modify}', '{file.path}', " - f"{year_val}, {month_val}, {state_val})" - ) - - new_columns = self._get_or_create_columns_raw( - columns_cursor, parquet_ext, dataset_id - ) - - for col in new_columns: - columns_cursor.execute( - f"INSERT INTO pysus.file_columns (file_id, column_id) " - f"VALUES ({file_id}, {col})" - ) - - catalog_conn.commit() - columns_conn.commit() - - catalog_cursor.execute("CHECKPOINT") - columns_cursor.execute("CHECKPOINT") - - if parquet_ext.path.exists(): - parquet_ext.path.unlink() - await self.pysus._delete_record(str(parquet_ext.path)) - - except BaseException as rollback_err: # noqa - try: - catalog_conn.rollback() - except Exception as inner_err: # noqa - error(f"Catalog rollback failed: {inner_err}") - try: - columns_conn.rollback() - except Exception as inner_err: # noqa - error(f"Columns rollback failed: {inner_err}") - raise rollback_err - - async def _upload_to_s3( - self, - local_path: Path, - s3_path: str, - callback: Callable[[int, int], None] | None = None, - ): - await upload_s3( - local_path=local_path, - access_key=str(self.access_key), - secret_key=str(self.secret_key), - remote_path=s3_path, - callback=callback, - ) - - async def _download_with_retry( - self, - file: FTPFile | APIFile, - callback: Callable[[int, int], None] | None = None, - max_retries: int = 3, - ) -> Parquet: - errors = (ConnectionResetError, ConnectionRefusedError, TimeoutError) - last_error = None - - for attempt in range(max_retries): - try: - return await self.pysus.download_to_parquet( - file=file, - token=self.dadosgov_token, - callback=callback, - ) - except errors as e: - last_error = e - wait_time = 2**attempt - error( - f"Download attempt {attempt + 1}/{max_retries} failed " - f"for {file.basename}: {e}. Retrying in {wait_time}s..." - ) - await asyncio.sleep(wait_time) - - raise RuntimeError( - f"Failed to download {file.basename} after {max_retries} " - f"attempts: {last_error}" - ) from last_error - - def _should_upload_raw( - self, - cursor, - file_id: int, - file: BaseRemoteFile, force: bool = False, ) -> bool: - if force: - return True + """Download, convert, upload and catalog a single remote file. - cursor.execute( - "SELECT origin_modified FROM pysus.files WHERE id = ?", - (file_id,), + Returns True if the file was (re)processed, False if the catalog + already holds an equally recent artifact. + """ + return await self.engine.upload_file( + file, callback=callback, force=force ) - row = cursor.fetchone() - if not row: - return True - - origin_modified = row[0] - if origin_modified is None: - return True - - file_mod = getattr(file, "modify", None) - if file_mod is None: - return True - - return str(file_mod) > str(origin_modified) - - def _get_or_create_columns_raw( - self, cursor, file: Parquet, dataset_id: int - ) -> list[int]: - schema = file.schema - type_map = { - "int64": "BIGINT", - "int32": "INTEGER", - "double": "DOUBLE", - "float": "FLOAT", - "bool": "BOOLEAN", - "timestamp[us]": "TIMESTAMP", - "string": "VARCHAR", - "binary": "BLOB", - } - - result = [] - for col_name in schema.names: - field = schema.field(col_name) - arrow_type = str(field.type) - sql_type = type_map.get(arrow_type, "VARCHAR") - - cursor.execute( - "SELECT id FROM pysus.dataset_columns " - "WHERE name = ? AND dataset_id = ?", - (col_name, dataset_id), - ) - existing = cursor.fetchone() - - if existing: - result.append(existing[0]) - else: - cursor.execute("SELECT MAX(id) FROM pysus.dataset_columns") - max_id_row = cursor.fetchone() - max_id = max_id_row[0] if max_id_row else None - new_id = (max_id or 0) + 1 - cursor.execute( - "INSERT INTO pysus.dataset_columns (id, dataset_id, name, " - "type, nullable) VALUES (?, ?, ?, ?, true)", - (new_id, dataset_id, col_name, sql_type), - ) - result.append(new_id) - - return result diff --git a/pysus/management/compare.py b/pysus/management/compare.py new file mode 100644 index 00000000..b81f0d3c --- /dev/null +++ b/pysus/management/compare.py @@ -0,0 +1,203 @@ +"""Cross-client comparison of inventoried files. + +Files with the same logical identity (dataset + group + year/month/state + +format-independent stem) are grouped into :class:`FileComparison` objects +that expose which origins carry the file, in which formats, and which +record should be preferred for download. + +Content-level equivalence (same data, different bytes/format) is confirmed +by :func:`content_fingerprint`, which operates on decompressed/parsed +data — raw sizes can never be compared across formats (a ``csv.zip`` is +not byte-comparable to a ``parquet``). +""" + +from __future__ import annotations + +import hashlib +from collections.abc import Iterable + +import pandas as pd + +from .records import DOWNLOAD_PRIORITY, FileComparison, FileRecord, IdentityKey + + +class Comparator: + """Group and compare :class:`FileRecord` objects across origins.""" + + def __init__( + self, + priorities: tuple[str, ...] = DOWNLOAD_PRIORITY, + ): + self.priorities = priorities + + def compare( + self, + records: Iterable[FileRecord], + ) -> list[FileComparison]: + """Group records by logical identity. + + The grouping key is ``(dataset, year, stem)``. Group, month and + state are deliberately excluded: they are encoded in the stem for + state/month/group-level files (e.g. ``PAAC2408``), and the same + logical file is known to carry different values per origin — + FTP leaves ``state``/``group`` null on national files while + DadosGov sets ``"BR"``, and legacy catalog rows have + ``group_id NULL``. Those attributes are preserved per record and + resolved onto the comparison key from the richest record. This is + deliberately permissive: content fingerprints can later veto a + wrong grouping. + """ + groups: dict[tuple, list[FileRecord]] = {} + order: list[tuple] = [] + for record in records: + key = record.identity_key() + group_key = (key.dataset, key.year, key.stem) + if group_key not in groups: + groups[group_key] = [] + order.append(group_key) + groups[group_key].append(record) + + comparisons: list[FileComparison] = [] + + def _pick(records: list[FileRecord], attr: str): + return next( + ( + getattr(r, attr) + for r in records + if getattr(r, attr) is not None + ), + None, + ) + + for group_key in order: + items = self._dedup_origin_formats(groups[group_key]) + dataset, year, stem = group_key + + comparisons.append( + FileComparison( + key=IdentityKey( + dataset=dataset, + group=_pick(items, "group"), + year=year, + month=_pick(items, "month"), + state=_pick(items, "state"), + stem=stem, + ), + records=items, + ) + ) + return comparisons + + @staticmethod + def _dedup_origin_formats( + records: list[FileRecord], + ) -> list[FileRecord]: + """Keep only one record per (origin, logical key). + + DadosGov publishes the same data as csv/json/xml triplets: when + several records from the same origin share size-agnostic identity, + prefer the format highest in ``FORMAT_PREFERENCE`` (csv first). + """ + preferred: dict[str, FileRecord] = {} + for record in records: + fmt = (record.format or "").lower() + current = preferred.get(record.origin) + if current is None or _format_rank(fmt) < _format_rank( + current.format or "" + ): + preferred[record.origin] = record + return list(preferred.values()) + + def pick( + self, + comparison: FileComparison, + ) -> FileRecord | None: + """Resolve the download source for a logical file. + + S3 (ducklake) first, then FTP, then DadosGov (which needs the API + token). Within an origin the most recently modified record wins. + """ + return comparison.best_record(priorities=self.priorities) + + +FORMAT_PREFERENCE: tuple[str, ...] = ( + "csv", + "csv.zip", + "json", + "xml", + "xlsx", + "dbf", + "dbc", + "parquet", +) + + +def _format_rank(fmt: str) -> int: + fmt = fmt.strip().lower() + if not fmt or fmt == "unknown": + return len(FORMAT_PREFERENCE) + 1 + try: + return FORMAT_PREFERENCE.index(fmt) + except ValueError: + return len(FORMAT_PREFERENCE) + + +def content_fingerprint( + frame: pd.DataFrame, + sample_size: int = 1000, + even_spacing: int = 1000, +) -> str: + """Compute a format-independent fingerprint of tabular content. + + The fingerprint combines the (sorted) schema, the row count, and a + hash over sampled rows with stringified values. It is stable across + formats (dbc, parquet, csv, json) as long as the data and columns are + equivalent, and it is computed only on parsed/decompressed content. + + Parameters + ---------- + frame : pd.DataFrame + The parsed/decompressed content. + sample_size : int, optional + Number of rows to sample from the head. + even_spacing : int, optional + Take one row every *even_spacing* rows to cover the whole file. + + Returns + ------- + str + Hex digest of the content fingerprint. + """ + schema = tuple( + sorted((str(col), str(dtype)) for col, dtype in frame.dtypes.items()) + ) + total_rows = len(frame) + + ordered = frame[sorted(frame.columns, key=str)] + sampled = list( + ordered.head(sample_size).astype(str).itertuples(index=False, name=None) + ) + if total_rows > sample_size: + sampled.extend( + ordered.iloc[::even_spacing] + .astype(str) + .itertuples(index=False, name=None) + ) + + digest = hashlib.sha256() + digest.update(repr(schema).encode()) + digest.update(str(total_rows).encode()) + for row in sampled: + digest.update(repr(row).encode()) + return digest.hexdigest() + + +def byte_hash( + path, algorithm: str = "sha256", chunk_size: int = 1024 * 1024 +) -> str: + """Compute the byte-level hash of a local file (exact identity only).""" + hash_obj = hashlib.new(algorithm) + with open(path, "rb") as f: + while chunk := f.read(chunk_size): + hash_obj.update(chunk) + return hash_obj.hexdigest() diff --git a/pysus/management/inventory.py b/pysus/management/inventory.py new file mode 100644 index 00000000..3024cd64 --- /dev/null +++ b/pysus/management/inventory.py @@ -0,0 +1,257 @@ +"""Inventory collection: snapshot every file visible on each client. + +The collectors reduce FTP, DadosGov and DuckLake listings into +:class:`~pysus.management.records.FileRecord` objects. Snapshots are +persisted locally (JSON) so consecutive runs can diff against the previous +state without re-listing. +""" + +from __future__ import annotations + +import json +from datetime import datetime +from pathlib import Path +from typing import TYPE_CHECKING, Any + +from pysus import CACHEPATH + +from .records import FileRecord, SnapshotDiff + +if TYPE_CHECKING: # pragma: no cover + from pysus.api.client import PySUS + from pysus.api.models import BaseRemoteFile + +SNAPSHOT_DIR: Path = Path(CACHEPATH) / "management" / "inventory" + +_ORIGIN_TO_CLIENT = { + "ftp": "ftp", + "dadosgov": "dadosgov", + "ducklake": "ducklake", +} + + +class Inventory: + """Collect and persist file listings from all three clients.""" + + def __init__(self, pysus: PySUS, snapshot_dir: Path | None = None): + self.pysus = pysus + self.snapshot_dir = snapshot_dir or SNAPSHOT_DIR + self.snapshot_dir.mkdir(parents=True, exist_ok=True) + + # ------------------------------------------------------------------ + # collection + # ------------------------------------------------------------------ + async def collect( + self, + origin: str, + datasets: list[str] | None = None, + dadosgov_token: str | None = None, + ) -> list[FileRecord]: + """Collect the file listing of *origin* (``ftp``, ``dadosgov`` or + ``ducklake``), optionally restricted to *datasets* (canonical + uppercase names).""" + origin = origin.strip().lower() + if origin == "ftp": + return await self._collect_ftp(datasets) + if origin == "dadosgov": + return await self._collect_dadosgov(datasets, token=dadosgov_token) + if origin == "ducklake": + return await self._collect_ducklake(datasets) + raise ValueError(f"Unknown origin: {origin!r}") + + async def collect_all( + self, + datasets: list[str] | None = None, + dadosgov_token: str | None = None, + ) -> dict[str, list[FileRecord]]: + """Collect from every client, returning ``{origin: records}``.""" + return { + "ducklake": await self._collect_ducklake(datasets), + "ftp": await self._collect_ftp(datasets), + "dadosgov": await self._collect_dadosgov( + datasets, token=dadosgov_token + ), + } + + async def _collect_ftp( + self, datasets: list[str] | None = None + ) -> list[FileRecord]: + client = await self.pysus.get_ftp() + records: list[FileRecord] = [] + for dataset in await client.datasets(): + if datasets and dataset.name.upper() not in datasets: + continue + for item in await dataset.content: + records.extend(await self._walk_ftp_item(item)) + return records + + async def _walk_ftp_item(self, item: Any) -> list[FileRecord]: + from pysus.api.ftp.models import Directory + from pysus.api.ftp.models import File as FTPFile + from pysus.api.models import BaseRemoteGroup + + if isinstance(item, FTPFile): + record = FileRecord( + origin="ftp", + dataset=item.dataset.name, + name=item.basename, + path=str(item.path), + size=item.size, + modified=_safe_modify(item), + group=getattr(item.group, "name", None), + year=item.year, + month=item.month, + state=item.state, + file=item, + ) + return [record] + + if isinstance(item, BaseRemoteGroup): + records: list[FileRecord] = [] + for file in await item.files: + records.extend(await self._walk_ftp_item(file)) + return records + + if isinstance(item, Directory): + dir_records: list[FileRecord] = [] + for child in await item.content: + dir_records.extend(await self._walk_ftp_item(child)) + return dir_records + + return [] + + async def _collect_dadosgov( + self, + datasets: list[str] | None = None, + token: str | None = None, + ) -> list[FileRecord]: + from pysus.api.models import BaseRemoteGroup + + client = await self.pysus.get_dadosgov(token) + records: list[FileRecord] = [] + for dataset in await client.datasets(): + if datasets and dataset.name.upper() not in datasets: + continue + for group in await dataset.content: + if not isinstance(group, BaseRemoteGroup): + continue + for file in await group.files: + records.append( + FileRecord( + origin="dadosgov", + dataset=dataset.name, + name=file.basename, + path=str(file.path), + size=file.size, + modified=_safe_modify(file), + group=getattr(group, "name", None), + year=file.year, + month=file.month, + state=file.state, + file=file, + ) + ) + return records + + async def _collect_ducklake( + self, datasets: list[str] | None = None + ) -> list[FileRecord]: + client = await self.pysus.get_ducklake() + records: list[FileRecord] = [] + for dataset in await client.datasets(): + if datasets and dataset.name.upper() not in datasets: + continue + for file in await dataset.query(): + record = file.record + records.append( + FileRecord( + origin="ducklake", + dataset=dataset.name, + name=file.basename, + path=str(file.path), + size=file.size, + modified=record.modified, + group=(record.group.name if record.group else None), + year=record.year, + month=record.month, + state=record.state, + sha256=record.sha256, + rows=record.rows, + source_path=record.origin_path, + source_size=record.origin_size, + source_modified=record.origin_modified, + file=file, + ) + ) + return records + + # ------------------------------------------------------------------ + # snapshot persistence + # ------------------------------------------------------------------ + def _snapshot_path(self, origin: str) -> Path: + return self.snapshot_dir / f"{origin.lower()}.json" + + def save_snapshot(self, origin: str, records: list[FileRecord]) -> Path: + """Persist *records* as the latest snapshot for *origin*.""" + path = self._snapshot_path(origin) + payload = { + "origin": origin.lower(), + "captured_at": datetime.now().isoformat(timespec="seconds"), + "count": len(records), + "records": [r.to_dict() for r in records], + } + path.write_text(json.dumps(payload, indent=2, default=str)) + return path + + def load_snapshot(self, origin: str) -> list[FileRecord] | None: + """Load the previous snapshot for *origin*, if any.""" + path = self._snapshot_path(origin) + if not path.exists(): + return None + try: + payload: dict[str, Any] = json.loads(path.read_text()) + return [FileRecord.from_dict(r) for r in payload["records"]] + except (json.JSONDecodeError, KeyError, TypeError, ValueError): + return None + + def diff( + self, + previous: list[FileRecord] | None, + current: list[FileRecord], + origin: str, + ) -> SnapshotDiff: + """Compare a previous snapshot with the current listing.""" + prev_by_path = {r.path: r for r in (previous or [])} + curr_by_path = {r.path: r for r in current} + + added = [r for p, r in curr_by_path.items() if p not in prev_by_path] + removed = [r for p, r in prev_by_path.items() if p not in curr_by_path] + changed = [ + (prev_by_path[p], curr_by_path[p]) + for p in prev_by_path.keys() & curr_by_path.keys() + if _record_changed(prev_by_path[p], curr_by_path[p]) + ] + + return SnapshotDiff( + origin=origin, + added=added, + removed=removed, + changed=changed, + ) + + +def _safe_modify(file: BaseRemoteFile) -> datetime | None: + try: + return file.modify + except (ValueError, AttributeError): + return None + + +def _record_changed(previous: FileRecord, current: FileRecord) -> bool: + return previous.size != current.size or _safe_iso( + previous.modified + ) != _safe_iso(current.modified) + + +def _safe_iso(value: datetime | None) -> str | None: + return value.isoformat() if value else None diff --git a/pysus/management/normalize.py b/pysus/management/normalize.py new file mode 100644 index 00000000..eec9a1f6 --- /dev/null +++ b/pysus/management/normalize.py @@ -0,0 +1,623 @@ +"""Bucket normalization: canonicalize S3 parquet keys and fix catalog paths. + +Ensures every parquet object is stored under the hierarchical key +convention (``public/data////// +/.parquet``) and that the DuckLake catalog ``files.path`` +rows point to the objects that actually exist. + +Renames are copy+delete (S3 has no move). Attribute gaps are enriched +with the per-dataset formatters already shipped with the clients, so each +dataset's specific filename conventions drive the migration. +""" + +from __future__ import annotations + +from collections import defaultdict +from collections.abc import Callable +from dataclasses import dataclass, field +from pathlib import Path + +import boto3 +from botocore.config import Config + +from .records import compose_s3_key, parquet_key + +_BUCKET = "pysus" +_ENDPOINT = "nbg1.your-objectstorage.com" +_REGION = "nbg1" + +_SCAN_PREFIXES = ( + "public/data/dadosgov/", + "public/data/ftp/ciha/", + "data/ftp/", +) + +_FORMAT_RANK = {"csv": 0, "json": 1, "xml": 2, "xlsx": 3} + + +def _format_of(name: str) -> str | None: + """Return the format token of a name, if present (csv/json/xml).""" + lower = name.lower() + for fmt in _FORMAT_RANK: + if f".{fmt}." in lower or lower.endswith(f"_{fmt}"): + return fmt + return None + + +@dataclass +class ObjectRename: + old: str + new: str + size: int = 0 + + +@dataclass +class CatalogPathFix: + catalog: str + old_path: str + new_path: str + + +@dataclass +class CatalogRowDelete: + catalog: str + path: str + reason: str = "" + + +@dataclass +class NormalizePlan: + object_renames: list[ObjectRename] = field(default_factory=list) + object_deletes: list[str] = field(default_factory=list) + catalog_fixes: list[CatalogPathFix] = field(default_factory=list) + catalog_row_deletes: list[CatalogRowDelete] = field(default_factory=list) + broken_rows: list[tuple[str, str]] = field(default_factory=list) + raw_objects: list[str] = field(default_factory=list) + + def summary(self) -> dict: + return { + "object_renames": len(self.object_renames), + "object_deletes": len(self.object_deletes), + "catalog_fixes": len(self.catalog_fixes), + "catalog_row_deletes": len(self.catalog_row_deletes), + "broken_rows": len(self.broken_rows), + "raw_objects": len(self.raw_objects), + } + + +_FORMATTER_CACHE: dict[tuple[str, str], Callable | None] = {} + + +def formatter_for(origin: str, dataset: str) -> Callable | None: + """Return the filename formatter for *origin*/*dataset*, if any. + + Formatters are the per-dataset parsers shipped with each client; they + encode each dataset's specific filename conventions (group codes, + state/month/year positions), so the migration stays data-driven. + """ + key = (origin.strip().lower(), dataset.strip().upper()) + if key in _FORMATTER_CACHE: + return _FORMATTER_CACHE[key] + + formatter: Callable | None = None + try: + if key[0] == "ftp": + from pysus.api.ftp.databases import ( + AVAILABLE_DATABASES as FTP_DATABASES, + ) + + for ftp_class in FTP_DATABASES: + if ftp_class.__name__.upper() == key[1]: + formatter = ftp_class.model_construct().formatter + break + elif key[0] == "dadosgov": + from pysus.api.dadosgov.databases import ( + AVAILABLE_DATABASES as DADOSGOV_DATABASES, + ) + + for gov_class in DADOSGOV_DATABASES: + if gov_class.__name__.upper() == key[1]: + formatter = gov_class.model_construct().formatter + break + except Exception: # noqa + formatter = None + + _FORMATTER_CACHE[key] = formatter + return formatter + + +class BucketNormalizer: + """Survey and normalize parquet object keys and catalog paths on S3.""" + + def __init__(self, access_key: str, secret_key: str): + self.client = boto3.client( + "s3", + endpoint_url=f"https://{_ENDPOINT}", + region_name=_REGION, + aws_access_key_id=access_key, + aws_secret_access_key=secret_key, + config=Config(signature_version="s3v4"), + ) + self.raw_objects: list[str] = [] + self.broken_rows: list[tuple[str, str]] = [] + + # ------------------------------------------------------------------ + # survey + # ------------------------------------------------------------------ + def _list_objects(self, prefix: str) -> list[tuple[str, int]]: + paginator = self.client.get_paginator("list_objects_v2") + objects: list[tuple[str, int]] = [] + for page in paginator.paginate(Bucket=_BUCKET, Prefix=prefix): + for obj in page.get("Contents", []): + objects.append((obj["Key"], obj["Size"])) + return objects + + def _object_exists(self, key: str) -> bool: + try: + self.client.head_object(Bucket=_BUCKET, Key=key) + return True + except Exception: # noqa + return False + + def survey_objects(self) -> tuple[list[ObjectRename], list[str]]: + """Find non-canonical *parquet* object keys and resolve collisions. + + Only objects whose key ends in ``.parquet`` are candidates — raw + source objects (``.dbc``, ``.dbf``, ``.zip``...) are never renamed + (that would mislabel content; converting them is ETL work, not a + rename). + + When csv/json/xml parquet variants normalize to the same key, the + format highest in ``_FORMAT_RANK`` (csv first) is kept and the + others are scheduled for deletion. + """ + renames: list[ObjectRename] = [] + deletes: list[str] = [] + by_canonical: dict[str, list[tuple[str, int]]] = defaultdict(list) + + for prefix in _SCAN_PREFIXES: + for key, size in self._list_objects(prefix): + if not key.endswith(".parquet"): + self.raw_objects.append(key) + continue + base = key.rsplit("/", 1)[-1] + canonical = key.rsplit("/", 1)[0] + "/" + parquet_key(base) + by_canonical[canonical].append((key, size)) + + for canonical, items in by_canonical.items(): + if len(items) == 1: + key, size = items[0] + if key != canonical: + renames.append( + ObjectRename(old=key, new=canonical, size=size) + ) + continue + + ranked = sorted( + items, + key=lambda kv: ( + _FORMAT_RANK.get(_format_of(kv[0]) or "", 99), + kv[1], + ), + ) + winner_key, winner_size = ranked[0] + if winner_key != canonical: + renames.append( + ObjectRename( + old=winner_key, new=canonical, size=winner_size + ) + ) + for loser_key, _ in ranked[1:]: + deletes.append(loser_key) + + return renames, deletes + + def survey_catalog( + self, catalog_dir: Path + ) -> tuple[list[CatalogPathFix], list[CatalogRowDelete]]: + """Build catalog path fixes and duplicate-row deletions. + + * Rows whose path ends ``.dbc`` and whose object exists on S3 are + raw artifacts: kept as-is (converting is ETL work). + * Rows whose path ends ``.dbc`` and whose object is missing while a + sibling ``.parquet`` row exists are duplicate stale rows: deleted. + * Other non-canonical rows are fixed when the canonical object + exists (e.g. ``*.csv.parquet`` -> ``*.parquet`` after the object + rename). + """ + import duckdb + + fixes: list[CatalogPathFix] = [] + row_deletes: list[CatalogRowDelete] = [] + catalog = catalog_dir.name.removesuffix(".duckdb").removeprefix( + "catalog_" + ) + con = duckdb.connect(str(catalog_dir), read_only=True) + try: + rows = con.execute("SELECT path FROM pysus.files").fetchall() + finally: + con.close() + + parquet_paths = {p for (p,) in rows if p.endswith(".parquet")} + + for (path,) in rows: + base = path.rsplit("/", 1)[-1] + canonical = path.rsplit("/", 1)[0] + "/" + parquet_key(base) + if canonical == path: + continue + + if path.endswith(".dbc"): + if self._object_exists(path): + self.raw_objects.append(path) + continue + sibling = path.rsplit("/", 1)[0] + "/" + parquet_key(base) + if sibling in parquet_paths: + row_deletes.append( + CatalogRowDelete( + catalog=catalog, + path=path, + reason=f"stale duplicate of {sibling}", + ) + ) + else: + self.broken_rows.append((catalog, path)) + continue + + if self._object_exists(canonical): + fixes.append( + CatalogPathFix( + catalog=catalog, old_path=path, new_path=canonical + ) + ) + else: + self.broken_rows.append((catalog, path)) + + return fixes, row_deletes + + # ------------------------------------------------------------------ + # apply + # ------------------------------------------------------------------ + def _copy_source(self, key: str, size: int | None = None) -> str: + """Return the actual content key for *key*, following alias markers. + + Objects larger than an alias marker (a few hundred bytes) are + never markers, so the HEAD can be skipped when *size* is known. + """ + if size is not None and size > 4096: + return key + for _ in range(5): + try: + head = self.client.head_object(Bucket=_BUCKET, Key=key) + except Exception: # noqa + return key + target = head.get("Metadata", {}).get("pysus-alias") + if not target: + return key + key = str(target) + raise RuntimeError(f"Too many alias hops resolving {key}") + + def _do_relocate(self, rename: ObjectRename, sizes: dict[str, int]) -> None: + from pysus.api.ducklake.functional import alias_marker + + source = self._copy_source(rename.old, sizes.get(rename.old)) + if source == rename.new: + # already relocated by a previous run (old key holds an alias + # marker pointing at the very same target) + return + self.client.copy_object( + Bucket=_BUCKET, + CopySource={"Bucket": _BUCKET, "Key": source}, + Key=rename.new, + ) + self.client.put_object( + Bucket=_BUCKET, + Key=rename.old, + Body=alias_marker(rename.new).encode(), + Metadata={"pysus-alias": rename.new}, + ) + + def apply_objects( + self, + renames: list[ObjectRename], + deletes: list[str], + dry_run: bool = True, + ) -> None: + """Apply renames (with aliases) and deletions on the bucket.""" + if renames: + self.apply_renames_with_aliases(renames, dry_run=dry_run) + + for key in deletes: + print( + f"{'DRY ' if dry_run else ''}" + f"delete {key} (duplicate format)" + ) + if dry_run: + continue + self.client.delete_object(Bucket=_BUCKET, Key=key) + + def apply_renames_with_aliases( + self, + renames: list[ObjectRename], + dry_run: bool = True, + object_sizes: dict[str, int] | None = None, + workers: int = 16, + ) -> dict[str, str]: + """Copy objects to their new keys and leave alias markers behind. + + Instead of deleting the old key, a tiny pointer object + (``{"pysus-alias": ""}`` content + ``pysus-alias`` custom + metadata) is written there, keeping old paths resolvable for + backwards compatibility. Sources that are themselves aliases are + copied from their target, so re-runs never propagate markers. + Copy+marker writes run in parallel (``workers`` threads); S3 + objects are immutable inputs, so ordering does not matter. + """ + from concurrent.futures import ThreadPoolExecutor + + sizes = object_sizes or {} + aliases: dict[str, str] = {} + failures: dict[str, str] = {} + total = len(renames) + done = 0 + + def apply_one(rename: ObjectRename) -> tuple[ObjectRename, str | None]: + if not dry_run: + try: + self._do_relocate(rename, sizes) + except Exception as exc: # noqa + return rename, str(exc) + return rename, None + + if dry_run: + for rename in renames: + print( + f"DRY relocate {rename.old} " + f"-> {rename.new} (alias kept)" + ) + return {} + + with ThreadPoolExecutor(max_workers=workers) as pool: + for rename, error in pool.map(apply_one, renames): + done += 1 + if error: + failures[rename.old] = error + else: + aliases[rename.old] = rename.new + if done % 250 == 0: + print(f" progress: {done}/{total}", flush=True) + + if failures: + print(f" failures: {len(failures)}", flush=True) + for old, error in list(failures.items())[:10]: + print(f" {old}: {error[:120]}", flush=True) + + return aliases + + def apply_catalog( + self, + catalog_dir: Path, + fixes: list[CatalogPathFix], + row_deletes: list[CatalogRowDelete] | None = None, + dry_run: bool = True, + ) -> None: + """Update/delete ``pysus.files`` rows in-place in a local duckdb.""" + row_deletes = row_deletes or [] + if not fixes and not row_deletes: + return + import duckdb + + for fix in fixes: + print( + f"{'DRY ' if dry_run else ''}catalog[{fix.catalog}] " + f"{fix.old_path} -> {fix.new_path}" + ) + for delete in row_deletes: + print( + f"{'DRY ' if dry_run else ''}catalog[{delete.catalog}] " + f"DELETE {delete.path} ({delete.reason})" + ) + if dry_run: + return + + con = duckdb.connect(str(catalog_dir)) + try: + for fix in fixes: + con.execute( + "UPDATE pysus.files SET path = ? WHERE path = ?", + (fix.new_path, fix.old_path), + ) + for delete in row_deletes: + con.execute( + "DELETE FROM pysus.file_columns WHERE file_id IN " + "(SELECT id FROM pysus.files WHERE path = ?)", + (delete.path,), + ) + con.execute( + "DELETE FROM pysus.files WHERE path = ?", + (delete.path,), + ) + con.execute("CHECKPOINT") + finally: + con.close() + + # ------------------------------------------------------------------ + # hierarchical relayout + # ------------------------------------------------------------------ + def _enrich( + self, + origin: str, + dataset: str, + name: str, + group: str | None, + year: int | None, + month: int | None, + state: str | None, + ) -> dict: + """Fill attribute gaps using the dataset formatter. + + Catalog values win; formatter output fills missing values and + replaces legacy directory names (e.g. group ``"Dados"``) with the + parsed group code. + """ + enriched = { + "group": group, + "year": year, + "month": month, + "state": state, + } + formatter = formatter_for(origin, dataset) + if formatter is None: + return enriched + try: + parsed = formatter(name) + except Exception: # noqa + parsed = {} + + parsed_group = parsed.get("group") + if parsed_group and isinstance(parsed_group, dict): + parsed_group = parsed_group.get("name") + + if parsed_group and str(parsed_group) != enriched["group"]: + # formatters are curated; catalog groups may be legacy + # (e.g. directory names) or NULL + enriched["group"] = str(parsed_group) + if enriched["year"] is None and parsed.get("year"): + enriched["year"] = int(parsed["year"]) + if enriched["month"] is None and parsed.get("month"): + enriched["month"] = int(parsed["month"]) + if enriched["state"] is None and parsed.get("state"): + enriched["state"] = str(parsed["state"]) + return enriched + + def survey_relayout( + self, + catalog_dir: Path, + object_keys: set[str], + ) -> NormalizePlan: + """Plan the hierarchical relayout for one per-dataset catalog.""" + import duckdb + + plan = NormalizePlan() + catalog = catalog_dir.name.removesuffix(".duckdb").removeprefix( + "catalog_" + ) + con = duckdb.connect(str(catalog_dir), read_only=True) + try: + rows = con.execute( + "SELECT f.path, f.year, f.month, f.state, g.name, " + "f.origin_path FROM pysus.files f " + "LEFT JOIN pysus.dataset_groups g ON f.group_id = g.id" + ).fetchall() + finally: + con.close() + + by_new: dict[str, list[tuple[str, str]]] = defaultdict(list) + for path, year, month, state, group, origin_path in rows: + origin, dataset = self._split_key(path) + if origin is None or dataset is None: + plan.broken_rows.append((catalog, path)) + continue + + source_name = ( + Path(origin_path).name if origin_path else Path(path).name + ) + enriched = self._enrich( + origin, dataset, source_name, group, year, month, state + ) + new_key = compose_s3_key( + origin=origin, + dataset=dataset, + name=source_name, + group=enriched["group"], + year=enriched["year"], + month=enriched["month"], + state=enriched["state"], + ) + if new_key == path: + continue + by_new[new_key].append((path, source_name)) + + for new_key, candidates in by_new.items(): + existing = [ + (old, src) for old, src in candidates if old in object_keys + ] + if not existing: + for old, _ in candidates: + plan.catalog_row_deletes.append( + CatalogRowDelete( + catalog=catalog, + path=old, + reason="object missing (stale duplicate)", + ) + ) + continue + + winner, winner_src = existing[0] + plan.object_renames.append(ObjectRename(old=winner, new=new_key)) + plan.catalog_fixes.append( + CatalogPathFix( + catalog=catalog, old_path=winner, new_path=new_key + ) + ) + for old, _src in existing[1:]: + plan.object_deletes.append(old) + plan.catalog_row_deletes.append( + CatalogRowDelete( + catalog=catalog, + path=old, + reason=f"duplicate of {new_key}", + ) + ) + for old, _src in candidates: + if old in object_keys: + continue + plan.catalog_row_deletes.append( + CatalogRowDelete( + catalog=catalog, + path=old, + reason="object missing (stale duplicate)", + ) + ) + + return plan + + @staticmethod + def _split_key(path: str) -> tuple[str | None, str | None]: + """Return ``(origin, dataset)`` from a ``public/data/...`` key.""" + parts = path.split("/") + if len(parts) < 4 or parts[0] != "public" or parts[1] != "data": + return None, None + return parts[2].lower(), parts[3] + + def relocate_uncataloged( + self, + object_keys: set[str], + cataloged_paths: set[str], + ) -> NormalizePlan: + """Plan hierarchical keys for objects missing from the catalogs.""" + plan = NormalizePlan() + for key in sorted(object_keys): + if key in cataloged_paths: + continue + origin, dataset = self._split_key(key) + if origin is None or dataset is None: + plan.raw_objects.append(key) + continue + name = Path(key).name + enriched = self._enrich( + origin, dataset, name, None, None, None, None + ) + new_key = compose_s3_key( + origin=origin, + dataset=dataset, + name=name, + group=enriched["group"], + year=enriched["year"], + month=enriched["month"], + state=enriched["state"], + ) + if new_key == key: + plan.raw_objects.append(key) + continue + plan.object_renames.append(ObjectRename(old=key, new=new_key)) + return plan diff --git a/pysus/management/records.py b/pysus/management/records.py new file mode 100644 index 00000000..d6f4f07d --- /dev/null +++ b/pysus/management/records.py @@ -0,0 +1,396 @@ +"""Normalized records shared by the management workflow. + +These dataclasses are origin-agnostic: every client (FTP, DadosGov, +DuckLake/S3) is reduced to the same :class:`FileRecord` shape so that +tracking, comparison and catalog persistence operate uniformly. +""" + +from __future__ import annotations + +import re +from dataclasses import dataclass, field +from datetime import datetime +from pathlib import Path +from typing import Any + +ORIGINS: tuple[str, ...] = ("ducklake", "ftp", "dadosgov") + +#: Download resolution order: S3 parquet first, FTP second, DadosGov last +#: (only used when a file is not on S3 and requires the API token). +DOWNLOAD_PRIORITY: tuple[str, ...] = ORIGINS + +_FORMAT_SUFFIXES = (".csv", ".json", ".xml", ".xlsx", ".xls", ".dbx") +_COMPRESSION_SUFFIXES = (".zip", ".gz", ".bz2", ".7z", ".rar") + +_NON_ALNUM_RE = re.compile(r"[^a-z0-9]+") + + +def canonical_dataset(name: str) -> str: + """Return the canonical uppercase dataset name.""" + return name.strip().upper() + + +def canonical_group(name: str | None) -> str | None: + """Return the canonical uppercase group code, if any.""" + if not name: + return None + return name.strip().upper() or None + + +def stem_of(name: str) -> str: + """Extract a format-independent file stem. + + ``DENGBR25.csv.zip``, ``DENGBR25.json`` and ``DENGBR25.dbc`` all map to + ``dengbr25``, making them comparable across clients despite different + formats and compression. + """ + stem = base_stem(name) + return _NON_ALNUM_RE.sub("_", stem.lower()).strip("_") + + +def base_stem(name: str) -> str: + """Return the format-independent, case-preserving stem of *name*. + + ``CHIKBR15.csv.zip`` → ``CHIKBR15``; ``DENGBR25.dbc`` → ``DENGBR25``; + ``Mortalidade_Geral_2022_csv.zip`` → ``Mortalidade_Geral_2022``. + + This is the canonical base used to build S3 parquet keys, so format + tokens (dotted or underscored) never leak into catalog paths. + """ + stem = name.strip() + + while stem.lower().endswith(_COMPRESSION_SUFFIXES): + stem = Path(stem).stem + stem = Path(stem).stem + + for suffix in ( + ".csv", + "_csv", + ".json", + "_json", + ".xml", + "_xml", + ".xlsx", + ".xls", + ): + if stem.lower().endswith(suffix): + stem = stem[: -len(suffix)] + break + + return stem + + +def parquet_key(name: str) -> str: + """Return the canonical parquet filename for *name*. + + ``CHIKBR15.csv.zip`` → ``CHIKBR15.parquet`` (not ``CHIKBR15.csv.parquet``). + """ + return f"{base_stem(name)}.parquet" + + +#: Placeholder used in S3 keys for attributes that are not present. +KEY_MISSING = "_" + +#: State used in S3 keys for national (country-wide) files. +NATIONAL_STATE = "BR" + +_KEY_SEGMENT_ORDER = ("group", "year", "month", "state") + + +def compose_s3_key( + origin: str, + dataset: str, + name: str, + group: str | None = None, + year: int | None = None, + month: int | None = None, + state: str | None = None, +) -> str: + """Build the hierarchical S3 key for a parquet artifact. + + The directory structure composes the file characteristics, making the + bucket navigable by prefix:: + + public/data///////.parquet + + Missing attributes use the ``_`` placeholder; a missing state is + interpreted as national (``BR``). Datasets with different metadata + shapes (SINAN has no month, SIA has all four, PNI has no month/state) + therefore keep a stable, predictable layout. + """ + segments = { + "group": canonical_group(group), + "year": str(year) if year is not None else None, + "month": f"{month:02d}" if month is not None else None, + "state": ( + (state.strip().upper() or NATIONAL_STATE) + if state + else NATIONAL_STATE + ), + } + dirs = [origin.strip().lower(), canonical_dataset(dataset).lower()] + dirs.extend(segments[key] or KEY_MISSING for key in _KEY_SEGMENT_ORDER) + return "/".join(["public/data", *dirs, parquet_key(name)]) + + +def format_of(name: str) -> str: + """Return the file format label (e.g. ``csv.zip``, ``dbc``, ``parquet``).""" + lower = name.strip().lower() + compression: list[str] = [] + while lower.endswith(_COMPRESSION_SUFFIXES): + for suffix in _COMPRESSION_SUFFIXES: + if lower.endswith(suffix): + compression.insert(0, suffix.lstrip(".")) + lower = lower[: -len(suffix)] + break + suffix = Path(lower).suffix.lstrip(".") + parts = ([suffix] if suffix else []) + compression + return ".".join(parts) if parts else "unknown" + + +@dataclass(frozen=True) +class IdentityKey: + """The logical identity of a file, independent of client and format.""" + + dataset: str + group: str | None + year: int | None + month: int | None + state: str | None + stem: str + + def as_tuple(self) -> tuple: + return ( + self.dataset, + self.group, + self.year, + self.month, + self.state, + self.stem, + ) + + +@dataclass +class FileRecord: + """One physical artifact of a file on a specific client.""" + + origin: str + dataset: str + name: str + path: str + size: int = 0 + modified: datetime | None = None + group: str | None = None + year: int | None = None + month: int | None = None + state: str | None = None + format: str | None = None + sha256: str | None = None + rows: int | None = None + source_path: str | None = None + source_size: int | None = None + source_modified: datetime | None = None + file: Any = field(default=None, repr=False, compare=False) + + def __post_init__(self) -> None: + self.dataset = canonical_dataset(self.dataset) + self.group = canonical_group(self.group) + self.origin = self.origin.strip().lower() + if self.format is None: + self.format = format_of(self.name) + if self.state is not None: + self.state = self.state.strip().upper() or None + + @property + def stem(self) -> str: + return stem_of(self.name) + + def identity_key(self) -> IdentityKey: + """Build the logical identity key for this record.""" + return IdentityKey( + dataset=self.dataset, + group=self.group, + year=self.year, + month=self.month, + state=self.state, + stem=self.stem, + ) + + def to_dict(self) -> dict[str, Any]: + """Serialize the record for snapshot persistence.""" + return { + "origin": self.origin, + "dataset": self.dataset, + "name": self.name, + "path": self.path, + "size": self.size, + "modified": self.modified.isoformat() if self.modified else None, + "group": self.group, + "year": self.year, + "month": self.month, + "state": self.state, + "format": self.format, + "sha256": self.sha256, + "rows": self.rows, + "source_path": self.source_path, + "source_size": self.source_size, + "source_modified": ( + self.source_modified.isoformat() + if self.source_modified + else None + ), + } + + @classmethod + def from_dict(cls, data: dict[str, Any]) -> FileRecord: + """Rehydrate a record from a serialized snapshot.""" + modified = data.get("modified") + source_modified = data.get("source_modified") + return cls( + origin=data["origin"], + dataset=data["dataset"], + name=data["name"], + path=data["path"], + size=data.get("size", 0), + modified=datetime.fromisoformat(modified) if modified else None, + group=data.get("group"), + year=data.get("year"), + month=data.get("month"), + state=data.get("state"), + format=data.get("format"), + sha256=data.get("sha256"), + rows=data.get("rows"), + source_path=data.get("source_path"), + source_size=data.get("source_size"), + source_modified=( + datetime.fromisoformat(source_modified) + if source_modified + else None + ), + ) + + +@dataclass +class FileComparison: + """Comparison of one logical file across all clients.""" + + key: IdentityKey + records: list[FileRecord] = field(default_factory=list) + + @property + def origins(self) -> set[str]: + return {r.origin for r in self.records} + + @property + def formats(self) -> set[str]: + return {r.format or "unknown" for r in self.records} + + def by_origin(self, origin: str) -> list[FileRecord]: + return [r for r in self.records if r.origin == origin.lower()] + + def _pick(self, origin: str) -> FileRecord | None: + records = self.by_origin(origin) + if not records: + return None + return max( + records, + key=lambda r: (r.modified or datetime.min, r.size), + ) + + def best_record( + self, priorities: tuple[str, ...] = DOWNLOAD_PRIORITY + ) -> FileRecord | None: + """Return the first available record following the priority order.""" + for origin in priorities: + record = self._pick(origin) + if record: + return record + return None + + @property + def is_on_s3(self) -> bool: + return "ducklake" in self.origins + + @property + def only_on_dadosgov(self) -> bool: + return self.origins == {"dadosgov"} + + @property + def needs_token(self) -> bool: + return self.only_on_dadosgov and "ftp" not in self.origins + + def to_dict(self) -> dict[str, Any]: + return { + "key": dict( + zip( + ("dataset", "group", "year", "month", "state", "stem"), + self.key.as_tuple(), + ) + ), + "origins": sorted(self.origins), + "formats": sorted(self.formats), + "records": [r.to_dict() for r in self.records], + } + + +@dataclass +class SnapshotDiff: + """Difference between two snapshots of the same origin.""" + + origin: str + added: list[FileRecord] = field(default_factory=list) + removed: list[FileRecord] = field(default_factory=list) + changed: list[tuple[FileRecord, FileRecord]] = field( + default_factory=list + ) # (previous, current) + + @property + def has_changes(self) -> bool: + return bool(self.added or self.removed or self.changed) + + @property + def changed_count(self) -> int: + return len(self.added) + len(self.removed) + len(self.changed) + + +@dataclass +class SyncOutcome: + """Outcome of a single file during a sync run.""" + + key: IdentityKey + origin: str + status: str # "skipped" | "uploaded" | "failed" | "needs_token" + detail: str = "" + + +@dataclass +class SyncReport: + """Aggregated result of a sync run.""" + + outcomes: list[SyncOutcome] = field(default_factory=list) + dataset: str | None = None + + @property + def uploaded(self) -> list[SyncOutcome]: + return [o for o in self.outcomes if o.status == "uploaded"] + + @property + def skipped(self) -> list[SyncOutcome]: + return [o for o in self.outcomes if o.status == "skipped"] + + @property + def failed(self) -> list[SyncOutcome]: + return [o for o in self.outcomes if o.status == "failed"] + + @property + def needs_token(self) -> list[SyncOutcome]: + return [o for o in self.outcomes if o.status == "needs_token"] + + def summary(self) -> dict[str, int]: + return { + "total": len(self.outcomes), + "uploaded": len(self.uploaded), + "skipped": len(self.skipped), + "failed": len(self.failed), + "needs_token": len(self.needs_token), + } diff --git a/pysus/management/report.py b/pysus/management/report.py new file mode 100644 index 00000000..7d289412 --- /dev/null +++ b/pysus/management/report.py @@ -0,0 +1,111 @@ +"""Cross-client comparison reporting. + +Given inventoried records from the three clients, produce a per-dataset +breakdown of logical files: how many exist on all three clients, on +exactly two, or on a single client (ftp-only, dadosgov-only, s3-only). +""" + +from __future__ import annotations + +from collections import defaultdict +from collections.abc import Iterable +from dataclasses import dataclass, field + +from .compare import Comparator +from .records import FileRecord + + +@dataclass +class DatasetComparison: + dataset: str + total: int = 0 + on_all_three: int = 0 + on_ftp_dadosgov: int = 0 + on_ftp_s3: int = 0 + on_dadosgov_s3: int = 0 + ftp_only: int = 0 + dadosgov_only: int = 0 + s3_only: int = 0 + examples: dict[str, list[str]] = field(default_factory=dict) + + def to_dict(self) -> dict: + return { + "dataset": self.dataset, + "total": self.total, + "on_all_three": self.on_all_three, + "on_ftp_dadosgov": self.on_ftp_dadosgov, + "on_ftp_s3": self.on_ftp_s3, + "on_dadosgov_s3": self.on_dadosgov_s3, + "ftp_only": self.ftp_only, + "dadosgov_only": self.dadosgov_only, + "s3_only": self.s3_only, + "examples": self.examples, + } + + +class ComparisonReporter: + """Build cross-client presence reports from inventoried records.""" + + def __init__(self, comparator: Comparator | None = None): + self.comparator = comparator or Comparator() + + def report( + self, + records: Iterable[FileRecord], + example_limit: int = 3, + ) -> list[DatasetComparison]: + """Return one :class:`DatasetComparison` per dataset in *records*.""" + by_dataset: dict[str, list] = defaultdict(list) + for comparison in self.comparator.compare(records): + by_dataset[comparison.key.dataset].append(comparison) + + reports: list[DatasetComparison] = [] + for dataset, comparisons in sorted(by_dataset.items()): + item = DatasetComparison(dataset=dataset) + item.total = len(comparisons) + + for comparison in comparisons: + origins = comparison.origins + label = self._label(comparison) + if origins == {"ftp", "dadosgov", "ducklake"}: + item.on_all_three += 1 + self._example(item, "all_three", label, example_limit) + elif origins == {"ftp", "dadosgov"}: + item.on_ftp_dadosgov += 1 + self._example(item, "ftp_dadosgov", label, example_limit) + elif origins == {"ftp", "ducklake"}: + item.on_ftp_s3 += 1 + elif origins == {"dadosgov", "ducklake"}: + item.on_dadosgov_s3 += 1 + self._example(item, "dadosgov_s3", label, example_limit) + elif origins == {"ftp"}: + item.ftp_only += 1 + self._example(item, "ftp_only", label, example_limit) + elif origins == {"dadosgov"}: + item.dadosgov_only += 1 + self._example(item, "dadosgov_only", label, example_limit) + elif origins == {"ducklake"}: + item.s3_only += 1 + self._example(item, "s3_only", label, example_limit) + + reports.append(item) + return reports + + @staticmethod + def _label(comparison) -> str: + key = comparison.key + return ( + f"{key.dataset}/{key.group or '-'}/{key.year or '-'}/" + f"{key.month or '-'}/{key.stem}" + ) + + @staticmethod + def _example( + item: DatasetComparison, + category: str, + label: str, + limit: int, + ) -> None: + examples = item.examples.setdefault(category, []) + if len(examples) < limit: + examples.append(label) diff --git a/pysus/management/scripts/__init__.py b/pysus/management/scripts/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/pysus/management/scripts/compare_clients.py b/pysus/management/scripts/compare_clients.py new file mode 100644 index 00000000..881b39f5 --- /dev/null +++ b/pysus/management/scripts/compare_clients.py @@ -0,0 +1,141 @@ +"""Compare datasets across the three clients (FTP, DadosGov, S3/DuckLake). + +Usage: + python -m pysus.management.scripts.compare_clients \ + [--datasets SINAN SIM] + python -m pysus.management.scripts.compare_clients \ + --json --output /tmp/report.json + +Requires ``.env`` (or environment) with ``ACCESS_KEY``, ``SECRET_KEY`` and +optionally ``DADOSGOV_TOKEN`` (DadosGov is skipped without the token). +""" + +from __future__ import annotations + +import argparse +import asyncio +import json +import sys +from pathlib import Path + +from pysus.management.report import ComparisonReporter +from pysus.management.sync import SyncEngine + + +def load_env(path: str = ".env") -> dict[str, str]: + env: dict[str, str] = {} + for line in Path(path).read_text().splitlines(): + line = line.strip() + if not line or line.startswith("#") or "=" not in line: + continue + key, value = line.split("=", 1) + env[key.strip()] = value.strip().strip('"').strip("'") + return env + + +async def run(datasets: list[str] | None) -> dict: + env = load_env() + + engine = SyncEngine( + access_key=env.get("ACCESS_KEY"), + secret_key=env.get("SECRET_KEY"), + dadosgov_token=env.get("DADOSGOV_TOKEN"), + ) + + async with engine: + records = { + "ducklake": await engine.inventory.collect("ducklake", datasets), + "ftp": await engine.inventory.collect("ftp", datasets), + } + records["dadosgov"] = [] + if env.get("DADOSGOV_TOKEN"): + records["dadosgov"] = await engine.inventory.collect( + "dadosgov", datasets, dadosgov_token=env["DADOSGOV_TOKEN"] + ) + + reporter = ComparisonReporter() + reports = reporter.report( + records["ducklake"] + records["ftp"] + records["dadosgov"] + ) + + return { + "origin_counts": { + origin: len(items) for origin, items in records.items() + }, + "reports": [r.to_dict() for r in reports], + } + + +def print_table(result: dict) -> None: + header = ( + f"{'dataset':<10} {'total':>7} {'all3':>6} {'ftp+dg':>7} " + f"{'ftp+s3':>7} {'dg+s3':>6} {'ftp':>6} {'dg':>6} {'s3':>6}" + ) + print(header) + print("-" * len(header)) + for report in result["reports"]: + print( + f"{report['dataset']:<10} {report['total']:>7} " + f"{report['on_all_three']:>6} {report['on_ftp_dadosgov']:>7} " + f"{report['on_ftp_s3']:>7} {report['on_dadosgov_s3']:>6} " + f"{report['ftp_only']:>6} {report['dadosgov_only']:>6} " + f"{report['s3_only']:>6}" + ) + print() + print( + "origin record counts: " + + ", ".join(f"{k}={v}" for k, v in result["origin_counts"].items()) + ) + + for report in result["reports"]: + examples = report.get("examples") or {} + interesting = { + k: v + for k, v in examples.items() + if k in ("ftp_only", "dadosgov_only", "s3_only", "all_three") + } + if interesting: + print(f"\n[{report['dataset']}] examples:") + for category, labels in interesting.items(): + for label in labels: + print(f" {category:<15} {label}") + + +def main() -> int: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument( + "--datasets", + nargs="*", + default=None, + help="Restrict to datasets (e.g. SINAN SIM)", + ) + parser.add_argument( + "--json", + action="store_true", + help="Print the full report as JSON instead of a table", + ) + parser.add_argument( + "--output", + default=None, + help="Also write the JSON report to this file", + ) + args = parser.parse_args() + + datasets = [d.upper() for d in args.datasets] if args.datasets else None + + result = asyncio.run(run(datasets)) + + if args.output: + Path(args.output).write_text(json.dumps(result, indent=2)) + print(f"report written to {args.output}", file=sys.stderr) + + if args.json: + print(json.dumps(result, indent=2)) + else: + print_table(result) + + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/pysus/management/scripts/relayout_bucket.py b/pysus/management/scripts/relayout_bucket.py new file mode 100644 index 00000000..298a5170 --- /dev/null +++ b/pysus/management/scripts/relayout_bucket.py @@ -0,0 +1,180 @@ +"""Relayout the bucket: hierarchical directories compose file attributes. + +Moves every cataloged parquet to:: + + public/data///////.parquet + +Missing attributes use ``_`` (state falls back to ``BR`` for national +files). Catalog ``files.path`` rows are updated accordingly and the +catalogs are re-uploaded to S3. Uncataloged objects are relocated using +the per-dataset formatters when possible. + +Usage: + python -m pysus.management.scripts.relayout_bucket --dry-run + python -m pysus.management.scripts.relayout_bucket +""" + +from __future__ import annotations + +import argparse +from pathlib import Path + +import httpx +from pysus.management.normalize import BucketNormalizer + +CATALOGS = ( + "catalog_ciha", + "catalog_cnes", + "catalog_covid19", + "catalog_ibge", + "catalog_pni", + "catalog_sia", + "catalog_sih", + "catalog_sim", + "catalog_sinan", + "catalog_sinasc", +) + +SCAN_PREFIXES = ("public/data/ftp/", "public/data/dadosgov/", "data/ftp/") + + +def load_env(path: str = ".env") -> dict[str, str]: + env: dict[str, str] = {} + for line in Path(path).read_text().splitlines(): + line = line.strip() + if not line or line.startswith("#") or "=" not in line: + continue + key, value = line.split("=", 1) + env[key.strip()] = value.strip().strip('"').strip("'") + return env + + +def main() -> int: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument( + "--dry-run", + action="store_true", + help="Only print the plan without touching S3", + ) + parser.add_argument( + "--workdir", + default="/tmp/opencode/relayout", + help="Working directory for downloaded catalogs", + ) + args = parser.parse_args() + + env = load_env() + normalizer = BucketNormalizer( + access_key=env["ACCESS_KEY"], + secret_key=env["SECRET_KEY"], + ) + + workdir = Path(args.workdir) + workdir.mkdir(parents=True, exist_ok=True) + + print("listing objects...", flush=True) + object_keys: set[str] = set() + object_sizes: dict[str, int] = {} + for prefix in SCAN_PREFIXES: + for key, size in normalizer._list_objects(prefix): + object_keys.add(key) + object_sizes[key] = size + print(f"objects: {len(object_keys)}", flush=True) + + print("downloading catalogs...", flush=True) + with httpx.Client(follow_redirects=True, timeout=600) as client: + for name in CATALOGS: + url = ( + f"https://nbg1.your-objectstorage.com/pysus/" + f"public/{name}.duckdb" + ) + response = client.get(url) + response.raise_for_status() + (workdir / f"{name}.duckdb").write_bytes(response.content) + print(f" {name} ({len(response.content)} bytes)", flush=True) + + all_old_paths: set[str] = set() + all_new_paths: set[str] = set() + plans = {} + + for name in CATALOGS: + catalog_dir = workdir / f"{name}.duckdb" + print(f"surveying {name}...", flush=True) + plan = normalizer.survey_relayout(catalog_dir, object_keys) + plans[name] = plan + all_old_paths.update(fix.old_path for fix in plan.catalog_fixes) + all_old_paths.update(delete.path for delete in plan.catalog_row_deletes) + all_new_paths.update(fix.new_path for fix in plan.catalog_fixes) + print(f" {name}: {plan.summary()}", flush=True) + + print("relocating uncataloged objects...", flush=True) + orphans = object_keys - all_old_paths + orphan_plan = normalizer.relocate_uncataloged(orphans, set()) + plans["__orphans__"] = orphan_plan + print(f" orphans: {orphan_plan.summary()}", flush=True) + + if args.dry_run: + print("DRY RUN — no changes applied", flush=True) + return 0 + + print("relocating objects (aliases kept at old keys)...", flush=True) + aliases: dict[str, str] = {} + for name, plan in plans.items(): + print(f" applying {name}...", flush=True) + aliases.update( + normalizer.apply_renames_with_aliases( + plan.object_renames, dry_run=False + ) + ) + normalizer.apply_objects([], plan.object_deletes, dry_run=False) + + print("applying catalog updates...", flush=True) + for name in CATALOGS: + plan = plans[name] + normalizer.apply_catalog( + workdir / f"{name}.duckdb", + plan.catalog_fixes, + plan.catalog_row_deletes, + dry_run=False, + ) + + print("uploading catalogs...", flush=True) + import asyncio + import json + + from pysus.api.ducklake.functional import upload_s3 + + async def upload_all(): + for name in CATALOGS: + await upload_s3( + local_path=workdir / f"{name}.duckdb", + remote_path=f"public/{name}.duckdb", + access_key=env["ACCESS_KEY"], + secret_key=env["SECRET_KEY"], + ) + print(f" uploaded {name}", flush=True) + + asyncio.run(upload_all()) + + print("writing alias registry...", flush=True) + registry_key = "public/data/.aliases.json" + try: + registry_obj = normalizer.client.get_object( + Bucket="pysus", Key=registry_key + ) + registry = json.loads(registry_obj["Body"].read()) + except Exception: # noqa + registry = {} + registry.update(aliases) + normalizer.client.put_object( + Bucket="pysus", + Key=registry_key, + Body=json.dumps(registry, indent=2, sort_keys=True).encode(), + ) + print(f"alias registry: {len(registry)} entries", flush=True) + print("done", flush=True) + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/pysus/management/scripts/sync_clients.py b/pysus/management/scripts/sync_clients.py new file mode 100644 index 00000000..9d5b1b69 --- /dev/null +++ b/pysus/management/scripts/sync_clients.py @@ -0,0 +1,90 @@ +"""Sync S3 with every file available on FTP and DadosGov. + +Downloads missing/outdated files, converts them to parquet, uploads to S3 +and updates the DuckLake catalogs. Resumable: files already cataloged with +an equally recent origin are skipped, and the catalogs are checkpointed to +S3 every N uploads. + +Usage: + python -m pysus.management.scripts.sync_clients --datasets SINAN SIM + python -m pysus.management.scripts.sync_clients --checkpoint-every 500 +""" + +from __future__ import annotations + +import argparse +import asyncio +from pathlib import Path + +from pysus.management.sync import SyncEngine + + +def load_env(path: str = ".env") -> dict[str, str]: + env: dict[str, str] = {} + for line in Path(path).read_text().splitlines(): + line = line.strip() + if not line or line.startswith("#") or "=" not in line: + continue + key, value = line.split("=", 1) + env[key.strip()] = value.strip().strip('"').strip("'") + return env + + +async def run(datasets, checkpoint_every, force) -> dict: + env = load_env() + engine = SyncEngine( + access_key=env.get("ACCESS_KEY"), + secret_key=env.get("SECRET_KEY"), + dadosgov_token=env.get("DADOSGOV_TOKEN"), + ) + + counts: dict[str, int] = {} + + def on_outcome(outcome) -> None: + counts[outcome.status] = counts.get(outcome.status, 0) + 1 + if outcome.status in ("uploaded", "failed", "needs_token"): + print(f"[{outcome.status}] {outcome.detail}", flush=True) + total = sum(counts.values()) + if total % 500 == 0: + print(f"progress: {counts}", flush=True) + + async with engine: + report = await engine.run( + datasets=datasets, + force=force, + checkpoint_every=checkpoint_every, + on_outcome=on_outcome, + ) + return report.summary() + + +def main() -> int: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument( + "--datasets", + nargs="*", + default=None, + help="Restrict to datasets (e.g. SINAN SIM)", + ) + parser.add_argument( + "--checkpoint-every", + type=int, + default=500, + help="Upload catalogs to S3 every N successful uploads", + ) + parser.add_argument( + "--force", + action="store_true", + help="Reprocess files even when the catalog is current", + ) + args = parser.parse_args() + + datasets = [d.upper() for d in args.datasets] if args.datasets else None + + summary = asyncio.run(run(datasets, args.checkpoint_every, args.force)) + print(summary) + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/pysus/management/sync.py b/pysus/management/sync.py new file mode 100644 index 00000000..9e9dea06 --- /dev/null +++ b/pysus/management/sync.py @@ -0,0 +1,538 @@ +"""Sync engine: the end-to-end workflow across FTP, DadosGov and S3. + +Pipeline per dataset: + +1. inventory — snapshot every file visible on each client + (:class:`~pysus.management.inventory.Inventory`); +2. compare — group artifacts into logical files + (:class:`~pysus.management.compare.Comparator`); +3. resolve — pick the download source following the fixed priority + S3 → FTP → DadosGov (token required only for DadosGov-only files); +4. load — download, convert to parquet, upload to S3, and persist the + metadata in the DuckLake catalog (:class:`~pysus.management.catalog + .CatalogWriter`). + +Single-writer assumption: the DuckLake catalog duckdbs are uploaded as +whole files on close; concurrent sync runs would clobber each other. +""" + +from __future__ import annotations + +import asyncio +from collections.abc import Callable +from datetime import datetime +from logging import error +from typing import TYPE_CHECKING, Any, cast + +from pysus.api.dadosgov.models import File as APIFile +from pysus.api.ducklake.functional import upload_s3 +from pysus.api.errors import AuthenticationError, ConnectionError +from pysus.api.extensions import Parquet +from pysus.api.ftp.models import File as FTPFile +from pysus.api.models import BaseRemoteFile + +from .catalog import CatalogWriter, sha256_of +from .compare import Comparator +from .inventory import Inventory +from .records import ( + DOWNLOAD_PRIORITY, + FileComparison, + FileRecord, + SyncOutcome, + SyncReport, + compose_s3_key, +) + +if TYPE_CHECKING: # pragma: no cover + from pysus.api.client import PySUS + from pysus.api.ducklake.client import DuckLake + +_RETRYABLE = (ConnectionResetError, ConnectionRefusedError, TimeoutError) + + +class _DatasetStub: + """Registry entry for a per-dataset adapter without a full DuckDataset. + + ``DuckLake.close`` iterates ``_datasets`` and calls ``ds.close``; + the stub delegates to its adapter so the catalog is uploaded on close. + """ + + def __init__(self, name: str, adapter): + self.name = name + self.adapter = adapter + + async def close(self, update_catalog: bool | None = None) -> None: + await self.adapter.close(update=bool(update_catalog)) + + +class SyncEngine: + """Orchestrates inventory → compare → download → parquet → catalog.""" + + pysus: PySUS | None + _ducklake: DuckLake | None + + def __init__( + self, + access_key: str | None = None, + secret_key: str | None = None, + dadosgov_token: str | None = None, + pysus: PySUS | None = None, + ): + self.access_key = access_key + self.secret_key = secret_key + self.dadosgov_token = dadosgov_token + self.pysus = pysus + self._ducklake = None + self._changed_catalog = False + + def _require_pysus(self) -> PySUS: + if self.pysus is None: + raise ConnectionError("PySUS orchestrator is not connected") + return self.pysus + + def _require_ducklake(self) -> DuckLake: + if self._ducklake is None: + raise ConnectionError("DuckLake is not connected") + return self._ducklake + + @property + def inventory(self) -> Inventory: + return Inventory(self._require_pysus()) + + @property + def comparator(self) -> Comparator: + return Comparator() + + @property + def writer(self) -> CatalogWriter: + return CatalogWriter(self._require_ducklake()) + + # ------------------------------------------------------------------ + # lifecycle + # ------------------------------------------------------------------ + async def __aenter__(self) -> SyncEngine: + if self.pysus is None: + from pysus.api.client import PySUS + + self.pysus = PySUS() + await self.pysus.__aenter__() + + self._ducklake = await self.pysus.get_ducklake() + if self.access_key and self.secret_key: + await self._ducklake.login( + access_key=self.access_key, + secret_key=self.secret_key, + ) + return self + + async def __aexit__(self, exc_type, exc_val, exc_tb) -> None: + try: + if not exc_type and self._ducklake: + await self._ducklake.close(update_catalog=self._changed_catalog) + finally: + if self.pysus is not None: + await self.pysus.__aexit__(exc_type, exc_val, exc_tb) + + # ------------------------------------------------------------------ + # single-file upload (used by CatalogManager) + # ------------------------------------------------------------------ + def s3_key_for(self, file: BaseRemoteFile) -> str: + """Return the hierarchical S3 key for *file*'s parquet artifact.""" + group = getattr(file, "group", None) + return compose_s3_key( + origin=file.client.name, + dataset=file.dataset.name, + name=file.basename, + group=getattr(group, "name", None) if group else None, + year=file.year, + month=file.month, + state=file.state, + ) + + async def upload_file( + self, + file: FTPFile | APIFile, + callback: Callable[[int, int], None] | None = None, + force: bool = False, + ) -> bool: + """Download *file*, convert to parquet, upload and catalog it. + + Writes metadata into the correct catalogs: the dataset registry + row into ``catalog.duckdb``, the file/group rows into the + per-dataset ``catalog_.duckdb``, and the column definitions + into ``catalog_columns.duckdb``. + + Returns True if the file was processed; False when the catalog + already holds an equally recent artifact (skip). + """ + if self._ducklake is None: + raise ConnectionError("DuckLake is not connected") + + s3_key = self.s3_key_for(file) + writer = self.writer + + dataset_adapter = self._dataset_adapter(file) + central_adapter = self._ducklake._catalog_adap + columns_adapter = self._ducklake._columns_adap + + await central_adapter.connect() + await columns_adapter.connect() + await dataset_adapter.connect() + + central_conn = central_adapter._engine.raw_connection() + dataset_conn = dataset_adapter._engine.raw_connection() + columns_conn = columns_adapter._engine.raw_connection() + + connections = (central_conn, dataset_conn, columns_conn) + try: + with central_conn, dataset_conn, columns_conn: + central_cursor = central_conn.cursor() + dataset_cursor = dataset_conn.cursor() + columns_cursor = columns_conn.cursor() + + writer._ensure_management_columns(dataset_cursor) + + existing = writer.get_file(dataset_cursor, s3_key) + if existing and not force: + _, origin_modified = existing + if self._is_current(file, origin_modified): + return False + + dataset_id = writer.ensure_dataset( + central_cursor, + file.dataset.name, + file.dataset.long_name, + getattr(file.dataset, "description", None), + ) + + group = getattr(file, "group", None) + group_name = ( + getattr(group, "name", None) if group is not None else None + ) + group_name = str(group_name) if group_name else None + group_id = writer.ensure_group( + dataset_cursor, + dataset_id, + group_name, + getattr(group, "long_name", None) if group else None, + getattr(group, "description", None) if group else None, + ) + + parquet_file = await self._download_with_retry(file, callback) + await upload_s3( + local_path=parquet_file.path, + remote_path=s3_key, + access_key=str(self.access_key), + secret_key=str(self.secret_key), + callback=callback, + ) + + digest = sha256_of(parquet_file.path) + writer.upsert_file( + dataset_cursor, + dataset_id=dataset_id, + group_id=group_id, + path=s3_key, + size=parquet_file.size, + rows=parquet_file.rows, + modified=datetime.now(), + origin_modified=self._safe_modify(file), + origin_size=self._safe_size(file), + origin_path=str(file.path), + year=file.year, + month=file.month, + state=file.state, + origin=file.client.name.lower(), + format="parquet", + sha256=digest, + file_type="PARQUET", + ) + + inserted = writer.get_file(dataset_cursor, s3_key) + assert inserted is not None + file_id, _ = inserted + writer.link_columns( + dataset_cursor, + columns_cursor, + file_id, + parquet_file.schema, + dataset_id, + ) + + central_conn.commit() + dataset_conn.commit() + columns_conn.commit() + dataset_cursor.execute("CHECKPOINT") + columns_cursor.execute("CHECKPOINT") + + central_adapter._local_dirty = True + dataset_adapter._local_dirty = True + columns_adapter._local_dirty = True + self._changed_catalog = True + return True + except BaseException as exc: # noqa + # the connection context managers roll back on exit + for conn in connections: + try: + conn.close() + except Exception: # noqa + pass + raise exc + + def _dataset_adapter(self, file: BaseRemoteFile): + """Return (and register) the per-dataset adapter for *file*.""" + ducklake = self._require_ducklake() + dataset_name = file.dataset.name.lower() + for ds in ducklake._datasets: + if getattr(ds, "name", "").lower() == dataset_name: + return ds.adapter + + from pysus.api.ducklake.catalog.adapters import DatasetAdapter + + adapter = DatasetAdapter( + name=dataset_name, + dataset_id=0, + credentials=ducklake.credentials, + update_on_close=ducklake.update_on_close, + ) + ducklake._datasets.append( + cast(Any, _DatasetStub(dataset_name, adapter)) + ) + return adapter + + async def _download_with_retry( + self, + file: FTPFile | APIFile, + callback: Callable[[int, int], None] | None = None, + max_retries: int = 3, + ) -> Parquet: + last_error: Exception | None = None + token = ( + self.dadosgov_token + if file.client.name.lower() == "dadosgov" + else None + ) + for attempt in range(max_retries): + try: + return await self._require_pysus().download_to_parquet( + file=file, + token=token, + callback=callback, + ) + except _RETRYABLE as exc: + last_error = exc + wait_time = 2**attempt + error( + f"Download attempt {attempt + 1}/{max_retries} failed " + f"for {file.basename}: {exc}. Retrying in {wait_time}s..." + ) + await asyncio.sleep(wait_time) + + raise RuntimeError( + f"Failed to download {file.basename} after {max_retries} " + f"attempts: {last_error}" + ) from last_error + + @staticmethod + def _is_current( + file: BaseRemoteFile, + origin_modified: datetime | None, + ) -> bool: + if origin_modified is None: + return False + try: + file_mod = file.modify + except ValueError: + return False + return str(file_mod) <= str(origin_modified) + + @staticmethod + def _safe_modify(file: BaseRemoteFile) -> datetime | None: + try: + return file.modify + except ValueError: + return None + + @staticmethod + def _safe_size(file: BaseRemoteFile) -> int: + try: + return file.size + except (ValueError, AttributeError): + return 0 + + # ------------------------------------------------------------------ + # full sync run + # ------------------------------------------------------------------ + async def run( + self, + datasets: list[str] | None = None, + force: bool = False, + callback: Callable[[int, int], None] | None = None, + save_snapshots: bool = True, + checkpoint_every: int | None = None, + on_outcome: Callable[[SyncOutcome], None] | None = None, + ) -> SyncReport: + """Run the full pipeline and return a :class:`SyncReport`. + + Files already on S3 (ducklake artifacts) are skipped; FTP is + preferred over DadosGov, which requires ``dadosgov_token``. + Files whose non-S3 artifact is newer than the S3 copy are + reprocessed (most-updated policy). + + ``checkpoint_every`` uploads the modified catalogs to S3 every N + successful uploads, making long runs resumable (files already + cataloged are skipped on the next run). ``on_outcome`` is called + once per processed logical file (e.g. for progress logging). + """ + report = SyncReport(dataset=",".join(datasets) if datasets else None) + + records: dict[str, list[FileRecord]] = { + "ducklake": await self.inventory.collect("ducklake", datasets), + "ftp": await self.inventory.collect("ftp", datasets), + } + records["dadosgov"] = [] + if self.dadosgov_token: + records["dadosgov"] = await self.inventory.collect( + "dadosgov", datasets, dadosgov_token=self.dadosgov_token + ) + + comparisons = self.comparator.compare( + records["ducklake"] + records["ftp"] + records["dadosgov"] + ) + + uploaded_since_checkpoint = 0 + for comparison in comparisons: + outcome = await self._process_comparison( + comparison, force=force, callback=callback + ) + report.outcomes.append(outcome) + if on_outcome: + on_outcome(outcome) + + if ( + outcome.status == "uploaded" + and checkpoint_every + and self._changed_catalog + ): + uploaded_since_checkpoint += 1 + if uploaded_since_checkpoint >= checkpoint_every: + await self._checkpoint() + uploaded_since_checkpoint = 0 + + if self._changed_catalog and checkpoint_every is not None: + await self._checkpoint() + + if save_snapshots: + for origin, items in records.items(): + self.inventory.save_snapshot(origin, items) + + return report + + async def _checkpoint(self) -> None: + """Upload all dirty catalogs to S3 and reconnect the adapters.""" + ducklake = self._require_ducklake() + for ds in ducklake._datasets: + await ds.close(update_catalog=True) + await ducklake._catalog_adap.close(update=True) + await ducklake._columns_adap.close(update=True) + await ducklake._catalog_adap.connect() + await ducklake._columns_adap.connect() + self._changed_catalog = False + + async def _process_comparison( + self, + comparison: FileComparison, + force: bool = False, + callback: Callable[[int, int], None] | None = None, + ) -> SyncOutcome: + key = comparison.key + label = ( + f"{key.dataset}/{key.group or '-'}/" + f"{key.year or '-'}/{key.month or '-'}/{key.stem}" + ) + + if comparison.is_on_s3: + if not force and not self._s3_is_stale(comparison): + return SyncOutcome(key=key, origin="ducklake", status="skipped") + return await self._reprocess(comparison, callback, force, label) + + return await self._reprocess(comparison, callback, force, label) + + @staticmethod + def _s3_is_stale(comparison: FileComparison) -> bool: + """True when any non-S3 artifact is newer than the S3 copy.""" + s3_record = comparison._pick("ducklake") + if s3_record is None: + return False + + s3_source_modified = s3_record.source_modified or s3_record.modified + if s3_source_modified is None: + return False + + for record in comparison.records: + if record.origin == "ducklake": + continue + modified = record.modified + if modified and modified > s3_source_modified: + return True + return False + + async def _reprocess( + self, + comparison: FileComparison, + callback: Callable[[int, int], None] | None, + force: bool, + label: str, + ) -> SyncOutcome: + key = comparison.key + for origin in DOWNLOAD_PRIORITY[1:]: # ducklake already checked + record = comparison._pick(origin) + if record is None or record.file is None: + continue + if origin == "dadosgov" and not self.dadosgov_token: + return SyncOutcome( + key=key, + origin="dadosgov", + status="needs_token", + detail=f"only on DadosGov: {label}", + ) + try: + processed = await self.upload_file( + record.file, + callback=callback, + force=force, + ) + if processed: + return SyncOutcome( + key=key, + origin=origin, + status="uploaded", + detail=label, + ) + return SyncOutcome( + key=key, + origin=origin, + status="skipped", + detail=f"already current: {label}", + ) + except AuthenticationError as exc: + return SyncOutcome( + key=key, + origin=origin, + status="needs_token", + detail=str(exc), + ) + except Exception as exc: # noqa + return SyncOutcome( + key=key, + origin=origin, + status="failed", + detail=f"{label}: {exc}", + ) + + return SyncOutcome( + key=key, + origin="unknown", + status="failed", + detail=f"no downloadable artifact: {label}", + ) diff --git a/pysus/tests/api/ducklake/test_client.py b/pysus/tests/api/ducklake/test_client.py index 5976cc80..2e959d31 100644 --- a/pysus/tests/api/ducklake/test_client.py +++ b/pysus/tests/api/ducklake/test_client.py @@ -184,6 +184,7 @@ async def __anext__(self): mock_http = MagicMock() mock_http.__aenter__.return_value = mock_http + mock_http.head = AsyncMock() httpx_patcher = patch( "pysus.api.ducklake.functional.httpx.AsyncClient", return_value=mock_http, @@ -239,6 +240,7 @@ async def __anext__(self): mock_http = MagicMock() mock_http.__aenter__.return_value = mock_http + mock_http.head = AsyncMock() httpx_patcher = patch( "pysus.api.ducklake.functional.httpx.AsyncClient", return_value=mock_http, @@ -275,6 +277,7 @@ async def test_download_with_callback(self, tmp_path): mock_http = MagicMock() mock_http.__aenter__.return_value = mock_http + mock_http.head = AsyncMock() stream_cm = MagicMock() diff --git a/pysus/tests/api/ducklake/test_functional.py b/pysus/tests/api/ducklake/test_functional.py index a261a61d..3154d3cc 100644 --- a/pysus/tests/api/ducklake/test_functional.py +++ b/pysus/tests/api/ducklake/test_functional.py @@ -1,11 +1,25 @@ """Tests for pysus.api.ducklake.functional (HTTP/S3 download utilities).""" +import json from pathlib import Path from unittest.mock import AsyncMock, MagicMock, patch import httpx import pytest -from pysus.api.ducklake.functional import download_http, download_s3 +from pysus.api.ducklake.functional import ( + alias_marker, + download_http, + download_s3, +) + + +def _head_response(alias: str | None = None) -> MagicMock: + response = MagicMock() + response.status_code = 200 + response.headers = {} + if alias: + response.headers["x-amz-meta-pysus-alias"] = alias + return response @pytest.mark.asyncio @@ -18,6 +32,8 @@ async def test_download_http_success(tmp_path): mock_client = MagicMock() mock_client.__aenter__.return_value = mock_client + mock_client.head = AsyncMock(return_value=_head_response()) + mock_client.head = AsyncMock(return_value=_head_response()) async def fake_aiter_bytes(**kwargs): yield content @@ -34,6 +50,61 @@ async def fake_aiter_bytes(**kwargs): assert local.read_bytes() == content +@pytest.mark.asyncio +async def test_download_http_follows_alias(tmp_path): + local = tmp_path / "test.bin" + content = b"alias followed" + + mock_response = MagicMock() + mock_response.headers = {"Content-Length": str(len(content))} + mock_client = MagicMock() + mock_client.__aenter__.return_value = mock_client + alias_key = "public/data/ftp/sinan/DENG/2025/_/BR/DENGBR25.parquet" + mock_client.head = AsyncMock( + side_effect=[_head_response(alias=alias_key), _head_response()] + ) + + async def fake_aiter_bytes(**kwargs): + yield content + + mock_response.aiter_bytes = fake_aiter_bytes + + mock_stream = MagicMock() + mock_stream.__aenter__.return_value = mock_response + mock_client.stream.return_value = mock_stream + + with patch("httpx.AsyncClient", return_value=mock_client): + await download_http("public/data/ftp/sinan/DENGBR25.parquet", local) + + assert local.read_bytes() == content + calls = mock_client.head.call_args_list + assert len(calls) == 2 + stream_call = mock_client.stream.call_args_list[-1] + assert "public/data/ftp/sinan/DENG/2025/_/BR/DENGBR25.parquet" in ( + stream_call.args[1] if len(stream_call.args) > 1 else "" + ) + + +@pytest.mark.asyncio +async def test_download_http_alias_too_many_hops(tmp_path): + local = tmp_path / "test.bin" + mock_client = MagicMock() + mock_client.__aenter__.return_value = mock_client + mock_client.head = AsyncMock(return_value=_head_response()) + mock_client.head = AsyncMock( + return_value=_head_response(alias="public/data/x") + ) + + with patch("httpx.AsyncClient", return_value=mock_client): + with pytest.raises(RuntimeError, match="alias hops"): + await download_http("public/data/old", local) + + +def test_alias_marker_content(): + marker = json.loads(alias_marker("public/data/new/key.parquet")) + assert marker == {"pysus-alias": "public/data/new/key.parquet"} + + @pytest.mark.asyncio async def test_download_http_retry_on_remote_protocol_error(tmp_path): local = tmp_path / "test.bin" @@ -42,6 +113,7 @@ async def test_download_http_retry_on_remote_protocol_error(tmp_path): mock_client = MagicMock() mock_client.__aenter__.return_value = mock_client + mock_client.head = AsyncMock(return_value=_head_response()) def make_response(): call_count[0] += 1 @@ -86,6 +158,7 @@ async def test_download_http_cleanup_partial_on_error(tmp_path): mock_client = MagicMock() mock_client.__aenter__.return_value = mock_client + mock_client.head = AsyncMock(return_value=_head_response()) class ErrorCtx: async def __aenter__(self): @@ -117,6 +190,7 @@ async def test_download_http_retry_on_http_error(tmp_path): mock_client = MagicMock() mock_client.__aenter__.return_value = mock_client + mock_client.head = AsyncMock(return_value=_head_response()) def make_response(): call_count[0] += 1 @@ -165,6 +239,7 @@ async def test_download_http_callback(tmp_path): mock_response.headers = {"Content-Length": str(len(content))} mock_client = MagicMock() mock_client.__aenter__.return_value = mock_client + mock_client.head = AsyncMock(return_value=_head_response()) async def fake_aiter_bytes(**kwargs): yield content @@ -250,6 +325,7 @@ async def test_download_http_unlink_os_error_swallowed(tmp_path): mock_client = MagicMock() mock_client.__aenter__.return_value = mock_client + mock_client.head = AsyncMock(return_value=_head_response()) class ErrorCtx: async def __aenter__(self): diff --git a/pysus/tests/management/__init__.py b/pysus/tests/management/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/pysus/tests/management/test_compare.py b/pysus/tests/management/test_compare.py new file mode 100644 index 00000000..21ea2570 --- /dev/null +++ b/pysus/tests/management/test_compare.py @@ -0,0 +1,176 @@ +"""Tests for pysus.management.compare and inventory diffing.""" + +from datetime import datetime +from unittest.mock import MagicMock + +import pandas as pd +from pysus.management.compare import ( + FORMAT_PREFERENCE, + Comparator, + content_fingerprint, +) +from pysus.management.inventory import Inventory +from pysus.management.records import FileComparison, FileRecord + + +def _record(origin, name, dataset="SINAN", group="DENG", year=2025, **kw): + return FileRecord( + origin=origin, + dataset=dataset, + name=name, + path=f"{origin}/{dataset}/{name}", + group=group, + year=year, + size=kw.pop("size", 100), + modified=kw.pop("modified", datetime(2026, 1, 1)), + file=kw.pop("file", MagicMock()), + **kw, + ) + + +class TestComparator: + def test_same_logical_file_groups_across_origins(self): + records = [ + _record("ftp", "DENGBR25.dbc"), + _record("dadosgov", "DENGBR25.csv.zip"), + _record("ducklake", "DENGBR25.parquet", size=10), + ] + comparisons = Comparator().compare(records) + assert len(comparisons) == 1 + assert comparisons[0].origins == {"ftp", "dadosgov", "ducklake"} + + def test_different_years_do_not_group(self): + records = [ + _record("ftp", "DENGBR25.dbc", year=2025), + _record("ftp", "DENGBR24.dbc", year=2024), + ] + comparisons = Comparator().compare(records) + assert len(comparisons) == 2 + + def test_csv_json_xml_triplet_deduped(self): + records = [ + _record("dadosgov", "DENGBR25.csv.zip"), + _record("dadosgov", "DENGBR25.json.zip"), + _record("dadosgov", "DENGBR25.xml.zip"), + ] + comparisons = Comparator().compare(records) + assert len(comparisons) == 1 + assert len(comparisons[0].records) == 1 + assert comparisons[0].records[0].format == "csv.zip" + + def test_best_record_priority_s3_first(self): + ftp = _record("ftp", "DENGBR25.dbc") + api = _record("dadosgov", "DENGBR25.csv.zip") + s3 = _record("ducklake", "DENGBR25.parquet") + comparison = FileComparison( + key=ftp.identity_key(), records=[api, ftp, s3] + ) + assert comparison.best_record() is s3 + assert comparison.is_on_s3 + + def test_best_record_ftp_before_dadosgov(self): + ftp = _record("ftp", "DENGBR25.dbc") + api = _record("dadosgov", "DENGBR25.csv.zip") + comparison = FileComparison(key=ftp.identity_key(), records=[api, ftp]) + assert comparison.best_record() is ftp + assert not comparison.is_on_s3 + assert not comparison.only_on_dadosgov + + def test_only_on_dadosgov(self): + api = _record("dadosgov", "DENGBR25.csv.zip") + comparison = FileComparison(key=api.identity_key(), records=[api]) + assert comparison.only_on_dadosgov + assert comparison.best_record() is api + + def test_state_none_vs_br_does_not_split(self): + ftp = _record("ftp", "DENGBR25.dbc", state=None) + api = _record("dadosgov", "DENGBR25.csv.zip", state="BR") + comparisons = Comparator().compare([ftp, api]) + assert len(comparisons) == 1 + assert comparisons[0].origins == {"ftp", "dadosgov"} + assert comparisons[0].key.state == "BR" + + def test_state_level_files_keep_own_groups_via_stem(self): + a = _record("ftp", "DNRJ2401.dbc", state="RJ", year=2024) + b = _record("ftp", "DNRS2401.dbc", state="RS", year=2024) + comparisons = Comparator().compare([a, b]) + assert len(comparisons) == 2 + + def test_group_none_vs_group_merges(self): + ftp = _record("ftp", "ACBIBR07.dbc", group="ACBI") + s3 = _record("ducklake", "ACBIBR07.parquet", group=None) + comparisons = Comparator().compare([ftp, s3]) + assert len(comparisons) == 1 + assert comparisons[0].key.group == "ACBI" + + def test_month_none_vs_month_merges(self): + ftp = _record("ftp", "PAAC2408.dbc", group="PA", month=8) + api = _record("dadosgov", "PAAC2408.csv.zip", group="PA", month=None) + comparisons = Comparator().compare([ftp, api]) + assert len(comparisons) == 1 + assert comparisons[0].key.month == 8 + + def test_format_preference_csv_first(self): + assert FORMAT_PREFERENCE[0] == "csv" + + +class TestContentFingerprint: + def test_same_content_same_fingerprint(self): + a = pd.DataFrame({"col": [1, 2, 3]}) + b = pd.DataFrame({"col": [1, 2, 3]}) + assert content_fingerprint(a) == content_fingerprint(b) + + def test_different_rows_differ(self): + a = pd.DataFrame({"col": [1, 2, 3]}) + b = pd.DataFrame({"col": [1, 2, 4]}) + assert content_fingerprint(a) != content_fingerprint(b) + + def test_column_order_invariant(self): + a = pd.DataFrame({"x": [1], "y": [2]}) + b = pd.DataFrame({"y": [2], "x": [1]}) + assert content_fingerprint(a) == content_fingerprint(b) + + +class TestInventoryDiff: + def _inventory(self): + return Inventory(pysus=MagicMock()) + + def test_diff_empty_previous(self): + inventory = self._inventory() + current = [ + _record("ftp", "A.dbc"), + _record("ftp", "B.dbc"), + ] + diff = inventory.diff(None, current, "ftp") + assert len(diff.added) == 2 + assert diff.removed == [] + assert diff.changed == [] + assert diff.has_changes + + def test_diff_added_removed_changed(self): + inventory = self._inventory() + previous = [ + _record("ftp", "A.dbc", size=100), + _record("ftp", "B.dbc", size=100), + ] + current = [ + _record("ftp", "B.dbc", size=200), + _record("ftp", "C.dbc", size=100), + ] + diff = inventory.diff(previous, current, "ftp") + assert [r.name for r in diff.added] == ["C.dbc"] + assert [r.name for r in diff.removed] == ["A.dbc"] + assert len(diff.changed) == 1 + assert diff.changed_count == 3 + + def test_snapshot_roundtrip(self, tmp_path): + inventory = Inventory(pysus=MagicMock(), snapshot_dir=tmp_path) + records = [_record("ftp", "A.dbc", sha256="deadbeef")] + path = inventory.save_snapshot("ftp", records) + loaded = inventory.load_snapshot("ftp") + assert path.exists() + assert loaded == records + + def test_load_missing_snapshot(self, tmp_path): + inventory = Inventory(pysus=MagicMock(), snapshot_dir=tmp_path) + assert inventory.load_snapshot("ftp") is None diff --git a/pysus/tests/management/test_records.py b/pysus/tests/management/test_records.py new file mode 100644 index 00000000..e290d85a --- /dev/null +++ b/pysus/tests/management/test_records.py @@ -0,0 +1,209 @@ +"""Tests for pysus.management.records: identity keys and normalization.""" + +from datetime import datetime + +from pysus.management.records import ( + DOWNLOAD_PRIORITY, + FileRecord, + base_stem, + canonical_dataset, + canonical_group, + compose_s3_key, + format_of, + parquet_key, + stem_of, +) + + +class TestComposeS3Key: + def test_full_attributes(self): + key = compose_s3_key( + origin="ftp", + dataset="SIA", + name="PAAC2501.dbc", + group="PA", + year=2025, + month=1, + state="AC", + ) + assert key == "public/data/ftp/sia/PA/2025/01/AC/PAAC2501.parquet" + + def test_missing_month_state_defaults_br(self): + key = compose_s3_key( + origin="ftp", + dataset="SINAN", + name="DENGBR25.dbc", + group="DENG", + year=2025, + ) + assert key == "public/data/ftp/sinan/DENG/2025/_/BR/DENGBR25.parquet" + + def test_missing_group_and_year(self): + key = compose_s3_key( + origin="dadosgov", + dataset="SINAN", + name="dados_tuberculose.csv", + ) + assert ( + key == "public/data/dadosgov/sinan/_/_/_/BR/" + "dados_tuberculose.parquet" + ) + + def test_month_zero_padded(self): + key = compose_s3_key( + origin="ftp", + dataset="SIH", + name="RDAC2502.dbc", + group="RD", + year=2025, + month=2, + state="AC", + ) + assert key == "public/data/ftp/sih/RD/2025/02/AC/RDAC2502.parquet" + + def test_dataset_lowercased(self): + key = compose_s3_key(origin="FTP", dataset="SINAN", name="DENGBR25.dbc") + assert key.startswith("public/data/ftp/sinan/") + + def test_csv_zip_and_dbc_share_key(self): + a = compose_s3_key( + origin="ftp", + dataset="SINAN", + name="DENGBR25.dbc", + group="DENG", + year=2025, + ) + b = compose_s3_key( + origin="dadosgov", + dataset="SINAN", + name="DENGBR25.csv.zip", + group="DENG", + year=2025, + ) + assert a.split("/")[-1] == b.split("/")[-1] == "DENGBR25.parquet" + + +class TestStemOf: + def test_dotted_format_zip(self): + assert stem_of("DENGBR25.csv.zip") == "dengbr25" + + def test_single_extension(self): + assert stem_of("DENGBR25.dbc") == "dengbr25" + + def test_underscore_format_token(self): + assert stem_of("Mortalidade_Geral_2022_csv.zip") == ( + "mortalidade_geral_2022" + ) + + def test_json_variant_matches_csv_variant(self): + assert stem_of("Mortalidade_Geral_2022_json.zip") == stem_of( + "Mortalidade_Geral_2022_csv.zip" + ) + + def test_parquet_is_identity(self): + assert stem_of("DENGBR25.parquet") == "dengbr25" + + def test_plain_csv(self): + assert stem_of("dados_aids_hiv.csv") == "dados_aids_hiv" + + +class TestBaseStem: + def test_case_preserved(self): + assert base_stem("CHIKBR15.csv.zip") == "CHIKBR15" + + def test_dbc(self): + assert base_stem("PFMS0508.dbc") == "PFMS0508" + + def test_underscore_token(self): + assert base_stem("Mortalidade_Geral_2022_csv.zip") == ( + "Mortalidade_Geral_2022" + ) + + +class TestParquetKey: + def test_no_format_token(self): + assert parquet_key("CHIKBR15.csv.zip") == "CHIKBR15.parquet" + + def test_json_variant(self): + assert parquet_key("Mortalidade_Geral_2022_json.parquet") == ( + "Mortalidade_Geral_2022.parquet" + ) + + def test_already_parquet(self): + assert parquet_key("DENGBR25.parquet") == "DENGBR25.parquet" + + def test_dbc(self): + assert parquet_key("PFMS0508.dbc") == "PFMS0508.parquet" + + +class TestFormatOf: + def test_csv_zip(self): + assert format_of("DENGBR25.csv.zip") == "csv.zip" + + def test_dbc(self): + assert format_of("DENGBR25.dbc") == "dbc" + + def test_parquet(self): + assert format_of("DENGBR25.parquet") == "parquet" + + def test_unknown(self): + assert format_of("no_extension_name") == "unknown" + + +class TestCanonical: + def test_dataset(self): + assert canonical_dataset("sinan") == "SINAN" + + def test_group_none(self): + assert canonical_group(None) is None + + def test_group(self): + assert canonical_group("deng") == "DENG" + + +class TestFileRecord: + def _record(self, **kwargs): + defaults = { + "origin": "ftp", + "dataset": "SINAN", + "name": "DENGBR25.dbc", + "path": "/dissemin/publicos/SINAN/DADOS/PRELIM/DENGBR25.dbc", + "size": 100, + "modified": datetime(2026, 1, 1), + "group": "DENG", + "year": 2025, + "state": "BR", + } + defaults.update(kwargs) + return FileRecord(**defaults) + + def test_identity_key(self): + key = self._record().identity_key() + assert key.dataset == "SINAN" + assert key.group == "DENG" + assert key.year == 2025 + assert key.stem == "dengbr25" + + def test_format_inferred(self): + record = self._record() + assert record.format == "dbc" + + def test_csv_zip_matches_dbc_key(self): + ftp = self._record() + api = FileRecord( + origin="dadosgov", + dataset="sinan", + name="DENGBR25.csv.zip", + path="https://example.com/DENGBR25.csv.zip", + group="deng", + year=2025, + state="br", + ) + assert ftp.identity_key() == api.identity_key() + + def test_roundtrip(self): + record = self._record(sha256="abc", rows=10) + assert FileRecord.from_dict(record.to_dict()) == record + + def test_download_priority(self): + assert DOWNLOAD_PRIORITY == ("ducklake", "ftp", "dadosgov") diff --git a/pysus/tests/management/test_report.py b/pysus/tests/management/test_report.py new file mode 100644 index 00000000..943f1f86 --- /dev/null +++ b/pysus/tests/management/test_report.py @@ -0,0 +1,77 @@ +"""Tests for pysus.management.report.""" + +from unittest.mock import MagicMock + +from pysus.management.records import FileRecord +from pysus.management.report import ComparisonReporter + + +def _record(origin, name, dataset="SINAN", year=2025, group=None): + return FileRecord( + origin=origin, + dataset=dataset, + name=name, + path=f"{origin}/{dataset}/{name}", + group=group, + year=year, + file=MagicMock(), + ) + + +class TestComparisonReporter: + def test_all_three(self): + records = [ + _record("ftp", "DENGBR25.dbc", group="DENG"), + _record("dadosgov", "DENGBR25.csv.zip", group="DENG"), + _record("ducklake", "DENGBR25.parquet", group="DENG"), + ] + reports = ComparisonReporter().report(records) + assert len(reports) == 1 + assert reports[0].total == 1 + assert reports[0].on_all_three == 1 + + def test_ftp_only(self): + records = [ + _record("ftp", "DENGBR25.dbc", group="DENG"), + _record("ducklake", "CHIKBR25.parquet", group="CHIK"), + ] + reports = ComparisonReporter().report(records) + report = reports[0] + assert report.total == 2 + assert report.ftp_only == 1 + assert report.s3_only == 1 + assert report.examples["ftp_only"] == ["SINAN/DENG/2025/-/dengbr25"] + + def test_dadosgov_only(self): + records = [ + _record("dadosgov", "MPX_2024_OPENDATASUS.csv.zip"), + ] + reports = ComparisonReporter().report(records) + assert reports[0].dadosgov_only == 1 + assert reports[0].examples["dadosgov_only"] == [ + "SINAN/-/2025/-/mpx_2024_opendatasus" + ] + + def test_dadosgov_s3_pair(self): + records = [ + _record("dadosgov", "dados_tuberculose.csv"), + _record("ducklake", "dados_tuberculose.parquet"), + ] + reports = ComparisonReporter().report(records) + assert reports[0].on_dadosgov_s3 == 1 + + def test_per_dataset_split(self): + records = [ + _record("ftp", "DENGBR25.dbc", dataset="SINAN"), + _record("ftp", "DO25OPEN.dbc", dataset="SIM"), + ] + reports = ComparisonReporter().report(records) + datasets = {r.dataset for r in reports} + assert datasets == {"SIM", "SINAN"} + + def test_to_dict_serializable(self): + import json + + records = [_record("ftp", "DENGBR25.dbc", group="DENG")] + reports = ComparisonReporter().report(records) + json.dumps([r.to_dict() for r in reports]) diff --git a/pysus/tests/management/test_sync.py b/pysus/tests/management/test_sync.py new file mode 100644 index 00000000..f114954c --- /dev/null +++ b/pysus/tests/management/test_sync.py @@ -0,0 +1,149 @@ +"""Tests for pysus.management.sync key resolution (no network).""" + +from unittest.mock import MagicMock + +from pysus.management.sync import SyncEngine + + +class TestSyncEngine: + def _engine(self): + return SyncEngine(access_key="ak", secret_key="sk") + + def _file(self, client_name, dataset_name, basename, **attrs): + file = MagicMock() + file.client.name = client_name + file.dataset.name = dataset_name + file.basename = basename + group = MagicMock() + group.name = attrs.get("group") + file.group = group + file.year = attrs.get("year") + file.month = attrs.get("month") + file.state = attrs.get("state") + return file + + def test_s3_key_ftp_dbc(self): + engine = self._engine() + file = self._file( + "ftp", + "SINAN", + "DENGBR25.dbc", + group="DENG", + year=2025, + ) + assert ( + engine.s3_key_for(file) + == "public/data/ftp/sinan/DENG/2025/_/BR/DENGBR25.parquet" + ) + + def test_s3_key_dadosgov_csv_zip(self): + engine = self._engine() + file = self._file( + "DadosGov", + "SINAN", + "DENGBR25.csv.zip", + group="DENG", + year=2025, + ) + assert ( + engine.s3_key_for(file) + == "public/data/dadosgov/sinan/DENG/2025/_/BR/DENGBR25.parquet" + ) + + def test_s3_key_dadosgov_json_variant_collides_with_csv(self): + engine = self._engine() + csv = self._file( + "DadosGov", + "SIM", + "Mortalidade_Geral_2022_csv.zip", + group="DO", + year=2022, + ) + jsn = self._file( + "DadosGov", + "SIM", + "Mortalidade_Geral_2022.json.zip", + group="DO", + year=2022, + ) + assert engine.s3_key_for(csv) == engine.s3_key_for(jsn) + + def test_s3_key_full_attributes(self): + engine = self._engine() + file = self._file( + "ftp", + "SIA", + "PAAC2501.dbc", + group="PA", + year=2025, + month=1, + state="AC", + ) + assert ( + engine.s3_key_for(file) + == "public/data/ftp/sia/PA/2025/01/AC/PAAC2501.parquet" + ) + + def test_is_current(self): + engine = self._engine() + from datetime import datetime + + file = MagicMock() + file.modify = datetime(2026, 1, 2) + assert engine._is_current(file, datetime(2026, 1, 2)) + assert not engine._is_current(file, datetime(2026, 1, 1)) + assert not engine._is_current(file, None) + + def test_s3_is_stale_when_ftp_newer(self): + from datetime import datetime + + from pysus.management.records import FileComparison, FileRecord + + ftp = FileRecord( + origin="ftp", + dataset="SINAN", + name="DENGBR25.dbc", + path="ftp/x", + modified=datetime(2026, 6, 1), + group="DENG", + year=2025, + ) + s3 = FileRecord( + origin="ducklake", + dataset="SINAN", + name="DENGBR25.parquet", + path="s3/x", + modified=datetime(2026, 1, 1), + source_modified=datetime(2026, 1, 1), + group="DENG", + year=2025, + ) + comparison = FileComparison(key=ftp.identity_key(), records=[ftp, s3]) + assert SyncEngine._s3_is_stale(comparison) + + def test_s3_not_stale_when_equal(self): + from datetime import datetime + + from pysus.management.records import FileComparison, FileRecord + + ftp = FileRecord( + origin="ftp", + dataset="SINAN", + name="DENGBR25.dbc", + path="ftp/x", + modified=datetime(2026, 1, 1), + group="DENG", + year=2025, + ) + s3 = FileRecord( + origin="ducklake", + dataset="SINAN", + name="DENGBR25.parquet", + path="s3/x", + modified=datetime(2026, 1, 2), + source_modified=datetime(2026, 1, 1), + group="DENG", + year=2025, + ) + comparison = FileComparison(key=ftp.identity_key(), records=[ftp, s3]) + assert not SyncEngine._s3_is_stale(comparison) From 7d711decac16abde044ae8bbba9600dd5ab6a72c Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Lu=C3=A3=20Bida=20Vacaro?= Date: Fri, 14 Aug 2026 16:25:25 -0300 Subject: [PATCH 2/5] fix(ducklake): process-shared catalog engines and parallel sync pipeline DuckDB keeps a single database instance per file per process; disposing an adapter's engine tore that instance down for every other adapter attached to the same catalog, breaking concurrent writers with "Connection already closed" errors. - adapters share one engine per catalog file via a process-wide registry; engines are only released right before a file is replaced by a re-download - add public BaseAdapter API: connected, ensure_connected, reconnect, raw_connection, transaction (short-lived direct DuckDB connections, decoupled from engine lifecycles), mark_dirty - DuckLake: public catalog_adapter/columns_adapter, get_dataset_adapter and flush_catalogs for checkpoints - SyncEngine: staged parallel pipeline (serialized FTP downloads on a connection pool, worker pool for conversion/upload, serial catalog writer) with per-entry self-healing and quiescent checkpoints - trust-the-catalog freshness policy (size-based staleness) plus byte-exact content veto (sha256/source_sha256) - exclude unstable Fortaleza SIA dataset from DadosGov; skip xlsx dictionaries; longer API timeouts and inventory retries Tests: 734 passing. --- pysus/api/dadosgov/client.py | 2 +- pysus/api/dadosgov/databases.py | 4 +- pysus/api/dadosgov/models.py | 4 +- pysus/api/ducklake/catalog/adapters.py | 168 +++- pysus/api/ducklake/client.py | 54 ++ pysus/management/catalog.py | 49 ++ pysus/management/scripts/relayout_bucket.py | 5 +- pysus/management/scripts/sync_clients.py | 28 +- pysus/management/sync.py | 804 ++++++++++++++++---- pysus/tests/api/dadosgov/test_client.py | 2 +- pysus/tests/api/dadosgov/test_databases.py | 6 +- pysus/tests/api/ducklake/test_client.py | 23 +- pysus/tests/management/test_catalog.py | 215 ++++++ pysus/tests/management/test_sync.py | 69 +- 14 files changed, 1235 insertions(+), 198 deletions(-) create mode 100644 pysus/tests/management/test_catalog.py diff --git a/pysus/api/dadosgov/client.py b/pysus/api/dadosgov/client.py index 6e38311d..801332b5 100644 --- a/pysus/api/dadosgov/client.py +++ b/pysus/api/dadosgov/client.py @@ -151,7 +151,7 @@ async def connect(self, token: str | None = None) -> None: self._client = httpx.AsyncClient( base_url=self.base_url, headers=headers, - timeout=30.0, + timeout=httpx.Timeout(120.0, connect=30.0), follow_redirects=True, ) diff --git a/pysus/api/dadosgov/databases.py b/pysus/api/dadosgov/databases.py index 8000d5f9..3a797b74 100644 --- a/pysus/api/dadosgov/databases.py +++ b/pysus/api/dadosgov/databases.py @@ -626,7 +626,9 @@ def formatter(self, filename: str) -> dict[str, Any]: AVAILABLE_DATABASES: list[type[Dataset]] = [ CNES, PNI, - SIA, + # SIA excluded: its only dataset (Fortaleza municipal) is served from + # dados.fortaleza.ce.gov.br, which is unstable (frequent 504s) and is + # not national SIA data. See SIA class below. SIM, SINAN, SINASC, diff --git a/pysus/api/dadosgov/models.py b/pysus/api/dadosgov/models.py index ddb924c3..4baa12e9 100644 --- a/pysus/api/dadosgov/models.py +++ b/pysus/api/dadosgov/models.py @@ -322,7 +322,9 @@ async def _fetch_files(self) -> list[BaseRemoteFile]: filename = ( recurso.file_name or recurso.url.split("/")[-1].split("?")[0] ) - if filename.lower().endswith(".pdf") or filename.startswith("get_"): + if filename.lower().endswith( + (".pdf", ".xlsx") + ) or filename.startswith("get_"): continue metadata = {} if self._formatter: diff --git a/pysus/api/ducklake/catalog/adapters.py b/pysus/api/ducklake/catalog/adapters.py index f21abf0f..144af7ac 100644 --- a/pysus/api/ducklake/catalog/adapters.py +++ b/pysus/api/ducklake/catalog/adapters.py @@ -2,6 +2,7 @@ import os from abc import ABC from collections.abc import Callable +from contextlib import contextmanager from pathlib import Path import httpx @@ -23,6 +24,26 @@ class DuckLakeCredentials(BaseModel): secret_key: SecretStr +_SHARED_ENGINES: dict[str, Engine] = {} + + +def _dispose_shared(db_local: Path) -> None: + """Dispose and forget the shared engine for *db_local*, if any. + + The shared engine is also the process-lifetime anchor: its pooled + connection keeps the DuckDB instance alive, so every other adapter + attached to the same file stays valid. Dispose only happens right + before the file is replaced by a re-download. + """ + key = str(db_local.resolve()) + engine = _SHARED_ENGINES.pop(key, None) + if engine is not None: + try: + engine.dispose() + except Exception: # noqa + pass + + class BaseAdapter(ABC): cache_dir: Path = Path(CACHEPATH) / "ducklake" db_local: Path @@ -46,6 +67,92 @@ def __init__( def remote_url(self) -> str: return f"https://{types.S3_ENDPOINT}/{types.S3_BUCKET}/{self.db_remote}" + @property + def connected(self) -> bool: + """True when the database engine is initialized.""" + return self._engine is not None + + @property + def local_dirty(self) -> bool: + """True when the local database has unsaved changes.""" + return self._local_dirty + + def mark_dirty(self) -> None: + """Flag the local database as modified (upload on close).""" + self._local_dirty = True + + async def ensure_connected( + self, + callback: Callable[[int, int], None] | None = None, + ) -> None: + """Connect the engine if it is not already connected. + + The shared engine doubles as the process-lifetime anchor, so the + DuckDB instance can never be torn down by other lifecycle paths. + """ + if self._engine is not None: + return + await self.connect(callback=callback) + + def checkpoint(self) -> None: + """Force a WAL checkpoint, persisting writes to ``db_local``.""" + self.setup_engine().raw_connection().execute("CHECKPOINT") + + async def reconnect(self) -> None: + """Dispose the shared engine and reinitialize from scratch. + + Used after the database file is found broken or right before it + is replaced by a fresh download. + """ + _dispose_shared(self.db_local) + self._engine = None + self._session_factory = None + await self.connect(force=True) + + def raw_connection(self): + """Return a raw DuckDB connection to the catalog database. + + Raises + ------ + CatalogError + If the engine is not initialized (call + :meth:`ensure_connected` first). + """ + if self._engine is None: + raise CatalogError( + "Database engine not initialized. " + "Call ensure_connected() first." + ) + return self._engine.raw_connection() + + @contextmanager + def transaction(self): + """Yield ``(connection, cursor)`` for a committed transaction. + + Uses the shared engine's pooled connection — the only way to + guarantee every connection to a DuckDB file carries identical + configuration (DuckDB rejects mixed-configuration opens) and to + keep the process-lifetime instance anchored. The connection is + *not* closed (it is the shared anchor); the transaction is + committed on success and rolled back on error. + """ + engine = self.setup_engine() + conn = engine.raw_connection() + try: + probe = conn.cursor() + try: + probe.execute("SELECT 1") + probe.fetchone() + except Exception as exc: # noqa + raise CatalogError( + f"Catalog connection is broken: {exc}" + ) from exc + with conn: + cursor = conn.cursor() + yield conn, cursor + finally: + pass + def get_session(self) -> Session: if not self._session_factory: raise CatalogError( @@ -74,6 +181,7 @@ async def connect( return if force: + _dispose_shared(self.db_local) await self._download_catalog( self.db_local, str(self.db_remote), @@ -100,6 +208,7 @@ async def connect( self._engine = await to_thread.run_sync(self.setup_engine) self._session_factory = sessionmaker(bind=self._engine) except Exception: # noqa + _dispose_shared(self.db_local) if self.db_local.exists(): try: os.remove(self.db_local) @@ -118,7 +227,21 @@ async def connect( def setup_engine( self, access_key: str | None = None, secret_key: str | None = None ) -> Engine: - engine: Engine = create_engine( + """Return the shared engine for this adapter's database file. + + DuckDB keeps one database instance per file per process; opening + the file again merely attaches to it, and disposing one engine's + connection tears the shared instance down for everyone else. + Adapters therefore share one engine per file (process-wide) that + is only disposed via :func:`_dispose_shared` — right before the + file itself is replaced by a re-download. + """ + key = str(self.db_local.resolve()) + engine = _SHARED_ENGINES.get(key) + if engine is not None: + return engine + + engine = create_engine( f"duckdb:///{self.db_local}", poolclass=StaticPool, ) @@ -126,36 +249,10 @@ def setup_engine( with engine.connect() as conn: conn.exec_driver_sql("INSTALL ducklake; LOAD ducklake;") conn.exec_driver_sql("CREATE SCHEMA IF NOT EXISTS pysus;") - - has_pysus = conn.exec_driver_sql( - statement=( - "SELECT 1 FROM information_schema.schemata " - "WHERE schema_name = 'pysus'" - ) - ).fetchone() - - if has_pysus: - conn.exec_driver_sql("SET search_path='pysus,main';") - else: - conn.exec_driver_sql("SET search_path='main';") - - s3_cfg = { - "s3_endpoint": types.S3_ENDPOINT, - "s3_region": types.S3_REGION, - "s3_url_style": "path", - "s3_use_ssl": "true", - } - - if access_key and secret_key: - s3_cfg["s3_access_key_id"] = access_key - s3_cfg["s3_secret_access_key"] = secret_key - - for key, value in s3_cfg.items(): - conn.exec_driver_sql(f"SET {key}='{value}'") - conn.commit() DatasetBase.metadata.create_all(bind=engine) + _SHARED_ENGINES[key] = engine return engine async def _download_catalog( @@ -228,6 +325,9 @@ async def _upload_catalog(self) -> None: if not self.db_local.exists(): raise FileNotFoundError("catalog file not found") + # persist pending writes before uploading the file + self.checkpoint() + await upload_s3( local_path=self.db_local, remote_path=str(self.db_remote), @@ -240,10 +340,14 @@ async def close(self, update: bool = False) -> None: await self._upload_catalog() self._local_dirty = False - if self._engine: - await to_thread.run_sync(self._engine.dispose) - self._engine = None - self._session_factory = None + # The engine is shared process-wide per database file and is + # never disposed here: DuckDB tears down the in-process database + # instance when its last connection closes, which would break + # every other adapter attached to the same file. The shared + # registry releases engines only via _dispose_shared() right + # before a file is replaced. + self._engine = None + self._session_factory = None def __del__(self) -> None: if not hasattr(self, "_engine") or not self._engine: diff --git a/pysus/api/ducklake/client.py b/pysus/api/ducklake/client.py index 19aa0a53..89fceafd 100644 --- a/pysus/api/ducklake/client.py +++ b/pysus/api/ducklake/client.py @@ -72,6 +72,47 @@ def catalog_path(self) -> Path: def columns_path(self) -> Path: return self._columns_adap.db_local + @property + def catalog_adapter(self) -> CatalogAdapter: + """The central dataset-registry adapter (``catalog.duckdb``).""" + return self._catalog_adap + + @property + def columns_adapter(self) -> ColumnsAdapter: + """The column-definitions adapter (``catalog_columns.duckdb``).""" + return self._columns_adap + + def get_dataset_adapter(self, name: str) -> DatasetAdapter: + """Return (and register) the per-dataset adapter for *name*. + + Creates a fresh adapter for datasets not yet seen by this client. + """ + wanted = str(name).lower() + for dataset in self._datasets: + if getattr(dataset, "name", "").lower() == wanted: + return dataset.adapter + + adapter = DatasetAdapter( + name=wanted, + dataset_id=0, + credentials=self.credentials, + update_on_close=self.update_on_close, + ) + self._datasets.append( + type( + "_DatasetEntry", + (), + { + "name": wanted, + "adapter": adapter, + "close": lambda self_, update_catalog=None: ( + self_.adapter.close(update=bool(update_catalog)) + ), + }, + )() + ) + return adapter + async def datasets(self, **kwargs) -> list[DuckDataset]: def _fetch(): with self._catalog_adap.get_session() as session: @@ -140,6 +181,19 @@ async def close(self, update_catalog: bool | None = None) -> None: await self._catalog_adap.close(update=should_update) await self._columns_adap.close(update=should_update) + async def flush_catalogs(self, update: bool = True) -> None: + """Upload dirty catalogs (if *update*) and reopen the adapters. + + Long-running writers use this to checkpoint: modified local + databases are pushed to S3 and the adapters are reconnected. + """ + for ds in self._datasets: + await ds.close(update_catalog=update) + await self._catalog_adap.close(update=update) + await self._columns_adap.close(update=update) + await self._catalog_adap.connect() + await self._columns_adap.connect() + async def download( self, file: BaseRemoteFile, diff --git a/pysus/management/catalog.py b/pysus/management/catalog.py index 9ef98e98..904e3f93 100644 --- a/pysus/management/catalog.py +++ b/pysus/management/catalog.py @@ -108,6 +108,9 @@ def _ensure_column( def _ensure_management_columns(self, catalog_cursor) -> None: self._ensure_column(catalog_cursor, "files", "origin", "VARCHAR") self._ensure_column(catalog_cursor, "files", "format", "VARCHAR") + self._ensure_column( + catalog_cursor, "files", "source_sha256", "VARCHAR(64)" + ) # ------------------------------------------------------------------ # datasets & groups @@ -200,6 +203,45 @@ def get_file(self, cursor, path: str) -> tuple[int, datetime | None] | None: return None return int(row[0]), row[1] + def get_file_full( + self, + cursor, + path: str, + ) -> tuple[int, datetime | None, int, str | None, str | None] | None: + """Return ``(id, origin_modified, origin_size, sha256, + source_sha256)`` for the S3 *path*, if present.""" + cursor.execute( + "SELECT id, origin_modified, origin_size, sha256, " + "source_sha256 FROM pysus.files WHERE path = ?", + (path,), + ) + row = cursor.fetchone() + if not row: + return None + return int(row[0]), row[1], int(row[2] or 0), row[3], row[4] + + def touch_file( + self, + cursor, + file_id: int, + origin_modified: datetime | None, + origin_size: int, + source_sha256: str | None = None, + ) -> None: + """Update origin metadata without replacing the artifact.""" + if source_sha256 is not None: + cursor.execute( + "UPDATE pysus.files SET origin_modified = ?, " + "origin_size = ?, source_sha256 = ? WHERE id = ?", + (origin_modified, origin_size, source_sha256, file_id), + ) + else: + cursor.execute( + "UPDATE pysus.files SET origin_modified = ?, " + "origin_size = ? WHERE id = ?", + (origin_modified, origin_size, file_id), + ) + def delete_file(self, cursor, file_id: int) -> None: cursor.execute( "DELETE FROM pysus.file_columns WHERE file_id = ?", (file_id,) @@ -225,6 +267,7 @@ def upsert_file( origin: str | None = None, format: str | None = None, sha256: str | None = None, + source_sha256: str | None = None, file_type: str | None = None, ) -> tuple[int, bool]: """Insert or update the file row keyed on S3 *path*. @@ -265,6 +308,9 @@ def upsert_file( if sha256 is not None: sets.append("sha256 = ?") update_values.append(sha256) + if source_sha256 is not None: + sets.append("source_sha256 = ?") + update_values.append(source_sha256) if file_type is not None: sets.append("type = ?") update_values.append(file_type) @@ -303,6 +349,9 @@ def upsert_file( if sha256 is not None: columns.append("sha256") values.append(sha256) + if source_sha256 is not None: + columns.append("source_sha256") + values.append(source_sha256) if file_type is not None: columns.append("type") values.append(file_type) diff --git a/pysus/management/scripts/relayout_bucket.py b/pysus/management/scripts/relayout_bucket.py index 298a5170..36114ba3 100644 --- a/pysus/management/scripts/relayout_bucket.py +++ b/pysus/management/scripts/relayout_bucket.py @@ -123,7 +123,10 @@ def main() -> int: print(f" applying {name}...", flush=True) aliases.update( normalizer.apply_renames_with_aliases( - plan.object_renames, dry_run=False + plan.object_renames, + dry_run=False, + object_sizes=object_sizes, + workers=48, ) ) normalizer.apply_objects([], plan.object_deletes, dry_run=False) diff --git a/pysus/management/scripts/sync_clients.py b/pysus/management/scripts/sync_clients.py index 9d5b1b69..cde2d6b4 100644 --- a/pysus/management/scripts/sync_clients.py +++ b/pysus/management/scripts/sync_clients.py @@ -30,7 +30,9 @@ def load_env(path: str = ".env") -> dict[str, str]: return env -async def run(datasets, checkpoint_every, force) -> dict: +async def run( + datasets, checkpoint_every, force, workers, ftp_connections +) -> dict: env = load_env() engine = SyncEngine( access_key=env.get("ACCESS_KEY"), @@ -54,6 +56,8 @@ def on_outcome(outcome) -> None: force=force, checkpoint_every=checkpoint_every, on_outcome=on_outcome, + workers=workers, + ftp_connections=ftp_connections, ) return report.summary() @@ -77,11 +81,31 @@ def main() -> int: action="store_true", help="Reprocess files even when the catalog is current", ) + parser.add_argument( + "--workers", + type=int, + default=16, + help="Concurrent ingestion workers", + ) + parser.add_argument( + "--ftp-connections", + type=int, + default=6, + help="FTP connection pool size", + ) args = parser.parse_args() datasets = [d.upper() for d in args.datasets] if args.datasets else None - summary = asyncio.run(run(datasets, args.checkpoint_every, args.force)) + summary = asyncio.run( + run( + datasets, + args.checkpoint_every, + args.force, + args.workers, + args.ftp_connections, + ) + ) print(summary) return 0 diff --git a/pysus/management/sync.py b/pysus/management/sync.py index 9e9dea06..e2fd5547 100644 --- a/pysus/management/sync.py +++ b/pysus/management/sync.py @@ -22,12 +22,15 @@ from collections.abc import Callable from datetime import datetime from logging import error -from typing import TYPE_CHECKING, Any, cast +from pathlib import Path +from typing import TYPE_CHECKING, Any +from uuid import uuid4 +import httpx +from pysus import CACHEPATH from pysus.api.dadosgov.models import File as APIFile from pysus.api.ducklake.functional import upload_s3 from pysus.api.errors import AuthenticationError, ConnectionError -from pysus.api.extensions import Parquet from pysus.api.ftp.models import File as FTPFile from pysus.api.models import BaseRemoteFile @@ -47,22 +50,13 @@ from pysus.api.client import PySUS from pysus.api.ducklake.client import DuckLake -_RETRYABLE = (ConnectionResetError, ConnectionRefusedError, TimeoutError) - - -class _DatasetStub: - """Registry entry for a per-dataset adapter without a full DuckDataset. - - ``DuckLake.close`` iterates ``_datasets`` and calls ``ds.close``; - the stub delegates to its adapter so the catalog is uploaded on close. - """ - - def __init__(self, name: str, adapter): - self.name = name - self.adapter = adapter - - async def close(self, update_catalog: bool | None = None) -> None: - await self.adapter.close(update=bool(update_catalog)) +_RETRYABLE = ( + ConnectionResetError, + ConnectionRefusedError, + TimeoutError, + BrokenPipeError, + OSError, +) class SyncEngine: @@ -123,13 +117,58 @@ async def __aenter__(self) -> SyncEngine: access_key=self.access_key, secret_key=self.secret_key, ) + self._acquire_sync_lock() return self + def _acquire_sync_lock(self) -> None: + """Guarantee a single sync process owns the catalogs at a time. + + DuckDB catalog files allow one writer process; concurrent syncs + corrupt each other's state. The lock file carries the PID and is + stolen only when that PID is no longer alive. + """ + import os + + lock_path = Path(CACHEPATH) / "ducklake" / ".sync.lock" + lock_path.parent.mkdir(parents=True, exist_ok=True) + if lock_path.exists(): + try: + owner = int(lock_path.read_text().strip() or "0") + except ValueError: + owner = 0 + alive = owner > 0 and self._pid_alive(owner) + if alive: + raise ConnectionError( + "another sync process (PID " + f"{owner}) holds the catalog lock" + ) + lock_path.write_text(str(os.getpid())) + + @staticmethod + def _pid_alive(pid: int) -> bool: + import os + + try: + os.kill(pid, 0) + except ProcessLookupError: + return False + except PermissionError: + return True + return True + + def _release_sync_lock(self) -> None: + lock_path = Path(CACHEPATH) / "ducklake" / ".sync.lock" + try: + lock_path.unlink(missing_ok=True) + except OSError: + pass + async def __aexit__(self, exc_type, exc_val, exc_tb) -> None: try: if not exc_type and self._ducklake: await self._ducklake.close(update_catalog=self._changed_catalog) finally: + self._release_sync_lock() if self.pysus is not None: await self.pysus.__aexit__(exc_type, exc_val, exc_tb) @@ -154,6 +193,7 @@ async def upload_file( file: FTPFile | APIFile, callback: Callable[[int, int], None] | None = None, force: bool = False, + ftp_client: Any | None = None, ) -> bool: """Download *file*, convert to parquet, upload and catalog it. @@ -162,8 +202,23 @@ async def upload_file( per-dataset ``catalog_.duckdb``, and the column definitions into ``catalog_columns.duckdb``. - Returns True if the file was processed; False when the catalog - already holds an equally recent artifact (skip). + ``ftp_client`` (optional) is used instead of the file's client + for the raw download, allowing a pool of FTP connections to be + shared safely across parallel workers. + + Content-veto (no mismatch possible — byte-exact hashes): + + * if the source ``modified`` is not newer than the catalog's + ``origin_modified`` → skip without downloading; + * else the raw file is downloaded (cache bypassed) and hashed: + same raw ``sha256`` as the stored ``source_sha256`` → only the + origin metadata is touched, no conversion/upload; + * else the parquet is converted and hashed: same parquet + ``sha256`` as stored → only metadata is touched, no upload; + * otherwise the artifact is uploaded and both hashes stored. + + Returns True if the file was (re)processed; False when the + existing artifact is current or content-identical. """ if self._ducklake is None: raise ConnectionError("DuckLake is not connected") @@ -172,16 +227,16 @@ async def upload_file( writer = self.writer dataset_adapter = self._dataset_adapter(file) - central_adapter = self._ducklake._catalog_adap - columns_adapter = self._ducklake._columns_adap + central_adapter = self._ducklake.catalog_adapter + columns_adapter = self._ducklake.columns_adapter await central_adapter.connect() await columns_adapter.connect() await dataset_adapter.connect() - central_conn = central_adapter._engine.raw_connection() - dataset_conn = dataset_adapter._engine.raw_connection() - columns_conn = columns_adapter._engine.raw_connection() + central_conn = central_adapter.raw_connection() + dataset_conn = dataset_adapter.raw_connection() + columns_conn = columns_adapter.raw_connection() connections = (central_conn, dataset_conn, columns_conn) try: @@ -192,33 +247,65 @@ async def upload_file( writer._ensure_management_columns(dataset_cursor) - existing = writer.get_file(dataset_cursor, s3_key) + existing = writer.get_file_full(dataset_cursor, s3_key) if existing and not force: - _, origin_modified = existing + _, origin_modified, _, _, _ = existing if self._is_current(file, origin_modified): return False - dataset_id = writer.ensure_dataset( - central_cursor, - file.dataset.name, - file.dataset.long_name, - getattr(file.dataset, "description", None), + raw_path = await self._download_raw_with_retry( + file, ftp_client=ftp_client ) + raw_digest = sha256_of(raw_path) - group = getattr(file, "group", None) - group_name = ( - getattr(group, "name", None) if group is not None else None - ) - group_name = str(group_name) if group_name else None - group_id = writer.ensure_group( - dataset_cursor, - dataset_id, - group_name, - getattr(group, "long_name", None) if group else None, - getattr(group, "description", None) if group else None, + if existing and not force: + file_id = existing[0] + stored_source = existing[4] + if stored_source and raw_digest == stored_source: + writer.touch_file( + dataset_cursor, + file_id, + self._safe_modify(file), + self._safe_size(file), + ) + dataset_conn.commit() + dataset_cursor.execute("CHECKPOINT") + dataset_adapter.mark_dirty() + self._changed_catalog = True + self._cleanup_local(raw_path) + return False + + from pysus.api.extensions import ExtensionFactory + + local_file = await ExtensionFactory.instantiate(raw_path) + if not hasattr(local_file, "to_parquet"): + raise RuntimeError( + f"{file.basename}: cannot convert to parquet" + ) + parquet_file = await local_file.to_parquet( + callback=callback, ) + parquet_digest = sha256_of(parquet_file.path) + + if existing and not force: + file_id = existing[0] + stored_sha = existing[3] + if stored_sha and parquet_digest == stored_sha: + writer.touch_file( + dataset_cursor, + file_id, + self._safe_modify(file), + self._safe_size(file), + source_sha256=raw_digest, + ) + dataset_conn.commit() + dataset_cursor.execute("CHECKPOINT") + dataset_adapter.mark_dirty() + self._changed_catalog = True + self._cleanup_local(raw_path) + self._cleanup_local(parquet_file.path) + return False - parquet_file = await self._download_with_retry(file, callback) await upload_s3( local_path=parquet_file.path, remote_path=s3_key, @@ -227,36 +314,20 @@ async def upload_file( callback=callback, ) - digest = sha256_of(parquet_file.path) - writer.upsert_file( - dataset_cursor, - dataset_id=dataset_id, - group_id=group_id, - path=s3_key, - size=parquet_file.size, - rows=parquet_file.rows, - modified=datetime.now(), - origin_modified=self._safe_modify(file), - origin_size=self._safe_size(file), - origin_path=str(file.path), - year=file.year, - month=file.month, - state=file.state, - origin=file.client.name.lower(), - format="parquet", - sha256=digest, - file_type="PARQUET", - ) - - inserted = writer.get_file(dataset_cursor, s3_key) - assert inserted is not None - file_id, _ = inserted - writer.link_columns( + payload = { + "s3_key": s3_key, + "size": parquet_file.path.stat().st_size, + "rows": parquet_file.rows, + "schema": parquet_file.schema, + "raw_digest": raw_digest, + "parquet_digest": parquet_digest, + } + self._catalog_rows( + central_cursor, dataset_cursor, columns_cursor, - file_id, - parquet_file.schema, - dataset_id, + file, + payload, ) central_conn.commit() @@ -265,10 +336,12 @@ async def upload_file( dataset_cursor.execute("CHECKPOINT") columns_cursor.execute("CHECKPOINT") - central_adapter._local_dirty = True - dataset_adapter._local_dirty = True - columns_adapter._local_dirty = True + central_adapter.mark_dirty() + dataset_adapter.mark_dirty() + columns_adapter.mark_dirty() self._changed_catalog = True + self._cleanup_local(raw_path) + self._cleanup_local(parquet_file.path) return True except BaseException as exc: # noqa # the connection context managers roll back on exit @@ -279,49 +352,39 @@ async def upload_file( pass raise exc + def _dataset_adapter_by_name(self, dataset_name: str): + """Return (and register) the per-dataset adapter for *name*.""" + return self._require_ducklake().get_dataset_adapter(dataset_name) + def _dataset_adapter(self, file: BaseRemoteFile): """Return (and register) the per-dataset adapter for *file*.""" - ducklake = self._require_ducklake() - dataset_name = file.dataset.name.lower() - for ds in ducklake._datasets: - if getattr(ds, "name", "").lower() == dataset_name: - return ds.adapter - - from pysus.api.ducklake.catalog.adapters import DatasetAdapter - - adapter = DatasetAdapter( - name=dataset_name, - dataset_id=0, - credentials=ducklake.credentials, - update_on_close=ducklake.update_on_close, - ) - ducklake._datasets.append( - cast(Any, _DatasetStub(dataset_name, adapter)) - ) - return adapter + return self._dataset_adapter_by_name(file.dataset.name) - async def _download_with_retry( + async def _download_raw_with_retry( self, file: FTPFile | APIFile, - callback: Callable[[int, int], None] | None = None, - max_retries: int = 3, - ) -> Parquet: + max_retries: int = 5, + ftp_client: Any | None = None, + ) -> Path: + """Download the raw artifact bypassing the local cache. + + The PySUS local cache matches on size only, which could serve + stale bytes for a same-size update — the content veto must hash + exactly what the client serves now. ``ftp_client`` overrides the + file's own FTP connection (pooled clients are not shared). + """ + raw_dir = Path(CACHEPATH) / "management" / "tmp" + raw_dir.mkdir(parents=True, exist_ok=True) + output = raw_dir / f"{uuid4().hex[:8]}-{file.basename}" + last_error: Exception | None = None - token = ( - self.dadosgov_token - if file.client.name.lower() == "dadosgov" - else None - ) for attempt in range(max_retries): try: - return await self._require_pysus().download_to_parquet( - file=file, - token=token, - callback=callback, - ) + await self._download_once(file, output, ftp_client) + return output except _RETRYABLE as exc: last_error = exc - wait_time = 2**attempt + wait_time = 2**attempt + (attempt * 2) error( f"Download attempt {attempt + 1}/{max_retries} failed " f"for {file.basename}: {exc}. Retrying in {wait_time}s..." @@ -333,6 +396,66 @@ async def _download_with_retry( f"attempts: {last_error}" ) from last_error + async def _download_once( + self, + file: BaseRemoteFile, + output: Path, + ftp_client: Any | None = None, + ) -> None: + """Perform one raw download to *output*.""" + from anyio import to_thread + + client = ftp_client if ftp_client is not None else file.client + ftp = getattr(client, "ftp", None) + if ftp_client is not None: + # never fall back to the shared client: reconnect the pooled + # session instead + if ftp is None: + await client.connect() + ftp = getattr(client, "ftp", None) + assert ftp is not None + remote_path = str(file.path) + + def _retr(): + total = ftp.size(remote_path) or 0 + with open(output, "wb") as f: + ftp.retrbinary( + f"RETR {remote_path}", lambda chunk: f.write(chunk) + ) + return total + + try: + await to_thread.run_sync(_retr) + return + except Exception: # noqa + try: + ftp.quit() + except Exception: # noqa + pass + setattr( # noqa: B010 — reset pooled FTP session + client, "_ftp", None + ) + raise + if ftp is not None: + remote_path = str(file.path) + + def _direct_retr(): + with open(output, "wb") as f: + ftp.retrbinary( + f"RETR {remote_path}", lambda chunk: f.write(chunk) + ) + + await to_thread.run_sync(_direct_retr) + return + await file._download(output=output) + + @staticmethod + def _cleanup_local(path: Path) -> None: + try: + path.unlink(missing_ok=True) + except OSError: + pass + @staticmethod def _is_current( file: BaseRemoteFile, @@ -371,28 +494,43 @@ async def run( save_snapshots: bool = True, checkpoint_every: int | None = None, on_outcome: Callable[[SyncOutcome], None] | None = None, + workers: int = 16, + ftp_connections: int = 6, ) -> SyncReport: """Run the full pipeline and return a :class:`SyncReport`. - Files already on S3 (ducklake artifacts) are skipped; FTP is - preferred over DadosGov, which requires ``dadosgov_token``. - Files whose non-S3 artifact is newer than the S3 copy are - reprocessed (most-updated policy). + Files already on S3 (ducklake artifacts) are skipped unless + ``force`` or the source size proves a change (trust-the-catalog + policy); FTP is preferred over DadosGov, which requires + ``dadosgov_token``. + + Missing files are ingested in parallel: ``workers`` asyncio tasks + download (via a pool of ``ftp_connections`` FTP clients), convert + and upload concurrently; catalog writes stay serialized and + checkpoints only run when all workers are quiescent. ``checkpoint_every`` uploads the modified catalogs to S3 every N - successful uploads, making long runs resumable (files already - cataloged are skipped on the next run). ``on_outcome`` is called - once per processed logical file (e.g. for progress logging). + successful uploads, making long runs resumable. ``on_outcome`` is + called once per processed logical file. """ report = SyncReport(dataset=",".join(datasets) if datasets else None) + async def collect_with_retry(origin: str, datasets=None, **kwargs): + for attempt in range(3): + try: + return await self.inventory.collect(origin, **kwargs) + except (*_RETRYABLE, httpx.HTTPError): + if attempt == 2: + raise + await asyncio.sleep(2**attempt) + records: dict[str, list[FileRecord]] = { - "ducklake": await self.inventory.collect("ducklake", datasets), - "ftp": await self.inventory.collect("ftp", datasets), + "ducklake": await collect_with_retry("ducklake", datasets), + "ftp": await collect_with_retry("ftp", datasets), } records["dadosgov"] = [] if self.dadosgov_token: - records["dadosgov"] = await self.inventory.collect( + records["dadosgov"] = await collect_with_retry( "dadosgov", datasets, dadosgov_token=self.dadosgov_token ) @@ -400,24 +538,217 @@ async def run( records["ducklake"] + records["ftp"] + records["dadosgov"] ) - uploaded_since_checkpoint = 0 + # Pre-connect every adapter involved so concurrent workers never + # race the initial catalog download. + await self._preconnect_adapters(records) + + parallel: list[tuple[FileComparison, FileRecord]] = [] for comparison in comparisons: - outcome = await self._process_comparison( - comparison, force=force, callback=callback - ) - report.outcomes.append(outcome) - if on_outcome: - on_outcome(outcome) - - if ( - outcome.status == "uploaded" - and checkpoint_every - and self._changed_catalog + if comparison.is_on_s3 and not ( + force or self._s3_is_stale(comparison) ): - uploaded_since_checkpoint += 1 - if uploaded_since_checkpoint >= checkpoint_every: - await self._checkpoint() - uploaded_since_checkpoint = 0 + outcome = SyncOutcome( + key=comparison.key, + origin="ducklake", + status="skipped", + ) + report.outcomes.append(outcome) + if on_outcome: + on_outcome(outcome) + continue + record = self._pick_source(comparison) + if record is None: + outcome = await self._process_comparison( + comparison, force=force, callback=callback + ) + report.outcomes.append(outcome) + if on_outcome: + on_outcome(outcome) + continue + parallel.append((comparison, record)) + + 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"] + + ftp_pool: list[Any] = [] + if ftp_items: + from pysus.api.ftp.client import FTP + + for _ in range(ftp_connections): + client = FTP() + await client.connect() + ftp_pool.append(client) + + raw_queue: asyncio.Queue = asyncio.Queue(maxsize=workers * 2) + write_queue: asyncio.Queue = asyncio.Queue(maxsize=workers * 2) + + async def ftp_downloader( + client: Any, items: list[tuple[FileComparison, FileRecord]] + ) -> None: + for comparison, record in items: + try: + raw = await self._download_raw_with_retry( + record.file, ftp_client=client + ) + await raw_queue.put((comparison, record, raw, None)) + except Exception as exc: # noqa + await raw_queue.put((comparison, record, None, str(exc))) + await raw_queue.put(None) + + async def raw_processor() -> None: + while True: + entry = await raw_queue.get() + try: + if entry is None: + return + comparison, record, raw, err = entry + if err is not None: + await write_queue.put((comparison, record, None, err)) + continue + payload = await self._convert_and_upload( + record.file, raw, callback=callback + ) + await write_queue.put((comparison, record, payload, None)) + except Exception as exc: # noqa + comparison, record, _, _ = entry + await write_queue.put((comparison, record, None, str(exc))) + finally: + raw_queue.task_done() + + async def gov_worker() -> None: + while gov_items: + comparison, record = gov_items.pop() + try: + raw = await self._download_raw_with_retry(record.file) + payload = await self._convert_and_upload( + record.file, raw, callback=callback + ) + await write_queue.put((comparison, record, payload, None)) + except Exception as exc: # noqa + await write_queue.put((comparison, record, None, str(exc))) + + async def catalog_writer() -> None: + """Serial consumer: catalog rows + outcomes + checkpoints.""" + uploaded_since_checkpoint = 0 + + ducklake = self._require_ducklake() + central_adapter = ducklake.catalog_adapter + columns_adapter = ducklake.columns_adapter + await central_adapter.connect() + await columns_adapter.connect() + dataset_adapters: dict[str, Any] = {} + + done_writers = 0 + while True: + entry = await write_queue.get() + try: + if entry is None: + done_writers += 1 + if done_writers >= writers_total: + break + continue + + comparison, record, payload, err = entry + if err is not None: + outcome = SyncOutcome( + key=comparison.key, + origin=record.origin, + status="failed", + detail=f"{self._label(comparison)}: {err}", + ) + else: + try: + adapter = dataset_adapters.get( + record.dataset.lower() + ) + if adapter is None: + adapter = self._dataset_adapter_by_name( + record.dataset + ) + dataset_adapters[record.dataset.lower()] = adapter + self._catalog_write_entry( + adapter, + central_adapter, + columns_adapter, + record.file, + payload, + ) + adapter.mark_dirty() + central_adapter.mark_dirty() + columns_adapter.mark_dirty() + self._changed_catalog = True + outcome = SyncOutcome( + key=comparison.key, + origin=record.origin, + status="uploaded", + detail=self._label(comparison), + ) + except Exception as exc: # noqa + import traceback + + error( + "catalog write failed for " + f"{self._label(comparison)}: {exc}" + ) + error(traceback.format_exc()) + outcome = SyncOutcome( + key=comparison.key, + origin=record.origin, + status="failed", + detail=( + f"{self._label(comparison)}: " + f"catalog write: {exc}" + ), + ) + + report.outcomes.append(outcome) + if on_outcome: + on_outcome(outcome) + if ( + outcome.status == "uploaded" + and checkpoint_every + and self._changed_catalog + ): + uploaded_since_checkpoint += 1 + if uploaded_since_checkpoint >= checkpoint_every: + await self._checkpoint() + uploaded_since_checkpoint = 0 + finally: + write_queue.task_done() + + processor_tasks = [ + asyncio.create_task(raw_processor()) for _ in range(workers) + ] + # Heavy sources (DadosGov multi-million-row archives) are + # processed one at a time to bound peak memory; FTP files are + # small and flow through the processor pool. + gov_workers_tasks = [ + asyncio.create_task(gov_worker()) + for _ in range(2 if gov_items else 0) + ] + ftp_tasks = [] + if ftp_pool: + step = max(1, len(ftp_pool)) + ftp_tasks = [ + asyncio.create_task(ftp_downloader(client, ftp_items[i::step])) + for i, client in enumerate(ftp_pool) + ] + else: + for _ in range(workers): + await raw_queue.put(None) + if not ftp_tasks and not gov_workers_tasks: + for _ in range(workers): + await write_queue.put(None) + + writers_total = len(gov_workers_tasks) + writer_task = asyncio.create_task(catalog_writer()) + await asyncio.gather(*ftp_tasks, *gov_workers_tasks) + for _ in processor_tasks: + await raw_queue.put(None) + await asyncio.gather(*processor_tasks) + for _ in range(writers_total): + await write_queue.put(None) + await writer_task if self._changed_catalog and checkpoint_every is not None: await self._checkpoint() @@ -428,15 +759,187 @@ async def run( return report + async def _convert_and_upload( + self, + file: BaseRemoteFile, + raw_path: Path, + callback: Callable[[int, int], None] | None = None, + ) -> dict: + """Convert a downloaded raw file, upload it, return the payload. + + The payload carries everything the catalog writer needs after the + local files are removed. + """ + from anyio import to_thread + from pysus.api.extensions import ExtensionFactory + + s3_key = self.s3_key_for(file) + raw_digest = await to_thread.run_sync(sha256_of, raw_path) + + local_file = await ExtensionFactory.instantiate(raw_path) + if not hasattr(local_file, "to_parquet"): + raise RuntimeError(f"{file.basename}: cannot convert to parquet") + parquet_file = await local_file.to_parquet(callback=callback) + try: + parquet_digest = await to_thread.run_sync( + sha256_of, parquet_file.path + ) + payload = { + "s3_key": s3_key, + "size": parquet_file.path.stat().st_size, + "rows": parquet_file.rows, + "schema": parquet_file.schema, + "raw_digest": raw_digest, + "parquet_digest": parquet_digest, + } + await upload_s3( + local_path=parquet_file.path, + remote_path=s3_key, + access_key=str(self.access_key), + secret_key=str(self.secret_key), + callback=callback, + ) + return payload + finally: + self._cleanup_local(raw_path) + self._cleanup_local(parquet_file.path) + + def _catalog_write_entry( + self, + adapter, + central_adapter, + columns_adapter, + file: BaseRemoteFile, + payload: dict, + ) -> None: + """Write one artifact's catalog rows in short transactions. + + All three catalogs are touched through short-lived direct + DuckDB connections (``transaction()``), completely decoupled + from engine lifecycles — DuckDB tears down the shared in-process + database instance when its last connection closes, so no + connection is ever held across operations here. + """ + with central_adapter.transaction() as ( + central_conn, + central_cursor, + ): + with columns_adapter.transaction() as ( + columns_conn, + columns_cursor, + ): + with adapter.transaction() as (conn, dataset_cursor): + self.writer._ensure_management_columns(dataset_cursor) + self._catalog_rows( + central_cursor, + dataset_cursor, + columns_cursor, + file, + payload, + ) + conn.commit() + dataset_cursor.execute("CHECKPOINT") + columns_conn.commit() + central_conn.commit() + + def _catalog_rows( + self, + central_cursor, + dataset_cursor, + columns_cursor, + file: BaseRemoteFile, + payload: dict, + ) -> None: + """Write dataset/group/file/column rows for an uploaded artifact.""" + writer = self.writer + dataset_id = writer.ensure_dataset( + central_cursor, + file.dataset.name, + file.dataset.long_name, + getattr(file.dataset, "description", None), + ) + + group = getattr(file, "group", None) + group_name = getattr(group, "name", None) if group is not None else None + group_name = str(group_name) if group_name else None + group_id = writer.ensure_group( + dataset_cursor, + dataset_id, + group_name, + getattr(group, "long_name", None) if group else None, + getattr(group, "description", None) if group else None, + ) + + writer.upsert_file( + dataset_cursor, + dataset_id=dataset_id, + group_id=group_id, + path=payload["s3_key"], + size=payload["size"], + rows=payload["rows"], + modified=datetime.now(), + origin_modified=self._safe_modify(file), + origin_size=self._safe_size(file), + origin_path=str(file.path), + year=file.year, + month=file.month, + state=file.state, + origin=file.client.name.lower(), + format="parquet", + sha256=payload["parquet_digest"], + source_sha256=payload["raw_digest"], + file_type="PARQUET", + ) + + inserted = writer.get_file(dataset_cursor, payload["s3_key"]) + assert inserted is not None + file_id, _ = inserted + writer.link_columns( + dataset_cursor, + columns_cursor, + file_id, + payload["schema"], + dataset_id, + ) + + async def _preconnect_adapters( + self, records: dict[str, list[FileRecord]] + ) -> None: + """Ensure all adapters are connected before parallel ingestion.""" + ducklake = self._require_ducklake() + await ducklake.catalog_adapter.ensure_connected() + await ducklake.columns_adapter.ensure_connected() + datasets = { + r.dataset.lower() + for items in records.values() + for r in items + if r.origin in ("ftp", "dadosgov") + } + for name in datasets: + await self._dataset_adapter_by_name(name).ensure_connected() + + @staticmethod + def _pick_source(comparison: FileComparison) -> FileRecord | None: + """Return the artifact to ingest (ftp over dadosgov), if any.""" + record = comparison._pick("ftp") + if record is None or record.file is None: + record = comparison._pick("dadosgov") + if record is None or record.file is None: + return None + return record + + @staticmethod + def _label(comparison: FileComparison) -> str: + key = comparison.key + return ( + f"{key.dataset}/{key.group or '-'}/" + f"{key.year or '-'}/{key.month or '-'}/{key.stem}" + ) + async def _checkpoint(self) -> None: """Upload all dirty catalogs to S3 and reconnect the adapters.""" ducklake = self._require_ducklake() - for ds in ducklake._datasets: - await ds.close(update_catalog=True) - await ducklake._catalog_adap.close(update=True) - await ducklake._columns_adap.close(update=True) - await ducklake._catalog_adap.connect() - await ducklake._columns_adap.connect() + await ducklake.flush_catalogs(update=True) self._changed_catalog = False async def _process_comparison( @@ -460,20 +963,31 @@ async def _process_comparison( @staticmethod def _s3_is_stale(comparison: FileComparison) -> bool: - """True when any non-S3 artifact is newer than the S3 copy.""" + """True when a non-S3 artifact certainly differs from the S3 copy. + + Trust-the-catalog policy: legacy ``origin_modified`` values are + upload timestamps (unreliable for freshness), so modification + dates are ignored on this pass. A file on S3 is only re-checked + when the source *size* is known to differ from the recorded + ``origin_size`` — size inequality proves the content changed, + so no false skips are possible. Equal sizes (or unknown legacy + sizes) trust the catalog; the content veto still guards any + re-check that does happen. + """ s3_record = comparison._pick("ducklake") if s3_record is None: return False - s3_source_modified = s3_record.source_modified or s3_record.modified - if s3_source_modified is None: + # origin_size is the raw source size recorded at upload time; + # the parquet size can never equal the raw size. + s3_size = s3_record.source_size + if not s3_size: return False for record in comparison.records: if record.origin == "ducklake": continue - modified = record.modified - if modified and modified > s3_source_modified: + if record.size and record.size != s3_size: return True return False diff --git a/pysus/tests/api/dadosgov/test_client.py b/pysus/tests/api/dadosgov/test_client.py index 2d7b4ea6..b8658182 100644 --- a/pysus/tests/api/dadosgov/test_client.py +++ b/pysus/tests/api/dadosgov/test_client.py @@ -269,7 +269,7 @@ async def test_connect_with_token_creates_client(self): "User-Agent": f"PySUS/{__version__}", "chave-api-dados-abertos": "test-token-123", }, - timeout=30.0, + timeout=httpx.Timeout(120.0, connect=30.0), follow_redirects=True, ) diff --git a/pysus/tests/api/dadosgov/test_databases.py b/pysus/tests/api/dadosgov/test_databases.py index 698b723a..65ed9c7f 100644 --- a/pysus/tests/api/dadosgov/test_databases.py +++ b/pysus/tests/api/dadosgov/test_databases.py @@ -401,9 +401,13 @@ def test_formatter_xlsx_prefixed_with_get(self): class TestAVAILABLEDATABASES: def test_contains_all_databases(self): - expected = {CNES, PNI, SIA, SINAN, SIM, SINASC, COVID19} + expected = {CNES, PNI, SINAN, SIM, SINASC, COVID19} assert set(AVAILABLE_DATABASES) == expected + def test_sia_excluded(self): + # Fortaleza municipal SIA is unstable and excluded from sync + assert SIA not in AVAILABLE_DATABASES + def test_all_can_be_instantiated(self): for db_class in AVAILABLE_DATABASES: ds = db_class(client=DadosGov()) diff --git a/pysus/tests/api/ducklake/test_client.py b/pysus/tests/api/ducklake/test_client.py index 2e959d31..fb6a8aa6 100644 --- a/pysus/tests/api/ducklake/test_client.py +++ b/pysus/tests/api/ducklake/test_client.py @@ -721,7 +721,8 @@ def test_remote_url(self, tmp_path): class TestAdapterClose: @pytest.mark.asyncio - async def test_close_disposes_engine(self, tmp_path): + async def test_close_keeps_shared_engine(self, tmp_path): + """Engines are shared per file and must outlive adapters.""" from unittest.mock import MagicMock from pysus.api.ducklake.catalog.adapters import CatalogAdapter @@ -731,10 +732,28 @@ async def test_close_disposes_engine(self, tmp_path): mock_engine = MagicMock() adapter._engine = mock_engine await adapter.close() - mock_engine.dispose.assert_called_once() + mock_engine.dispose.assert_not_called() assert adapter._engine is None assert adapter._session_factory is None + @pytest.mark.asyncio + async def test_reconnect_disposes_shared_engine(self, tmp_path): + from unittest.mock import patch + + from pysus.api.ducklake.catalog.adapters import CatalogAdapter + + with patch("pysus.api.ducklake.catalog.adapters.CACHEPATH", tmp_path): + adapter = CatalogAdapter() + with patch.object( + adapter, "connect", new_callable=AsyncMock + ) as mock_connect: + with patch( + "pysus.api.ducklake.catalog.adapters._dispose_shared" + ) as mock_dispose: + await adapter.reconnect() + mock_dispose.assert_called_once_with(adapter.db_local) + assert mock_connect.called + def test_destructor_no_engine(self, tmp_path): from pysus.api.ducklake.catalog.adapters import CatalogAdapter diff --git a/pysus/tests/management/test_catalog.py b/pysus/tests/management/test_catalog.py new file mode 100644 index 00000000..884eee75 --- /dev/null +++ b/pysus/tests/management/test_catalog.py @@ -0,0 +1,215 @@ +"""Tests for pysus.management.catalog CatalogWriter DB operations.""" + +from datetime import datetime +from unittest.mock import MagicMock + +import duckdb +import pytest +from pysus.management.catalog import CatalogWriter + +_SCHEMA = """ +CREATE SCHEMA pysus; +CREATE TABLE pysus.files ( + id INTEGER, + dataset_id INTEGER, + group_id INTEGER, + path VARCHAR UNIQUE, + size BIGINT, + rows INTEGER, + type VARCHAR, + modified TIMESTAMP, + origin_modified TIMESTAMP, + origin_size BIGINT, + origin_path VARCHAR, + sha256 VARCHAR, + source_sha256 VARCHAR, + origin VARCHAR, + format VARCHAR, + year INTEGER, + month INTEGER, + state VARCHAR +); +""" + +_MINIMAL_SCHEMA = """ +CREATE SCHEMA pysus; +CREATE TABLE pysus.files ( + id INTEGER, + dataset_id INTEGER, + group_id INTEGER, + path VARCHAR UNIQUE, + size BIGINT, + rows INTEGER, + type VARCHAR, + modified TIMESTAMP, + origin_modified TIMESTAMP, + origin_size BIGINT, + origin_path VARCHAR, + sha256 VARCHAR, + year INTEGER, + month INTEGER, + state VARCHAR +); +""" + + +@pytest.fixture +def writer_and_cursor(): + writer = CatalogWriter(ducklake=MagicMock()) + con = duckdb.connect(":memory:") + con.execute(_SCHEMA) + return writer, con.cursor(), con + + +def _insert( + cursor, + path, + origin_modified=None, + origin_size=0, + sha256=None, + source_sha256=None, +): + cursor.execute( + "INSERT INTO pysus.files (id, dataset_id, group_id, path, size, " + "rows, modified, origin_modified, origin_size, origin_path, " + "sha256, source_sha256, year, month, state) " + "VALUES (?, ?, ?, ?, 0, 0, CURRENT_TIMESTAMP, ?, ?, ?, ?, ?, " + "NULL, NULL, NULL)", + ( + 1, + 8, + None, + path, + origin_modified, + origin_size, + "ftp/x", + sha256, + source_sha256, + ), + ) + + +class TestGetFileFull: + def test_missing(self, writer_and_cursor): + writer, cursor, _ = writer_and_cursor + assert writer.get_file_full(cursor, "public/data/x.parquet") is None + + def test_full_row(self, writer_and_cursor): + writer, cursor, _ = writer_and_cursor + _insert( + cursor, + "public/data/x.parquet", + origin_modified=datetime(2026, 1, 1), + origin_size=42, + sha256="aa" * 32, + source_sha256="bb" * 32, + ) + result = writer.get_file_full(cursor, "public/data/x.parquet") + assert result is not None + file_id, modified, size, sha, source = result + assert file_id == 1 + assert modified == datetime(2026, 1, 1) + assert size == 42 + assert sha == "aa" * 32 + assert source == "bb" * 32 + + +class TestTouchFile: + def test_updates_origin_metadata(self, writer_and_cursor): + writer, cursor, _ = writer_and_cursor + _insert(cursor, "public/data/x.parquet") + writer.touch_file( + cursor, 1, datetime(2026, 2, 2), 99, source_sha256="cc" * 32 + ) + result = writer.get_file_full(cursor, "public/data/x.parquet") + assert result is not None + assert result[1] == datetime(2026, 2, 2) + assert result[2] == 99 + assert result[4] == "cc" * 32 + assert result[3] is None + + def test_without_source_sha256(self, writer_and_cursor): + writer, cursor, _ = writer_and_cursor + _insert(cursor, "public/data/x.parquet") + writer.touch_file(cursor, 1, datetime(2026, 2, 2), 99) + result = writer.get_file_full(cursor, "public/data/x.parquet") + assert result is not None + assert result[1] == datetime(2026, 2, 2) + assert result[2] == 99 + + +class TestUpsertFile: + def test_insert_stores_hashes(self, writer_and_cursor): + writer, cursor, _ = writer_and_cursor + file_id, created = writer.upsert_file( + cursor, + dataset_id=8, + group_id=None, + path="public/data/x.parquet", + size=10, + rows=3, + modified=datetime(2026, 1, 1), + origin_modified=datetime(2026, 1, 1), + origin_size=11, + origin_path="ftp/x", + year=2026, + month=None, + state=None, + origin="ftp", + format="parquet", + sha256="aa" * 32, + source_sha256="bb" * 32, + ) + assert created is True + result = writer.get_file_full(cursor, "public/data/x.parquet") + assert result is not None + assert result[0] == file_id + assert result[3] == "aa" * 32 + assert result[4] == "bb" * 32 + + def test_update_stores_hashes(self, writer_and_cursor): + writer, cursor, _ = writer_and_cursor + _insert(cursor, "public/data/x.parquet") + file_id, created = writer.upsert_file( + cursor, + dataset_id=8, + group_id=None, + path="public/data/x.parquet", + size=99, + rows=9, + modified=datetime(2026, 1, 1), + origin_modified=datetime(2026, 1, 1), + origin_size=11, + origin_path="ftp/x", + year=None, + month=None, + state=None, + sha256="aa" * 32, + source_sha256="bb" * 32, + ) + assert created is False + assert file_id == 1 + result = writer.get_file_full(cursor, "public/data/x.parquet") + assert result is not None + assert result[2] == 11 + assert result[3] == "aa" * 32 + assert result[4] == "bb" * 32 + + +class TestEnsureManagementColumns: + def test_adds_source_sha256_column(self): + writer = CatalogWriter(ducklake=MagicMock()) + con = duckdb.connect(":memory:") + con.execute(_MINIMAL_SCHEMA) + cursor = con.cursor() + writer._ensure_management_columns(cursor) + cursor.execute( + "SELECT column_name FROM information_schema.columns " + "WHERE table_schema = 'pysus' AND table_name = 'files' " + "AND column_name IN ('origin', 'format', 'source_sha256')" + ) + assert {row[0] for row in cursor.fetchall()} == { + "origin", + "format", + "source_sha256", + } diff --git a/pysus/tests/management/test_sync.py b/pysus/tests/management/test_sync.py index f114954c..4378c984 100644 --- a/pysus/tests/management/test_sync.py +++ b/pysus/tests/management/test_sync.py @@ -94,9 +94,7 @@ def test_is_current(self): assert not engine._is_current(file, datetime(2026, 1, 1)) assert not engine._is_current(file, None) - def test_s3_is_stale_when_ftp_newer(self): - from datetime import datetime - + def test_s3_is_stale_when_source_size_differs(self): from pysus.management.records import FileComparison, FileRecord ftp = FileRecord( @@ -104,7 +102,8 @@ def test_s3_is_stale_when_ftp_newer(self): dataset="SINAN", name="DENGBR25.dbc", path="ftp/x", - modified=datetime(2026, 6, 1), + size=200, + modified=None, group="DENG", year=2025, ) @@ -113,17 +112,40 @@ def test_s3_is_stale_when_ftp_newer(self): dataset="SINAN", name="DENGBR25.parquet", path="s3/x", - modified=datetime(2026, 1, 1), - source_modified=datetime(2026, 1, 1), + size=5000, + source_size=100, group="DENG", year=2025, ) comparison = FileComparison(key=ftp.identity_key(), records=[ftp, s3]) assert SyncEngine._s3_is_stale(comparison) - def test_s3_not_stale_when_equal(self): - from datetime import datetime + def test_s3_not_stale_when_sizes_equal(self): + from pysus.management.records import FileComparison, FileRecord + ftp = FileRecord( + origin="ftp", + dataset="SINAN", + name="DENGBR25.dbc", + path="ftp/x", + size=100, + group="DENG", + year=2025, + ) + s3 = FileRecord( + origin="ducklake", + dataset="SINAN", + name="DENGBR25.parquet", + path="s3/x", + size=5000, + source_size=100, + group="DENG", + year=2025, + ) + comparison = FileComparison(key=ftp.identity_key(), records=[ftp, s3]) + assert not SyncEngine._s3_is_stale(comparison) + + def test_s3_not_stale_when_origin_size_unknown(self): from pysus.management.records import FileComparison, FileRecord ftp = FileRecord( @@ -131,7 +153,32 @@ def test_s3_not_stale_when_equal(self): dataset="SINAN", name="DENGBR25.dbc", path="ftp/x", - modified=datetime(2026, 1, 1), + size=100, + group="DENG", + year=2025, + ) + s3 = FileRecord( + origin="ducklake", + dataset="SINAN", + name="DENGBR25.parquet", + path="s3/x", + size=5000, + source_size=0, + group="DENG", + year=2025, + ) + comparison = FileComparison(key=ftp.identity_key(), records=[ftp, s3]) + assert not SyncEngine._s3_is_stale(comparison) + + def test_s3_not_stale_when_source_size_zero(self): + from pysus.management.records import FileComparison, FileRecord + + ftp = FileRecord( + origin="dadosgov", + dataset="SINAN", + name="DENGBR25.csv.zip", + path="api/x", + size=0, group="DENG", year=2025, ) @@ -140,8 +187,8 @@ def test_s3_not_stale_when_equal(self): dataset="SINAN", name="DENGBR25.parquet", path="s3/x", - modified=datetime(2026, 1, 2), - source_modified=datetime(2026, 1, 1), + size=5000, + source_size=100, group="DENG", year=2025, ) From 7f3abc57eb112385241eb72767717b014ae838ca Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Lu=C3=A3=20Bida=20Vacaro?= Date: Fri, 14 Aug 2026 17:24:58 -0300 Subject: [PATCH 3/5] test(management): cover sync internals, normalize, inventory, adapters and scripts Add mock-based suites for the parallel sync pipeline (locks, pooled FTP downloads, conversion/upload payloads, catalog row writes), bucket relayout planning, inventory collectors, adapter registry/transaction lifecycle, and the CLI scripts. 823 tests passing. --- pysus/tests/api/ducklake/test_adapters.py | 196 ++++++++++ pysus/tests/management/test_inventory.py | 300 +++++++++++++++ pysus/tests/management/test_normalize.py | 380 +++++++++++++++++++ pysus/tests/management/test_scripts.py | 194 ++++++++++ pysus/tests/management/test_sync_internal.py | 342 +++++++++++++++++ 5 files changed, 1412 insertions(+) create mode 100644 pysus/tests/api/ducklake/test_adapters.py create mode 100644 pysus/tests/management/test_inventory.py create mode 100644 pysus/tests/management/test_normalize.py create mode 100644 pysus/tests/management/test_scripts.py create mode 100644 pysus/tests/management/test_sync_internal.py diff --git a/pysus/tests/api/ducklake/test_adapters.py b/pysus/tests/api/ducklake/test_adapters.py new file mode 100644 index 00000000..b3dfc87e --- /dev/null +++ b/pysus/tests/api/ducklake/test_adapters.py @@ -0,0 +1,196 @@ +"""Tests for the shared-engine/transaction plumbing in adapters.""" + +from pathlib import Path +from unittest.mock import AsyncMock, MagicMock, patch + +import pytest +from pysus.api.ducklake.catalog import adapters as adapters_module +from pysus.api.ducklake.catalog.adapters import ( + CatalogAdapter, + ColumnsAdapter, + DatasetAdapter, +) +from pysus.api.errors import CatalogError + + +class TestRegistry: + def test_dispose_shared_missing(self): + adapters_module._dispose_shared(Path("/tmp/nope.duckdb")) + + def test_dispose_shared_existing(self, tmp_path): + path = tmp_path / "catalog.duckdb" + engine = MagicMock() + adapters_module._SHARED_ENGINES[str(path)] = engine + adapters_module._dispose_shared(path) + engine.dispose.assert_called_once() + assert str(path) not in adapters_module._SHARED_ENGINES + + +class TestBaseAdapterSurface: + def test_mark_dirty_and_flag(self): + adapter = CatalogAdapter() + assert not adapter.local_dirty + adapter.mark_dirty() + assert adapter.local_dirty + + def test_connected_false_initially(self): + assert not CatalogAdapter().connected + + def test_raw_connection_not_connected(self): + with pytest.raises(CatalogError, match="not initialized"): + CatalogAdapter().raw_connection() + + def test_checkpoint_uses_shared_engine(self, tmp_path, monkeypatch): + monkeypatch.setattr(adapters_module, "CACHEPATH", tmp_path) + adapter = CatalogAdapter() + engine = MagicMock() + engine.raw_connection.return_value.execute = MagicMock() + adapters_module._SHARED_ENGINES[str(adapter.db_local.resolve())] = ( + engine + ) + adapter.checkpoint() + engine.raw_connection.return_value.execute.assert_called_once_with( + "CHECKPOINT" + ) + + @pytest.mark.asyncio + async def test_ensure_connected_creates_engine(self, tmp_path, monkeypatch): + monkeypatch.setattr(adapters_module, "CACHEPATH", tmp_path) + adapter = CatalogAdapter() + engine = MagicMock() + with patch.object(adapter, "connect", new=AsyncMock()) as mock_connect: + with patch.object( + adapter, + "setup_engine", + return_value=engine, + ): + await adapter.ensure_connected() + mock_connect.assert_awaited_once() + + @pytest.mark.asyncio + async def test_ensure_connected_engine_exists(self, tmp_path, monkeypatch): + monkeypatch.setattr(adapters_module, "CACHEPATH", tmp_path) + adapter = CatalogAdapter() + adapter._engine = MagicMock() + with patch.object(adapter, "connect", new=AsyncMock()) as mock: + await adapter.ensure_connected() + mock.assert_not_awaited() + + @pytest.mark.asyncio + async def test_reconnect_disposes_and_reconnects( + self, tmp_path, monkeypatch + ): + monkeypatch.setattr(adapters_module, "CACHEPATH", tmp_path) + adapter = CatalogAdapter() + with patch.object(adapter, "connect", new=AsyncMock()) as mock_connect: + with patch.object( + adapters_module, "_dispose_shared" + ) as mock_dispose: + await adapter.reconnect() + mock_dispose.assert_called_once() + mock_connect.assert_awaited_once_with(force=True) + + +class TestTransaction: + def test_transaction_commits(self, tmp_path): + import duckdb + + path = tmp_path / "catalog.duckdb" + con = duckdb.connect(str(path)) + con.execute("CREATE TABLE t (x INTEGER)") + con.close() + + adapter = DatasetAdapter(name="x", dataset_id=1, engine=None) + adapter.db_local = path + engine = MagicMock() + raw_conn = duckdb.connect(str(path)) + engine.raw_connection.return_value = raw_conn + + with patch.object(adapter, "setup_engine", return_value=engine): + with adapter.transaction() as (conn, cursor): + cursor.execute("INSERT INTO t VALUES (42)") + conn.commit() + + check = duckdb.connect(str(path), read_only=True) + rows = check.execute("SELECT * FROM t").fetchall() + check.close() + assert rows == [(42,)] + + def test_transaction_broken_raises(self, tmp_path): + adapter = DatasetAdapter(name="x", dataset_id=1, engine=None) + adapter.db_local = tmp_path / "catalog.duckdb" + engine = MagicMock() + probe = MagicMock() + probe.execute.side_effect = RuntimeError("closed") + broken = MagicMock() + broken.cursor.return_value = probe + engine.raw_connection.return_value = broken + + with patch.object(adapter, "setup_engine", return_value=engine): + with pytest.raises(CatalogError, match="broken"): + with adapter.transaction(): + pass + + +class TestAdapterKinds: + def test_adapter_paths(self, tmp_path, monkeypatch): + monkeypatch.setattr(adapters_module, "CACHEPATH", tmp_path) + assert CatalogAdapter().db_remote == Path("public/catalog.duckdb") + assert ColumnsAdapter().db_remote == Path( + "public/catalog_columns.duckdb" + ) + ds = DatasetAdapter(name="sinan", dataset_id=8) + assert ds.db_remote == Path("public/catalog_sinan.duckdb") + assert ds.dataset_name == "sinan" + + def test_upload_catalog_missing_credentials(self, tmp_path, monkeypatch): + monkeypatch.setattr(adapters_module, "CACHEPATH", tmp_path) + adapter = CatalogAdapter() + adapter.db_local = tmp_path / "catalog.duckdb" + adapter.db_local.write_bytes(b"") + with pytest.raises(PermissionError, match="credentials"): + import asyncio + + asyncio.run(adapter._upload_catalog()) + + def test_upload_catalog_missing_file(self, tmp_path, monkeypatch): + monkeypatch.setattr(adapters_module, "CACHEPATH", tmp_path) + adapter = CatalogAdapter() + adapter.db_local = tmp_path / "missing.duckdb" + adapter.credentials = MagicMock() + adapter.checkpoint = MagicMock() + with pytest.raises(FileNotFoundError): + import asyncio + + asyncio.run(adapter._upload_catalog()) + + @pytest.mark.asyncio + async def test_upload_catalog_calls_s3(self, tmp_path, monkeypatch): + monkeypatch.setattr(adapters_module, "CACHEPATH", tmp_path) + adapter = CatalogAdapter() + adapter.db_local = tmp_path / "catalog.duckdb" + adapter.db_local.write_bytes(b"x") + creds = MagicMock() + creds.access_key.get_secret_value.return_value = "ak" + creds.secret_key.get_secret_value.return_value = "sk" + adapter.credentials = creds + adapter.checkpoint = MagicMock() + with patch.object( + adapters_module, "upload_s3", new=AsyncMock() + ) as mock_upload: + await adapter._upload_catalog() + mock_upload.assert_awaited_once() + + @pytest.mark.asyncio + async def test_close_clears_refs_only(self, tmp_path, monkeypatch): + monkeypatch.setattr(adapters_module, "CACHEPATH", tmp_path) + adapter = CatalogAdapter() + engine = MagicMock() + adapters_module._SHARED_ENGINES[str(adapter.db_local.resolve())] = ( + engine + ) + adapter._engine = engine + adapter._session_factory = MagicMock() + await adapter.close() + engine.dispose.assert_not_called() + assert adapter._engine is None diff --git a/pysus/tests/management/test_inventory.py b/pysus/tests/management/test_inventory.py new file mode 100644 index 00000000..20515af9 --- /dev/null +++ b/pysus/tests/management/test_inventory.py @@ -0,0 +1,300 @@ +"""Tests for pysus.management.inventory collectors (mocked clients).""" + +from datetime import datetime +from unittest.mock import AsyncMock, MagicMock + +import pytest +from pysus.management.inventory import Inventory, _record_changed, _safe_modify +from pysus.management.records import FileRecord + + +@pytest.fixture +def inventory(): + pysus = MagicMock() + return Inventory(pysus=pysus) + + +def _awaitable(value): + """An awaitable attribute (AsyncMock attrs cannot be awaited directly).""" + + class _Aw: + def __await__(self): + async def _inner(): + return value + + return _inner().__await__() + + return _Aw() + + +_DATASET = None + + +def _fake_dataset(): + from pysus.api.ftp.client import FTP + from pysus.api.ftp.models import Dataset + + global _DATASET + if _DATASET is None: + + class _FakeDataset(Dataset): + @property + def name(self): + return "SINAN" + + @property + def long_name(self): + return "Sistema de Informacao de Agravos" + + @property + def description(self): + return "" + + def formatter(self, filename): + return {} + + _DATASET = _FakeDataset(client=FTP()) + return _DATASET + + +def _ftp_file(name="DENGBR25.dbc"): + from pysus.api.ftp.models import File as FTPFile + + info = { + "name": name, + "path": f"/dissemin/publicos/SINAN/{name}", + "size": 100, + "modify": datetime(2026, 1, 1), + "type": "file", + "group": {"name": "DENG", "long_name": "Dengue"}, + "year": 2025, + "month": None, + "state": None, + } + return FTPFile( + path=info["path"], + dataset=_fake_dataset(), + type=info["type"], + _info=info, + ) + + +class TestSafeModify: + def test_modify_available(self): + file = MagicMock() + file.modify = datetime(2026, 1, 1) + assert _safe_modify(file) == datetime(2026, 1, 1) + + def test_modify_raises(self): + file = MagicMock() + type(file).modify = property( + lambda self: (_ for _ in ()).throw(ValueError("nope")) + ) + assert _safe_modify(file) is None + + +class TestRecordChanged: + def test_size_change(self): + a = FileRecord( + origin="ftp", dataset="X", name="a.dbc", path="p", size=1 + ) + b = FileRecord( + origin="ftp", dataset="X", name="a.dbc", path="p", size=2 + ) + assert _record_changed(a, b) + + def test_modify_change(self): + a = FileRecord( + origin="ftp", + dataset="X", + name="a.dbc", + path="p", + size=1, + modified=datetime(2026, 1, 1), + ) + b = FileRecord( + origin="ftp", + dataset="X", + name="a.dbc", + path="p", + size=1, + modified=datetime(2026, 1, 2), + ) + assert _record_changed(a, b) + + def test_identical(self): + a = FileRecord( + origin="ftp", + dataset="X", + name="a.dbc", + path="p", + size=1, + modified=datetime(2026, 1, 1), + ) + b = FileRecord( + origin="ftp", + dataset="X", + name="a.dbc", + path="p", + size=1, + modified=datetime(2026, 1, 1), + ) + assert not _record_changed(a, b) + + +class TestWalkFtpItem: + @pytest.mark.asyncio + async def test_ftp_file(self, inventory): + file = _ftp_file() + records = await inventory._walk_ftp_item(file) + assert len(records) == 1 + assert records[0].origin == "ftp" + assert records[0].group == "DENG" + assert records[0].year == 2025 + assert records[0].file is file + + @pytest.mark.asyncio + async def test_group_walk(self, inventory): + from pysus.api.models import BaseRemoteGroup + + group = MagicMock(spec=BaseRemoteGroup) + group.files = [AsyncMock()] + group.files = _awaitable((_ftp_file(), _ftp_file("DENGBR24.dbc"))) + records = await inventory._walk_ftp_item(group) + assert len(records) == 2 + + @pytest.mark.asyncio + async def test_directory_walk(self, inventory): + from pysus.api.ftp.models import Directory + + directory = Directory("/dissemin/publicos/SINAN") + directory.loaded = True + directory._content = [_ftp_file("A.dbc")] + records = await inventory._walk_ftp_item(directory) + assert len(records) == 1 + assert records[0].name == "A.dbc" + + @pytest.mark.asyncio + async def test_unknown_item(self, inventory): + records = await inventory._walk_ftp_item(object()) + assert records == [] + + +class TestCollectFtp: + @pytest.mark.asyncio + async def test_collect_ftp(self, inventory): + dataset = MagicMock() + dataset.name = "SINAN" + dataset.content = _awaitable([_ftp_file()]) + client = MagicMock() + client.datasets = AsyncMock(return_value=[dataset]) + inventory.pysus.get_ftp = AsyncMock(return_value=client) + + records = await inventory.collect("ftp") + assert len(records) == 1 + assert records[0].dataset == "SINAN" + + @pytest.mark.asyncio + async def test_collect_ftp_filtered(self, inventory): + dataset = MagicMock() + dataset.name = "SINAN" + dataset.content = _awaitable([_ftp_file()]) + other = MagicMock() + other.name = "SIM" + other.content = _awaitable([]) + client = MagicMock() + client.datasets = AsyncMock(return_value=[dataset, other]) + inventory.pysus.get_ftp = AsyncMock(return_value=client) + + records = await inventory.collect("ftp", ["SINAN"]) + assert len(records) == 1 + + @pytest.mark.asyncio + async def test_collect_unknown_origin(self, inventory): + with pytest.raises(ValueError, match="Unknown origin"): + await inventory.collect("nope") + + +class TestCollectDadosgov: + @pytest.mark.asyncio + async def test_collect_dadosgov(self, inventory): + from pathlib import Path + from unittest.mock import patch + + from pysus.api.models import BaseRemoteGroup + + class _FakeGroup(BaseRemoteGroup): + @property + def name(self): + return "DENG" + + @property + def long_name(self): + return "Dengue" + + @property + def description(self): + return "" + + async def _fetch_files(self): + return [] + + group = _FakeGroup(path=Path("x"), type="file", dataset=_fake_dataset()) + group._files = [_ftp_file("DENGBR25.csv.zip")] + dataset = MagicMock() + dataset.name = "SINAN" + dataset.content = _awaitable([group]) + client = MagicMock() + client.datasets = AsyncMock(return_value=[dataset]) + inventory.pysus.get_dadosgov = AsyncMock(return_value=client) + + record = FileRecord( + origin="dadosgov", + dataset="SINAN", + name="DENGBR25.csv.zip", + path="http://x", + group="DENG", + year=2025, + ) + with patch.object( + Inventory, + "_walk_ftp_item", + new=AsyncMock(return_value=[record]), + ): + records = await inventory.collect("dadosgov", dadosgov_token="tok") + assert len(records) == 1 + assert records[0].origin == "dadosgov" + + +class TestCollectDucklake: + @pytest.mark.asyncio + async def test_collect_ducklake(self, inventory): + record = MagicMock() + record.rows = 10 + record.sha256 = "a" * 64 + record.year = 2025 + record.month = None + record.state = None + record.modified = datetime(2026, 1, 1) + record.origin_path = "/ftp/x.dbc" + record.origin_size = 50 + record.origin_modified = datetime(2026, 1, 1) + record.group = None + + file = MagicMock() + file.basename = "X.parquet" + file.path = "public/data/ftp/sinan/X.parquet" + file.size = 100 + file.record = record + + dataset = MagicMock() + dataset.name = "sinan" + dataset.query = AsyncMock(return_value=[file]) + client = MagicMock() + client.datasets = AsyncMock(return_value=[dataset]) + inventory.pysus.get_ducklake = AsyncMock(return_value=client) + + records = await inventory.collect("ducklake") + assert len(records) == 1 + assert records[0].sha256 == "a" * 64 + assert records[0].source_path == "/ftp/x.dbc" diff --git a/pysus/tests/management/test_normalize.py b/pysus/tests/management/test_normalize.py new file mode 100644 index 00000000..35bcc21b --- /dev/null +++ b/pysus/tests/management/test_normalize.py @@ -0,0 +1,380 @@ +"""Tests for pysus.management.normalize.""" + +from unittest.mock import MagicMock + +import pytest +from pysus.management.normalize import ( + _SCAN_PREFIXES, + BucketNormalizer, + CatalogPathFix, + CatalogRowDelete, + ObjectRename, + formatter_for, +) + +_SCHEMA = """ +CREATE SCHEMA pysus; +CREATE TABLE pysus.dataset_groups ( + id INTEGER, dataset_id INTEGER, name VARCHAR, long_name VARCHAR, + description VARCHAR +); +CREATE TABLE pysus.files ( + id INTEGER, dataset_id INTEGER, group_id INTEGER, path VARCHAR UNIQUE, + size BIGINT, rows INTEGER, type VARCHAR, modified TIMESTAMP, + origin_modified TIMESTAMP, origin_size BIGINT, origin_path VARCHAR, + sha256 VARCHAR, source_sha256 VARCHAR, origin VARCHAR, format VARCHAR, + year INTEGER, month INTEGER, state VARCHAR +); +CREATE TABLE pysus.file_columns (file_id INTEGER, column_id INTEGER); +""" + + +@pytest.fixture +def normalizer(): + client = MagicMock() + normalizer = BucketNormalizer(access_key="ak", secret_key="sk") + normalizer.client = client + client.head_object.return_value = {"Metadata": {}} + return normalizer + + +class TestFormatterFor: + def test_ftp_formatter(self): + formatter = formatter_for("ftp", "SINAN") + assert formatter is not None + parsed = formatter("DENGBR25.dbc") + assert parsed["group"]["name"] == "DENG" + assert parsed["year"] == 2025 + + def test_dadosgov_formatter(self): + formatter = formatter_for("dadosgov", "SIM") + assert formatter is not None + parsed = formatter("Mortalidade_Geral_2022_csv.zip") + assert parsed["year"] == 2022 + + def test_unknown_dataset(self): + assert formatter_for("ftp", "NOPE") is None + + def test_cached(self): + first = formatter_for("ftp", "CIHA") + second = formatter_for("ftp", "CIHA") + assert first is second + + +class TestSplitKey: + def test_public_data_key(self): + assert BucketNormalizer._split_key( + "public/data/ftp/sinan/DENG/2025/_/BR/X.parquet" + ) == ("ftp", "sinan") + + def test_non_public_key(self): + assert BucketNormalizer._split_key("data/ftp/x") == (None, None) + + def test_short_key(self): + assert BucketNormalizer._split_key("public/data") == (None, None) + + +class TestEnrich: + def test_catalog_values_win(self, normalizer): + enriched = normalizer._enrich( + "ftp", "SINAN", "DENGBR25.dbc", "DENG", 2025, None, None + ) + assert enriched == { + "group": "DENG", + "year": 2025, + "month": None, + "state": None, + } + + def test_formatter_fills_gaps(self, normalizer): + enriched = normalizer._enrich( + "ftp", "SINAN", "DENGBR25.dbc", None, None, None, None + ) + assert enriched["group"] == "DENG" + assert enriched["year"] == 2025 + + def test_legacy_group_replaced_by_formatter(self, normalizer): + enriched = normalizer._enrich( + "ftp", "CIHA", "CIHAMA2209.parquet", "Dados", 2022, 9, "MA" + ) + assert enriched["group"] == "CIHA" + + def test_unknown_formatter_keeps_values(self, normalizer): + enriched = normalizer._enrich( + "ftp", "NOPE", "X.dbc", "G", 2020, 1, "AC" + ) + assert enriched == { + "group": "G", + "year": 2020, + "month": 1, + "state": "AC", + } + + +class TestSurveyRelayout: + def test_survey_renames_and_deletes(self, normalizer, tmp_path): + import duckdb as _duckdb + + con = _duckdb.connect(str(tmp_path / "catalog_sinan.duckdb")) + con.execute(_SCHEMA) + con.execute( + "INSERT INTO pysus.files (id, dataset_id, group_id, path, size, " + "rows, modified, origin_modified, origin_size, origin_path, " + "sha256, year, month, state) VALUES (1, 8, NULL, " + "'public/data/ftp/sinan/DENGBR25.parquet', 100, 5, " + "CURRENT_TIMESTAMP, CURRENT_TIMESTAMP, 50, " + "'/ftp/DENGBR25.dbc', NULL, 2025, NULL, NULL)" + ) + con.execute( + "INSERT INTO pysus.files (id, dataset_id, group_id, path, size, " + "rows, modified, origin_modified, origin_size, origin_path, " + "sha256, year, month, state) VALUES (2, 8, NULL, " + "'public/data/ftp/sinan/DENGBR25.dbc', 50, 5, " + "CURRENT_TIMESTAMP, CURRENT_TIMESTAMP, 50, " + "'/ftp/DENGBR25.dbc', NULL, 2025, NULL, NULL)" + ) + con.close() + + objects = { + "public/data/ftp/sinan/DENGBR25.parquet", + "public/data/ftp/sinan/DENGBR25.dbc", + } + plan = normalizer.survey_relayout( + tmp_path / "catalog_sinan.duckdb", objects + ) + assert len(plan.catalog_fixes) == 1 + fix = plan.catalog_fixes[0] + assert fix.new_path == ( + "public/data/ftp/sinan/DENG/2025/_/BR/DENGBR25.parquet" + ) + assert len(plan.object_renames) == 1 + assert len(plan.catalog_row_deletes) == 1 + + def test_missing_object_rows_deleted(self, normalizer, tmp_path): + import duckdb as _duckdb + + con = _duckdb.connect(str(tmp_path / "catalog_x.duckdb")) + con.execute(_SCHEMA) + con.execute( + "INSERT INTO pysus.files (id, dataset_id, group_id, path, size, " + "rows, modified, origin_modified, origin_size, origin_path, " + "sha256, year, month, state) VALUES (1, 8, NULL, " + "'public/data/ftp/sinan/DENGBR25.parquet', 100, 5, " + "CURRENT_TIMESTAMP, CURRENT_TIMESTAMP, 50, " + "'/ftp/DENGBR25.dbc', NULL, 2025, NULL, NULL)" + ) + con.close() + + plan = normalizer.survey_relayout(tmp_path / "catalog_x.duckdb", set()) + assert len(plan.catalog_fixes) == 0 + assert len(plan.catalog_row_deletes) == 1 + + +class TestRelocateUncataloged: + def test_relocates_with_formatter(self, normalizer): + keys = { + "public/data/ftp/sinan/DENGBR25.parquet", + "public/data/dadosgov/sim/Mortalidade_Geral_2022.parquet", + } + plan = normalizer.relocate_uncataloged(keys, set()) + assert len(plan.object_renames) == 2 + targets = {r.new for r in plan.object_renames} + assert "public/data/ftp/sinan/DENG/2025/_/BR/DENGBR25.parquet" in ( + targets + ) + + def test_unparsable_kept_with_placeholders(self, normalizer): + plan = normalizer.relocate_uncataloged( + {"public/data/ftp/sinan/weird-file-xyz.parquet"}, set() + ) + assert len(plan.object_renames) == 1 + assert plan.object_renames[0].new == ( + "public/data/ftp/sinan/_/_/_/BR/weird-file-xyz.parquet" + ) + + def test_non_public_key_kept_raw(self, normalizer): + plan = normalizer.relocate_uncataloged( + {"data/ftp/sih/ERAC1101.parquet"}, set() + ) + assert plan.object_renames == [] + assert len(plan.raw_objects) == 1 + + +class TestCopySource: + def test_large_object_skips_head(self, normalizer): + assert ( + normalizer._copy_source("public/data/k", size=100000) + == "public/data/k" + ) + normalizer.client.head_object.assert_not_called() + + def test_alias_followed(self, normalizer): + normalizer.client.head_object.side_effect = [ + {"Metadata": {"pysus-alias": "public/data/new"}}, + {"Metadata": {}}, + ] + assert ( + normalizer._copy_source("public/data/old", size=10) + == "public/data/new" + ) + + def test_head_error_returns_key(self, normalizer): + normalizer.client.head_object.side_effect = Exception("boom") + assert ( + normalizer._copy_source("public/data/k", size=10) == "public/data/k" + ) + + +class TestDoRelocate: + def test_self_copy_skipped(self, normalizer): + rename = ObjectRename(old="k", new="k") + normalizer._do_relocate(rename, {"k": 10}) + normalizer.client.copy_object.assert_not_called() + + def test_relocate_copies_and_aliases(self, normalizer): + normalizer.client.head_object.return_value = {"Metadata": {}} + rename = ObjectRename(old="old", new="new") + normalizer._do_relocate(rename, {"old": 100}) + normalizer.client.copy_object.assert_called_once() + normalizer.client.put_object.assert_called_once() + meta = normalizer.client.put_object.call_args.kwargs["Metadata"] + assert meta == {"pysus-alias": "new"} + + +class TestApplyRenamesWithAliases: + def test_dry_run_returns_empty(self, normalizer): + aliases = normalizer.apply_renames_with_aliases( + [ObjectRename(old="a", new="b")], dry_run=True + ) + assert aliases == {} + normalizer.client.copy_object.assert_not_called() + + def test_parallel_apply(self, normalizer): + normalizer.client.head_object.return_value = {"Metadata": {}} + renames = [ + ObjectRename(old=f"old{i}", new=f"new{i}") for i in range(10) + ] + aliases = normalizer.apply_renames_with_aliases( + renames, + dry_run=False, + object_sizes={"old0": 100}, + workers=4, + ) + assert set(aliases) == {f"old{i}" for i in range(10)} + assert normalizer.client.copy_object.call_count == 10 + + +class TestApplyObjects: + def test_deletes(self, normalizer): + normalizer.apply_objects([], ["k1", "k2"], dry_run=False) + assert normalizer.client.delete_object.call_count == 2 + + def test_dry_run(self, normalizer): + normalizer.apply_objects([], ["k1"], dry_run=True) + normalizer.client.delete_object.assert_not_called() + + +class TestSurveyCatalog: + def test_survey_broken_and_raw(self, normalizer, tmp_path): + import duckdb as _duckdb + + con = _duckdb.connect(str(tmp_path / "catalog_ciha.duckdb")) + con.execute(_SCHEMA) + con.execute( + "INSERT INTO pysus.files (id, dataset_id, group_id, path, size, " + "rows, modified, origin_modified, origin_size, origin_path, " + "sha256, year, month, state) VALUES (1, 8, NULL, " + "'public/data/ftp/ciha/CIHAAC2201.dbc', 100, 5, " + "CURRENT_TIMESTAMP, CURRENT_TIMESTAMP, 50, " + "'/ftp/CIHAAC2201.dbc', NULL, 2022, 1, 'AC')" + ) + con.close() + + normalizer._object_exists = lambda key: key.endswith(".dbc") + fixes, deletes = normalizer.survey_catalog( + tmp_path / "catalog_ciha.duckdb" + ) + assert fixes == [] + assert deletes == [] + assert normalizer.raw_objects == ["public/data/ftp/ciha/CIHAAC2201.dbc"] + + +class TestSurveyObjects: + def test_non_parquet_raw(self, normalizer): + normalizer._list_objects = MagicMock( + return_value=[("public/data/ftp/ciha/X.dbc", 100)] + ) + renames, deletes = normalizer.survey_objects() + assert renames == [] + assert normalizer.raw_objects.count( + "public/data/ftp/ciha/X.dbc" + ) == len(_SCAN_PREFIXES) + + def test_format_token_parquet(self, normalizer): + def _listing(prefix): + if prefix == "public/data/dadosgov/": + return [("public/data/dadosgov/sinan/X.csv.parquet", 100)] + return [] + + normalizer._list_objects = MagicMock(side_effect=_listing) + renames, deletes = normalizer.survey_objects() + assert len(renames) == 1 + assert renames[0].new == "public/data/dadosgov/sinan/X.parquet" + + def test_collision_keeps_csv(self, normalizer): + def _listing(prefix): + if prefix == "public/data/dadosgov/": + return [ + ("public/data/dadosgov/sim/M_2022.csv.parquet", 200), + ("public/data/dadosgov/sim/M_2022.json.parquet", 100), + ] + return [] + + normalizer._list_objects = MagicMock(side_effect=_listing) + renames, deletes = normalizer.survey_objects() + assert len(renames) == 1 + assert renames[0].old.endswith("csv.parquet") + assert deletes == ["public/data/dadosgov/sim/M_2022.json.parquet"] + + +class TestApplyCatalog: + def test_apply_fixes_and_deletes(self, normalizer, tmp_path): + import duckdb as _duckdb + + path = tmp_path / "catalog_sinan.duckdb" + con = _duckdb.connect(str(path)) + con.execute(_SCHEMA) + con.execute( + "INSERT INTO pysus.files (id, dataset_id, group_id, path, size, " + "rows, modified, origin_modified, origin_size, origin_path, " + "sha256, year, month, state) VALUES (1, 8, NULL, " + "'public/data/ftp/sinan/DENGBR25.parquet', 100, 5, " + "CURRENT_TIMESTAMP, CURRENT_TIMESTAMP, 50, " + "'/ftp/DENGBR25.dbc', NULL, 2025, NULL, NULL)" + ) + con.close() + + normalizer.apply_catalog( + path, + [ + CatalogPathFix( + catalog="sinan", + old_path="public/data/ftp/sinan/DENGBR25.parquet", + new_path="public/data/ftp/sinan/DENG/2025/_/BR/" + "DENGBR25.parquet", + ) + ], + [ + CatalogRowDelete( + catalog="sinan", + path="public/data/ftp/sinan/DENGBR25.parquet", + ) + ], + dry_run=False, + ) + + con = _duckdb.connect(str(path), read_only=True) + rows = con.execute("SELECT path FROM pysus.files").fetchall() + con.close() + expected = "public/data/ftp/sinan/DENG/2025/_/BR/DENGBR25.parquet" + assert rows == [(expected,)] diff --git a/pysus/tests/management/test_scripts.py b/pysus/tests/management/test_scripts.py new file mode 100644 index 00000000..f76fc49b --- /dev/null +++ b/pysus/tests/management/test_scripts.py @@ -0,0 +1,194 @@ +"""Tests for pysus.management.scripts (mocked engines).""" + +from unittest.mock import AsyncMock, MagicMock, patch + +import pytest + + +class TestSyncClientsScript: + def test_load_env(self, tmp_path): + from pysus.management.scripts.sync_clients import load_env + + env_file = tmp_path / ".env" + env_file.write_text( + "ACCESS_KEY=ak\nSECRET_KEY=sk\nDADOSGOV_TOKEN=tok\n" + ) + env = load_env(str(env_file)) + assert env == { + "ACCESS_KEY": "ak", + "SECRET_KEY": "sk", + "DADOSGOV_TOKEN": "tok", + } + + @pytest.mark.asyncio + async def test_run(self, tmp_path, capsys): + from pysus.management.scripts.sync_clients import run + + with patch( + "pysus.management.scripts.sync_clients.load_env", + return_value={ + "ACCESS_KEY": "ak", + "SECRET_KEY": "sk", + "DADOSGOV_TOKEN": "tok", + }, + ): + with patch( + "pysus.management.scripts.sync_clients.SyncEngine" + ) as mock_cls: + engine = mock_cls.return_value + engine.__aenter__ = AsyncMock(return_value=engine) + engine.__aexit__ = AsyncMock(return_value=None) + engine.run = AsyncMock( + return_value=MagicMock(summary=lambda: {"total": 1}) + ) + summary = await run(["SINAN"], 500, False, 4, 2) + assert summary == {"total": 1} + + def test_main_runs(self, tmp_path, capsys): + from pysus.management.scripts import sync_clients + + with patch.object( + sync_clients, "run", new=AsyncMock(return_value={"uploaded": 1}) + ): + with patch.object(sync_clients, "load_env", return_value={}): + with patch( + "sys.argv", + ["sync_clients", "--datasets", "SINAN"], + ): + assert sync_clients.main() == 0 + + +class TestCompareClientsScript: + def test_load_env(self, tmp_path): + from pysus.management.scripts.compare_clients import load_env + + env_file = tmp_path / ".env" + env_file.write_text("ACCESS_KEY=ak\n") + assert load_env(str(env_file)) == {"ACCESS_KEY": "ak"} + + @pytest.mark.asyncio + async def test_run(self): + from pysus.management.scripts.compare_clients import run + + engine = MagicMock() + engine.__aenter__ = AsyncMock(return_value=engine) + engine.__aexit__ = AsyncMock(return_value=None) + engine.inventory.collect = AsyncMock(return_value=[]) + with patch( + "pysus.management.scripts.compare_clients.load_env", + return_value={"ACCESS_KEY": "ak"}, + ): + with patch( + "pysus.management.scripts.compare_clients.SyncEngine", + return_value=engine, + ): + result = await run(["SINAN"]) + assert result["origin_counts"] == { + "ducklake": 0, + "ftp": 0, + "dadosgov": 0, + } + assert result["reports"] == [] + + def test_print_table(self, capsys): + from pysus.management.scripts.compare_clients import print_table + + result = { + "origin_counts": {"ftp": 1, "dadosgov": 0, "ducklake": 0}, + "reports": [ + { + "dataset": "SINAN", + "total": 1, + "on_all_three": 0, + "on_ftp_dadosgov": 0, + "on_ftp_s3": 0, + "on_dadosgov_s3": 0, + "ftp_only": 1, + "dadosgov_only": 0, + "s3_only": 0, + "examples": {"ftp_only": ["SINAN/DENG/2025/-/dengbr25"]}, + } + ], + } + print_table(result) + out = capsys.readouterr().out + assert "SINAN" in out + assert "ftp_only" in out + + def test_main_json(self, tmp_path, capsys): + from pysus.management.scripts import compare_clients + + with patch.object( + compare_clients, + "run", + new=AsyncMock( + return_value={ + "origin_counts": {}, + "reports": [], + } + ), + ): + with patch( + "sys.argv", + ["compare_clients", "--json", "--datasets", "SINAN"], + ): + assert compare_clients.main() == 0 + assert '"origin_counts"' in capsys.readouterr().out + + +class TestRelayoutBucketScript: + def test_load_env(self, tmp_path): + from pysus.management.scripts.relayout_bucket import load_env + + env_file = tmp_path / ".env" + env_file.write_text("ACCESS_KEY=ak\nSECRET_KEY=sk\n") + assert load_env(str(env_file)) == { + "ACCESS_KEY": "ak", + "SECRET_KEY": "sk", + } + + def test_main_dry_run(self, tmp_path, capsys): + from pysus.management.scripts import relayout_bucket + + normalizer = MagicMock() + normalizer._list_objects.return_value = [] + normalizer.survey_relayout.return_value = MagicMock( + object_renames=[], + object_deletes=[], + catalog_fixes=[], + catalog_row_deletes=[], + broken_rows=[], + raw_objects=[], + summary=lambda: {}, + ) + normalizer.relocate_uncataloged.return_value = MagicMock( + object_renames=[], + object_deletes=[], + catalog_fixes=[], + catalog_row_deletes=[], + broken_rows=[], + raw_objects=[], + summary=lambda: {}, + ) + with patch( + "pysus.management.scripts.relayout_bucket.load_env", + return_value={"ACCESS_KEY": "ak", "SECRET_KEY": "sk"}, + ): + with patch( + "pysus.management.scripts.relayout_bucket.BucketNormalizer", + return_value=normalizer, + ): + with patch( + "pysus.management.scripts.relayout_bucket.httpx.Client" + ) as mock_http: + response = MagicMock() + response.content = b"" + response.raise_for_status = MagicMock() + http_get = mock_http.return_value.__enter__.return_value.get + http_get.return_value = response + with patch( + "sys.argv", + ["relayout_bucket", "--dry-run"], + ): + assert relayout_bucket.main() == 0 + assert "DRY RUN" in capsys.readouterr().out diff --git a/pysus/tests/management/test_sync_internal.py b/pysus/tests/management/test_sync_internal.py new file mode 100644 index 00000000..1ac69be3 --- /dev/null +++ b/pysus/tests/management/test_sync_internal.py @@ -0,0 +1,342 @@ +"""Tests for pysus.management.sync connection and helper paths.""" + +from unittest.mock import AsyncMock, MagicMock, PropertyMock, patch + +import pytest +from pysus.management.records import FileComparison, FileRecord +from pysus.management.sync import SyncEngine + + +@pytest.fixture +def engine(): + return SyncEngine(access_key="ak", secret_key="sk") + + +def _record(origin, name, dataset="SINAN", year=2025, **kw): + return FileRecord( + origin=origin, + dataset=dataset, + name=name, + path=f"{origin}/{dataset}/{name}", + year=year, + file=kw.pop("file", MagicMock()), + **kw, + ) + + +class TestPickSourceAndLabel: + def test_pick_source_prefers_ftp(self): + ftp = _record("ftp", "DENGBR25.dbc", file=MagicMock()) + gov = _record("dadosgov", "DENGBR25.csv.zip", file=MagicMock()) + comparison = FileComparison(key=ftp.identity_key(), records=[ftp, gov]) + assert SyncEngine._pick_source(comparison) is ftp + + def test_pick_source_falls_back_to_dadosgov(self): + gov = _record("dadosgov", "DENGBR25.csv.zip", file=MagicMock()) + comparison = FileComparison(key=gov.identity_key(), records=[gov]) + assert SyncEngine._pick_source(comparison) is gov + + def test_pick_source_none_without_files(self): + gov = _record("dadosgov", "DENGBR25.csv.zip", file=None) + comparison = FileComparison(key=gov.identity_key(), records=[gov]) + assert SyncEngine._pick_source(comparison) is None + + def test_label(self): + ftp = _record( + "ftp", + "PAAC2501.dbc", + dataset="SIA", + group="PA", + month=1, + state="AC", + year=2025, + ) + comparison = FileComparison(key=ftp.identity_key(), records=[ftp]) + assert SyncEngine._label(comparison) == "SIA/PA/2025/1/paac2501" + + +class TestSyncLock: + def test_acquire_creates_lock(self, engine, tmp_path, monkeypatch): + monkeypatch.setattr("pysus.management.sync.CACHEPATH", tmp_path) + engine._acquire_sync_lock() + lock = tmp_path / "ducklake" / ".sync.lock" + assert lock.exists() + assert int(lock.read_text()) > 0 + + def test_acquire_conflict_raises(self, engine, tmp_path, monkeypatch): + import os + + monkeypatch.setattr("pysus.management.sync.CACHEPATH", tmp_path) + lock = tmp_path / "ducklake" / ".sync.lock" + lock.parent.mkdir(parents=True, exist_ok=True) + lock.write_text(str(os.getpid())) + with pytest.raises(Exception, match="holds the catalog lock"): + engine._acquire_sync_lock() + + def test_acquire_steals_dead_lock(self, engine, tmp_path, monkeypatch): + monkeypatch.setattr("pysus.management.sync.CACHEPATH", tmp_path) + lock = tmp_path / "ducklake" / ".sync.lock" + lock.parent.mkdir(parents=True, exist_ok=True) + lock.write_text("99999999") + engine._acquire_sync_lock() + assert int(lock.read_text()) != 99999999 + + def test_release_removes_lock(self, engine, tmp_path, monkeypatch): + monkeypatch.setattr("pysus.management.sync.CACHEPATH", tmp_path) + engine._acquire_sync_lock() + engine._release_sync_lock() + assert not (tmp_path / "ducklake" / ".sync.lock").exists() + + def test_pid_alive(self): + import os + + assert SyncEngine._pid_alive(os.getpid()) + assert not SyncEngine._pid_alive(99999999) + + +class TestDownloadOnce: + @pytest.mark.asyncio + async def test_download_once_ftp_pooled(self, engine, tmp_path): + ftp = MagicMock() + ftp.size.return_value = 5 + file = MagicMock() + file.path = "/remote/X.dbc" + client = MagicMock() + client.ftp = ftp + + def _retrbinary(cmd, cb): + cb(b"hello") + + ftp.retrbinary = _retrbinary + out = tmp_path / "x.dbc" + + with patch( + "anyio.to_thread.run_sync", + new=AsyncMock(side_effect=lambda fn, *a, **kw: fn(*a, **kw)), + ): + await engine._download_once(file, out, ftp_client=client) + + assert out.read_bytes() == b"hello" + + @pytest.mark.asyncio + async def test_download_once_reconnects_broken_pool(self, engine, tmp_path): + file = MagicMock() + file.path = "/remote/X.dbc" + client = MagicMock() + broken_ftp = MagicMock() + broken_ftp.size.side_effect = OSError("broken") + client.ftp = None + + async def _connect(): + client.ftp = broken_ftp + + client.connect = _connect + out = tmp_path / "x.dbc" + + with pytest.raises(OSError): + await engine._download_once(file, out, ftp_client=client) + assert client._ftp is None # reset after failure + + @pytest.mark.asyncio + async def test_download_once_falls_back_to_file(self, engine, tmp_path): + file = MagicMock() + file.client = MagicMock() + file.client.ftp = None + file._download = AsyncMock() + out = tmp_path / "x.dbc" + await engine._download_once(file, out) + file._download.assert_awaited_once_with(output=out) + + +class TestDownloadRawWithRetry: + @pytest.mark.asyncio + async def test_retries_then_succeeds(self, engine, tmp_path): + file = MagicMock() + file.basename = "X.dbc" + attempts = [] + + async def _once(f, output, ftp_client=None): + attempts.append(1) + if len(attempts) < 3: + raise ConnectionResetError("drop") + output.write_bytes(b"data") + return output + + with patch.object( + engine, "_download_once", new=AsyncMock(side_effect=_once) + ): + with patch( + "pysus.management.sync.asyncio.sleep", + new=AsyncMock(), + ): + raw = await engine._download_raw_with_retry(file) + assert raw.read_bytes() == b"data" + assert len(attempts) == 3 + + @pytest.mark.asyncio + async def test_exhausts_retries(self, engine): + file = MagicMock() + file.basename = "X.dbc" + with patch.object( + engine, + "_download_once", + new=AsyncMock(side_effect=TimeoutError("slow")), + ): + with patch( + "pysus.management.sync.asyncio.sleep", + new=AsyncMock(), + ): + with pytest.raises(RuntimeError, match="Failed to download"): + await engine._download_raw_with_retry(file, max_retries=2) + + +class TestConvertAndUpload: + @pytest.mark.asyncio + async def test_convert_and_upload(self, engine, tmp_path): + raw = tmp_path / "X.dbc" + raw.write_bytes(b"raw") + parquet = tmp_path / "X.parquet" + parquet.write_bytes(b"pq") + + fake_parquet = MagicMock() + fake_parquet.path = parquet + fake_parquet.rows = 10 + fake_parquet.schema = "schema" + + local_file = MagicMock() + local_file.to_parquet = AsyncMock(return_value=fake_parquet) + + engine.s3_key_for = MagicMock(return_value="public/data/k") + with patch( + "pysus.api.extensions.ExtensionFactory", + MagicMock(instantiate=AsyncMock(return_value=local_file)), + ): + with patch( + "pysus.management.sync.upload_s3", new=AsyncMock() + ) as mock_upload: + with patch( + "anyio.to_thread.run_sync", + new=AsyncMock( + side_effect=lambda fn, *a, **kw: fn(*a, **kw) + ), + ): + payload = await engine._convert_and_upload(MagicMock(), raw) + assert payload["s3_key"] == "public/data/k" + assert payload["rows"] == 10 + assert payload["schema"] == "schema" + assert len(payload["raw_digest"]) == 64 + mock_upload.assert_awaited_once() + assert not raw.exists() + assert not parquet.exists() + + @pytest.mark.asyncio + async def test_convert_and_upload_non_tabular(self, engine, tmp_path): + raw = tmp_path / "X.xyz" + raw.write_bytes(b"raw") + + class NotTabular: + pass + + fake = MagicMock() + fake.basename = "X.xyz" + engine.s3_key_for = MagicMock(return_value="public/data/k") + with patch( + "pysus.api.extensions.ExtensionFactory", + MagicMock(instantiate=AsyncMock(return_value=NotTabular())), + ): + with pytest.raises(RuntimeError, match="cannot convert"): + await engine._convert_and_upload(fake, raw) + + +class TestCatalogRows: + def test_catalog_rows(self, engine): + writer = MagicMock() + writer.ensure_dataset.return_value = 8 + writer.ensure_group.return_value = 3 + writer.get_file.return_value = (99, None) + + central_cursor = MagicMock() + dataset_cursor = MagicMock() + columns_cursor = MagicMock() + group = MagicMock() + group.name = "DENG" + file = MagicMock() + file.dataset.name = "SINAN" + file.dataset.long_name = "Sistema..." + file.dataset.description = "desc" + file.group = group + file.path = "/ftp/x" + file.year = 2025 + file.month = None + file.state = None + file.client.name = "ftp" + + payload = { + "s3_key": "public/data/ftp/sinan/DENG/2025/_/BR/X.parquet", + "size": 100, + "rows": 5, + "schema": MagicMock(), + "raw_digest": "a" * 64, + "parquet_digest": "b" * 64, + } + + with patch.object( + SyncEngine, "writer", new_callable=PropertyMock + ) as mock_writer_prop: + mock_writer_prop.return_value = writer + engine._catalog_rows( + central_cursor, dataset_cursor, columns_cursor, file, payload + ) + + writer.ensure_dataset.assert_called_once() + writer.upsert_file.assert_called_once() + writer.link_columns.assert_called_once() + kwargs = writer.upsert_file.call_args.kwargs + assert kwargs["path"] == payload["s3_key"] + assert kwargs["sha256"] == payload["parquet_digest"] + assert kwargs["source_sha256"] == payload["raw_digest"] + assert kwargs["origin"] == "ftp" + + +class TestCatalogWriteEntry: + def test_catalog_write_entry(self, engine): + adapter = MagicMock() + central = MagicMock() + columns = MagicMock() + writer = MagicMock() + writer.get_file.return_value = (42, True) + file = MagicMock() + payload = { + "s3_key": "k", + "size": 10, + "rows": 1, + "schema": MagicMock(), + "raw_digest": "a", + "parquet_digest": "b", + } + + # nested transactions: use real context managers via MagicMock + class _Ctx: + def __init__(self, cursor): + self.cursor = cursor + + def __enter__(self): + return MagicMock(), self.cursor + + def __exit__(self, *a): + return False + + adapter.transaction = MagicMock(return_value=_Ctx(MagicMock())) + central.transaction = MagicMock(return_value=_Ctx(MagicMock())) + columns.transaction = MagicMock(return_value=_Ctx(MagicMock())) + + with patch.object( + SyncEngine, "writer", new_callable=PropertyMock + ) as mock_writer_prop: + mock_writer_prop.return_value = writer + engine._catalog_write_entry( + adapter, central, columns, file, payload + ) + + adapter.transaction.assert_called_once() + writer._ensure_management_columns.assert_called_once() From 58eef1e1606fe3f132613a6db9153357b49d8afd Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Lu=C3=A3=20Bida=20Vacaro?= Date: Fri, 14 Aug 2026 17:35:16 -0300 Subject: [PATCH 4/5] fix(sync): cross-platform PID liveness check for the catalog lock Windows' os.kill raises OSError (WinError 87) for invalid PIDs instead of ProcessLookupError, breaking lock stealing. Use OpenProcess on Windows and treat OSError as not-alive elsewhere. --- pysus/management/sync.py | 18 ++++++++++++++++++ 1 file changed, 18 insertions(+) diff --git a/pysus/management/sync.py b/pysus/management/sync.py index e2fd5547..9972646e 100644 --- a/pysus/management/sync.py +++ b/pysus/management/sync.py @@ -148,12 +148,30 @@ def _acquire_sync_lock(self) -> None: def _pid_alive(pid: int) -> bool: import os + if os.name == "nt": + try: + import ctypes + + process_query_limited = 0x1000 + handle = ctypes.windll.kernel32.OpenProcess( + process_query_limited, False, pid + ) + if not handle: + return False + ctypes.windll.kernel32.CloseHandle(handle) + return True + except Exception: # noqa + return False + try: os.kill(pid, 0) except ProcessLookupError: return False except PermissionError: return True + except OSError: + # invalid/out-of-range pid on some platforms + return False return True def _release_sync_lock(self) -> None: From affde5cc65b101130bdbaa01220d5ac0835f97bf Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Lu=C3=A3=20Bida=20Vacaro?= Date: Fri, 14 Aug 2026 17:49:19 -0300 Subject: [PATCH 5/5] fix(sync): mypy-safe windll access in PID liveness check --- pysus/management/sync.py | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/pysus/management/sync.py b/pysus/management/sync.py index 9972646e..72a396c8 100644 --- a/pysus/management/sync.py +++ b/pysus/management/sync.py @@ -152,13 +152,16 @@ def _pid_alive(pid: int) -> bool: try: import ctypes + windll = getattr(ctypes, "windll", None) + if windll is None: + return False process_query_limited = 0x1000 - handle = ctypes.windll.kernel32.OpenProcess( + handle = windll.kernel32.OpenProcess( process_query_limited, False, pid ) if not handle: return False - ctypes.windll.kernel32.CloseHandle(handle) + windll.kernel32.CloseHandle(handle) return True except Exception: # noqa return False