hotdata-marimo 0.1.0__py3-none-any.whl

hotdata_marimo/sql_engine.py

@@ -0,0 +1,371 @@
+"""Marimo ``mo.sql`` engine integration for :class:`~hotdata_runtime.HotdataClient`."""
+
+from __future__ import annotations
+
+from collections import defaultdict
+from typing import Any, Literal
+
+from hotdata_runtime import HotdataClient
+
+from marimo import _loggers
+from marimo._data.models import (
+    Database,
+    DataSourceConnection,
+    DataTable,
+    DataTableColumn,
+    Schema,
+)
+from marimo._sql.engines.types import InferenceConfig, SQLConnection
+from marimo._sql.utils import convert_to_output, sql_type_to_data_type
+from marimo._types.ids import VariableName
+
+LOGGER = _loggers.marimo_logger()
+
+
+def _table_schema_name(t: Any) -> str:
+    return str(t.var_schema)
+
+
+class HotdataMarimoEngine(SQLConnection[HotdataClient]):
+    """Marimo :class:`~marimo._sql.engines.types.SQLConnection` backed by Hotdata.
+
+    Catalog methods support Marimo's Data Sources panel. ``execute()`` only runs SQL
+    via :meth:`~hotdata_runtime.HotdataClient.execute_sql` (no catalog calls in that path).
+    """
+
+    def __init__(
+        self,
+        connection: HotdataClient,
+        engine_name: VariableName | None = None,
+    ) -> None:
+        super().__init__(connection, engine_name)
+        self._connections_cache: list[Any] | None = None
+
+    @property
+    def source(self) -> str:
+        return "hotdata"
+
+    @property
+    def dialect(self) -> str:
+        # Marimo labels engines as ``{dialect} ({variable_name})``; display_name is patched to "Hotdata".
+        return "hotdata"
+
+    @staticmethod
+    def is_compatible(var: Any) -> bool:
+        return isinstance(var, HotdataClient)
+
+    @property
+    def inference_config(self) -> InferenceConfig:
+        return InferenceConfig(
+            auto_discover_schemas=True,
+            auto_discover_tables="auto",
+            auto_discover_columns="auto",
+        )
+
+    def _resolve_should_auto_discover(
+        self,
+        value: bool | Literal["auto"],
+    ) -> bool:
+        if value == "auto":
+            return True
+        return value
+
+    def _connection_id(self, connection_name: str) -> str | None:
+        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:
+            self._connections_cache = list(
+                self._connection.connections().list_connections().connections
+            )
+        return self._connections_cache
+
+    def _iter_grouped(
+        self,
+        *,
+        connection_id: str | None,
+        include_columns: bool,
+    ) -> dict[str, dict[str, list[Any]]]:
+        grouped: dict[str, dict[str, list[Any]]] = defaultdict(
+            lambda: defaultdict(list)
+        )
+        for t in self._connection.iter_tables(
+            connection_id=connection_id,
+            include_columns=include_columns,
+        ):
+            grouped[str(t.connection)][_table_schema_name(t)].append(t)
+        return grouped
+
+    def get_default_database(self) -> str | None:
+        listing = self._connections()
+        if not listing:
+            return None
+        return str(listing[0].name)
+
+    def get_default_schema(self) -> str | None:
+        return None
+
+    def get_databases(
+        self,
+        *,
+        include_schemas: bool | Literal["auto"],
+        include_tables: bool | Literal["auto"],
+        include_table_details: bool | Literal["auto"],
+    ) -> list[Database]:
+        databases: list[Database] = []
+        for c in self._connections():
+            name = str(c.name)
+            if self._resolve_should_auto_discover(include_schemas):
+                schemas = self.get_schemas(
+                    database=name,
+                    include_tables=self._resolve_should_auto_discover(
+                        include_tables
+                    ),
+                    include_table_details=self._resolve_should_auto_discover(
+                        include_table_details
+                    ),
+                )
+            else:
+                schemas = []
+            databases.append(
+                Database(
+                    name=name,
+                    dialect=self.dialect,
+                    schemas=schemas,
+                    engine=self._engine_name,
+                )
+            )
+        return databases
+
+    def get_schemas(
+        self,
+        *,
+        database: str | None,
+        include_tables: bool,
+        include_table_details: bool,
+    ) -> list[Schema]:
+        if not database:
+            return []
+        conn_id = self._connection_id(database)
+        if conn_id is None:
+            LOGGER.warning("Unknown Hotdata connection name %r", database)
+            return []
+        grouped = self._iter_grouped(
+            connection_id=conn_id,
+            include_columns=include_table_details,
+        )
+        inner = grouped.get(database, {})
+        schemas: list[Schema] = []
+        for schema_name in sorted(inner.keys()):
+            tables: list[DataTable] = []
+            if include_tables:
+                tables = self.get_tables_in_schema(
+                    schema=schema_name,
+                    database=database,
+                    include_table_details=include_table_details,
+                )
+                if not tables:
+                    continue
+            schemas.append(Schema(name=schema_name, tables=tables))
+        return schemas
+
+    def _data_table_from_table_info(self, t: Any) -> DataTable:
+        cols: list[DataTableColumn] = []
+        for col in t.columns or []:
+            cols.append(
+                DataTableColumn(
+                    name=str(col.name),
+                    type=sql_type_to_data_type(str(col.data_type)),
+                    external_type=str(col.data_type),
+                    sample_values=[],
+                )
+            )
+        return DataTable(
+            source_type="connection",
+            source=self.source,
+            name=str(t.table),
+            num_rows=None,
+            num_columns=len(cols) if cols else None,
+            variable_name=None,
+            engine=self._engine_name,
+            type="table",
+            columns=cols,
+            primary_keys=None,
+            indexes=None,
+        )
+
+    def get_tables_in_schema(
+        self,
+        *,
+        schema: str,
+        database: str,
+        include_table_details: bool,
+    ) -> list[DataTable]:
+        conn_id = self._connection_id(database)
+        if conn_id is None:
+            return []
+        grouped = self._iter_grouped(
+            connection_id=conn_id,
+            include_columns=include_table_details,
+        )
+        tables_info = grouped.get(database, {}).get(schema, [])
+        out: list[DataTable] = []
+        for t in sorted(tables_info, key=lambda x: str(x.table)):
+            if include_table_details:
+                if t.columns:
+                    out.append(self._data_table_from_table_info(t))
+                    continue
+                dt = self.get_table_details(
+                    table_name=str(t.table),
+                    schema_name=schema,
+                    database_name=database,
+                )
+                if dt is not None:
+                    out.append(dt)
+            else:
+                out.append(
+                    DataTable(
+                        source_type="connection",
+                        source=self.source,
+                        name=str(t.table),
+                        num_rows=None,
+                        num_columns=len(t.columns or []) if t.columns else None,
+                        variable_name=None,
+                        engine=self._engine_name,
+                        type="table",
+                        columns=[],
+                        primary_keys=None,
+                        indexes=None,
+                    )
+                )
+        return out
+
+    def get_table_details(
+        self,
+        *,
+        table_name: str,
+        schema_name: str,
+        database_name: str,
+    ) -> DataTable | None:
+        conn_id = self._connection_id(database_name)
+        if conn_id is None:
+            return None
+        qualified = f"{database_name}.{schema_name}.{table_name}"
+        try:
+            cols_raw = self._connection.columns_for_qualified(
+                qualified, connection_id=conn_id
+            )
+        except Exception:
+            LOGGER.warning(
+                "Failed to load columns for %s",
+                qualified,
+                exc_info=True,
+            )
+            return None
+        cols: list[DataTableColumn] = []
+        for col in cols_raw:
+            cols.append(
+                DataTableColumn(
+                    name=str(col.name),
+                    type=sql_type_to_data_type(str(col.data_type)),
+                    external_type=str(col.data_type),
+                    sample_values=[],
+                )
+            )
+        return DataTable(
+            source_type="connection",
+            source=self.source,
+            name=table_name,
+            num_rows=None,
+            num_columns=len(cols),
+            variable_name=None,
+            engine=self._engine_name,
+            type="table",
+            columns=cols,
+            primary_keys=None,
+            indexes=None,
+        )
+
+    def execute(self, query: str) -> Any:
+        qr = self._connection.execute_sql(query)
+        fmt = self.sql_output_format()
+
+        def to_polars() -> Any:
+            import polars as pl
+
+            if not qr.columns:
+                return pl.DataFrame()
+            return pl.DataFrame(qr.rows, schema=qr.columns, orient="row")
+
+        return convert_to_output(
+            sql_output_format=fmt,
+            to_polars=to_polars,
+            to_pandas=qr.to_pandas,
+            to_native=to_polars,
+        )
+
+
+_HOTDATA_ENGINE_DISPLAY_NAME = "Hotdata"
+_ORIGINAL_ENGINE_TO_CONNECTION = None
+
+
+def _install_hotdata_engine_display_name() -> None:
+    """Show ``Hotdata`` in Marimo's SQL engine / Data Sources UI (not ``sql (client)``)."""
+    global _ORIGINAL_ENGINE_TO_CONNECTION
+    if _ORIGINAL_ENGINE_TO_CONNECTION is not None:
+        return
+
+    import marimo._sql.get_engines as ge
+
+    _ORIGINAL_ENGINE_TO_CONNECTION = ge.engine_to_data_source_connection
+
+    def engine_to_data_source_connection(
+        variable_name: VariableName, engine: object
+    ) -> DataSourceConnection:
+        conn = _ORIGINAL_ENGINE_TO_CONNECTION(variable_name, engine)  # type: ignore[arg-type]
+        if not isinstance(engine, HotdataMarimoEngine):
+            return conn
+        return DataSourceConnection(
+            source=conn.source,
+            dialect=conn.dialect,
+            name=conn.name,
+            display_name=_HOTDATA_ENGINE_DISPLAY_NAME,
+            databases=conn.databases,
+            default_database=conn.default_database,
+            default_schema=conn.default_schema,
+        )
+
+    _set_engine_to_data_source_connection(engine_to_data_source_connection)
+
+
+def _set_engine_to_data_source_connection(fn: object) -> None:
+    """Marimo imports this helper in multiple modules; patch all bindings."""
+    import marimo._runtime.runner.hooks_post_execution as hpe
+    import marimo._runtime.runtime as rt
+    import marimo._sql.get_engines as ge
+
+    ge.engine_to_data_source_connection = fn  # type: ignore[assignment]
+    hpe.engine_to_data_source_connection = fn  # type: ignore[assignment]
+    rt.engine_to_data_source_connection = fn  # type: ignore[assignment]
+
+
+def register_hotdata_sql_engine() -> None:
+    """Register :class:`HotdataMarimoEngine` with Marimo's SQL engine registry (idempotent)."""
+    _install_hotdata_engine_display_name()
+    from marimo._sql.get_engines import SUPPORTED_ENGINES
+
+    if HotdataMarimoEngine in SUPPORTED_ENGINES:
+        return
+    SUPPORTED_ENGINES.insert(0, HotdataMarimoEngine)
+
+
+def unregister_hotdata_sql_engine() -> None:
+    """Remove :class:`HotdataMarimoEngine` from Marimo's registry (mostly for tests)."""
+    from marimo._sql.get_engines import SUPPORTED_ENGINES
+
+    while HotdataMarimoEngine in SUPPORTED_ENGINES:
+        SUPPORTED_ENGINES.remove(HotdataMarimoEngine)

hotdata_marimo/table_browser.py

@@ -0,0 +1,221 @@
+from __future__ import annotations
+
+from typing import Any
+
+import marimo as mo
+
+from hotdata_runtime import HotdataClient
+
+from hotdata_marimo._options import (
+    connection_picker,
+    empty_dropdown,
+    resolve_connection_picker,
+)
+
+
+class TableBrowser:
+    """Pick a fully qualified `connection.schema.table` and inspect columns.
+
+    Marimo does not allow reading ``.value`` on UI elements in the same cell that
+    constructs them. Build ``TableBrowser`` in one cell and use ``.ui`` in another.
+
+    The table dropdown is not recreated on every render (that would make it
+    "born" in the layout cell). It is only rebuilt when the active connection
+    changes: after a rebuild, ``.value`` is not read until a later run.
+    """
+
+    def __init__(
+        self,
+        client: HotdataClient,
+        *,
+        table_limit: int = 5000,
+        connection_id: str | None = None,
+    ) -> None:
+        self._client = client
+        self._table_limit = table_limit
+        self._override_connection_id = connection_id
+        self._conn_pick: Any | None = None
+        self._implicit_connection_id: str | None = None
+
+        if self._override_connection_id is None:
+            self._conn_pick, self._implicit_connection_id = resolve_connection_picker(
+                client
+            )
+
+        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._implicit_connection_id or None
+
+    def _names_for_active_connection(self) -> list[str]:
+        cid = self._active_connection_id()
+        if cid is None or 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:
+        self._all_names = names
+        if not names:
+            self._empty_catalog = True
+            self.table_pick = empty_dropdown(
+                label="Table",
+                message="(no tables in catalog)",
+            )
+        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._rebuilt_table_pick_this_run = True
+
+    @property
+    def selected_connection_id(self) -> str | None:
+        return self._active_connection_id()
+
+    @property
+    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:
+            _ = self.table_pick.value
+
+        sel = None if self._rebuilt_table_pick_this_run else self.selected_table
+        cid = self._active_connection_id()
+        conn_header = (
+            mo.md(f"**Connection** `{self._active_connection_id()}`")
+            if self._active_connection_id()
+            else None
+        )
+        if not sel:
+            if self._conn_pick is not None and not cid:
+                hint = "Choose a connection above to load tables."
+            elif self._empty_catalog:
+                hint = (
+                    "_No tables returned from the information schema. "
+                    "Try refreshing a connection in Hotdata._"
+                )
+            else:
+                hint = "Choose a table below (search inside the dropdown when needed)."
+            stack = [
+                mo.md(
+                    f"**Workspace** `{self._client.workspace_id}` — {hint}"
+                ),
+            ]
+            if conn_header is not None:
+                stack.append(conn_header)
+            if self._conn_pick is not None:
+                stack.append(self._conn_pick)
+            stack.append(self.table_pick)
+            return mo.vstack(stack, gap=1)
+
+        cols = self._client.columns_for_qualified(
+            sel,
+            connection_id=self.selected_connection_id,
+        )
+        if not cols:
+            body = mo.md("_No column metadata returned (check catalog sync)._")
+        else:
+            lines = [
+                f"| `{c.name}` | {c.data_type} | "
+                f"{'NULL' if c.nullable else 'NOT NULL'} |"
+                for c in cols
+            ]
+            table = (
+                "| column | type | null |\n| --- | --- | --- |\n"
+                + "\n".join(lines)
+            )
+            body = mo.md(table)
+        starter = f"```sql\nSELECT *\nFROM {sel}\nLIMIT 100\n```"
+        stack2 = [
+            mo.md(
+                f"**Workspace** `{self._client.workspace_id}` — "
+                f"**selected** `{sel}`"
+            ),
+        ]
+        if conn_header is not None:
+            stack2.append(conn_header)
+        if self._conn_pick is not None:
+            stack2.append(self._conn_pick)
+        stack2.extend(
+            [
+                self.table_pick,
+                mo.md("### Columns"),
+                body,
+                mo.md("### Starter query"),
+                mo.md(starter),
+            ]
+        )
+        return mo.vstack(stack2, gap=1)
+
+
+def table_browser(
+    client: HotdataClient,
+    *,
+    table_limit: int = 5000,
+    connection_id: str | None = None,
+) -> TableBrowser:
+    return TableBrowser(
+        client, table_limit=table_limit, connection_id=connection_id
+    )

hotdata_marimo/workspace_selector.py

@@ -0,0 +1,93 @@
+from __future__ import annotations
+
+import marimo as mo
+from hotdata_runtime import (
+    HotdataClient,
+    default_api_key,
+    default_host,
+    default_session_id,
+    resolve_workspace_selection,
+)
+
+from hotdata_marimo._options import unique_label_options
+
+
+class WorkspaceSelector:
+    """Workspace picker that rebuilds `HotdataClient` as selection changes."""
+
+    def __init__(
+        self,
+        *,
+        api_key: str,
+        host: str | None = None,
+        session_id: str | None = None,
+        label: str = "Workspace",
+    ) -> None:
+        self._api_key = api_key
+        self._host = host or default_host()
+        self._session_id = session_id
+        selection = resolve_workspace_selection(api_key, self._host, session_id)
+        self._explicit = selection.source == "explicit_env"
+        if self._explicit:
+            self._pick = None
+            self._workspace_id = selection.workspace_id
+            return
+
+        workspaces = selection.workspaces
+        if len(workspaces) == 1:
+            self._pick = None
+            self._workspace_id = workspaces[0].public_id
+            return
+
+        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
+    def workspace_id(self) -> str:
+        if self._pick is None:
+            return self._workspace_id
+        v = self._pick.value
+        return v if v else self._workspace_id
+
+    @property
+    def client(self) -> HotdataClient:
+        return HotdataClient(
+            self._api_key,
+            self.workspace_id,
+            host=self._host,
+            session_id=self._session_id,
+        )
+
+    @property
+    def ui(self):
+        if self._pick is None:
+            return mo.md(f"**Workspace** `{self.workspace_id}`")
+        _ = self._pick.value
+        return mo.vstack([self._pick], gap=1)
+
+
+def workspace_selector_from_env(*, label: str = "Workspace") -> WorkspaceSelector:
+    api_key = default_api_key()
+    if not api_key:
+        raise RuntimeError("HOTDATA_API_KEY must be set.")
+    host = default_host()
+    session = default_session_id()
+    return WorkspaceSelector(
+        api_key=api_key,
+        host=host,
+        session_id=session,
+        label=label,
+    )