planar 0.10.0__py3-none-any.whl → 0.12.0__py3-none-any.whl

planar/app.py

@@ -28,7 +28,6 @@ from planar.routers import (
     create_workflow_router,
 )
 from planar.routers.agents_router import create_agent_router
-from planar.routers.dataset_router import create_dataset_router
 from planar.routers.entity_router import create_entities_router
 from planar.routers.object_config_router import create_object_config_router
 from planar.routers.rule import create_rule_router
@@ -130,14 +129,27 @@ class PlanarApp:
             create_human_task_routes(),
             prefix="/human-tasks",
         )
-        self.router_v1.include_router(
-            create_dataset_router(),
-            prefix="/datasets",
-        )
+
+        if self.config.data:
+            try:
+                from planar.routers.dataset_router import create_dataset_router
+
+                self.router_v1.include_router(
+                    create_dataset_router(),
+                    prefix="/datasets",
+                )
+            except ImportError:
+                logger.error(
+                    "Data dependencies not installed. Ensure you install the `data` optional dependency in your project (planar[data])"
+                )
+                raise
 
         self.router_v1.include_router(
             create_info_router(
-                title=title or "Planar API", description=description or "Planar API"
+                title=title or "Planar API",
+                description=description or "Planar API",
+                config=self.config,
+                registry=self._object_registry,
             ),
             prefix="",
         )
@@ -262,6 +274,14 @@ class PlanarApp:
             # Reset the config in the context
             config_var.reset(config_tok)
 
+            if self.config.data:
+                try:
+                    from planar.data.connection import reset_connection_cache
+                except ImportError as exc:  # pragma: no cover - optional dependency
+                    logger.debug("skipping data connection cleanup", error=str(exc))
+                else:
+                    await reset_connection_cache()
+
         await self.db_manager.disconnect()
 
         if self.storage:

planar/cli.py

@@ -2,6 +2,7 @@ import os
 import subprocess
 import sys
 from pathlib import Path
+from typing import Annotated
 
 import typer
 import uvicorn
@@ -9,10 +10,35 @@ from jinja2 import Environment as JinjaEnvironment
 from jinja2 import FileSystemLoader
 
 from planar.config import Environment
+from planar.version import get_version
 
 app = typer.Typer(help="Planar CLI tool")
 
 
+def version_callback(value: bool) -> bool:
+    if value:
+        typer.echo(f"planar {get_version()}")
+        raise typer.Exit()
+    return value
+
+
+@app.callback()
+def root_callback(
+    version: Annotated[
+        bool | None,
+        typer.Option(
+            "--version",
+            "-v",
+            help="Show Planar version and exit.",
+            callback=version_callback,
+            is_flag=True,
+            is_eager=True,
+        ),
+    ] = None,
+) -> None:
+    """Entry point for Planar CLI with shared options."""
+
+
 def find_default_app_path() -> Path:
     """Checks for default app file paths (app.py, then main.py)."""
     for filename in ["app.py", "main.py"]:

planar/data/__init__.py

@@ -6,6 +6,7 @@ lazy_exports(
     __name__,
     {
         "PlanarDataset": (".dataset", "PlanarDataset"),
+        "reset_connection_cache": (".connection", "reset_connection_cache"),
     },
 )
 

planar/data/config.py

@@ -2,7 +2,7 @@
 
 from typing import Annotated, Literal
 
-from pydantic import BaseModel, Field
+from pydantic import BaseModel, ConfigDict, Field
 
 from planar.files.storage.config import StorageConfig
 
@@ -13,6 +13,8 @@ class DuckDBCatalogConfig(BaseModel):
     type: Literal["duckdb"]
     path: str  # Path to .ducklake file
 
+    model_config = ConfigDict(frozen=True)
+
 
 class PostgresCatalogConfig(BaseModel):
     """Configuration for PostgreSQL catalog backend."""
@@ -24,6 +26,8 @@ class PostgresCatalogConfig(BaseModel):
     password: str | None = None
     db: str
 
+    model_config = ConfigDict(frozen=True)
+
 
 class SQLiteCatalogConfig(BaseModel):
     """Configuration for SQLite catalog backend."""
@@ -31,6 +35,8 @@ class SQLiteCatalogConfig(BaseModel):
     type: Literal["sqlite"]
     path: str  # Path to .sqlite file
 
+    model_config = ConfigDict(frozen=True)
+
 
 # Discriminated union for catalog configurations
 CatalogConfig = Annotated[
@@ -47,3 +53,8 @@ class DataConfig(BaseModel):
 
     # Optional settings
     catalog_name: str = "planar_data"  # Default catalog name in Ducklake
+
+    model_config = ConfigDict(frozen=True)
+
+    def is_sqlite_catalog(self) -> bool:
+        return self.catalog.type == "sqlite"

planar/data/connection.py

@@ -1,4 +1,5 @@
 import asyncio
+from dataclasses import dataclass
 
 import ibis
 from ibis.backends.duckdb import Backend as DuckDBBackend
@@ -17,6 +18,43 @@ 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))
+
+
 async def _create_connection(config: PlanarConfig) -> DuckDBBackend:
     """Create Ibis DuckDB connection with Ducklake."""
     data_config = config.data
@@ -95,8 +133,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 +155,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():