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

planar/data/utils.py

@@ -0,0 +1,89 @@
+import asyncio
+from typing import TypedDict
+
+import ibis.expr.datatypes as dt
+from ibis.common.exceptions import TableNotFound
+
+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
+
+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)
+    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",
+    ]
+
+    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]
+
+
+async def get_dataset(dataset_name: str, schema_name: str = "main") -> PlanarDataset:
+    # TODO: add schema_name as a parameter
+
+    dataset = PlanarDataset(name=dataset_name)
+
+    if not await dataset.exists():
+        raise DatasetNotFoundError(f"Dataset {dataset_name} not found")
+
+    return dataset
+
+
+async def get_dataset_row_count(dataset_name: str) -> int:
+    conn = await _get_connection()
+
+    try:
+        value = await asyncio.to_thread(
+            lambda conn, dataset_name: conn.table(dataset_name).count().to_polars(),
+            conn,
+            dataset_name,
+        )
+
+        assert isinstance(value, int), "Scalar must be an integer"
+
+        return value
+    except TableNotFound:
+        raise  # re-raise the exception and allow the caller to handle it
+
+
+class DatasetMetadata(TypedDict):
+    schema: dict[str, dt.DataType]
+    row_count: int
+
+
+async def get_dataset_metadata(
+    dataset_name: str, schema_name: str
+) -> DatasetMetadata | None:
+    conn = await _get_connection()
+
+    try:
+        schema, row_count = await asyncio.gather(
+            asyncio.to_thread(conn.get_schema, dataset_name, database=schema_name),
+            get_dataset_row_count(dataset_name),
+        )
+
+        return DatasetMetadata(schema=schema.fields, row_count=row_count)
+
+    except TableNotFound:
+        return None

planar/db/alembic/env.py

@@ -1,5 +1,7 @@
+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
 
@@ -48,6 +50,28 @@ def include_name(name, type_, _):
         return True
 
 
+sqlite_schema_translate_map = {PLANAR_SCHEMA: None}
+
+
+def schema_translate_wrapper(f):
+    @wraps(f)
+    def format_table_name_with_schema(compiler, name, schema):
+        # when on sqlite, we need to translate the schema to None
+        is_sqlite = compiler.dialect.name == "sqlite"
+        if is_sqlite:
+            translated_schema = sqlite_schema_translate_map.get(schema, schema)
+        else:
+            translated_schema = schema
+        return f(compiler, name, translated_schema)
+
+    return format_table_name_with_schema
+
+
+alembic_base.format_table_name = schema_translate_wrapper(
+    alembic_base.format_table_name
+)
+
+
 def run_migrations_online() -> None:
     """Run migrations in 'online' mode.
 
@@ -102,7 +126,7 @@ def run_migrations_online() -> None:
         config_dict = config.get_section(config.config_ini_section, {})
         url = config_dict["sqlalchemy.url"]
         is_sqlite = url.startswith("sqlite://")
-        translate_map = {PLANAR_SCHEMA: None} if is_sqlite else {}
+        translate_map = sqlite_schema_translate_map if is_sqlite else {}
         connectable = engine_from_config(
             config_dict,
             prefix="sqlalchemy.",

planar/files/storage/azure_blob.py

@@ -278,7 +278,7 @@ class AzureBlobStorage(Storage):
 
         elif self.auth_method.name == "AZURE_AD":
             # Generate a User Delegation SAS signed with a user delegation key
-            start_time = datetime.utcnow()
+            start_time = datetime.now(UTC)
             user_delegation_key = await self.client.get_user_delegation_key(
                 key_start_time=start_time, key_expiry_time=expiry_time
             )

planar/registry_items.py

@@ -47,6 +47,7 @@ class RegisteredWorkflow:
     input_schema: dict[str, Any]
     output_schema: dict[str, Any]
     pydantic_model: Type[BaseModel]
+    is_interactive: bool
 
     @staticmethod
     def from_workflow(workflow: "WorkflowWrapper") -> "RegisteredWorkflow":
@@ -63,4 +64,5 @@ class RegisteredWorkflow:
                 workflow.original_fn
             ),
             pydantic_model=create_pydantic_model_for_workflow(workflow),
+            is_interactive=workflow.is_interactive,
         )

planar/routers/dataset_router.py

@@ -0,0 +1,213 @@
+import io
+from typing import AsyncGenerator
+
+import pyarrow as pa
+import pyarrow.parquet as pq
+from fastapi import APIRouter, HTTPException, Query
+from fastapi.responses import StreamingResponse
+from ibis.common.exceptions import TableNotFound
+from pydantic import BaseModel
+
+from planar.data.exceptions import DatasetNotFoundError
+from planar.data.utils import (
+    get_dataset,
+    get_dataset_metadata,
+    list_datasets,
+    list_schemas,
+)
+from planar.logging import get_logger
+from planar.security.authorization import (
+    DatasetAction,
+    DatasetResource,
+    validate_authorization_for,
+)
+
+logger = get_logger(__name__)
+
+
+class DatasetMetadata(BaseModel):
+    name: str
+    table_schema: dict
+    row_count: int
+
+
+def create_dataset_router() -> APIRouter:
+    router = APIRouter(tags=["Planar Datasets"])
+
+    @router.get("/schemas", response_model=list[str])
+    async def get_schemas():
+        validate_authorization_for(
+            DatasetResource(), DatasetAction.DATASET_LIST_SCHEMAS
+        )
+        schemas = await list_schemas()
+        return schemas
+
+    @router.get("/metadata", response_model=list[DatasetMetadata])
+    async def list_planar_datasets(
+        limit: int = Query(100, ge=1, le=1000),
+        offset: int = Query(0, ge=0),
+        schema_name: str = Query("main"),
+    ):
+        validate_authorization_for(DatasetResource(), DatasetAction.DATASET_LIST)
+        datasets = await list_datasets(limit, offset)
+
+        response = []
+        for dataset in datasets:
+            metadata = await get_dataset_metadata(dataset.name, schema_name)
+
+            if not metadata:
+                continue
+
+            schema = metadata["schema"]
+            row_count = metadata["row_count"]
+
+            response.append(
+                DatasetMetadata(
+                    name=dataset.name,
+                    row_count=row_count,
+                    table_schema={
+                        field_name: str(field_type)
+                        for field_name, field_type in schema.items()
+                    },
+                )
+            )
+
+        return response
+
+    @router.get("/metadata/{dataset_name}", response_model=DatasetMetadata)
+    async def get_planar_dataset(dataset_name: str, schema_name: str = "main"):
+        validate_authorization_for(
+            DatasetResource(dataset_name=dataset_name),
+            DatasetAction.DATASET_VIEW_DETAILS,
+        )
+        try:
+            metadata = await get_dataset_metadata(dataset_name, schema_name)
+
+            if not metadata:
+                raise HTTPException(
+                    status_code=404, detail=f"Dataset {dataset_name} not found"
+                )
+
+            schema = metadata["schema"]
+            row_count = metadata["row_count"]
+
+            return DatasetMetadata(
+                name=dataset_name,
+                row_count=row_count,
+                table_schema={
+                    field_name: str(field_type)
+                    for field_name, field_type in schema.items()
+                },
+            )
+        except (DatasetNotFoundError, TableNotFound):
+            raise HTTPException(
+                status_code=404, detail=f"Dataset {dataset_name} not found"
+            )
+
+    @router.get(
+        "/content/{dataset_name}/arrow-stream", response_class=StreamingResponse
+    )
+    async def stream_dataset_content(
+        dataset_name: str,
+        batch_size: int = Query(100, ge=1, le=1000),
+        limit: int | None = Query(None, ge=1),
+    ):
+        validate_authorization_for(
+            DatasetResource(dataset_name=dataset_name),
+            DatasetAction.DATASET_STREAM_CONTENT,
+        )
+        try:
+            dataset = await get_dataset(dataset_name)
+
+            # Apply limit parameter if specified
+            table = await dataset.read(limit=limit)
+
+            schema = table.schema().to_pyarrow()
+
+            async def stream_content() -> AsyncGenerator[bytes, None]:
+                sink = io.BytesIO()
+
+                try:
+                    with pa.ipc.new_stream(sink, schema) as writer:
+                        yield sink.getvalue()  # yield the schema
+
+                        batch_count = 0
+                        for batch in table.to_pyarrow_batches(chunk_size=batch_size):
+                            # reset the sink to only stream
+                            # the current batch
+                            # we don't want to stream the schema or previous
+                            # batches again
+                            sink.seek(0)
+                            sink.truncate(0)
+
+                            writer.write_batch(batch)
+                            yield sink.getvalue()
+                            batch_count += 1
+
+                        # For empty datasets, ensure we have a complete stream
+                        if batch_count == 0:
+                            # Write an empty batch to ensure valid Arrow stream format
+                            empty_batch = pa.RecordBatch.from_arrays(
+                                [pa.array([], type=field.type) for field in schema],
+                                schema=schema,
+                            )
+                            sink.seek(0)
+                            sink.truncate(0)
+                            writer.write_batch(empty_batch)
+                            yield sink.getvalue()
+                finally:
+                    # Explicit BytesIO cleanup for memory safety
+                    sink.close()
+
+            return StreamingResponse(
+                stream_content(),
+                media_type="application/vnd.apache.arrow.stream",
+                headers={
+                    "Content-Disposition": f"attachment; filename={dataset_name}.arrow",
+                    "X-Batch-Size": str(batch_size),
+                    "X-Row-Limit": str(limit) if limit else "unlimited",
+                },
+            )
+        except (DatasetNotFoundError, TableNotFound):
+            raise HTTPException(
+                status_code=404, detail=f"Dataset {dataset_name} not found"
+            )
+
+    @router.get("/content/{dataset_name}/download")
+    async def download_dataset(dataset_name: str, schema_name: str = "main"):
+        validate_authorization_for(
+            DatasetResource(dataset_name=dataset_name),
+            DatasetAction.DATASET_DOWNLOAD,
+        )
+        try:
+            arrow_buffer = pa.BufferOutputStream()
+            dataset = await get_dataset(dataset_name, schema_name)
+
+            pyarrow_table = await dataset.to_pyarrow()
+
+            pq.write_table(pyarrow_table, arrow_buffer)
+
+            if arrow_buffer.tell() == 0:
+                logger.warning(
+                    "Dataset is empty",
+                    dataset_name=dataset_name,
+                    schema_name=schema_name,
+                )
+
+            buffer = arrow_buffer.getvalue()
+            parquet_bytes = buffer.to_pybytes()
+            bytes_io = io.BytesIO(parquet_bytes)
+
+            return StreamingResponse(
+                bytes_io,
+                media_type="application/x-parquet",
+                headers={
+                    "Content-Disposition": f"attachment; filename={dataset_name}.parquet"
+                },
+            )
+        except (DatasetNotFoundError, TableNotFound):
+            raise HTTPException(
+                status_code=404, detail=f"Dataset {dataset_name} not found"
+            )
+
+    return router

planar/routers/models.py

@@ -79,6 +79,7 @@ class WorkflowDefinition(BaseModel):
     total_runs: int
     run_statuses: WorkflowRunStatusCounts
     durations: DurationStats | None = None
+    is_interactive: bool
 
 
 class StepStats(BaseModel):