refactor: dedupe marimo dropdown and connection helpers

Extract shared option builders into _options.py, reuse resolve_connection_picker
in TableBrowser, and resolve SQL engine connection ids via runtime.

Author: eddietejeda

hotdata_marimo/_options.py

@@ -0,0 +1,89 @@
+"""Shared dropdown option helpers for Marimo UI widgets."""
+
+from __future__ import annotations
+
+from collections.abc import Callable
+from typing import Any
+
+import marimo as mo
+
+from hotdata_runtime import HotdataClient
+
+
+def unique_label_options(
+    pairs: list[tuple[str, str]],
+    *,
+    disambiguate: Callable[[str, str, int], str] | None = None,
+) -> dict[str, str]:
+    """Build a label→value map, suffixing repeated labels when needed."""
+    counts: dict[str, int] = {}
+    options: dict[str, str] = {}
+    for label, value in pairs:
+        count = counts.get(label, 0)
+        counts[label] = count + 1
+        if count == 0:
+            key = label
+        elif disambiguate is not None:
+            key = disambiguate(label, value, count)
+        else:
+            key = f"{label} ({count + 1})"
+        options[key] = value
+    return options
+
+
+def empty_dropdown(
+    *,
+    label: str,
+    message: str,
+    full_width: bool = True,
+):
+    return mo.ui.dropdown(
+        options={message: ""},
+        label=label,
+        full_width=full_width,
+    )
+
+
+def connection_options(conns: list[Any]) -> dict[str, str]:
+    pairs = [(str(c.name), str(c.id)) for c in conns]
+    return unique_label_options(
+        pairs,
+        disambiguate=lambda label, value, count: f"{label} ({value})",
+    )
+
+
+def connection_picker(
+    client: HotdataClient,
+    *,
+    label: str = "Connection",
+    full_width: bool = True,
+):
+    listing = client.connections().list_connections()
+    conns = listing.connections
+    if not conns:
+        return empty_dropdown(
+            label=label,
+            message="(no connections)",
+            full_width=full_width,
+        )
+    return mo.ui.dropdown(
+        options=connection_options(conns),
+        label=label,
+        full_width=full_width,
+    )
+
+
+def resolve_connection_picker(
+    client: HotdataClient,
+    *,
+    label: str = "Connection",
+    full_width: bool = True,
+) -> tuple[Any | None, str | None]:
+    """Return ``(dropdown_or_none, implicit_connection_id)`` for table browsers."""
+    listing = client.connections().list_connections()
+    conns = listing.connections
+    if not conns:
+        return None, ""
+    if len(conns) == 1:
+        return None, conns[0].id
+    return connection_picker(client, label=label, full_width=full_width), None

hotdata_marimo/display.py

@@ -6,18 +6,7 @@
 
 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
+from hotdata_marimo._options import empty_dropdown, unique_label_options
 
 
 def query_result(
@@ -76,11 +65,15 @@ def __init__(self, client: HotdataClient, *, limit: int = 50) -> None:
             (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",
-            full_width=True,
+        options = unique_label_options(option_pairs)
+        self.pick = (
+            empty_dropdown(label="Recent results", message="(no results)")
+            if not options
+            else mo.ui.dropdown(
+                options=options,
+                label="Recent results",
+                full_width=True,
+            )
         )
 
     @property

hotdata_marimo/sql_engine.py

@@ -40,7 +40,6 @@ def __init__(
     ) -> None:
         super().__init__(connection, engine_name)
         self._connections_cache: list[Any] | None = None
-        self._connection_id_cache: dict[str, str] | None = None
 
     @property
     def source(self) -> str:
@@ -71,15 +70,12 @@ def _resolve_should_auto_discover(
             return True
         return value
 
-    def _connection_ids(self) -> dict[str, str]:
-        if self._connection_id_cache is None:
-            self._connection_id_cache = {
-                str(c.name): str(c.id) for c in self._connections()
-            }
-        return self._connection_id_cache
-
     def _connection_id(self, connection_name: str) -> str | None:
-        return self._connection_ids().get(connection_name)
+        try:
+            return self._connection.connection_id_by_name().get(connection_name)
+        except RuntimeError as e:
+            LOGGER.warning("%s", e)
+            return None
 
     def _connections(self) -> list[Any]:
         if self._connections_cache is None:

hotdata_marimo/table_browser.py

@@ -6,39 +6,11 @@
 
 from hotdata_runtime import HotdataClient
 
-
-def _connection_options(conns: list[Any]) -> dict[str, str]:
-    counts: dict[str, int] = {}
-    options: dict[str, str] = {}
-    for c in conns:
-        label = c.name
-        count = counts.get(label, 0)
-        counts[label] = count + 1
-        key = label if count == 0 else f"{label} ({c.id})"
-        options[key] = c.id
-    return options
-
-
-def connection_picker(
-    client: HotdataClient,
-    *,
-    label: str = "Connection",
-    full_width: bool = True,
-):
-    listing = client.connections().list_connections()
-    conns = listing.connections
-    if not conns:
-        return mo.ui.dropdown(
-            options={"(no connections)": ""},
-            label=label,
-            full_width=full_width,
-        )
-    options = _connection_options(conns)
-    return mo.ui.dropdown(
-        options=options,
-        label=label,
-        full_width=full_width,
-    )
+from hotdata_marimo._options import (
+    connection_picker,
+    empty_dropdown,
+    resolve_connection_picker,
+)
 
 
 class TableBrowser:
@@ -66,44 +38,40 @@ def __init__(
         self._implicit_connection_id: str | None = None
 
         if self._override_connection_id is None:
-            listing = client.connections().list_connections()
-            conns = listing.connections
-            if len(conns) > 1:
-                self._conn_pick = connection_picker(client)
-            elif len(conns) == 1:
-                self._implicit_connection_id = conns[0].id
-            else:
-                self._implicit_connection_id = ""
+            self._conn_pick, self._implicit_connection_id = resolve_connection_picker(
+                client
+            )
 
         self._table_pick_ctx: str | None = None
+        self._init_table_pick()
 
+    def _init_table_pick(self) -> None:
         if self._conn_pick is not None:
-            self.table_pick = mo.ui.dropdown(
-                options={"(select connection above)": ""},
+            self.table_pick = empty_dropdown(
                 label="Table",
-                full_width=True,
+                message="(select connection above)",
+            )
+            self._empty_catalog = True
+            self._all_names = []
+            return
+
+        names = self._names_for_active_connection()
+        self._all_names = names
+        if not names:
+            self.table_pick = empty_dropdown(
+                label="Table",
+                message="(no tables in catalog)",
             )
             self._empty_catalog = True
-            self._all_names: list[str] = []
         else:
-            names = self._names_for_active_connection()
-            self._all_names = names
-            if not names:
-                self.table_pick = mo.ui.dropdown(
-                    options={"(no tables in catalog)": ""},
-                    label="Table",
-                    full_width=True,
-                )
-                self._empty_catalog = True
-            else:
-                self._empty_catalog = False
-                self.table_pick = mo.ui.dropdown(
-                    options={n: n for n in names},
-                    label="Table",
-                    full_width=True,
-                    searchable=True,
-                )
-            self._table_pick_ctx = self._active_connection_id()
+            self._empty_catalog = False
+            self.table_pick = mo.ui.dropdown(
+                options={n: n for n in names},
+                label="Table",
+                full_width=True,
+                searchable=True,
+            )
+        self._table_pick_ctx = self._active_connection_id()
 
     def _active_connection_id(self) -> str | None:
         if self._override_connection_id is not None:
@@ -128,10 +96,9 @@ def _rebuild_table_pick(self, names: list[str]) -> None:
         self._all_names = names
         if not names:
             self._empty_catalog = True
-            self.table_pick = mo.ui.dropdown(
-                options={"(no tables in catalog)": ""},
+            self.table_pick = empty_dropdown(
                 label="Table",
-                full_width=True,
+                message="(no tables in catalog)",
             )
         else:
             self._empty_catalog = False

hotdata_marimo/workspace_selector.py

@@ -9,6 +9,8 @@
     resolve_workspace_selection,
 )
 
+from hotdata_marimo._options import unique_label_options
+
 
 class WorkspaceSelector:
     """Workspace picker that rebuilds `HotdataClient` as selection changes."""
@@ -37,17 +39,20 @@ def __init__(
             self._workspace_id = workspaces[0].public_id
             return
 
-        labels: list[tuple[str, str]] = []
-        seen: set[str] = set()
-        for w in workspaces:
-            base = w.name
-            label_text = base if base not in seen else f"{base} ({w.public_id})"
-            seen.add(base)
-            labels.append((label_text, w.public_id))
-
-        labels.sort(key=lambda t: 0 if t[1] == selection.workspace_id else 1)
-        options = {k: v for k, v in labels}
-        self._pick = mo.ui.dropdown(options=options, label=label, full_width=True)
+        pairs = [(w.name, w.public_id) for w in workspaces]
+        options = unique_label_options(
+            pairs,
+            disambiguate=lambda name, public_id, count: f"{name} ({public_id})",
+        )
+        items = sorted(
+            options.items(),
+            key=lambda item: 0 if item[1] == selection.workspace_id else 1,
+        )
+        self._pick = mo.ui.dropdown(
+            options=dict(items),
+            label=label,
+            full_width=True,
+        )
         self._workspace_id = selection.workspace_id
 
     @property

tests/test_hotdata_marimo.py

@@ -7,9 +7,8 @@
 
 import hotdata_marimo as hm
 from hotdata_runtime import HotdataClient
-from hotdata_marimo.display import _option_map_with_unique_labels
+from hotdata_marimo._options import connection_options, unique_label_options
 from hotdata_marimo.sql_engine import HotdataMarimoEngine
-from hotdata_marimo.table_browser import _connection_options
 from hotdata_marimo.workspace_selector import WorkspaceSelector, workspace_selector_from_env
 from marimo._types.ids import VariableName
 
@@ -39,8 +38,8 @@ def _workspace_row(name: str, public_id: str, *, active: bool = True):
         ),
     ],
 )
-def test_option_map_with_unique_labels(labels, expected):
-    assert _option_map_with_unique_labels(labels) == expected
+def test_unique_label_options(labels, expected):
+    assert unique_label_options(labels) == expected
 
 
 def test_connection_options_disambiguates_duplicate_names():
@@ -49,7 +48,7 @@ def test_connection_options_disambiguates_duplicate_names():
         SimpleNamespace(name="Warehouse", id="conn_2"),
         SimpleNamespace(name="Analytics", id="conn_3"),
     ]
-    assert _connection_options(conns) == {
+    assert connection_options(conns) == {
         "Warehouse": "conn_1",
         "Warehouse (conn_2)": "conn_2",
         "Analytics": "conn_3",