chore: implement context managers to every duckdb interaction

Author: luabida

pysus/api/ducklake/README.md

@@ -0,0 +1,37 @@
+# DuckLake Catalog Component
+
+This module provides the application-level models and data adapters required to interface with remote DuckLake resources over S3. It wraps low-level database operations into unified interfaces (`File`, `DuckDataset`, and `DuckGroup`) and implements resilient database connection management via async context managers.
+
+## Features
+
+* **Deterministic Resource Management**: Implements asynchronous context managers (`async with`) across clients and database adapters to prevent DuckDB file locks and connection leaks.
+* **Fail-Safe Cleanups**: Features fallback `__del__` destructors to safely terminate remaining active engines during garbage collection or interpreter shutdown.
+* **Intelligent Syncing**: Employs an `update_on_close` mechanism to optionally push state changes back to S3 automatically upon exiting a context.
+* **SQLAlchemy Eager Loading Fixes**: Optimizes attribute mappings using isolated path strategies (`joinedload` vs. `contains_eager`) based on query parameters.
+* **Resilient S3 Verifications**: Gracefully intercepts 404 responses during S3 download handshakes to stop failing connection retry loops early.
+
+---
+
+## Architecture Overview
+
+The system separates the raw database management layer (Adapters) from the client wrapper layer (Client Models). 
+
+1. **Adapters (`BaseAdapter`)**: Track local and remote `.duckdb` target states, manage connections, handle S3 transfers, and expose scoped SQLAlchemy transaction sessions.
+2. **Client Components (`DuckLake`)**: Coordinate high-level actions, parse credential models, route queries, and handle collection loops.
+
+---
+
+## Lifecycle & Connection Handling
+
+### Using Context Managers (Recommended)
+
+Using `async with` blocks guarantees deterministic resource teardown. The moment execution exits the context layout block—even due to a runtime crash—all engines are disposed of cleanly.
+
+```python
+from pysus.api.ducklake.client import DuckLake
+
+async with DuckLake() as dl:
+    datasets = await dl.datasets()
+
+sia = datasets[4] # e.g., SIA
+files = await sia.query(state="SP", year=2026)

pysus/api/ducklake/__init__.py

@@ -1,7 +0,0 @@
-"""DuckLake subpackage for interacting with the PySUS S3 catalog.
-
-Provides a DuckDB-based client for querying and downloading
-public health datasets stored in object storage.
-"""
-
-from .client import DuckLake as DuckLakeClient  # noqa

pysus/api/ducklake/catalog/adapters.py

@@ -1,5 +1,6 @@
 from abc import ABC
 from pathlib import Path
+import asyncio
 
 import httpx
 from anyio import to_thread
@@ -12,6 +13,7 @@
 from pysus import CACHEPATH
 from pysus.api import types
 from pysus.api.ducklake.functional import download_s3, upload_s3
+from pysus.api.ducklake.catalog.orm.dataset import DatasetBase
 
 
 class DuckLakeCredentials(BaseModel):
@@ -25,12 +27,17 @@ class BaseAdapter(ABC):
     db_remote: Path
 
     def __init__(
-        self, engine=None, credentials: DuckLakeCredentials | None = None, **data
+        self,
+        engine=None,
+        credentials: DuckLakeCredentials | None = None,
+        update_on_close: bool = False,
+        **data,
     ) -> None:
         self._engine = engine
         self._session_factory = None
         self.cache_dir.mkdir(parents=True, exist_ok=True)
         self.credentials = credentials
+        self.update_on_close = update_on_close
 
     @property
     def remote_url(self) -> str:
@@ -64,6 +71,7 @@ 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(
                 "SELECT 1 FROM information_schema.schemata WHERE schema_name = 'pysus'"
@@ -90,6 +98,7 @@ def setup_engine(
 
             conn.commit()
 
+        DatasetBase.metadata.create_all(bind=engine)
         return engine
 
     async def _download_catalog(self, local_path: Path, remote_path: str) -> None:
@@ -106,8 +115,14 @@ async def _download_catalog(self, local_path: Path, remote_path: str) -> None:
         async with httpx.AsyncClient(follow_redirects=True) as client:
             try:
                 head = await client.head(url)
+
+                if head.status_code == 404:
+                    return
+
                 head.raise_for_status()
                 remote_size = int(head.headers.get("content-length", 0))
+            except httpx.HTTPStatusError:
+                return
             except Exception:
                 remote_size = 0
 
@@ -153,6 +168,28 @@ async def close(self, update: bool = False) -> None:
             self._engine = None
             self._session_factory = None
 
+    def __del__(self) -> None:
+        if not hasattr(self, "_engine") or not self._engine:
+            return
+        try:
+            loop = asyncio.get_running_loop()
+            if loop.is_running():
+                loop.create_task(self.close(update=False))
+        except RuntimeError:
+            try:
+                asyncio.run(self.close(update=False))
+            except Exception:  # noqa
+                pass
+        except Exception:  # noqa
+            pass
+
+    async def __aenter__(self):
+        await self.connect()
+        return self
+
+    async def __aexit__(self, exc_type, exc_val, exc_tb) -> None:
+        await self.close(update=self.update_on_close)
+
 
 class CatalogAdapter(BaseAdapter):
     def __init__(self, engine=None, **data) -> None:
@@ -174,3 +211,4 @@ def __init__(self, engine=None, **data) -> None:
         super().__init__(engine=engine, **data)
         self.db_local: Path = self.cache_dir / "columns.duckdb"
         self.db_remote: str = "public/columns.duckdb"
+

pysus/api/ducklake/client.py

@@ -4,11 +4,12 @@
 capabilities backed by per-dataset DuckDB engines.
 """
 
+import asyncio
 from collections.abc import Callable
 from pathlib import Path
 
 from anyio import to_thread
-from pydantic import SecretStr, PrivateAttr
+from pydantic import SecretStr, PrivateAttr, Field
 from pysus.api.models import BaseRemoteClient
 from pysus.api.types import DUCKLAKE
 
@@ -21,13 +22,17 @@
 
 class DuckLake(BaseRemoteClient):
     credentials: DuckLakeCredentials | None = None
+    update_on_close: bool = Field(default=False, exclude=True)
     _datasets: list[DuckDataset] = PrivateAttr(default_factory=list)
+    _catalog_adap: CatalogAdapter = PrivateAttr()
 
-    def __init__(self, engine=None, **data) -> None:
+    def __init__(self, engine=None, update_on_close: bool = False, **data) -> None:
         super().__init__(**data)
-        self.catalog_adap = CatalogAdapter(
+        self.update_on_close = update_on_close
+        self._catalog_adap = CatalogAdapter(
             engine=engine,
             credentials=self.credentials,
+            update_on_close=self.update_on_close,
         )
 
     @property
@@ -43,24 +48,31 @@ def description(self) -> str:
         return ""
 
     async def datasets(self, **kwargs) -> list[DuckDataset]:
-        await self.catalog_adap.connect()
-
         def _fetch():
-            with self.catalog_adap.get_session() as session:
+            with self._catalog_adap.get_session() as session:
                 results = session.query(Dataset).all()
                 session.expunge_all()
                 return results
 
-        records = await to_thread.run_sync(_fetch)
-
         duck_datasets: list[DuckDataset] = []
-        for rec in records:
-            dataset_adapter = DatasetAdapter(
-                name=str(rec.name), credentials=self.credentials
-            )
-            duck_datasets.append(
-                DuckDataset(record=rec, client=self, adapter=dataset_adapter)
-            )
+
+        async with self._catalog_adap:
+            records = await to_thread.run_sync(_fetch)
+
+            for rec in records:
+                dataset_adapter = DatasetAdapter(
+                    name=str(rec.name),
+                    credentials=self.credentials,
+                    update_on_close=self.update_on_close,
+                )
+                duck_datasets.append(
+                    DuckDataset(
+                        record=rec,
+                        client=self,
+                        adapter=dataset_adapter,
+                        update_on_close=self.update_on_close,
+                    )
+                )
 
         self._datasets = duck_datasets
         return duck_datasets
@@ -75,18 +87,21 @@ async def login(
             access_key=SecretStr(access_key),
             secret_key=SecretStr(secret_key),
         )
-        self.catalog_adap.credentials = self.credentials
-        await self.catalog_adap.connect(force=True)
+        self._catalog_adap.credentials = self.credentials
+        await self._catalog_adap.connect(force=True)
 
     async def connect(self, force: bool = False) -> None:
-        await self.catalog_adap.connect(force=force)
+        await self._catalog_adap.connect(force=force)
+
+    async def close(self, update_catalog: bool | None = None) -> None:
+        should_update = (
+            self.update_on_close if update_catalog is None else update_catalog
+        )
 
-    async def close(self, update_catalog: bool = False) -> None:
         for ds in self._datasets:
-            await ds.close(update_catalog=update_catalog)
+            await ds.close(update_catalog=should_update)
 
-        await self.catalog_adap.close(update=update_catalog)
-        self._datasets.clear()
+        await self._catalog_adap.close(update=should_update)
 
     async def download(
         self,
@@ -95,7 +110,7 @@ async def download(
         callback: Callable[[int, int], None] | None = None,
     ) -> Path:
         if not isinstance(file, File):
-            raise ValueError("FTP File was not properly instantiated")
+            raise ValueError("DuckLake File was not properly instantiated")
 
         access_key = (
             self.credentials.access_key.get_secret_value() if self.credentials else None
@@ -113,5 +128,27 @@ async def download(
         )
         return output
 
+    async def __aenter__(self):
+        await self.connect()
+        return self
+
+    async def __aexit__(self, exc_type, exc_val, exc_tb) -> None:
+        await self.close(update_catalog=None)
+
+    def __del__(self) -> None:
+        if not hasattr(self, "_catalog_adap"):
+            return
+        try:
+            loop = asyncio.get_running_loop()
+            if loop.is_running():
+                loop.create_task(self.close(update_catalog=False))
+        except RuntimeError:
+            try:
+                asyncio.run(self.close(update_catalog=False))
+            except Exception:
+                pass
+        except Exception:
+            pass
+
 
 DuckDataset.model_rebuild(_types_namespace={"DuckLake": DuckLake})