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

planar/routers/test_dataset_router.py

@@ -0,0 +1,429 @@
+import math
+
+import polars as pl
+import pyarrow as pa
+import pytest
+
+from planar.data.dataset import PlanarDataset
+from planar.testing.planar_test_client import PlanarTestClient
+
+
+@pytest.fixture(name="app")
+def app_fixture(app_with_data):
+    """Use the shared app_with_data fixture as 'app' for this test module."""
+    return app_with_data
+
+
+async def test_stream_arrow_chunks(
+    client: PlanarTestClient,
+):
+    dataset_name = "test_streaming"
+    dataset_size = 10_000
+    batch_size = 1000
+
+    dataset = await PlanarDataset.create(dataset_name)
+
+    df = pl.DataFrame({"id": range(dataset_size)}).with_columns(
+        pl.format("value_{}", pl.col("id")).alias("value")
+    )
+
+    await dataset.write(df)
+
+    response = await client.get(
+        f"/planar/v1/datasets/content/{dataset_name}/arrow-stream",
+        params={"batch_size": batch_size, "limit": dataset_size},
+    )
+
+    assert response.status_code == 200
+    assert response.headers["content-type"] == "application/vnd.apache.arrow.stream"
+    assert "test_streaming.arrow" in response.headers.get("content-disposition", "")
+    assert response.headers.get("x-batch-size") == str(batch_size)
+
+    content = await response.aread()
+    buffer = pa.py_buffer(content)
+    reader = pa.ipc.open_stream(buffer)
+
+    batch_info = []
+    total_rows_received = 0
+    all_ids = []
+
+    try:
+        while True:
+            arrow_batch = reader.read_next_batch()
+            batch_info.append(
+                {
+                    "rows": arrow_batch.num_rows,
+                    "columns": arrow_batch.num_columns,
+                }
+            )
+            total_rows_received += arrow_batch.num_rows
+
+            id_column = arrow_batch.column("id")
+            batch_ids = id_column.to_pylist()
+            all_ids.extend(batch_ids)
+    except StopIteration:
+        pass
+
+    expected_batches = math.ceil(dataset_size / batch_size)
+
+    assert len(batch_info) == expected_batches
+    assert total_rows_received == dataset_size
+
+    # Verify data integrity - check that we received all expected IDs
+    assert len(all_ids) == dataset_size
+    assert set(all_ids) == set(range(dataset_size))
+    assert sum(all_ids) == sum(range(dataset_size))
+
+
+async def test_stream_arrow_with_limit(
+    client: PlanarTestClient,
+):
+    """Test that the limit parameter properly restricts the number of rows streamed."""
+    dataset_name = "test_streaming_limit"
+    dataset_size = 1000
+    batch_size = 100
+    row_limit = 250  # Should get 3 batches (100 + 100 + 50)
+
+    dataset = await PlanarDataset.create(dataset_name)
+
+    # Create test data
+    df = pl.DataFrame({"id": range(dataset_size)}).with_columns(
+        pl.format("value_{}", pl.col("id")).alias("value")
+    )
+
+    await dataset.write(df)
+
+    response = await client.get(
+        f"/planar/v1/datasets/content/{dataset_name}/arrow-stream",
+        params={"batch_size": batch_size, "limit": row_limit},
+    )
+
+    assert response.status_code == 200
+    assert response.headers["x-row-limit"] == str(row_limit)
+
+    content = await response.aread()
+    buffer = pa.py_buffer(content)
+    reader = pa.ipc.open_stream(buffer)
+
+    total_rows_received = 0
+    batch_count = 0
+
+    try:
+        while True:
+            arrow_batch = reader.read_next_batch()
+            total_rows_received += arrow_batch.num_rows
+            batch_count += 1
+    except StopIteration:
+        pass
+
+    # Should receive exactly the limited number of rows
+    assert total_rows_received == row_limit
+    # Should receive expected number of batches (3: 100, 100, 50)
+    expected_batches = math.ceil(row_limit / batch_size)
+    assert batch_count == expected_batches
+
+
+async def test_stream_arrow_empty_dataset(
+    client: PlanarTestClient,
+):
+    """Test streaming behavior with an empty dataset."""
+    dataset_name = "test_empty_stream"
+    batch_size = 100
+
+    dataset = await PlanarDataset.create(dataset_name)
+
+    # Create empty dataset
+    df = pl.DataFrame(
+        {"id": [], "value": []}, schema={"id": pl.Int64, "value": pl.Utf8}
+    )
+    await dataset.write(df)
+
+    response = await client.get(
+        f"/planar/v1/datasets/content/{dataset_name}/arrow-stream",
+        params={"batch_size": batch_size},
+    )
+
+    assert response.status_code == 200
+
+    content = await response.aread()
+    buffer = pa.py_buffer(content)
+    reader = pa.ipc.open_stream(buffer)
+
+    # Should be able to read the schema and get one empty batch
+    total_rows = 0
+    batch_count = 0
+
+    try:
+        while True:
+            arrow_batch = reader.read_next_batch()
+            total_rows += arrow_batch.num_rows
+            batch_count += 1
+    except StopIteration:
+        pass
+
+    # Should have exactly 1 empty batch (our fallback for empty datasets)
+    assert batch_count == 1
+    assert total_rows == 0
+
+
+async def test_stream_arrow_single_batch(
+    client: PlanarTestClient,
+):
+    """Test streaming when dataset size is smaller than batch size."""
+    dataset_name = "test_single_batch"
+    dataset_size = 50
+    batch_size = 100
+
+    dataset = await PlanarDataset.create(dataset_name)
+
+    df = pl.DataFrame({"id": range(dataset_size)}).with_columns(
+        pl.format("value_{}", pl.col("id")).alias("value")
+    )
+
+    await dataset.write(df)
+
+    response = await client.get(
+        f"/planar/v1/datasets/content/{dataset_name}/arrow-stream",
+        params={"batch_size": batch_size},
+    )
+
+    assert response.status_code == 200
+
+    content = await response.aread()
+    buffer = pa.py_buffer(content)
+    reader = pa.ipc.open_stream(buffer)
+
+    total_rows = 0
+    batch_count = 0
+
+    try:
+        while True:
+            arrow_batch = reader.read_next_batch()
+            total_rows += arrow_batch.num_rows
+            batch_count += 1
+    except StopIteration:
+        pass
+
+    assert batch_count == 1
+    assert total_rows == dataset_size
+
+
+async def test_get_schemas_endpoint(
+    client: PlanarTestClient,
+):
+    """Test the GET /schemas endpoint."""
+    response = await client.get("/planar/v1/datasets/schemas")
+
+    assert response.status_code == 200
+    schemas = response.json()
+    assert isinstance(schemas, list)
+    assert "main" in schemas  # Default schema should exist
+
+
+async def test_list_datasets_metadata_endpoint(
+    client: PlanarTestClient,
+):
+    """Test the GET /metadata endpoint (list all datasets)."""
+    # Create a test dataset first
+    dataset_name = "test_list_datasets"
+    dataset = await PlanarDataset.create(dataset_name)
+
+    df = pl.DataFrame({"id": [1, 2, 3], "name": ["a", "b", "c"]})
+    await dataset.write(df)
+
+    response = await client.get("/planar/v1/datasets/metadata")
+
+    assert response.status_code == 200
+    datasets = response.json()
+    assert isinstance(datasets, list)
+
+    # Find our test dataset
+    test_dataset = next((d for d in datasets if d["name"] == dataset_name), None)
+    assert test_dataset is not None
+    assert test_dataset["row_count"] == 3
+    assert "id" in test_dataset["table_schema"]
+    assert "name" in test_dataset["table_schema"]
+
+
+async def test_list_datasets_metadata_with_pagination(
+    client: PlanarTestClient,
+):
+    """Test the GET /metadata endpoint with pagination parameters."""
+    response = await client.get(
+        "/planar/v1/datasets/metadata",
+        params={"limit": 5, "offset": 0, "schema_name": "main"},
+    )
+
+    assert response.status_code == 200
+    datasets = response.json()
+    assert isinstance(datasets, list)
+    assert len(datasets) <= 5  # Should respect limit
+
+
+async def test_get_dataset_metadata_endpoint(
+    client: PlanarTestClient,
+):
+    """Test the GET /metadata/{dataset_name} endpoint."""
+    dataset_name = "test_single_metadata"
+    dataset = await PlanarDataset.create(dataset_name)
+
+    df = pl.DataFrame(
+        {
+            "id": [1, 2, 3, 4, 5],
+            "value": ["apple", "banana", "cherry", "date", "elderberry"],
+        }
+    )
+    await dataset.write(df)
+
+    response = await client.get(f"/planar/v1/datasets/metadata/{dataset_name}")
+
+    assert response.status_code == 200
+    metadata = response.json()
+    assert metadata["name"] == dataset_name
+    assert metadata["row_count"] == 5
+    assert "id" in metadata["table_schema"]
+    assert "value" in metadata["table_schema"]
+
+
+async def test_get_dataset_metadata_not_found(
+    client: PlanarTestClient,
+):
+    """Test the GET /metadata/{dataset_name} endpoint with non-existent dataset."""
+    response = await client.get("/planar/v1/datasets/metadata/nonexistent_dataset")
+
+    assert response.status_code == 404
+    error = response.json()
+    assert "not found" in error["detail"].lower()
+
+
+async def test_download_dataset_endpoint(
+    client: PlanarTestClient,
+):
+    """Test the GET /content/{dataset_name}/download endpoint."""
+    dataset_name = "test_download"
+    dataset = await PlanarDataset.create(dataset_name)
+
+    df = pl.DataFrame({"id": [1, 2, 3], "value": ["x", "y", "z"]})
+    await dataset.write(df)
+
+    response = await client.get(f"/planar/v1/datasets/content/{dataset_name}/download")
+
+    assert response.status_code == 200
+    assert response.headers["content-type"] == "application/x-parquet"
+    assert f"{dataset_name}.parquet" in response.headers.get("content-disposition", "")
+
+    # Verify we get valid parquet content
+    content = await response.aread()
+    assert len(content) > 0
+
+    # Verify it's valid parquet by reading it back
+    import pyarrow.parquet as pq
+
+    parquet_buffer = pa.py_buffer(content)
+    table = pq.read_table(parquet_buffer)
+    assert table.num_rows == 3
+    assert table.num_columns == 2
+
+
+async def test_download_dataset_not_found(
+    client: PlanarTestClient,
+):
+    """Test the GET /content/{dataset_name}/download endpoint with non-existent dataset."""
+    response = await client.get(
+        "/planar/v1/datasets/content/nonexistent_dataset/download"
+    )
+
+    assert response.status_code == 404
+    error = response.json()
+    assert "not found" in error["detail"].lower()
+
+
+async def test_stream_arrow_dataset_not_found(
+    client: PlanarTestClient,
+):
+    """Test the GET /content/{dataset_name}/arrow-stream endpoint with non-existent dataset."""
+    response = await client.get(
+        "/planar/v1/datasets/content/nonexistent_dataset/arrow-stream"
+    )
+
+    assert response.status_code == 404
+    error = response.json()
+    assert "not found" in error["detail"].lower()
+
+
+async def test_get_dataset_metadata_empty_dataset(
+    client: PlanarTestClient,
+):
+    """Test GET /metadata/{dataset_name} with empty dataset."""
+    dataset_name = "test_empty_metadata"
+    dataset = await PlanarDataset.create(dataset_name)
+
+    # Create empty dataset
+    df = pl.DataFrame(
+        {"id": [], "value": []}, schema={"id": pl.Int64, "value": pl.Utf8}
+    )
+    await dataset.write(df)
+
+    response = await client.get(f"/planar/v1/datasets/metadata/{dataset_name}")
+    assert response.status_code == 200
+
+    metadata = response.json()
+    assert metadata["name"] == dataset_name
+    assert metadata["row_count"] == 0
+    assert "id" in metadata["table_schema"]
+    assert "value" in metadata["table_schema"]
+
+
+async def test_list_datasets_metadata_empty_dataset(
+    client: PlanarTestClient,
+):
+    """Test GET /metadata with empty dataset in the list."""
+    dataset_name = "test_empty_in_list"
+    dataset = await PlanarDataset.create(dataset_name)
+
+    # Create empty dataset
+    df = pl.DataFrame(
+        {"id": [], "value": []}, schema={"id": pl.Int64, "value": pl.Utf8}
+    )
+    await dataset.write(df)
+
+    response = await client.get("/planar/v1/datasets/metadata")
+    assert response.status_code == 200
+
+    datasets = response.json()
+    empty_dataset = next((d for d in datasets if d["name"] == dataset_name), None)
+    assert empty_dataset is not None
+    assert empty_dataset["row_count"] == 0
+
+
+async def test_download_empty_dataset(
+    client: PlanarTestClient,
+):
+    """Test GET /content/{dataset_name}/download with empty dataset."""
+    dataset_name = "test_empty_download"
+    dataset = await PlanarDataset.create(dataset_name)
+
+    # Create empty dataset
+    df = pl.DataFrame(
+        {"id": [], "value": []}, schema={"id": pl.Int64, "value": pl.Utf8}
+    )
+    await dataset.write(df)
+
+    response = await client.get(f"/planar/v1/datasets/content/{dataset_name}/download")
+    assert response.status_code == 200
+    assert response.headers["content-type"] == "application/x-parquet"
+    assert f"{dataset_name}.parquet" in response.headers.get("content-disposition", "")
+
+    # Verify we get valid parquet content (even if empty)
+    content = await response.aread()
+    assert len(content) > 0  # Should have parquet metadata even for empty data
+
+    # Verify it's valid parquet by reading it back
+    import pyarrow.parquet as pq
+
+    parquet_buffer = pa.py_buffer(content)
+    table = pq.read_table(parquet_buffer)
+    assert table.num_rows == 0
+    assert table.num_columns == 2  # id and value columns
+    assert table.schema.field("id").type == pa.int64()
+    assert table.schema.field("value").type == pa.string()

planar/routers/test_workflow_router.py

@@ -102,7 +102,7 @@ class FileProcessingResult(BaseModel):
     file_id: UUID = Field(description="ID of the processed file")
 
 
-@workflow(name="test_file_processing_workflow")
+@workflow(name="test_file_processing_workflow", is_interactive=False)
 async def file_processing_workflow(file: PlanarFile):
     """
     Workflow that processes a text file and returns basic information about it.
@@ -284,6 +284,9 @@ async def test_list_workflows(client: PlanarTestClient):
     # Verify that the file workflow input schema includes file parameter
     assert "file" in file_workflow["input_schema"]["properties"]
 
+    # Verify that we propagated the `is_interactive` flag
+    assert file_workflow["is_interactive"] is False
+
     # Verify run statistics are present
     assert "total_runs" in expense_workflow
     assert "run_statuses" in expense_workflow
@@ -537,3 +540,25 @@ async def test_get_compute_step(
     data = resp.json()
     assert "meta" in data
     assert data["meta"] is None
+
+
+async def test_list_interactive_workflow(app: PlanarApp, client: PlanarTestClient):
+    """
+    We propagate interactive workflows all the way to the `/workflows` endpoint.
+    """
+
+    # This is here rather than at the top bc it's not registered as part of the `app` fixture.
+    @workflow(name="interactive_workflow", is_interactive=True)
+    async def interactive_workflow():
+        pass
+
+    app.register_workflow(interactive_workflow)
+
+    response = await client.get("/planar/v1/workflows/")
+    assert response.status_code == 200
+
+    data = response.json()
+    expense_workflow = next(
+        item for item in data["items"] if item["name"] == "interactive_workflow"
+    )
+    assert expense_workflow["is_interactive"] is True

planar/routers/workflow.py

@@ -217,6 +217,7 @@ def create_workflow_router(
                         }
                     ),
                     durations=duration_stats,
+                    is_interactive=workflow.is_interactive,
                 )
             )
 
@@ -257,6 +258,7 @@ def create_workflow_router(
                 **{status.value: count for status, count in run_statuses.items()}
             ),
             durations=duration_stats,
+            is_interactive=wf.is_interactive,
         )
 
     @router.get("/{workflow_name}/runs", response_model=WorkflowRunList)

planar/security/authorization.py

@@ -87,12 +87,23 @@ class RuleAction(str, Enum):
     RULE_SIMULATE = "Rule::Simulate"
 
 
+class DatasetAction(str, Enum):
+    """Actions that can be performed on datasets."""
+
+    DATASET_LIST_SCHEMAS = "Dataset::ListSchemas"
+    DATASET_LIST = "Dataset::List"
+    DATASET_VIEW_DETAILS = "Dataset::ViewDetails"
+    DATASET_STREAM_CONTENT = "Dataset::StreamContent"
+    DATASET_DOWNLOAD = "Dataset::Download"
+
+
 class ResourceType(str, Enum):
     PRINCIPAL = "Principal"
     WORKFLOW = "Workflow"
     ENTITY = "Entity"
     AGENT = "Agent"
     Rule = "Rule"
+    DATASET = "Dataset"
 
 
 class EntityIdentifier(TypedDict):
@@ -129,7 +140,12 @@ class RuleResource:
     rule_name: str | None = None
 
 
-ResourceDescriptor = AgentResource | WorkflowResource | RuleResource
+@dataclass(frozen=True, slots=True)
+class DatasetResource:
+    dataset_name: str | None = None
+
+
+ResourceDescriptor = AgentResource | WorkflowResource | RuleResource | DatasetResource
 
 
 class CedarEntity(BaseModel):
@@ -209,6 +225,15 @@ class CedarEntity(BaseModel):
             resource_attributes={"rule_name": rule_name},
         )
 
+    @staticmethod
+    def from_dataset(dataset_name: str | None) -> "CedarEntity":
+        """Create a CedarEntity instance from dataset data"""
+        return CedarEntity(
+            resource_type=ResourceType.DATASET,
+            resource_key="dataset_name",
+            resource_attributes={"dataset_name": dataset_name},
+        )
+
 
 class PolicyService:
     """Service for managing and evaluating Authorization policies."""
@@ -272,7 +297,7 @@ class PolicyService:
     def is_allowed(
         self,
         principal: CedarEntity,
-        action: str | WorkflowAction | AgentAction | RuleAction,
+        action: str | WorkflowAction | AgentAction | RuleAction | DatasetAction,
         resource: CedarEntity,
     ) -> bool:
         """Check if the principal is permitted to perform the action on the resource.
@@ -294,6 +319,7 @@ class PolicyService:
             isinstance(action, WorkflowAction)
             or isinstance(action, AgentAction)
             or isinstance(action, RuleAction)
+            or isinstance(action, DatasetAction)
         ):
             action = f'Action::"{action.value}"'
         else:
@@ -346,7 +372,7 @@ class PolicyService:
 
 def validate_authorization_for(
     resource_descriptor: ResourceDescriptor,
-    action: WorkflowAction | AgentAction | RuleAction,
+    action: WorkflowAction | AgentAction | RuleAction | DatasetAction,
 ):
     authz_service = get_policy_service()
 
@@ -363,6 +389,8 @@ def validate_authorization_for(
             entity = CedarEntity.from_agent(resource_descriptor.id)
         case RuleAction() if isinstance(resource_descriptor, RuleResource):
             entity = CedarEntity.from_rule(resource_descriptor.rule_name)
+        case DatasetAction() if isinstance(resource_descriptor, DatasetResource):
+            entity = CedarEntity.from_dataset(resource_descriptor.dataset_name)
         case _:
             raise ValueError(
                 f"Invalid resource descriptor {type(resource_descriptor).__name__} for action {action}"

planar/security/default_policies.cedar

@@ -74,4 +74,29 @@ permit (
     principal,
     action == Action::"Rule::Simulate",
     resource
+);
+permit (
+    principal,
+    action == Action::"Dataset::ListSchemas",
+    resource
+);
+permit (
+    principal,
+    action == Action::"Dataset::List",
+    resource
+);
+permit (
+    principal,
+    action == Action::"Dataset::ViewDetails",
+    resource
+);
+permit (
+    principal,
+    action == Action::"Dataset::StreamContent",
+    resource
+);
+permit (
+    principal,
+    action == Action::"Dataset::Download",
+    resource
 );

planar/testing/fixtures.py

@@ -16,6 +16,8 @@ Usage in external projects:
 
 Available fixtures:
 - storage: In-memory file storage for tests
+- data_config: Test data configuration with SQLite catalog and local storage
+- app_with_data: PlanarApp instance with data configuration
 - tmp_db_url: Parametrized database URL (SQLite/PostgreSQL)
 - session: Database session
 - client: Planar test client
@@ -33,8 +35,11 @@ from pathlib import Path
 
 import pytest
 
+from planar.app import PlanarApp
 from planar.config import load_config
+from planar.data.config import DataConfig, SQLiteCatalogConfig
 from planar.db import DatabaseManager, new_session
+from planar.files.storage.config import LocalDirectoryConfig
 from planar.files.storage.context import set_storage
 from planar.logging import set_context_metadata
 from planar.object_registry import ObjectRegistry
@@ -114,6 +119,31 @@ async def storage():
     yield storage
 
 
+@pytest.fixture()
+def data_config(tmp_path):
+    """Create a test data configuration."""
+    data_dir = tmp_path / "data"
+    data_dir.mkdir(exist_ok=True)
+
+    catalog_path = data_dir / "test.sqlite"
+    storage_path = data_dir / "ducklake_files"
+    storage_path.mkdir(exist_ok=True)
+
+    return DataConfig(
+        catalog=SQLiteCatalogConfig(type="sqlite", path=str(catalog_path)),
+        storage=LocalDirectoryConfig(backend="localdir", directory=str(storage_path)),
+    )
+
+
+@pytest.fixture(name="app_with_data")
+def app_with_data_fixture(data_config):
+    """Create a PlanarApp with data configuration."""
+    app = PlanarApp()
+    # Add data config to the app's config
+    app.config.data = data_config
+    return app
+
+
 @pytest.fixture()
 def tmp_sqlite_url(tmp_db_path: str):
     return f"sqlite+aiosqlite:///{tmp_db_path}"

planar/testing/planar_test_client.py

@@ -53,5 +53,5 @@ async def wait_all_event_loop_tasks():
             break
         try:
             await asyncio.gather(*other_tasks)
-        except asyncio.CancelledError:
+        except (asyncio.CancelledError, Exception):
             pass

planar/workflows/decorators.py

@@ -78,7 +78,7 @@ def step(
     return decorator
 
 
-def workflow(*, name: str | None = None):
+def workflow(*, name: str | None = None, is_interactive: bool = False):
     """
     Decorator to define a workflow.
 
@@ -177,6 +177,7 @@ def workflow(*, name: str | None = None):
             start_step=start_workflow_step,
             wait_for_completion=wait_for_completion,
             wrapped_fn=run_workflow,
+            is_interactive=is_interactive,
         )
 
         return wf_wrapper

planar/workflows/wrappers.py

@@ -33,6 +33,7 @@ class WorkflowWrapper(Wrapper[P, T, U, R]):
     start: Callable[P, Coroutine[T, U, Workflow]]
     start_step: Callable[P, Coroutine[T, U, UUID]]
     wait_for_completion: Callable[[UUID], Coroutine[T, U, R]]
+    is_interactive: bool
 
 
 @dataclass(kw_only=True)