chore: make ducklake client be a duckdataset manager

Author: luabida

pysus/api/README.md

@@ -1,2 +1 @@
 ## Roadmap
-

pysus/api/_impl/databases.py

@@ -24,8 +24,8 @@
 from typing import Literal
 
 import pandas as pd
-from pysus.api.client import PySUS
 from pysus.api import types
+from pysus.api.client import PySUS
 from tqdm import tqdm
 
 

pysus/api/client.py

@@ -11,12 +11,11 @@
 from typing import TYPE_CHECKING, Literal
 
 import anyio
-
-from pysus.api.types import Origin
 import duckdb
-from duckdb import func
 import pandas as pd
+from duckdb import func
 from pysus import CACHEPATH
+from pysus.api.types import Origin
 from sqlalchemy import DateTime, Enum, Integer, String, create_engine
 from sqlalchemy.orm import DeclarativeBase, Mapped, mapped_column, sessionmaker
 from sqlalchemy.pool import NullPool
@@ -518,7 +517,9 @@ async def query(
         all_datasets = await self._ducklake.datasets()
 
         if dataset:
-            matching = [d for d in all_datasets if d.name.lower() == dataset.lower()]
+            matching = [
+                d for d in all_datasets if d.name.lower() == dataset.lower()
+            ]
             if not matching:
                 return []
             target = matching[0]
@@ -617,7 +618,9 @@ def get_columns(path: Path) -> set[tuple[str, str]]:
 
         else:
             paths_str = ", ".join(f"'{p}'" for p in paths)
-            query = f"SELECT * FROM read_parquet([{paths_str}], union_by_name=True)"
+            query = (
+                f"SELECT * FROM read_parquet([{paths_str}], union_by_name=True)"
+            )
 
         if sql:
             if sql.upper().startswith("SELECT"):
@@ -630,7 +633,9 @@ def get_columns(path: Path) -> set[tuple[str, str]]:
         if not add_dv:
             return base
 
-        geocode_cols = [col[0] for col in base.description if is_geocode_column(col[0])]
+        geocode_cols = [
+            col[0] for col in base.description if is_geocode_column(col[0])
+        ]
         if not geocode_cols:
             return base
 

pysus/api/dadosgov/client.py

@@ -277,7 +277,11 @@ async def _download_file(
                 "Client not connected. Call login(token=...) first.",
             )
 
-        url = str(file.path).replace("https:/", "https://").replace("http:/", "http://")
+        url = (
+            str(file.path)
+            .replace("https:/", "https://")
+            .replace("http:/", "http://")
+        )
 
         async with self._client.stream("GET", url) as response:
             response.raise_for_status()

pysus/api/dadosgov/models.py

@@ -30,9 +30,13 @@ def _dedup_entries(
         if m:
             stem = filename[: m.start()]
             fmt = m.group(1).lower()
-            grouped.setdefault(stem, []).append((fmt, filename, recurso, metadata))
+            grouped.setdefault(stem, []).append(
+                (fmt, filename, recurso, metadata)
+            )
         else:
-            grouped.setdefault(filename, []).append(("", filename, recurso, metadata))
+            grouped.setdefault(filename, []).append(
+                ("", filename, recurso, metadata)
+            )
 
     result: list[tuple[str, Any, dict]] = []
     for _, items in grouped.items():
@@ -245,7 +249,9 @@ class Group(BaseRemoteGroup):
     """A group of files within a dataset."""
 
     record: ConjuntoDados
-    _formatter: Callable[[str], dict[str, Any]] | None = PrivateAttr(default=None)
+    _formatter: Callable[[str], dict[str, Any]] | None = PrivateAttr(
+        default=None
+    )
 
     def __init__(
         self,
@@ -313,7 +319,9 @@ async def _fetch_files(self) -> list[BaseRemoteFile]:
         """Build File objects from the underlying resources."""
         entries: list[tuple[str, Any, dict]] = []
         for recurso in self.record.resources:
-            filename = recurso.file_name or recurso.url.split("/")[-1].split("?")[0]
+            filename = (
+                recurso.file_name or recurso.url.split("/")[-1].split("?")[0]
+            )
             if filename.lower().endswith(".pdf") or filename.startswith("get_"):
                 continue
             metadata = {}

pysus/api/ducklake/catalog/columns.py

@@ -6388,7 +6388,7 @@
 
 TIPPRE = {"sia": ""}
 
-TIPPRE   = {"sia": ""}
+TIPPRE = {"sia": ""}
 
 TIP_DIARRE = {"sinan": ""}
 

pysus/api/ducklake/catalog/parsers.py

@@ -1 +0,0 @@
-

pysus/api/ducklake/client.py

@@ -319,7 +319,9 @@ async def _download(
                 else:
                     raise e
 
-    async def _download_catalog(self, local_path: Path, remote_path: str) -> None:
+    async def _download_catalog(
+        self, local_path: Path, remote_path: str
+    ) -> None:
         """Download a catalog database from remote storage with retries.
 
         Parameters
@@ -373,7 +375,9 @@ def _get_s3_client(self):
             "s3",
             endpoint_url=f"https://{self.endpoint}",
             aws_access_key_id=self.credentials.access_key.get_secret_value(),
-            aws_secret_access_key=(self.credentials.secret_key.get_secret_value()),
+            aws_secret_access_key=(
+                self.credentials.secret_key.get_secret_value()
+            ),
             region_name=self.region,
             config=Config(signature_version="s3v4"),
         )

pysus/api/ducklake/models.py

@@ -8,16 +8,12 @@
 from collections.abc import Callable
 from datetime import datetime
 from pathlib import Path
-from typing import TYPE_CHECKING, Any, Union, Optional
+from typing import TYPE_CHECKING, Any, Optional, Union
 
 from anyio import to_thread
 from pydantic import Field, PrivateAttr
 from pysus import CACHEPATH
-from pysus.api.models import (
-    BaseRemoteDataset,
-    BaseRemoteFile,
-    BaseRemoteGroup,
-)
+from pysus.api.models import BaseRemoteDataset, BaseRemoteFile, BaseRemoteGroup
 from sqlalchemy.orm import contains_eager, joinedload, sessionmaker
 
 from .catalog import CatalogDataset, CatalogFile, DatasetGroup
@@ -266,7 +262,9 @@ async def connect(
             return
 
         await self.client._download(
-            f"public/{self._catalog_name}", self._catalog_local, callback=callback
+            f"public/{self._catalog_name}",
+            self._catalog_local,
+            callback=callback,
         )
         self._engine = await to_thread.run_sync(
             lambda: self.client._setup_engine(self._catalog_local)

pysus/api/extensions.py

@@ -252,7 +252,9 @@ def columns(self) -> list["Column"]:
 
         schema = pq.read_schema(self.path)
         return [
-            Column.from_schema(name=field.name, dtype=_map_dtype(str(field.type)))
+            Column.from_schema(
+                name=field.name, dtype=_map_dtype(str(field.type))
+            )
             for field in schema
         ]
 
@@ -356,7 +358,9 @@ def columns(self) -> list["Column"]:
             "M": "VARCHAR",
         }
         return [
-            Column.from_schema(name=f.name, dtype=_DBF_DTYPE.get(f.type, "VARCHAR"))
+            Column.from_schema(
+                name=f.name, dtype=_DBF_DTYPE.get(f.type, "VARCHAR")
+            )
             for f in reader.fields
         ]
 
@@ -840,7 +844,9 @@ class DBCNotImported(BaseTabularFile):
 
     path: Path = Field(default_factory=lambda: Path("..."))
     type: str | FileType = Field(default="remote")
-    import_err: ClassVar[str] = """
+    import_err: ClassVar[
+        str
+    ] = """
         run "pip install pysus[dbc]" to handle DBC files.
         Make sure you also have libffi installed on the system. It may not work
         on Windows