stack/dev/scripts/setup_unity_catalog.py

"""Setup Unity Catalog structure from aco.table schemas.

This script creates the catalog, schemas, and table definitions in
Databricks Unity Catalog to match the canonical aco.table structure.

Usage::

    # Dry run to see what would be created
    uv run python dev/scripts/setup_unity_catalog.py --dry-run

    # Actually create the catalog structure
    uv run python dev/scripts/setup_unity_catalog.py --catalog aco --token <pat>

    # Create with custom workspace
    uv run python dev/scripts/setup_unity_catalog.py \
        --workspace-id 7474655000864661 \
        --catalog aco \
        --token <pat>
"""

from __future__ import annotations

import argparse
import os
import sys

from aco.lake.catalog import Catalog
from aco.lake.unity import UnityClient, setup_catalog_from_schemas


def main():
    """Setup Unity Catalog from command line arguments."""
    parser = argparse.ArgumentParser(
        description="Setup Unity Catalog structure from aco.table schemas"
    )

    parser.add_argument(
        "--workspace-id",
        default="7474655000864661",
        help="Databricks workspace ID (default: 7474655000864661)",
    )

    parser.add_argument(
        "--catalog",
        default="aco",
        help="Target catalog name (default: aco)",
    )

    parser.add_argument(
        "--token",
        help="Databricks Personal Access Token (or set DATABRICKS_TOKEN env var)",
    )

    parser.add_argument(
        "--dry-run",
        action="store_true",
        help="Show what would be created without making changes",
    )

    parser.add_argument(
        "--skip-existing",
        action="store_true",
        default=True,
        help="Skip schemas/tables that already exist (default: True)",
    )

    parser.add_argument(
        "--show-schemas",
        action="store_true",
        help="Show discovered schemas from aco.table and exit",
    )

    args = parser.parse_args()

    # Get token from args or environment
    token = args.token or os.getenv("DATABRICKS_TOKEN")
    if not token and not args.show_schemas:
        print("Error: --token required or set DATABRICKS_TOKEN environment variable")
        sys.exit(1)

    # Show schemas mode
    if args.show_schemas:
        print("Discovering schemas from aco.table...")
        print("=" * 70)

        catalog = Catalog()
        schemas = catalog.schemas()

        print(f"\nFound {len(schemas)} schemas:")
        for schema_name in schemas:
            tables = catalog.tables(schema_name)
            print(f"\n  {schema_name} ({len(tables)} tables)")
            for table_ref in tables[:3]:
                _, table = table_ref.split(".", 1)
                print(f"    - {table}")
            if len(tables) > 3:
                print(f"    ... and {len(tables) - 3} more tables")

        print("\n" + "=" * 70)
        print("\nTo create these in Unity Catalog, run without --show-schemas flag")
        return

    # Create Unity client
    print("Connecting to Unity Catalog...")
    print(f"  Workspace ID: {args.workspace_id}")
    print(f"  Catalog:      {args.catalog}")
    print()

    client = UnityClient(
        workspace_id=args.workspace_id,
        token=token,
    )

    # Test connection
    try:
        catalogs = client.list_catalogs()
        print("✓ Connected to Unity Catalog")
        print(f"  Found {len(catalogs)} existing catalogs:")
        for cat in catalogs:
            print(f"    - {cat.name}")
        print()
    except Exception as e:
        print(f"✗ Failed to connect to Unity Catalog: {e}")
        sys.exit(1)

    # Setup catalog from schemas
    if args.dry_run:
        print("=" * 70)
        print("DRY RUN MODE - No changes will be made")
        print("=" * 70)
        print()

    print("Setting up catalog structure...")
    print()

    report = setup_catalog_from_schemas(
        client=client,
        catalog_name=args.catalog,
        dry_run=args.dry_run,
        skip_existing=args.skip_existing,
    )

    # Print summary
    print()
    print("=" * 70)
    print("Setup Summary")
    print("=" * 70)

    if report["catalogs_created"]:
        print(f"\nCatalogs created: {len(report['catalogs_created'])}")
        for cat in report["catalogs_created"]:
            print(f"  ✓ {cat}")

    if report["schemas_created"]:
        print(f"\nSchemas created: {len(report['schemas_created'])}")
        for schema in report["schemas_created"]:
            print(f"  ✓ {schema}")

    if report["tables_created"]:
        print(f"\nTables created: {len(report['tables_created'])}")
        for table in report["tables_created"]:
            print(f"  ✓ {table}")

    if report["errors"]:
        print(f"\nErrors encountered: {len(report['errors'])}")
        for error in report["errors"]:
            print(f"  ✗ {error}")

    if args.dry_run:
        print("\n" + "=" * 70)
        print("DRY RUN COMPLETE - Run without --dry-run to apply changes")
        print("=" * 70)
    else:
        print("\n" + "=" * 70)
        print("✓ Setup complete")
        print("=" * 70)

        # Show how to use with IcebergContext
        print("\nNext steps:")
        print()
        print("  1. Use with IcebergContext:")
        print()
        print("     from aco.lake import IcebergContext")
        print()
        print("     ctx = IcebergContext(")
        print(
            f'         catalog_uri="https://dbc-{args.workspace_id}.cloud.databricks.com/api/2.1/unity-catalog/iceberg",'
        )
        print(f'         warehouse="{args.catalog}",')
        print('         properties={"token": "<your-token>"},')
        print("     )")
        print()
        print("  2. Run pipelines:")
        print()
        print("     from aco.lake.engine import execute")
        print("     from aco.pipe import readmissions")
        print()
        print("     results = execute(readmissions.pipeline, ctx)")


if __name__ == "__main__":
    main()