planar 0.11.0__py3-none-any.whl → 0.13.0__py3-none-any.whl

planar/data/connection.py

@@ -1,4 +1,6 @@
 import asyncio
+from dataclasses import dataclass
+from urllib.parse import urlparse
 
 import ibis
 from ibis.backends.duckdb import Backend as DuckDBBackend
@@ -17,6 +19,85 @@ from planar.session import get_config
 logger = get_logger(__name__)
 
 
+@dataclass
+class _ConnectionPool:
+    connections: list[DuckDBBackend]
+    cursor: int = 0
+
+
+# In production a Planar app typically runs with a single data config, so we only
+# ever have one signature, but we still rotate through a handful of cached
+# backends to reduce the risk of concurrent calls sharing the same DuckDB
+# connection. During testing we create many ephemeral configs (temp dirs, sqlite
+# files, etc.), so the cache also avoids paying the attachment cost on every
+# request. We keep up to `_MAX_CONNECTIONS_PER_SIGNATURE` backends per signature
+# and hand them out in round-robin order; concurrency safety ultimately depends on
+# DuckDB tolerating overlapping use of an individual backend.
+_connection_cache: dict[int, _ConnectionPool] = {}
+_cache_lock: asyncio.Lock | None = None
+
+# Maximum number of cached connections per configuration signature.
+_MAX_CONNECTIONS_PER_SIGNATURE = 10
+
+
+def _config_signature(config: PlanarConfig) -> int:
+    """Create a stable signature for caching connections."""
+
+    assert config.data is not None, "data configuration must be set"
+    return hash(config.data)
+
+
+async def _close_backend(connection: DuckDBBackend) -> None:
+    close_fn = getattr(connection, "close", None)
+    try:
+        if callable(close_fn):
+            await asyncio.to_thread(close_fn)
+    except Exception as exc:
+        logger.warning("failed to close DuckDB connection", error=str(exc))
+
+
+def _make_aws_s3_secret_query(config: S3Config) -> str:
+    """
+    https://duckdb.org/docs/stable/core_extensions/httpfs/s3api
+    """
+    columns = ["TYPE s3"]
+
+    if config.region:
+        columns.append(
+            f"REGION '{config.region}'",
+        )
+
+    if config.endpoint_url:
+        parsed_url = urlparse(config.endpoint_url)
+        endpoint_host = parsed_url.hostname
+        endpoint_port = f":{parsed_url.port}" if parsed_url.port else ""
+        endpoint = f"{endpoint_host}{endpoint_port}"
+
+        columns.append(f"ENDPOINT '{endpoint}'")
+
+    if config.access_key and config.secret_key:
+        columns.extend(
+            [
+                "PROVIDER config",
+                f"KEY_ID '{config.access_key}'",
+                f"SECRET '{config.secret_key}'",
+            ]
+        )
+    else:
+        columns.extend(
+            [
+                "PROVIDER credential_chain",
+                "CHAIN 'env;sts;instance'",
+            ]
+        )
+
+    return f"""
+        CREATE OR REPLACE SECRET secret (
+        {", ".join(columns)}
+        );
+    """
+
+
 async def _create_connection(config: PlanarConfig) -> DuckDBBackend:
     """Create Ibis DuckDB connection with Ducklake."""
     data_config = config.data
@@ -67,6 +148,14 @@ async def _create_connection(config: PlanarConfig) -> DuckDBBackend:
     if isinstance(storage, LocalDirectoryConfig):
         data_path = storage.directory
     elif isinstance(storage, S3Config):
+        await asyncio.to_thread(con.raw_sql, "INSTALL httpfs;")
+        await asyncio.to_thread(con.raw_sql, "LOAD httpfs;")
+
+        await asyncio.to_thread(
+            con.raw_sql,
+            _make_aws_s3_secret_query(storage),
+        )
+
         data_path = f"s3://{storage.bucket_name}/"
     else:
         # Generic fallback
@@ -95,8 +184,21 @@ async def _create_connection(config: PlanarConfig) -> DuckDBBackend:
     return con
 
 
-async def _get_connection() -> DuckDBBackend:
-    """Get Ibis connection to Ducklake."""
+def _get_cache_lock() -> asyncio.Lock:
+    # Create a lock on the first call to this function, or re-create it if the
+    # loop has changed (happens on tests).
+    global _cache_lock
+    loop = asyncio.get_running_loop()
+    lock = _cache_lock
+    if lock is None or getattr(lock, "_loop", None) is not loop:
+        lock = asyncio.Lock()
+        _cache_lock = lock
+    return lock
+
+
+async def get_connection() -> DuckDBBackend:
+    """Return a cached DuckDB connection using round-robin selection."""
+
     config = get_config()
 
     if not config.data:
@@ -104,5 +206,39 @@ async def _get_connection() -> DuckDBBackend:
             "Data configuration not found. Please configure 'data' in your planar.yaml"
         )
 
-    # TODO: Add cached connection pooling or memoize the connection
-    return await _create_connection(config)
+    signature = _config_signature(config)
+    lock = _get_cache_lock()
+
+    async with lock:
+        pool = _connection_cache.get(signature)
+
+        if pool is None:
+            connection = await _create_connection(config)
+            _connection_cache[signature] = _ConnectionPool(connections=[connection])
+            return connection
+
+        if len(pool.connections) < _MAX_CONNECTIONS_PER_SIGNATURE:
+            connection = await _create_connection(config)
+            pool.connections.append(connection)
+            return connection
+
+        connection = pool.connections[pool.cursor]
+        pool.cursor = (pool.cursor + 1) % len(pool.connections)
+        return connection
+
+
+async def reset_connection_cache() -> None:
+    """Reset the cached DuckDB connection, closing it if necessary."""
+
+    lock = _get_cache_lock()
+
+    async with lock:
+        pools = list(_connection_cache.values())
+        _connection_cache.clear()
+
+    for pool in pools:
+        for connection in pool.connections:
+            await _close_backend(connection)
+
+    global _cache_lock
+    _cache_lock = None

planar/data/dataset.py

@@ -6,10 +6,11 @@ from typing import Literal, Self
 import ibis
 import polars as pl
 import pyarrow as pa
+from ibis.backends.duckdb import Backend as DuckDBBackend
 from ibis.common.exceptions import TableNotFound
 from pydantic import BaseModel
 
-from planar.data.connection import _get_connection
+from planar.data.connection import get_connection
 from planar.logging import get_logger
 
 from .exceptions import DataError, DatasetAlreadyExistsError, DatasetNotFoundError
@@ -67,7 +68,11 @@ class PlanarDataset(BaseModel):
 
     async def exists(self) -> bool:
         """Check if the dataset exists in Ducklake."""
-        con = await _get_connection()
+        con = await get_connection()
+        return await self._table_exists(con)
+
+    async def _table_exists(self, con: DuckDBBackend) -> bool:
+        """Check for table existence using the provided connection."""
 
         try:
             # TODO: Query for the table name directly
@@ -88,11 +93,13 @@ class PlanarDataset(BaseModel):
             data: Data to write (Polars DataFrame/LazyFrame, PyArrow Table, or Ibis expression)
             mode: Write mode - "append" or "overwrite"
         """
-        con = await _get_connection()
         overwrite = mode == "overwrite"
 
         try:
-            if not await self.exists():
+            con = await get_connection()
+            table_exists = await self._table_exists(con)
+
+            if not table_exists:
                 await asyncio.to_thread(
                     con.create_table, self.name, data, overwrite=overwrite
                 )
@@ -133,9 +140,8 @@ class PlanarDataset(BaseModel):
         Returns:
             Ibis table expression that can be further filtered using Ibis methods
         """
-        con = await _get_connection()
-
         try:
+            con = await get_connection()
             table = await asyncio.to_thread(con.table, self.name)
 
             if columns:
@@ -162,8 +168,8 @@ class PlanarDataset(BaseModel):
 
     async def delete(self) -> None:
         """Delete the dataset."""
-        con = await _get_connection()
         try:
+            con = await get_connection()
             await asyncio.to_thread(con.drop_table, self.name, force=True)
             logger.info("deleted dataset", dataset_name=self.name)
         except Exception as e:

planar/data/utils.py

@@ -1,42 +1,45 @@
 import asyncio
-from typing import TypedDict
+from collections import defaultdict
+from typing import Sequence, TypedDict
 
+import ibis
 import ibis.expr.datatypes as dt
+import pyarrow as pa
+from ibis.backends.duckdb import Backend as DuckDBBackend
 from ibis.common.exceptions import TableNotFound
+from sqlglot import exp
 
-from planar.data.connection import _get_connection
+from planar.data.connection import get_connection
 from planar.data.dataset import PlanarDataset
 from planar.data.exceptions import DatasetNotFoundError
 from planar.logging import get_logger
+from planar.session import get_config
 
 logger = get_logger(__name__)
 
 
-# TODO: consider connection pooling or memoize the connection
-
-
 async def list_datasets(limit: int = 100, offset: int = 0) -> list[PlanarDataset]:
-    conn = await _get_connection()
-    tables = await asyncio.to_thread(conn.list_tables)
+    conn = await get_connection()
+    tables = sorted(await asyncio.to_thread(conn.list_tables))[offset : offset + limit]
     return [PlanarDataset(name=table) for table in tables]
 
 
 async def list_schemas() -> list[str]:
-    METADATA_SCHEMAS = [
-        "information_schema",
-        # FIXME: why is list_databases returning pg_catalog
-        # if the ducklake catalog is sqlite?
-        "pg_catalog",
-    ]
+    config = get_config()
+
+    if config.data is None:
+        return []
 
-    conn = await _get_connection()
+    METADATA_SCHEMAS = [config.data.catalog_name, "main"]
+
+    conn = await get_connection()
 
     # in ibis, "databases" are schemas in the traditional sense
     # e.g. psql: schema == ibis: database
     # https://ibis-project.org/concepts/backend-table-hierarchy
     schemas = await asyncio.to_thread(conn.list_databases)
 
-    return [schema for schema in schemas if schema not in METADATA_SCHEMAS]
+    return [schema for schema in schemas if schema in METADATA_SCHEMAS]
 
 
 async def get_dataset(dataset_name: str, schema_name: str = "main") -> PlanarDataset:
@@ -51,7 +54,7 @@ async def get_dataset(dataset_name: str, schema_name: str = "main") -> PlanarDat
 
 
 async def get_dataset_row_count(dataset_name: str) -> int:
-    conn = await _get_connection()
+    conn = await get_connection()
 
     try:
         value = await asyncio.to_thread(
@@ -72,18 +75,135 @@ class DatasetMetadata(TypedDict):
     row_count: int
 
 
-async def get_dataset_metadata(
-    dataset_name: str, schema_name: str
-) -> DatasetMetadata | None:
-    conn = await _get_connection()
+async def _fetch_column_schemas(
+    conn: DuckDBBackend,
+    dataset_names: Sequence[str],
+    schema_name: str,
+) -> dict[str, dict[str, dt.DataType]]:
+    columns = conn.table("columns", database="information_schema")
+    schema_literal = ibis.literal(schema_name)
+    dataset_literals = [ibis.literal(name) for name in dataset_names]
+    filtered = columns.filter(
+        (columns.table_schema == schema_literal)
+        & (columns.table_name.isin(dataset_literals))
+    )
+
+    selected = filtered.select(
+        columns.table_name.name("table_name"),
+        columns.column_name.name("column_name"),
+        columns.ordinal_position.name("ordinal_position"),
+        columns.data_type.name("data_type"),
+        columns.is_nullable.name("is_nullable"),
+    )
+
+    arrow_table: pa.Table = await asyncio.to_thread(selected.to_pyarrow)
+    rows = arrow_table.to_pylist()
+
+    schema_fields: dict[str, list[tuple[int, str, dt.DataType]]] = defaultdict(list)
+    type_mapper = conn.compiler.type_mapper
+
+    for row in rows:
+        table_name = row["table_name"]
+        column_name = row["column_name"]
+        ordinal_position = row["ordinal_position"]
+        dtype = type_mapper.from_string(
+            row["data_type"], nullable=row.get("is_nullable") == "YES"
+        )
 
-    try:
-        schema, row_count = await asyncio.gather(
-            asyncio.to_thread(conn.get_schema, dataset_name, database=schema_name),
-            get_dataset_row_count(dataset_name),
+        schema_fields[table_name].append((ordinal_position, column_name, dtype))
+
+    ordered_fields: dict[str, dict[str, dt.DataType]] = {}
+    for table_name, fields in schema_fields.items():
+        ordered_fields[table_name] = {
+            column_name: dtype
+            for _, column_name, dtype in sorted(fields, key=lambda entry: entry[0])
+        }
+
+    return ordered_fields
+
+
+async def _fetch_row_counts(
+    conn: DuckDBBackend,
+    dataset_names: Sequence[str],
+    schema_name: str,
+) -> dict[str, int]:
+    if not dataset_names:
+        return {}
+
+    quoted = conn.compiler.quoted
+    count_queries: list[exp.Select] = []
+
+    for dataset_name in dataset_names:
+        table_expr = exp.Table(
+            this=exp.Identifier(this=dataset_name, quoted=quoted),
+            db=(
+                exp.Identifier(this=schema_name, quoted=quoted) if schema_name else None
+            ),
         )
+        select_expr = (
+            exp.Select()
+            .select(
+                exp.Literal.string(dataset_name).as_("dataset_name"),
+                exp.Count(this=exp.Star()).as_("row_count"),
+            )
+            .from_(table_expr)
+        )
+        count_queries.append(select_expr)
+
+    if not count_queries:
+        return {}
+
+    union_query: exp.Expression = count_queries[0]
+    for query in count_queries[1:]:
+        union_query = exp.Union(this=union_query, expression=query, distinct=False)
+
+    def _execute() -> dict[str, int]:
+        with conn._safe_raw_sql(union_query) as cursor:  # type: ignore[attr-defined]
+            rows = cursor.fetchall()
 
-        return DatasetMetadata(schema=schema.fields, row_count=row_count)
+        return {str(dataset_name): int(row_count) for dataset_name, row_count in rows}
 
+    return await asyncio.to_thread(_execute)
+
+
+async def get_datasets_metadata(
+    dataset_names: Sequence[str], schema_name: str
+) -> dict[str, DatasetMetadata]:
+    if not dataset_names:
+        return {}
+
+    dataset_list = list(dict.fromkeys(dataset_names))
+    if not dataset_list:
+        return {}
+
+    conn = await get_connection()
+
+    schemas = await _fetch_column_schemas(conn, dataset_list, schema_name)
+    row_counts = await _fetch_row_counts(conn, list(schemas.keys()), schema_name)
+
+    metadata: dict[str, DatasetMetadata] = {}
+
+    for dataset_name in dataset_list:
+        schema = schemas.get(dataset_name)
+        row_count = row_counts.get(dataset_name)
+
+        if not schema or row_count is None:
+            continue
+
+        metadata[dataset_name] = DatasetMetadata(
+            schema=schema,
+            row_count=row_count,
+        )
+
+    return metadata
+
+
+async def get_dataset_metadata(
+    dataset_name: str, schema_name: str
+) -> DatasetMetadata | None:
+    try:
+        metadata = await get_datasets_metadata([dataset_name], schema_name)
     except TableNotFound:
         return None
+
+    return metadata.get(dataset_name)

planar/db/alembic/env.py

@@ -1,9 +1,11 @@
+import asyncio
 from functools import wraps
 from logging.config import fileConfig
 
 import alembic.ddl.base as alembic_base
 from alembic import context
-from sqlalchemy import Connection, engine_from_config, pool
+from sqlalchemy import Connection, pool
+from sqlalchemy.ext.asyncio import create_async_engine
 
 from planar.db import PLANAR_FRAMEWORK_METADATA, PLANAR_SCHEMA
 
@@ -72,6 +74,69 @@ alembic_base.format_table_name = schema_translate_wrapper(
 )
 
 
+async def run_migrations_online_async() -> None:
+    """Run migrations in 'online' mode using async engine for development."""
+    # Import models to ensure they're registered with PLANAR_FRAMEWORK_METADATA
+    try:
+        from planar.files.models import PlanarFileMetadata  # noqa: F401, PLC0415
+        from planar.human.models import HumanTask  # noqa: F401, PLC0415
+        from planar.object_config.models import (  # noqa: F401, PLC0415
+            ObjectConfiguration,
+        )
+        from planar.workflows.models import (  # noqa: PLC0415
+            LockedResource,  # noqa: F401
+            Workflow,  # noqa: F401
+            WorkflowEvent,  # noqa: F401
+            WorkflowStep,  # noqa: F401
+        )
+    except ImportError as e:
+        raise RuntimeError(
+            f"Failed to import system models for migration generation: {e}"
+        )
+
+    config_dict = config.get_section(config.config_ini_section, {})
+    url = config_dict["sqlalchemy.url"]
+    is_sqlite = url.startswith("sqlite://")
+
+    # Create async engine
+    connectable = create_async_engine(
+        url,
+        poolclass=pool.NullPool,
+        execution_options={
+            # SQLite doesn't support schemas, so we need to translate the planar schema
+            # name to None in order to ignore it.
+            "schema_translate_map": sqlite_schema_translate_map if is_sqlite else {},
+        },
+    )
+
+    async with connectable.connect() as connection:
+        is_sqlite = connection.dialect.name == "sqlite"
+        if is_sqlite:
+            connection.dialect.default_schema_name = PLANAR_SCHEMA
+
+        def do_run_migrations(sync_conn):
+            context.configure(
+                connection=sync_conn,
+                target_metadata=target_metadata,
+                # For SQLite, don't use schema since it's not supported
+                version_table_schema=None if is_sqlite else PLANAR_SCHEMA,
+                include_schemas=True,
+                include_name=include_name,
+                # SQLite doesn't support alter table, so we need to use render_as_batch
+                # to create the tables in a single transaction. For other databases,
+                # the batch op is no-op.
+                # https://alembic.sqlalchemy.org/en/latest/batch.html#running-batch-migrations-for-sqlite-and-other-databases
+                render_as_batch=True,
+            )
+
+            with context.begin_transaction():
+                context.run_migrations()
+
+        await connection.run_sync(do_run_migrations)
+
+    await connectable.dispose()
+
+
 def run_migrations_online() -> None:
     """Run migrations in 'online' mode.
 
@@ -103,62 +168,8 @@ def run_migrations_online() -> None:
         with context.begin_transaction():
             context.run_migrations()
     else:
-        # Development mode: create engine from alembic.ini
-        # Used for alembic to generate migrations
-        # Import models to ensure they're registered with PLANAR_FRAMEWORK_METADATA
-        try:
-            from planar.files.models import PlanarFileMetadata  # noqa: F401, PLC0415
-            from planar.human.models import HumanTask  # noqa: F401, PLC0415
-            from planar.object_config.models import (  # noqa: F401, PLC0415
-                ObjectConfiguration,
-            )
-            from planar.workflows.models import (  # noqa: PLC0415
-                LockedResource,  # noqa: F401
-                Workflow,  # noqa: F401
-                WorkflowEvent,  # noqa: F401
-                WorkflowStep,  # noqa: F401
-            )
-        except ImportError as e:
-            raise RuntimeError(
-                f"Failed to import system models for migration generation: {e}"
-            )
-
-        config_dict = config.get_section(config.config_ini_section, {})
-        url = config_dict["sqlalchemy.url"]
-        is_sqlite = url.startswith("sqlite://")
-        translate_map = sqlite_schema_translate_map if is_sqlite else {}
-        connectable = engine_from_config(
-            config_dict,
-            prefix="sqlalchemy.",
-            poolclass=pool.NullPool,
-            execution_options={
-                # SQLite doesn't support schemas, so we need to translate the planar schema
-                # name to None in order to ignore it.
-                "schema_translate_map": translate_map,
-            },
-        )
-
-        with connectable.connect() as connection:
-            is_sqlite = connection.dialect.name == "sqlite"
-            if is_sqlite:
-                connection.dialect.default_schema_name = PLANAR_SCHEMA
-
-            context.configure(
-                connection=connection,
-                target_metadata=target_metadata,
-                # For SQLite, don't use schema since it's not supported
-                version_table_schema=None if is_sqlite else PLANAR_SCHEMA,
-                include_schemas=True,
-                include_name=include_name,
-                # SQLite doesn't support alter table, so we need to use render_as_batch
-                # to create the tables in a single transaction. For other databases,
-                # the batch op is no-op.
-                # https://alembic.sqlalchemy.org/en/latest/batch.html#running-batch-migrations-for-sqlite-and-other-databases
-                render_as_batch=True,
-            )
-
-            with context.begin_transaction():
-                context.run_migrations()
+        # Development mode: run migrations asynchronously
+        asyncio.run(run_migrations_online_async())
 
 
 if context.is_offline_mode():

planar/db/alembic.ini

@@ -70,7 +70,7 @@ version_path_separator = os
 # so it sometimes incorrectly thinks it needs to re-generate things (like indices) that already
 # exist in the database from a prior migration. Using postgres obviates that issue.
 # https://github.com/sqlalchemy/alembic/issues/555
-sqlalchemy.url = postgresql+psycopg2://postgres:postgres@localhost:5432/postgres
+sqlalchemy.url = postgresql+asyncpg://postgres:postgres@localhost:5432/postgres
 
 # [post_write_hooks]
 # This section defines scripts or Python functions that are run

planar/files/storage/config.py

@@ -2,7 +2,7 @@ from __future__ import annotations
 
 from typing import TYPE_CHECKING, Annotated, Literal
 
-from pydantic import BaseModel, Field, model_validator
+from pydantic import BaseModel, ConfigDict, Field, model_validator
 
 from .local_directory import LocalDirectoryStorage
 from .s3 import S3Storage
@@ -15,6 +15,8 @@ class LocalDirectoryConfig(BaseModel):
     backend: Literal["localdir"]
     directory: str
 
+    model_config = ConfigDict(frozen=True)
+
 
 class S3Config(BaseModel):
     backend: Literal["s3"]
@@ -25,6 +27,8 @@ class S3Config(BaseModel):
     endpoint_url: str | None = None
     presigned_url_ttl: int = 3600
 
+    model_config = ConfigDict(frozen=True)
+
 
 class AzureBlobConfig(BaseModel):
     backend: Literal["azure_blob"]
@@ -39,6 +43,8 @@ class AzureBlobConfig(BaseModel):
     # Common settings
     sas_ttl: int = 3600  # SAS URL expiry time in seconds
 
+    model_config = ConfigDict(frozen=True)
+
     @model_validator(mode="after")
     def validate_auth_config(self):
         """Ensure exactly one valid authentication configuration."""

planar/human/models.py

@@ -7,7 +7,6 @@ from pydantic import BaseModel
 from sqlmodel import JSON, Column, Field
 
 from planar.db import PlanarInternalBase
-from planar.modeling.field_helpers import JsonSchema
 from planar.modeling.mixins import TimestampMixin
 from planar.modeling.mixins.auditable import AuditableMixin
 from planar.modeling.mixins.uuid_primary_key import UUIDPrimaryKeyMixin
@@ -48,12 +47,12 @@ class HumanTask(
     workflow_name: str
 
     # Input data for context
-    input_schema: Optional[JsonSchema] = Field(default=None, sa_column=Column(JSON))
+    input_schema: Optional[dict[str, Any]] = Field(default=None, sa_column=Column(JSON))
     input_data: Optional[dict[str, Any]] = Field(default=None, sa_column=Column(JSON))
     message: Optional[str] = Field(default=None)
 
     # Schema for expected output
-    output_schema: JsonSchema = Field(sa_column=Column(JSON))
+    output_schema: dict[str, Any] = Field(sa_column=Column(JSON))
     output_data: Optional[dict[str, Any]] = Field(default=None, sa_column=Column(JSON))
 
     # Suggested data for the form (optional)

planar/logging/attributes.py

@@ -18,6 +18,7 @@ def _in_event_loop_task() -> bool:
         return (
             threading.main_thread() == threading.current_thread()
             and asyncio.get_running_loop() is not None
+            and asyncio.current_task() is not None
         )
     except RuntimeError:
         return False