Merge pull request #1 from hotdata-dev/feat/use-runtime-workspace-selection

Use runtime workspace resolver in Marimo

Author: eddietejeda

README.md

@@ -2,6 +2,16 @@
 
 Marimo UI helpers for [Hotdata](https://hotdata.dev): run SQL from a notebook, browse catalog metadata, and render results as tables.
 
+## Features
+
+- **Workspace-aware setup** — build a `HotdataClient` from environment variables, or use `workspace_selector_from_env()` to choose a workspace interactively when no workspace is pinned.
+- **Connection health** — show a compact status callout with API, workspace, and optional sandbox context.
+- **Catalog browsing** — browse Hotdata connections, schemas, tables, and columns from Marimo UI controls.
+- **SQL editor widget** — run SQL against Hotdata, cache the latest successful result, and render results in downstream reactive cells.
+- **Native `mo.sql` engine** — register `HotdataMarimoEngine` so Marimo SQL cells can execute through a live `HotdataClient` with `engine=client`.
+- **Result display helpers** — render query results, recent results, and run history as notebook-friendly UI.
+- **Marimo UI aliases** — importing `hotdata_marimo` attaches helpers such as `mo.ui.hotdata_sql_editor` and `mo.ui.hotdata_table_browser` for discoverability.
+
 ## Install
 
 ```bash
@@ -39,13 +49,45 @@ Importing `hotdata_marimo` registers discoverability aliases on Marimo’s UI na
 
 Use `hm.connection_status(client)` (or `mo.ui.hotdata_connection_status(client)`) for a small API/workspace health callout.
 
+## Marimo SQL Cells
+
+Register the Hotdata SQL engine once during setup, then pass a `HotdataClient` to Marimo SQL cells:
+
+```python
+import hotdata_marimo as hm
+
+hm.register_hotdata_sql_engine()
+client = hm.from_env()
+```
+
+```python
+_df = mo.sql(
+    """
+    SELECT 1 AS example_value
+    """,
+    engine=client,
+)
+```
+
+The engine also exposes Hotdata catalog metadata to Marimo's data-source UI. Hotdata connections are labeled **Hotdata** in the SQL connection picker.
+
 ## Two-cell pattern
 
 Keep the editor in one cell and consume `editor.result` in another. The editor caches the last successful run so downstream cells do not re-query the API on every refresh; click **Run on Hotdata** again after you change SQL. While a query is running, a Marimo status spinner is shown.
 
 Marimo only shows **what you `return` from a cell**. Calling `mo.vstack(...)` or `hm.query_result(...)` without returning it produces no visible output.
 
-See `examples/hotdata_basic.py` for a full notebook: five Python cells (`mo.vstack` for **controls only**, then a separate cell `return hm.query_result(editor.result)` so results show immediately — **avoid** `mo.lazy` here: it only renders after the block scrolls into view, which looks like an empty cell). If Marimo shows **empty cells**, quit and remove `examples/__marimo__/` so the UI reloads from the `.py` file only.
+See `examples/demo.py` for a full runnable notebook flow.
+
+## Examples
+
+- `examples/demo.py` — end-to-end browser + editor + result rendering flow.
+
+Run:
+
+```bash
+uv run marimo edit examples/demo.py --no-token
+```
 
 ## Layout
 
@@ -58,7 +100,7 @@ This package depends on [**hotdata-runtime**](https://github.com/hotdata-dev/hot
 ```bash
 uv sync --locked
 uv run pytest
-marimo edit examples/hotdata_basic.py --no-token
+marimo edit examples/demo.py --no-token
 ```
 
 To pin **hotdata-runtime** from Git instead of the sibling path, remove the `[tool.uv.sources]` block, set the dependency line as needed, and run `uv lock` again.

examples/demo.py

@@ -0,0 +1,95 @@
+import marimo
+
+__generated_with = "0.23.5"
+app = marimo.App()
+
+
+@app.cell
+def _():
+    import os
+
+    import marimo as mo
+
+    import hotdata_marimo as hm
+
+    hm.register_hotdata_sql_engine()
+    return hm, mo, os
+
+
+@app.cell
+def _(hm, mo, os):
+    mo.stop(
+        not os.environ.get("HOTDATA_API_KEY"),
+        mo.callout(
+            mo.md(
+                "Add **HOTDATA_API_KEY** to your environment "
+                "to run this example."
+            ),
+            kind="warn",
+        ),
+    )
+    workspace = hm.workspace_selector_from_env()
+    return (workspace,)
+
+
+@app.cell
+def _(hm, workspace):
+    client = workspace.client
+    status = hm.connection_status(client)
+    browser = hm.table_browser(client)
+    editor = hm.sql_editor(
+        client,
+        default_sql="SELECT 1 AS ok",
+    )
+    recent = hm.recent_results(client, limit=20)
+    history = hm.run_history(client, limit=10)
+    return browser, client, editor, history, recent, status, workspace
+
+
+@app.cell
+def _(browser, editor, mo, recent, status, workspace):
+    return mo.vstack(
+        [
+            workspace.ui,
+            status,
+            browser.ui,
+            editor.ui,
+            recent.ui,
+        ],
+        gap=2,
+    )
+
+
+@app.cell
+def _(history):
+    return history
+
+
+@app.cell
+def _(editor, hm):
+    # Explicitly touch nested widget values so Marimo reruns this cell on clicks.
+    _run = editor.run.value
+    _rerun = editor.rerun.value
+    _clear = editor.clear.value
+    return hm.query_result(editor.result), _clear, _rerun, _run
+
+
+@app.cell
+def _(hm, recent):
+    _selected = recent.pick.value
+    return hm.query_result(recent.result, label="Recent result"), _selected
+
+
+@app.cell
+def _(client, mo):
+    _df = mo.sql(
+        """
+        SELECT 1 AS example_value
+        """,
+        engine=client,
+    )
+    return
+
+
+if __name__ == "__main__":
+    app.run()

examples/hotdata_basic.py

@@ -16,13 +16,19 @@
     recent_results,
     run_history,
 )
+from hotdata_marimo.sql_engine import (
+    HotdataMarimoEngine,
+    register_hotdata_sql_engine,
+    unregister_hotdata_sql_engine,
+)
 from hotdata_marimo.sql_editor import SqlEditor, sql_editor
 from hotdata_marimo.table_browser import TableBrowser, connection_picker, table_browser
 from hotdata_marimo.workspace_selector import WorkspaceSelector, workspace_selector_from_env
 
 __all__ = [
     "__version__",
     "HotdataClient",
+    "HotdataMarimoEngine",
     "QueryResult",
     "RecentResults",
     "SqlEditor",
@@ -39,11 +45,13 @@
     "hotdata_workspace_selector",
     "query_result",
     "recent_results",
+    "register_hotdata_sql_engine",
+    "register_mo_ui_hotdata_aliases",
     "run_history",
     "sql_editor",
     "table_browser",
+    "unregister_hotdata_sql_engine",
     "workspace_selector_from_env",
-    "register_mo_ui_hotdata_aliases",
 ]
 
 hotdata_sql_editor = sql_editor

hotdata_marimo/__init__.py

@@ -4,9 +4,20 @@
 
 import marimo as mo
 
-from hotdata_runtime.client import HotdataClient
-from hotdata_runtime.health import workspace_health_lines
-from hotdata_runtime.result import QueryResult
+from hotdata_runtime import HotdataClient, QueryResult, workspace_health_lines
+
+
+def _option_map_with_unique_labels(
+    pairs: list[tuple[str, str]],
+) -> dict[str, str]:
+    counts: dict[str, int] = {}
+    options: dict[str, str] = {}
+    for label, value in pairs:
+        count = counts.get(label, 0)
+        counts[label] = count + 1
+        key = label if count == 0 else f"{label} ({count + 1})"
+        options[key] = value
+    return options
 
 
 def query_result(
@@ -27,17 +38,18 @@ def query_result(
         )
     else:
         trunc = None
+    meta = result.metadata_dict()
     meta_bits = []
-    if result.result_id:
-        meta_bits.append(f"**result_id** `{result.result_id}`")
-    if result.query_run_id:
-        meta_bits.append(f"**query_run_id** `{result.query_run_id}`")
-    if result.execution_time_ms is not None:
-        meta_bits.append(f"**execution_time_ms** {result.execution_time_ms}")
-    if result.warning:
-        meta_bits.append(f"**warning** {result.warning}")
-    if result.error_message:
-        meta_bits.append(f"**error** {result.error_message}")
+    if meta["result_id"]:
+        meta_bits.append(f"**result_id** `{meta['result_id']}`")
+    if meta["query_run_id"]:
+        meta_bits.append(f"**query_run_id** `{meta['query_run_id']}`")
+    if meta["execution_time_ms"] is not None:
+        meta_bits.append(f"**execution_time_ms** {meta['execution_time_ms']}")
+    if meta["warning"]:
+        meta_bits.append(f"**warning** {meta['warning']}")
+    if meta["error_message"]:
+        meta_bits.append(f"**error** {meta['error_message']}")
     header = mo.md(" · ".join(meta_bits) if meta_bits else "_No metadata._")
     df = result.to_pandas()
     tbl = mo.ui.table(
@@ -59,11 +71,12 @@ def query_result(
 class RecentResults:
     def __init__(self, client: HotdataClient, *, limit: int = 50) -> None:
         self._client = client
-        listing = client.results().list_results(limit=limit, offset=0)
-        self._results = listing.results
-        options = {
-            f"{r.created_at} · {r.status} · {r.id}": r.id for r in self._results
-        }
+        self._results = client.list_recent_results(limit=limit, offset=0)
+        option_pairs = [
+            (f"{r.created_at} · {r.status} · {r.result_id}", r.result_id)
+            for r in self._results
+        ]
+        options = _option_map_with_unique_labels(option_pairs)
         self.pick = mo.ui.dropdown(
             options=options or {"(no results)": ""},
             label="Recent results",
@@ -97,20 +110,19 @@ def run_history(
     limit: int = 20,
     label: str = "Run history",
 ):
-    runs = client.query_runs().list_query_runs(limit=limit).query_runs
+    runs = client.list_run_history(limit=limit)
     if not runs:
         return mo.md("_No query runs returned._")
 
     rows: list[dict[str, object]] = []
     for r in runs:
         rows.append(
             {
-                "created_at": getattr(r, "created_at", None),
-                "status": getattr(r, "status", None),
-                "execution_time_ms": getattr(r, "execution_time_ms", None),
-                "result_id": getattr(r, "result_id", None),
-                "query_run_id": getattr(r, "id", None)
-                or getattr(r, "query_run_id", None),
+                "created_at": r.created_at,
+                "status": r.status,
+                "execution_time_ms": r.execution_time_ms,
+                "result_id": r.result_id,
+                "query_run_id": r.query_run_id,
             }
         )
 

hotdata_marimo/display.py

@@ -2,8 +2,7 @@
 
 import marimo as mo
 
-from hotdata_runtime.client import HotdataClient
-from hotdata_runtime.result import QueryResult
+from hotdata_runtime import HotdataClient, QueryResult
 
 
 class SqlEditor: