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 62edc56e..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 @@ -36,15 +57,102 @@ 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: 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( @@ -73,12 +181,19 @@ async def connect( return if force: + _dispose_shared(self.db_local) await self._download_catalog( self.db_local, str(self.db_remote), 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 @@ -93,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) @@ -111,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, ) @@ -119,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( @@ -221,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), @@ -229,13 +336,18 @@ 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() - - if self._engine: - await to_thread.run_sync(self._engine.dispose) - self._engine = None - self._session_factory = None + self._local_dirty = False + + # 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/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..904e3f93 --- /dev/null +++ b/pysus/management/catalog.py @@ -0,0 +1,416 @@ +"""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") + self._ensure_column( + catalog_cursor, "files", "source_sha256", "VARCHAR(64)" + ) + + # ------------------------------------------------------------------ + # 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 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,) + ) + 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, + source_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 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) + 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 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) + 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..36114ba3 --- /dev/null +++ b/pysus/management/scripts/relayout_bucket.py @@ -0,0 +1,183 @@ +"""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, + object_sizes=object_sizes, + workers=48, + ) + ) + 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..cde2d6b4 --- /dev/null +++ b/pysus/management/scripts/sync_clients.py @@ -0,0 +1,114 @@ +"""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, workers, ftp_connections +) -> 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, + workers=workers, + ftp_connections=ftp_connections, + ) + 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", + ) + 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, + args.workers, + args.ftp_connections, + ) + ) + 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..72a396c8 --- /dev/null +++ b/pysus/management/sync.py @@ -0,0 +1,1073 @@ +"""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 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.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, + BrokenPipeError, + OSError, +) + + +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, + ) + 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 + + if os.name == "nt": + try: + import ctypes + + windll = getattr(ctypes, "windll", None) + if windll is None: + return False + process_query_limited = 0x1000 + handle = windll.kernel32.OpenProcess( + process_query_limited, False, pid + ) + if not handle: + return False + 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: + 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) + + # ------------------------------------------------------------------ + # 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, + ftp_client: Any | None = None, + ) -> 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``. + + ``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") + + s3_key = self.s3_key_for(file) + writer = self.writer + + dataset_adapter = self._dataset_adapter(file) + 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.raw_connection() + dataset_conn = dataset_adapter.raw_connection() + columns_conn = columns_adapter.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_full(dataset_cursor, s3_key) + if existing and not force: + _, origin_modified, _, _, _ = existing + if self._is_current(file, origin_modified): + return False + + raw_path = await self._download_raw_with_retry( + file, ftp_client=ftp_client + ) + raw_digest = sha256_of(raw_path) + + 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 + + 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, + ) + + 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, + payload, + ) + + central_conn.commit() + dataset_conn.commit() + columns_conn.commit() + dataset_cursor.execute("CHECKPOINT") + columns_cursor.execute("CHECKPOINT") + + 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 + for conn in connections: + try: + conn.close() + except Exception: # noqa + 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*.""" + return self._dataset_adapter_by_name(file.dataset.name) + + async def _download_raw_with_retry( + self, + file: FTPFile | APIFile, + 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 + for attempt in range(max_retries): + try: + await self._download_once(file, output, ftp_client) + return output + except _RETRYABLE as exc: + last_error = exc + 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..." + ) + await asyncio.sleep(wait_time) + + raise RuntimeError( + f"Failed to download {file.basename} after {max_retries} " + 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, + 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, + 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 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. ``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 collect_with_retry("ducklake", datasets), + "ftp": await collect_with_retry("ftp", datasets), + } + records["dadosgov"] = [] + if self.dadosgov_token: + records["dadosgov"] = await collect_with_retry( + "dadosgov", datasets, dadosgov_token=self.dadosgov_token + ) + + comparisons = self.comparator.compare( + records["ducklake"] + records["ftp"] + records["dadosgov"] + ) + + # 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: + if comparison.is_on_s3 and not ( + force or self._s3_is_stale(comparison) + ): + 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() + + if save_snapshots: + for origin, items in records.items(): + self.inventory.save_snapshot(origin, items) + + 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() + await ducklake.flush_catalogs(update=True) + 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 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 + + # 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 + if record.size and record.size != s3_size: + 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/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_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/api/ducklake/test_client.py b/pysus/tests/api/ducklake/test_client.py index 5976cc80..fb6a8aa6 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() @@ -718,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 @@ -728,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/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_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_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_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_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_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.py b/pysus/tests/management/test_sync.py new file mode 100644 index 00000000..4378c984 --- /dev/null +++ b/pysus/tests/management/test_sync.py @@ -0,0 +1,196 @@ +"""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_source_size_differs(self): + from pysus.management.records import FileComparison, FileRecord + + ftp = FileRecord( + origin="ftp", + dataset="SINAN", + name="DENGBR25.dbc", + path="ftp/x", + size=200, + modified=None, + 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 SyncEngine._s3_is_stale(comparison) + + 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( + 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=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, + ) + 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) 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()