stack/notebooks/nessie_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("""
    # Nessie Tutorial - Git for Your Data Lake

    Nessie provides Git-like version control for Iceberg tables: branches, commits, merges, and time travel.

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


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

    import requests

    from conf import cfg

    NESSIE_API = f"{cfg.services.nessie}/api/v2"
    return NESSIE_API, json, requests


@app.cell(hide_code=True)
def _(mo):
    mo.md("""
    ## 1. Check Nessie Configuration
    """)
    return


@app.cell(hide_code=True)
def _(NESSIE_API, requests):
    config = requests.get(f"{NESSIE_API}/config").json()
    print(f"Nessie API Version: {config['specVersion']}")
    print(f"Default Branch: {config['defaultBranch']}")
    return


@app.cell(hide_code=True)
def _(mo):
    mo.md("""
    ## 2. List Branches and Tags
    """)
    return


@app.cell(hide_code=True)
def _(NESSIE_API, requests):
    refs = requests.get(f"{NESSIE_API}/trees").json()
    print("References:")
    for _r in refs.get("references", []):
        print(f"  {_r['type']:6} {_r['name']:20} @ {_r['hash'][:12]}...")
    return


@app.cell(hide_code=True)
def _(mo):
    mo.md("""
    ## 3. Get Main Branch Details
    """)
    return


@app.cell(hide_code=True)
def _(NESSIE_API, requests):
    main_resp = requests.get(f"{NESSIE_API}/trees/main").json()
    main_ref = main_resp.get("reference", main_resp)
    print(f"Branch: {main_ref['name']}")
    print(f"Hash:   {main_ref['hash']}")
    return (main_ref,)


@app.cell(hide_code=True)
def _(mo):
    mo.md("""
    ## 4. Create an Experiment Branch

    The v2 API uses:
    - Query params: `name` (new branch) and `type` (branch/tag)
    - Body: Source reference object `{"type": "BRANCH", "name": "main", "hash": "..."}`
    """)
    return


@app.cell(hide_code=True)
def _(NESSIE_API, json, main_ref, requests):
    # Create experiment branch from main
    new_branch = "experiment"

    resp = requests.post(
        f"{NESSIE_API}/trees?name={new_branch}&type=branch",
        headers={"Content-Type": "application/json"},
        data=json.dumps(
            {"type": "BRANCH", "name": main_ref["name"], "hash": main_ref["hash"]}
        ),
    )

    if resp.status_code == 200:
        print(f"Created branch: {new_branch}")
        print(resp.json())
    elif resp.status_code == 409:
        print(f"Branch '{new_branch}' already exists")
    else:
        print(f"Error {resp.status_code}: {resp.text}")
    return


@app.cell(hide_code=True)
def _(mo):
    mo.md("""
    ## 5. List Tables on a Branch
    """)
    return


@app.cell(hide_code=True)
def _(NESSIE_API, requests):
    entries = requests.get(f"{NESSIE_API}/trees/main/entries").json()
    print("Tables on main:")
    for _e in entries.get("entries", []):
        _key = ".".join(_e["name"]["elements"])
        print(f"  {_e['type']:15} {_key}")
    if not entries.get("entries"):
        print("  (none yet)")
    return


@app.cell(hide_code=True)
def _(mo):
    mo.md("""
    ## 6. View Commit History
    """)
    return


@app.cell(hide_code=True)
def _(NESSIE_API, requests):
    log = requests.get(f"{NESSIE_API}/trees/main/history").json()
    print("Recent commits:")
    for _e in log.get("logEntries", [])[:5]:
        _c = _e["commitMeta"]
        _author = _c.get("authors", ["?"])[0]
        print(f"  {_c['hash'][:12]} {_c.get('message', '')[:40]} ({_author})")
    return


@app.cell(hide_code=True)
def _(mo):
    mo.md("""
    ## 7. Query with Trino
    """)
    return


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

    conn = trino.dbapi.connect(
        host="trino",
        port=8080,
        user="kert",
        catalog="iceberg",
    )

    def query(sql):
        cur = conn.cursor()
        cur.execute(sql)
        return cur.fetchall()

    print("Connected to Trino")
    return (query,)


@app.cell(hide_code=True)
def _(query):
    schemas = query("SHOW SCHEMAS IN iceberg")
    print("Schemas:", [s[0] for s in schemas])
    return


@app.cell(hide_code=True)
def _(mo):
    mo.md("""
    ## 8. Create Sample Data
    """)
    return


@app.cell(hide_code=True)
def _(query):
    query("CREATE SCHEMA IF NOT EXISTS iceberg.tutorial")
    print("Schema ready")
    return


@app.cell(hide_code=True)
def _(query):
    query("""
        CREATE TABLE IF NOT EXISTS iceberg.tutorial.events (
            id VARCHAR,
            ts TIMESTAMP(6),
            event VARCHAR
        ) WITH (format = 'PARQUET')
    """)
    print("Table ready")
    return


@app.cell(hide_code=True)
def _(query):
    query("""
        INSERT INTO iceberg.tutorial.events VALUES
        ('e1', CURRENT_TIMESTAMP, 'click'),
        ('e2', CURRENT_TIMESTAMP, 'view')
    """)
    print("Data inserted")
    return


@app.cell(hide_code=True)
def _(query):
    rows = query("SELECT * FROM iceberg.tutorial.events")
    print("Events:")
    for _r in rows:
        print(f"  {_r}")
    return


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

    Query specific branches with `FOR VERSION AS OF 'branch'`
    """)
    return


@app.cell(hide_code=True)
def _(query):
    count = query(
        "SELECT count(*) FROM iceberg.tutorial.events FOR VERSION AS OF 'main'"
    )
    print(f"Events on main: {count[0][0]}")
    return


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


@app.cell(hide_code=True)
def _(NESSIE_API, requests):
    def delete_branch(name):
        ref = requests.get(f"{NESSIE_API}/trees/{name}").json()
        _hash = ref.get("reference", ref).get("hash")
        if not _hash:
            print(f"Branch {name} not found")
            return
        resp = requests.delete(
            f"{NESSIE_API}/trees/{name}", headers={"Expected-Hash": _hash}
        )
        if resp.status_code == 204:
            print(f"Deleted: {name}")
        else:
            print(f"Error: {resp.status_code}")

    # Uncomment to delete:
    # delete_branch("experiment")
    print("delete_branch() ready")
    return


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

    | Operation | API |
    |-----------|-----|
    | List refs | `GET /trees` |
    | Get ref | `GET /trees/{name}` |
    | Create branch | `POST /trees?name=X&type=branch` + body |
    | Delete ref | `DELETE /trees/{name}` + Expected-Hash header |
    | List tables | `GET /trees/{ref}/entries` |
    | History | `GET /trees/{ref}/history` |
    """)
    return


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