stack/dev/scripts/unity_catalog_example.py

"""Example usage of Unity Catalog integration.

This demonstrates the full workflow:
1. Connect to Unity Catalog
2. Create catalogs, schemas, tables
3. Use with IcebergContext for data access
4. Run pipelines against Unity Catalog

Prerequisites:
- Databricks workspace with Unity Catalog enabled
- Personal Access Token (PAT) with catalog permissions
- Set DATABRICKS_TOKEN environment variable

Usage::

    export DATABRICKS_TOKEN="<your-pat>"
    uv run python dev/scripts/unity_catalog_example.py
"""

from __future__ import annotations

import os

from aco.lake import IcebergContext, UnityClient
from aco.lake.catalog import Catalog


def example_unity_client_basics():
    """Example 1: Unity Client basics - catalogs, schemas, tables."""
    print("=" * 70)
    print("Example 1: Unity Client Basics")
    print("=" * 70)
    print()

    # Connect to Unity Catalog
    token = os.getenv("DATABRICKS_TOKEN")
    if not token:
        print("⚠ Set DATABRICKS_TOKEN environment variable to run this example")
        return

    client = UnityClient(
        workspace_id="7474655000864661",
        token=token,
    )

    # List existing catalogs
    print("1. Listing existing catalogs")
    print("-" * 70)
    catalogs = client.list_catalogs()
    print(f"Found {len(catalogs)} catalogs:")
    for cat in catalogs:
        print(f"  - {cat.name}: {cat.comment or '(no description)'}")
    print()

    # Create a new catalog
    print("2. Creating new catalog")
    print("-" * 70)
    try:
        catalog = client.create_catalog(
            name="aco_dev",
            comment="ACO analytics development catalog",
        )
        print(f"✓ Created catalog: {catalog.name}")
        print(f"  Owner: {catalog.owner}")
        print(f"  Created: {catalog.created_at}")
    except Exception as e:
        print(f"Catalog may already exist: {e}")
    print()

    # Create a schema
    print("3. Creating schema")
    print("-" * 70)
    try:
        schema = client.create_schema(
            catalog_name="aco_dev",
            schema_name="core",
            comment="Core healthcare data tables",
        )
        print(f"✓ Created schema: {schema.full_name}")
    except Exception as e:
        print(f"Schema may already exist: {e}")
    print()

    # List schemas in catalog
    print("4. Listing schemas")
    print("-" * 70)
    schemas = client.list_schemas("aco_dev")
    print(f"Found {len(schemas)} schemas in aco_dev:")
    for schema in schemas:
        print(f"  - {schema.name}: {schema.comment or '(no description)'}")
    print()


def example_schema_mapping():
    """Example 2: Schema mapping for enterprise tables."""
    print("=" * 70)
    print("Example 2: Schema Mapping (Enterprise Tables)")
    print("=" * 70)
    print()

    # Create catalog with schema mapping
    print("1. Setting up catalog with enterprise schema mapping")
    print("-" * 70)

    catalog = Catalog(
        schema_map={
            # Map Unity Catalog names to canonical aco.table names
            "aco.encounters": "core.encounter",
            "aco.patients": "core.patient",
            "aco.claims": "core.medical_claim",
        },
        column_map={
            "aco.encounters": {
                # Map Unity Catalog column names to canonical names
                "encntr_id": "encounter_id",
                "mbr_id": "person_id",
                "encntr_type": "encounter_type",
            },
            "aco.patients": {
                "mbr_id": "patient_id",
                "birth_dt": "birth_date",
                "death_dt": "death_date",
            },
        },
    )

    # Test mapping
    print("✓ Catalog configured with schema mappings")
    print()

    print("2. Testing table reference mapping")
    print("-" * 70)
    test_refs = [
        "aco.encounters",
        "aco.patients",
        "core.condition",  # Unmapped, passes through
    ]

    for ref in test_refs:
        physical = catalog.physical_table_ref(ref)
        print(f"  {ref:25s} → {physical}")
    print()

    print("3. Testing column name mapping")
    print("-" * 70)
    test_cols = [
        ("aco.encounters", "encntr_id"),
        ("aco.encounters", "mbr_id"),
        ("aco.encounters", "unmapped_col"),
        ("aco.patients", "birth_dt"),
    ]

    for table_ref, col in test_cols:
        physical_col = catalog.physical_column_name(table_ref, col)
        print(f"  {table_ref}.{col:20s} → {physical_col}")
    print()


def example_iceberg_context():
    """Example 3: Using Unity Catalog with IcebergContext."""
    print("=" * 70)
    print("Example 3: IcebergContext with Unity Catalog")
    print("=" * 70)
    print()

    token = os.getenv("DATABRICKS_TOKEN")
    if not token:
        print("⚠ Set DATABRICKS_TOKEN environment variable to run this example")
        return

    # Create IcebergContext pointing to Unity Catalog
    print("1. Creating IcebergContext for Unity Catalog")
    print("-" * 70)

    workspace_id = "7474655000864661"

    ctx = IcebergContext(
        catalog_uri=f"https://dbc-{workspace_id}.cloud.databricks.com/api/2.1/unity-catalog/iceberg",
        warehouse="aco_dev",
        properties={
            "token": token,
            # Unity Catalog specific properties
            "header.X-Databricks-Cluster-Id": "auto",
        },
    )

    print("✓ IcebergContext configured")
    print(f"  Catalog URI: {ctx.catalog_uri}")
    print(f"  Warehouse: {ctx.warehouse}")
    print()

    print("2. Using context to load tables")
    print("-" * 70)
    print("  (Would load tables from Unity Catalog if they exist)")
    print("  Example: df = ctx.load('core.encounter')")
    print()


def example_pipeline_execution():
    """Example 4: Running pipelines against Unity Catalog."""
    print("=" * 70)
    print("Example 4: Pipeline Execution")
    print("=" * 70)
    print()

    token = os.getenv("DATABRICKS_TOKEN")
    if not token:
        print("⚠ Set DATABRICKS_TOKEN environment variable to run this example")
        return

    workspace_id = "7474655000864661"

    # Create context with schema mapping
    print("1. Setting up IcebergContext with schema mapping")
    print("-" * 70)

    catalog = Catalog(
        schema_map={
            "aco_prod.encounters": "core.encounter",
            "aco_prod.patients": "core.patient",
        }
    )

    IcebergContext(
        catalog_uri=f"https://dbc-{workspace_id}.cloud.databricks.com/api/2.1/unity-catalog/iceberg",
        warehouse="aco_dev",
        catalog=catalog,
        properties={"token": token},
    )

    print("✓ Context configured with schema mapping")
    print()

    print("2. Executing pipeline (example - would run if tables exist)")
    print("-" * 70)
    print("  from aco.pipe import readmissions")
    print("  results = execute(readmissions.pipeline, ctx)")
    print()
    print("  This would:")
    print("  - Read input tables from Unity Catalog")
    print("  - Execute narwhals transforms locally")
    print("  - Optionally write results back to Unity Catalog")
    print()


def main():
    """Run all examples."""
    examples = [
        example_unity_client_basics,
        example_schema_mapping,
        example_iceberg_context,
        example_pipeline_execution,
    ]

    for i, example_fn in enumerate(examples, 1):
        try:
            example_fn()
        except Exception as e:
            print(f"✗ Example {i} failed: {e}")
            print()

        if i < len(examples):
            print("\n" + "=" * 70)
            print()


if __name__ == "__main__":
    main()