refactor: eliminate flag-based side-effect tracking, fix unregister, remove dead code

- table_browser: extract _set_table_pick() replacing duplicate _init/_rebuild methods;
  _sync_table_catalog returns bool so ui drops _rebuilt_table_pick_this_run flag;
  standardize _active_connection_id to use 'or None' consistently
- sql_engine: unregister now restores original engine_to_data_source_connection and
  resets sentinel so register/unregister/register round-trip works correctly
- sql_editor: remove dead 'or ""' on _cached_sql (already guarded by None check above)
- workspace_selector: cache HotdataClient, only reconstruct when workspace_id changes

Author: eddietejeda

hotdata_marimo/sql_editor.py

@@ -105,7 +105,7 @@ def _execute_or_cached(self) -> QueryResult | None:
                 title="Running on Hotdata",
                 subtitle="Re-running last query and waiting for results…",
             ):
-                result = self._client.execute_sql(self._cached_sql or "", database=self._database)
+                result = self._client.execute_sql(self._cached_sql, database=self._database)
             self._result_cache = result
             self._last_rerun_n = rerun_n
             return result

hotdata_marimo/sql_engine.py

@@ -368,7 +368,11 @@ def register_hotdata_sql_engine() -> None:
 
 def unregister_hotdata_sql_engine() -> None:
     """Remove :class:`HotdataMarimoEngine` from Marimo's registry (mostly for tests)."""
+    global _ORIGINAL_ENGINE_TO_CONNECTION
     from marimo._sql.get_engines import SUPPORTED_ENGINES
 
     while HotdataMarimoEngine in SUPPORTED_ENGINES:
         SUPPORTED_ENGINES.remove(HotdataMarimoEngine)
+    if _ORIGINAL_ENGINE_TO_CONNECTION is not None:
+        _set_engine_to_data_source_connection(_ORIGINAL_ENGINE_TO_CONNECTION)
+        _ORIGINAL_ENGINE_TO_CONNECTION = None

hotdata_marimo/table_browser.py

@@ -12,6 +12,8 @@
     resolve_connection_picker,
 )
 
+__all__ = ["TableBrowser", "connection_picker", "table_browser"]
+
 
 class TableBrowser:
     """Pick a fully qualified `connection.schema.table` and inspect columns.
@@ -43,58 +45,26 @@ def __init__(
             )
 
         self._table_pick_ctx: str | None = None
-        self._rebuilt_table_pick_this_run = False
         self._init_table_pick()
 
-    def _init_table_pick(self) -> None:
-        if self._conn_pick is not None:
-            self.table_pick = empty_dropdown(
-                label="Table",
-                message="(select connection above)",
-            )
-            self._empty_catalog = True
-            self._all_names = []
-            self._table_pick_ctx = ""
-            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
-        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()
-
     def _active_connection_id(self) -> str | None:
         if self._override_connection_id is not None:
             return self._override_connection_id or None
         if self._conn_pick is not None:
-            v = self._conn_pick.value  # type: ignore[attr-defined]
-            return v if v else None
-        if self._implicit_connection_id is None:
-            return None
+            return self._conn_pick.value or None  # type: ignore[attr-defined]
         return self._implicit_connection_id or None
 
     def _names_for_active_connection(self) -> list[str]:
         cid = self._active_connection_id()
-        if cid is None or cid == "":
+        if not cid:
             return []
         return self._client.list_qualified_table_names(
             limit=self._table_limit,
             connection_id=cid,
         )
 
-    def _rebuild_table_pick(self, names: list[str]) -> None:
+    def _set_table_pick(self, names: list[str]) -> None:
+        """Create or replace the table dropdown for the given names list."""
         self._all_names = names
         if not names:
             self._empty_catalog = True
@@ -111,7 +81,32 @@ def _rebuild_table_pick(self, names: list[str]) -> None:
                 searchable=True,
             )
         self._table_pick_ctx = self._active_connection_id()
-        self._rebuilt_table_pick_this_run = True
+
+    def _init_table_pick(self) -> None:
+        if self._conn_pick is not None:
+            self._all_names = []
+            self._empty_catalog = True
+            self.table_pick = empty_dropdown(
+                label="Table",
+                message="(select connection above)",
+            )
+            self._table_pick_ctx = ""
+            return
+        self._set_table_pick(self._names_for_active_connection())
+
+    def _sync_table_catalog(self) -> bool:
+        """Refresh the table dropdown when the active connection changes.
+
+        Returns True if the dropdown was rebuilt (so the caller knows not to
+        read ``.value`` on the new widget in the same Marimo run).
+        """
+        if self._conn_pick is not None:
+            _ = self._conn_pick.value  # type: ignore[attr-defined]
+        cid = self._active_connection_id()
+        if not cid or cid == self._table_pick_ctx:
+            return False
+        self._set_table_pick(self._names_for_active_connection())
+        return True
 
     @property
     def selected_connection_id(self) -> str | None:
@@ -122,30 +117,17 @@ def selected_table(self) -> str | None:
         v = self.table_pick.value
         return v if v else None
 
-    def _sync_table_catalog(self) -> None:
-        """Refresh the table dropdown when the active connection changes."""
-        if self._conn_pick is not None:
-            _ = self._conn_pick.value  # type: ignore[attr-defined]
-        cid = self._active_connection_id()
-        if not cid:
-            return
-        if cid == self._table_pick_ctx:
-            return
-        self._rebuild_table_pick(self._names_for_active_connection())
-
     @property
     def ui(self):
-        self._rebuilt_table_pick_this_run = False
-        self._sync_table_catalog()
-
-        if not self._rebuilt_table_pick_this_run:
+        rebuilt = self._sync_table_catalog()
+        if not rebuilt:
             _ = self.table_pick.value
 
-        sel = None if self._rebuilt_table_pick_this_run else self.selected_table
+        sel = None if rebuilt else self.selected_table
         cid = self._active_connection_id()
         conn_header = (
-            mo.md(f"**Connection** `{self._active_connection_id()}`")
-            if self._active_connection_id()
+            mo.md(f"**Connection** `{cid}`")
+            if cid
             else None
         )
         if not sel:

hotdata_marimo/workspace_selector.py

@@ -26,6 +26,8 @@ def __init__(
         self._api_key = api_key
         self._host = host or default_host()
         self._session_id = session_id
+        self._client_cache: HotdataClient | None = None
+        self._client_cache_wid: str | None = None
         selection = resolve_workspace_selection(api_key, self._host, session_id)
         self._explicit = selection.source == "explicit_env"
         if self._explicit:
@@ -64,12 +66,16 @@ def workspace_id(self) -> str:
 
     @property
     def client(self) -> HotdataClient:
-        return HotdataClient(
-            self._api_key,
-            self.workspace_id,
-            host=self._host,
-            session_id=self._session_id,
-        )
+        wid = self.workspace_id
+        if self._client_cache is None or self._client_cache_wid != wid:
+            self._client_cache = HotdataClient(
+                self._api_key,
+                wid,
+                host=self._host,
+                session_id=self._session_id,
+            )
+            self._client_cache_wid = wid
+        return self._client_cache
 
     @property
     def ui(self):