create a PySUS orchestrator that will be a bridge between remote clients and local files

Author: luabida

.pre-commit-config.yaml

@@ -34,6 +34,19 @@ repos:
             'pydocstyle>=6.3.0',
         ]
 
+  - repo: https://github.com/pre-commit/mirrors-mypy
+    rev: v1.8.0
+    hooks:
+      - id: mypy
+        additional_dependencies: [
+            'types-python-dateutil',
+            'types-requests',
+            'types-setuptools',
+            'pandas-stubs',
+            'pydantic>=2.0.0',
+        ]
+        args: [--ignore-missing-imports, --explicit-package-bases]
+
   - repo: https://github.com/asottile/pyupgrade
     rev: v3.15.0
     hooks:

pysus/__init__.py

@@ -5,9 +5,11 @@
 from importlib import metadata as importlib_metadata
 from typing import Final
 
-CACHEPATH: Final[str] = os.getenv(
-    "PYSUS_CACHEPATH",
-    os.path.join(str(pathlib.Path.home()), "pysus"),
+CACHEPATH: Final[pathlib.Path] = pathlib.Path(
+    os.getenv(
+        "PYSUS_CACHEPATH",
+        os.path.join(str(pathlib.Path.home()), "pysus"),
+    )
 )
 
 # from pysus.api.ftp.databases import *  # noqa

pysus/api/client.py

@@ -0,0 +1,225 @@
+import enum
+from collections.abc import Callable
+from datetime import datetime
+from pathlib import Path
+
+from pysus import CACHEPATH
+from sqlalchemy import Column, Integer, DateTime, Enum, String, create_engine
+from sqlalchemy.orm import declarative_base, sessionmaker
+
+from .dadosgov import DadosGovClient
+from .ducklake import DuckLakeClient
+from .ftp import FTPClient
+from .models import BaseLocalFile, BaseRemoteFile
+
+Base = declarative_base()
+
+
+class DownloadStatus(enum.Enum):
+    PENDING = "pending"
+    DOWNLOADING = "downloading"
+    COMPLETED = "completed"
+    FAILED = "failed"
+    MISSING = "missing"
+
+
+class LocalFileState(Base):
+    __tablename__ = "local_file_state"
+    path = Column(String, primary_key=True)
+    remote_path = Column(String, nullable=False)
+    client_name = Column(String, nullable=False)
+    year = Column(Integer, nullable=True)
+    month = Column(Integer, nullable=True)
+    state = Column(String, nullable=True)
+    status = Column(Enum(DownloadStatus), default=DownloadStatus.PENDING)
+    sha256 = Column(String, nullable=True)
+    last_synced = Column(DateTime, default=datetime.utcnow)
+
+
+class PySUS:
+    def __init__(self, db_path: str = CACHEPATH / "config.db"):
+        db_path = Path(db_path)
+        db_path.parent.mkdir(parents=True, exist_ok=True)
+
+        self.engine = create_engine(f"duckdb:///{db_path}")
+        Base.metadata.create_all(self.engine)
+        self.Session = sessionmaker(bind=self.engine)
+
+        self._ducklake: DuckLakeClient | None = None
+        self._ftp: FTPClient | None = None
+        self._dadosgov: DadosGovClient | None = None
+
+    async def __aenter__(self):
+        self._ducklake = DuckLakeClient(engine=self.engine)
+        await self._ducklake._load_catalog()
+        self._attach_client_catalog("ducklake", self._ducklake.catalog_path)
+        return self
+
+    async def get_dadosgov(self, access_token: str) -> DadosGovClient:
+        if self._dadosgov is None:
+            self._dadosgov = DadosGovClient()
+            await self._dadosgov.connect(token=access_token)
+        return self._dadosgov
+
+    async def get_ftp(self) -> FTPClient:
+        if self._ftp is None:
+            self._ftp = FTPClient()
+            await self._ftp.connect()
+        return self._ftp
+
+    async def get_local_file(
+        self,
+        file: BaseRemoteFile,
+    ) -> BaseLocalFile | None:
+        from pysus.api.extensions import ExtensionFactory
+
+        client_name = file.client.name.lower()
+        remote_path = file.path
+
+        with self.Session() as session:
+            records = (
+                session.query(LocalFileState)
+                .filter_by(
+                    remote_path=remote_path,
+                    client_name=client_name,
+                    status=DownloadStatus.COMPLETED,
+                )
+                .all()
+            )
+
+            if not records:
+                return None
+
+            parquet_version = next(
+                (r for r in records if r.path.endswith(".parquet")), None
+            )
+            file = parquet_version or records[0]
+
+            return await ExtensionFactory.instantiate(file.path)
+
+    def _attach_client_catalog(self, name: str, path: str):
+        abs_path = str(Path(path).absolute())
+        with self.engine.connect() as conn:
+            q = "SELECT database_name FROM duckdb_databases() WHERE path = ?"
+            existing = conn.exec_driver_sql(q, (abs_path,)).fetchone()
+
+            if not existing:
+                conn.exec_driver_sql(f"ATTACH '{abs_path}' AS {
+                                     name} (READ_ONLY)")
+
+    async def __aexit__(self, exc_type, exc_val, exc_tb):
+        if self._ducklake:
+            await self._ducklake.close()
+        if self._ftp:
+            await self._ftp.close()
+        if self._dadosgov:
+            await self._dadosgov.close()
+        self.engine.dispose()
+
+    def _get_dest_path(self, client_name: str, remote_path: str) -> Path:
+        return CACHEPATH / "downloads" / client_name / remote_path.lstrip("/")
+
+    async def _update_state(
+        self,
+        local_path: Path,
+        remote_path: str,
+        client_name: str,
+        status: DownloadStatus,
+        year: int = None,
+        month: int = None,
+        state: str = None,
+    ):
+        with self.Session() as session:
+            record = (
+                session.query(LocalFileState).filter_by(
+                    path=str(local_path)).first()
+            )
+            if not record:
+                record = LocalFileState(
+                    path=str(local_path),
+                    remote_path=remote_path,
+                    client_name=client_name,
+                    year=year,
+                    month=month,
+                    state=state,
+                )
+                session.add(record)
+
+            record.status = status
+            record.last_synced = datetime.utcnow()
+            session.commit()
+
+    async def download(
+        self,
+        file: BaseRemoteFile,
+        token: str = None,
+        callback: Callable = None,
+    ):
+        from pysus.api.extensions import ExtensionFactory
+
+        existing_local = await self.get_local_file(file)
+        if existing_local and existing_local.path.exists():
+            return existing_local
+
+        client_name = file.client.name.lower()
+        remote_path = file.path
+        local_path = self._get_dest_path(client_name, remote_path)
+
+        local_path.parent.mkdir(parents=True, exist_ok=True)
+
+        await self._update_state(
+            local_path, remote_path, client_name, DownloadStatus.DOWNLOADING
+        )
+
+        try:
+            if client_name == "ducklake":
+                await self._ducklake._download_file(file, local_path, callback)
+            elif client_name == "ftp":
+                client = await self.get_ftp()
+                await client._download_file(file, local_path, callback)
+            elif client_name == "dadosgov":
+                client = await self.get_dadosgov(token)
+                await client._download_file(file, local_path, callback)
+            else:
+                raise ValueError(f"No download logic for client: {client_name}")
+
+            await self._update_state(
+                local_path=local_path,
+                remote_path=remote_path,
+                client_name=client_name,
+                status=DownloadStatus.DOWNLOADING,
+                year=file.year,
+                month=file.month,
+                state=file.state,
+            )
+            return await ExtensionFactory.instantiate(local_path)
+
+        except Exception:
+            await self._update_state(
+                local_path, remote_path, client_name, DownloadStatus.FAILED
+            )
+            raise
+
+    async def download_to_parquet(
+        self,
+        file: BaseRemoteFile,
+        token: str = None,
+        callback: Callable = None,
+    ):
+        local_file = await self.download(
+            file=file,
+            token=token,
+            callback=callback,
+        )
+
+        if hasattr(local_file, "to_parquet"):
+            parquet_file = await local_file.to_parquet()
+
+            await self._update_state(
+                local_path=parquet_file.path,
+                remote_path=file.path,
+                client_name=file.client.name.lower(),
+                status=DownloadStatus.COMPLETED,
+            )
+            return parquet_file
+        return local_file

pysus/api/dadosgov/__init__.py

@@ -0,0 +1 @@
+from .client import DadosGov as DadosGovClient  # noqa

pysus/api/dadosgov/client.py

@@ -76,7 +76,7 @@ async def connect(self, token: str | None = None) -> None:
         self._client = httpx.AsyncClient(
             base_url=self.base_url,
             headers=headers,
-            timeout=60.0,
+            timeout=30.0,
             follow_redirects=True,
         )
 

pysus/api/dadosgov/databases.py

@@ -1,5 +1,3 @@
-from typing import List
-
 from .models import Dataset
 
 

pysus/api/dadosgov/models.py

@@ -1,11 +1,12 @@
 import asyncio
 import pathlib
-from abc import ABC
+from abc import ABC, abstractmethod
 from collections.abc import Callable
 from datetime import datetime as dt
-from typing import Any, List, Optional, Union
+from typing import Any
 
 import httpx
+from pydantic import PrivateAttr
 from pysus.api.models import (
     BaseRemoteClient,
     BaseRemoteDataset,
@@ -19,6 +20,12 @@
 class File(BaseRemoteFile):
     record: Recurso
     type: str | None = "remote"
+    _metadata: dict[str, Any] = PrivateAttr(default_factory=dict)
+
+    def __init__(self, **data):
+        metadata = data.pop("_metadata", {})
+        super().__init__(**data)
+        self._metadata = metadata
 
     def __repr__(self):
         return self.basename
@@ -51,6 +58,18 @@ def size(self) -> int:
     def modify(self) -> dt:
         return self.record.last_modified
 
+    @property
+    def year(self) -> int | None:
+        return self._metadata.get("year")
+
+    @property
+    def month(self) -> int | None:
+        return self._metadata.get("month")
+
+    @property
+    def state(self) -> str | None:
+        return self._metadata.get("state")
+
     async def _download(
         self,
         output: pathlib.Path | None = None,
@@ -64,7 +83,7 @@ async def fetch_size(self) -> int:
         try:
             async with httpx.AsyncClient(
                 follow_redirects=True,
-                timeout=1,
+                timeout=3,
             ) as client:
                 response = await client.head(self.path)
 
@@ -85,6 +104,19 @@ async def fetch_size(self) -> int:
 
 class Group(BaseRemoteGroup):
     record: ConjuntoDados
+    _formatter: Callable[[Recurso, "Group"], dict[str, Any]] | None = (
+        PrivateAttr(default=None)
+    )
+
+    def __init__(
+        self,
+        record: ConjuntoDados,
+        dataset: BaseRemoteDataset,
+        formatter: Callable | None = None,
+    ):
+        super().__init__(dataset=dataset)
+        self.record = record
+        self._formatter = formatter
 
     def __repr__(self):
         return self.name
@@ -99,16 +131,15 @@ def long_name(self) -> str:
 
     @property
     def description(self) -> str:
-        return ""  # TODO:
+        return ""
 
     async def _fetch_files(self) -> list[File]:
         files = []
         for recurso in self.record.resources:
-            file = File(
-                record=recurso,
-                parent=self,
+            metadata = (
+                self._formatter(recurso, self) if self._formatter else {}
             )
-            await file.fetch_size()
+            file = File(record=recurso, parent=self, _metadata=metadata)
             files.append(file)
         return files
 
@@ -119,13 +150,20 @@ class Dataset(BaseRemoteDataset, ABC):
     def __repr__(self):
         return self.name
 
-    async def _fetch_content(
-        self,
-    ) -> list[Group]:
+    @property
+    @abstractmethod
+    def formatter(self) -> Callable[[Recurso, Group], dict[str, Any]]:
+        pass
+
+    async def _fetch_content(self) -> list[Group]:
         items: list[Group] = []
         client: BaseRemoteClient = self.client
         if self.ids:
             for group_id in self.ids:
                 record = await client.get_dataset(group_id)
-                items.append(Group(record=record, dataset=self))
+                items.append(
+                    Group(
+                        record=record, dataset=self, formatter=self.formatter
+                    )
+                )
         return items