feat: improve ducklake api struture

feat: global variables usage

feat: include more orm models to manage .duckdb files

Author: luabida

docker/Dockerfile Dockerfiledocker/Dockerfile renamed to Dockerfile

@@ -25,12 +25,10 @@ RUN useradd -ms /bin/bash pysus \
 
 COPY pyproject.toml poetry.lock LICENSE README.md /usr/src/
 COPY pysus /usr/src/pysus
-COPY docker/scripts/entrypoint.sh /entrypoint.sh
-COPY docker/notebooks/ /home/pysus/Notebooks/
+COPY entrypoint.sh /entrypoint.sh
 
 RUN pip install poetry \
   && cd /usr/src && poetry config virtualenvs.create false && poetry install --with docs \
-  && pip install 'httpx<0.28' \
   && chown -R pysus:pysus /home/pysus
 
 USER pysus

docker/docker-compose.yaml docker-compose.yamldocker/docker-compose.yaml renamed to docker-compose.yaml

@@ -1,8 +1,8 @@
 services:
   jupyter:
     build:
-      context: ".."
-      dockerfile: docker/Dockerfile
+      context: "."
+      dockerfile: Dockerfile
     hostname: pysus-jupyter
     container_name: pysus-jupyter
     ports:

docker/notebooks/Welcome.ipynb

@@ -24,3 +24,4 @@ def get_version() -> str:
 __version__: str = version
 
 from pysus.api._impl.databases import *  # noqa
+from pysus.api import types  # noqa

docker/scripts/poetry-install.sh

@@ -0,0 +1,173 @@
+from abc import ABC
+from pathlib import Path
+
+import httpx
+from anyio import to_thread
+from pydantic import BaseModel, SecretStr
+from sqlalchemy.engine import Engine
+from sqlalchemy import create_engine
+from sqlalchemy.pool import StaticPool
+from sqlalchemy.orm import sessionmaker, Session
+
+from pysus import CACHEPATH, types
+
+
+class DuckLakeCredentials(BaseModel):
+    """Credentials for authenticating with the S3-compatible object storage.
+
+    Parameters
+    ----------
+    access_key : SecretStr
+        The S3 access key ID.
+    secret_key : SecretStr
+        The S3 secret access key.
+
+    :meta private:
+    """
+
+    access_key: SecretStr
+    secret_key: SecretStr
+
+
+class BaseAdapter(ABC):
+    cache_dir: Path = Path(CACHEPATH) / "ducklake"
+    db_local: Path
+    db_remote: Path
+
+    def __init__(
+        self, engine=None, credentials: DuckLakeCredentials | None = None, **data
+    ) -> None:
+        self._engine = engine
+        self._session_factory = None
+        self.cache_dir.mkdir(parents=True, exist_ok=True)
+        self.credentials = credentials
+
+    @property
+    def remote_url(self) -> str:
+        return f"https://{types.S3_ENDPOINT}/{types.S3_BUCKET}/{self.db_remote}"
+
+    def get_session(self) -> Session:
+        if not self._session_factory:
+            raise RuntimeError("Database engine not initialized. Call connect() first.")
+        return self._session_factory()
+
+    async def connect(self, force: bool = False) -> None:
+        if self._engine and not force:
+            if not self._Session:
+                self._Session = sessionmaker(bind=self._engine)
+            return
+
+        await self._download_catalog(
+            self._catalog_local,
+            self._catalog_remote,
+        )
+        self._engine = await to_thread.run_sync(self.setup_engine)
+        self._Session = sessionmaker(bind=self._engine)
+
+    def setup_engine(
+        self, access_key: str | None = None, secret_key: str | None = None
+    ) -> Engine:
+        engine: Engine = create_engine(
+            f"duckdb:///{self.db_local}",
+            poolclass=StaticPool,
+        )
+
+        with engine.connect() as conn:
+            conn.exec_driver_sql("INSTALL ducklake; LOAD ducklake;")
+
+            has_pysus = conn.exec_driver_sql(
+                "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()
+
+        return engine
+
+    async def _download_catalog(self, local_path: Path, remote_path: str) -> None:
+        url = f"https://{types.S3_ENDPOINT}/{types.S3_BUCKET}/{remote_path}"
+
+        if local_path.exists():
+            try:
+                local_size = local_path.stat().st_size
+            except OSError:
+                local_size = -1
+        else:
+            local_size = -1
+
+        async with httpx.AsyncClient(follow_redirects=True) as client:
+            try:
+                head = await client.head(url)
+                head.raise_for_status()
+                remote_size = int(head.headers.get("content-length", 0))
+            except Exception:  # noqa: B902
+                remote_size = 0
+
+        if remote_size == local_size:
+            return
+
+        await self._download(remote_path, local_path)
+
+    async def _upload_catalog(self) -> None:
+        """Upload all per-dataset catalogs to remote storage.
+
+        Requires authenticated credentials.
+
+        :meta private:
+        """
+        if not self.credentials:
+            raise PermissionError(
+                "Admin credentials required to upload catalog.",
+            )
+
+        if not self.db_local.exists():
+            raise FileNotFoundError("catalog file not found")
+
+        def _upload():
+            self._s3_client.upload_file(
+                self.db_local.absolute(),
+                types.S3_BUCKET,
+                self.db_remote,
+            )
+
+        await to_thread.run_sync(_upload)
+
+
+class CatalogAdapter(BaseAdapter):
+    def __init__(self, engine=None, **data) -> None:
+        super().__init__(engine=engine, **data)
+        self.db_local: Path = self.cache_dir / "catalog.duckdb"
+        self.db_remote: str = "public/catalog.duckdb"
+
+
+class DatasetAdapter(BaseAdapter):
+    def __init__(self, name: str, engine=None, **data) -> None:
+        super().__init__(engine=engine, **data)
+        self.dataset_name: str = name
+        self.db_local: Path = self.cache_dir / f"catalog_{name}.duckdb"
+        self.db_remote: str = f"datasets/catalog_{name}.duckdb"
+
+
+class ColumnsAdapter(BaseAdapter):
+    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"