stack/notebooks/polaris_tutorial.py

import marimo

__generated_with = "0.19.8"
app = marimo.App(width="medium")


@app.cell(hide_code=True)
def _():
    import marimo as mo

    mo.md("""
    # Apache Polaris Tutorial - Iceberg Catalog with Governance

    Polaris is an open-source Iceberg catalog with Unity Catalog-like governance features:
    RBAC, multi-catalog management, and fine-grained access control.

    This notebook walks through the key concepts and operations.
    """)
    return (mo,)


@app.cell(hide_code=True)
def _():
    import json
    import os

    import requests

    from conf import cfg

    POLARIS_API = cfg.services.polaris
    CLIENT_ID = "root"
    CLIENT_SECRET = os.environ.get("POLARIS_ROOT_SECRET", "polaris_root_secret")
    return POLARIS_API, CLIENT_ID, CLIENT_SECRET, json, requests, os


@app.cell(hide_code=True)
def _(mo):
    mo.md("""
    ## 1. Authenticate with Polaris

    Polaris uses OAuth2 client credentials flow. Get an access token first.
    """)
    return


@app.cell(hide_code=True)
def _(POLARIS_API, CLIENT_ID, CLIENT_SECRET, requests):
    token_resp = requests.post(
        f"{POLARIS_API}/api/catalog/v1/oauth/tokens",
        headers={"Content-Type": "application/x-www-form-urlencoded"},
        data={
            "grant_type": "client_credentials",
            "client_id": CLIENT_ID,
            "client_secret": CLIENT_SECRET,
            "scope": "PRINCIPAL_ROLE:ALL",
        },
    )

    if token_resp.status_code == 200:
        _token_data = token_resp.json()
        access_token = _token_data["access_token"]
        print(f"Token obtained (expires in {_token_data['expires_in']}s)")
        print(f"Token prefix: {access_token[:50]}...")
    else:
        print(f"Auth failed: {token_resp.status_code}")
        print(token_resp.text)
        access_token = None
    return (access_token,)


@app.cell(hide_code=True)
def _(mo):
    mo.md("""
    ## 2. List Catalogs

    Catalogs are top-level containers for namespaces and tables.
    """)
    return


@app.cell(hide_code=True)
def _(POLARIS_API, access_token, requests):
    headers = {"Authorization": f"Bearer {access_token}"}
    catalogs_resp = requests.get(
        f"{POLARIS_API}/api/management/v1/catalogs", headers=headers
    )

    if catalogs_resp.status_code == 200:
        _catalogs = catalogs_resp.json().get("catalogs", [])
        print("Catalogs:")
        for _c in _catalogs:
            print(f"  {_c['name']} ({_c['type']})")
        if not _catalogs:
            print("  (none yet)")
    else:
        print(f"Error: {catalogs_resp.status_code}")
    return (headers,)


@app.cell(hide_code=True)
def _(mo):
    mo.md("""
    ## 3. Create a Catalog

    Create a catalog with S3 storage configuration.
    """)
    return


@app.cell(hide_code=True)
def _(POLARIS_API, headers, json, requests):
    catalog_name = "analytics"

    create_resp = requests.post(
        f"{POLARIS_API}/api/management/v1/catalogs",
        headers={**headers, "Content-Type": "application/json"},
        data=json.dumps(
            {
                "name": catalog_name,
                "type": "INTERNAL",
                "properties": {
                    "default-base-location": f"s3://polaris/{catalog_name}/"
                },
                "storageConfigInfo": {
                    "storageType": "S3",
                    "allowedLocations": ["s3://polaris/", "s3://lakehouse/"],
                    "s3": {
                        "region": "us-east-1",
                        "endpoint": cfg.services.rustfs,
                        "pathStyleAccess": True,
                    },
                },
            }
        ),
    )

    if create_resp.status_code == 200:
        print(f"Created catalog: {catalog_name}")
        print(create_resp.json())
    elif create_resp.status_code == 409:
        print(f"Catalog '{catalog_name}' already exists")
    else:
        print(f"Error {create_resp.status_code}: {create_resp.text}")
    return (catalog_name,)


@app.cell(hide_code=True)
def _(mo):
    mo.md("""
    ## 4. Get Catalog Details
    """)
    return


@app.cell(hide_code=True)
def _(POLARIS_API, catalog_name, headers, requests):
    cat_resp = requests.get(
        f"{POLARIS_API}/api/management/v1/catalogs/{catalog_name}", headers=headers
    )

    if cat_resp.status_code == 200:
        _cat = cat_resp.json()
        print(f"Catalog: {_cat['name']}")
        print(f"Type: {_cat['type']}")
        print(f"Properties: {_cat.get('properties', {})}")
    else:
        print(f"Error: {cat_resp.status_code}")
    return


@app.cell(hide_code=True)
def _(mo):
    mo.md("""
    ## 5. List Principals

    Principals are users or service accounts that can access catalogs.
    """)
    return


@app.cell(hide_code=True)
def _(POLARIS_API, headers, requests):
    principals_resp = requests.get(
        f"{POLARIS_API}/api/management/v1/principals", headers=headers
    )

    if principals_resp.status_code == 200:
        _principals = principals_resp.json().get("principals", [])
        print("Principals:")
        for _p in _principals:
            print(f"  {_p['name']} ({_p['type']})")
    else:
        print(f"Error: {principals_resp.status_code}")
    return


@app.cell(hide_code=True)
def _(mo):
    mo.md("""
    ## 6. List Principal Roles

    Principal roles group permissions that can be assigned to principals.
    """)
    return


@app.cell(hide_code=True)
def _(POLARIS_API, headers, requests):
    roles_resp = requests.get(
        f"{POLARIS_API}/api/management/v1/principal-roles", headers=headers
    )

    if roles_resp.status_code == 200:
        _roles = roles_resp.json().get("roles", [])
        print("Principal Roles:")
        for _r in _roles:
            print(f"  {_r['name']}")
        if not _roles:
            print("  (none yet)")
    else:
        print(f"Error: {roles_resp.status_code}")
    return


@app.cell(hide_code=True)
def _(mo):
    mo.md("""
    ## 7. Iceberg REST Catalog API

    Polaris implements the Iceberg REST Catalog specification.
    """)
    return


@app.cell(hide_code=True)
def _(POLARIS_API, catalog_name, headers, requests):
    # Get catalog config via Iceberg REST API
    config_resp = requests.get(
        f"{POLARIS_API}/api/catalog/v1/config",
        headers=headers,
        params={"warehouse": catalog_name},
    )

    if config_resp.status_code == 200:
        _config = config_resp.json()
        print("Iceberg Catalog Config:")
        print(f"  Overrides: {_config.get('overrides', {})}")
        print(f"  Defaults: {_config.get('defaults', {})}")
    else:
        print(f"Error: {config_resp.status_code}")
        print(config_resp.text)
    return


@app.cell(hide_code=True)
def _(mo):
    mo.md("""
    ## 8. List Namespaces

    Namespaces are like schemas/databases in the catalog.
    """)
    return


@app.cell(hide_code=True)
def _(POLARIS_API, catalog_name, headers, requests):
    ns_resp = requests.get(
        f"{POLARIS_API}/api/catalog/v1/{catalog_name}/namespaces", headers=headers
    )

    if ns_resp.status_code == 200:
        _namespaces = ns_resp.json().get("namespaces", [])
        print(f"Namespaces in {catalog_name}:")
        for _ns in _namespaces:
            print(f"  {'.'.join(_ns)}")
        if not _namespaces:
            print("  (none yet)")
    else:
        print(f"Error: {ns_resp.status_code}")
        print(ns_resp.text)
    return


@app.cell(hide_code=True)
def _(mo):
    mo.md("""
    ## 9. Create a Namespace
    """)
    return


@app.cell(hide_code=True)
def _(POLARIS_API, catalog_name, headers, json, requests):
    ns_name = "tutorial"

    create_ns_resp = requests.post(
        f"{POLARIS_API}/api/catalog/v1/{catalog_name}/namespaces",
        headers={**headers, "Content-Type": "application/json"},
        data=json.dumps(
            {
                "namespace": [ns_name],
                "properties": {"description": "Tutorial namespace"},
            }
        ),
    )

    if create_ns_resp.status_code == 200:
        print(f"Created namespace: {ns_name}")
    elif create_ns_resp.status_code == 409:
        print(f"Namespace '{ns_name}' already exists")
    else:
        print(f"Error {create_ns_resp.status_code}: {create_ns_resp.text}")
    return (ns_name,)


@app.cell(hide_code=True)
def _(mo):
    mo.md("""
    ## 10. Delete a Catalog

    To clean up, you can delete catalogs (must be empty first).
    """)
    return


@app.cell(hide_code=True)
def _(POLARIS_API, headers, requests):
    def delete_catalog(name):
        resp = requests.delete(
            f"{POLARIS_API}/api/management/v1/catalogs/{name}", headers=headers
        )
        if resp.status_code == 204:
            print(f"Deleted: {name}")
        elif resp.status_code == 409:
            print(f"Catalog '{name}' is not empty")
        else:
            print(f"Error: {resp.status_code} - {resp.text}")

    # Uncomment to delete:
    # delete_catalog("analytics")
    print("delete_catalog() ready")
    return


@app.cell(hide_code=True)
def _(mo):
    mo.md("""
    ## Summary

    | Operation | API |
    |-----------|-----|
    | Get token | `POST /api/catalog/v1/oauth/tokens` |
    | List catalogs | `GET /api/management/v1/catalogs` |
    | Create catalog | `POST /api/management/v1/catalogs` |
    | Delete catalog | `DELETE /api/management/v1/catalogs/{name}` |
    | List principals | `GET /api/management/v1/principals` |
    | List roles | `GET /api/management/v1/principal-roles` |
    | List namespaces | `GET /api/catalog/v1/{catalog}/namespaces` |
    | Create namespace | `POST /api/catalog/v1/{catalog}/namespaces` |

    ## Polaris vs Nessie

    | Feature | Nessie | Polaris |
    |---------|--------|---------|
    | Git-like branching | Yes | No |
    | Time travel | Yes | Yes |
    | RBAC/governance | Basic | Full |
    | Multi-catalog | No | Yes |
    | Unity Catalog-like | No | Yes |
    """)
    return


if __name__ == "__main__":
    app.run()