sqlrooms 0.1.0__py3-none-any.whl

sqlrooms/web/db_bridge/connectors/postgres.py

@@ -0,0 +1,118 @@
+from __future__ import annotations
+
+import importlib.util
+from dataclasses import dataclass
+from typing import Any
+from urllib.parse import quote
+
+from .base import BaseSqlBridgeConnector
+
+
+@dataclass(frozen=True)
+class PostgresConnectorSettings:
+    connection_id: str = "postgres-default"
+    title: str = "Postgres"
+    host: str = "localhost"
+    port: str = "5432"
+    database: str = ""
+    user: str = ""
+    password: str | None = None
+
+    def resolve_dsn(self) -> str:
+        pw = f":{quote(self.password, safe='')}" if self.password else ""
+        return f"postgresql://{quote(self.user, safe='')}{pw}@{self.host}:{self.port}/{self.database}"
+
+
+@dataclass(frozen=True)
+class PostgresBridgeConnector(BaseSqlBridgeConnector):
+    settings: PostgresConnectorSettings
+    engine_id: str = "postgres"
+
+    @property
+    def connection_id(self) -> str:
+        return self.settings.connection_id
+
+    @property
+    def title(self) -> str:
+        return self.settings.title
+
+    def config_dict(self) -> dict[str, Any]:
+        s = self.settings
+        d: dict[str, Any] = {}
+        for k in ("host", "port", "database", "user", "password"):
+            v = getattr(s, k)
+            if v:
+                d[k] = v
+        return d
+
+    def _connect(self):
+        try:
+            import psycopg  # type: ignore
+        except ImportError as exc:
+            raise RuntimeError(
+                "Postgres bridge requires `psycopg`. Install it to enable Postgres."
+            ) from exc
+        return psycopg.connect(self.settings.resolve_dsn())
+
+    def dependency_diagnostics(self) -> dict[str, Any]:
+        available = importlib.util.find_spec("psycopg") is not None
+        if available:
+            return {"available": True}
+        return {
+            "available": False,
+            "error": "Missing Python dependency: psycopg",
+            "requiredPackages": ["psycopg[binary]>=3.2.0"],
+            "installCommands": {
+                "uvProject": "uv sync --extra postgres",
+                "uvxRelaunch": (
+                    'uvx --from "sqlrooms[postgres]" sqlrooms --db-path :memory:'
+                ),
+                "uvxWith": (
+                    'uvx --from sqlrooms --with "psycopg[binary]>=3.2.0" '
+                    "sqlrooms --db-path :memory:"
+                ),
+            },
+        }
+
+    def list_catalog(self) -> dict[str, list[dict[str, Any]]]:
+        with self._connect() as conn:
+            with conn.cursor() as cur:
+                cur.execute("SELECT current_database()")
+                result = cur.fetchone()
+                if result is None:
+                    raise RuntimeError(
+                        "Postgres bridge could not resolve current database."
+                    )
+                current_db = result[0]
+                cur.execute(
+                    """
+                    SELECT schema_name
+                    FROM information_schema.schemata
+                    ORDER BY schema_name
+                    """
+                )
+                schemas = [
+                    {"database": current_db, "schema": row[0]} for row in cur.fetchall()
+                ]
+                cur.execute(
+                    """
+                    SELECT table_schema, table_name, table_type
+                    FROM information_schema.tables
+                    WHERE table_schema NOT IN ('pg_catalog', 'information_schema')
+                    ORDER BY table_schema, table_name
+                    """
+                )
+                tables = [
+                    {
+                        "database": current_db,
+                        "schema": row[0],
+                        "table": row[1],
+                        "isView": str(row[2]).upper().endswith("VIEW"),
+                    }
+                    for row in cur.fetchall()
+                ]
+        return {
+            "databases": [{"database": current_db}],
+            "schemas": schemas,
+            "tables": tables,
+        }

sqlrooms/web/db_bridge/connectors/snowflake.py

@@ -0,0 +1,161 @@
+from __future__ import annotations
+
+import importlib.util
+from dataclasses import dataclass
+from typing import Any
+
+from ..utils import cursor_columns, quoted_ident, rows_to_json_rows
+from .base import BaseSqlBridgeConnector
+
+SNOWFLAKE_INSTALL_COMMANDS = {
+    "uvProject": "uv sync --extra snowflake",
+    "uvxRelaunch": 'uvx --from "sqlrooms[snowflake]" sqlrooms --db-path :memory:',
+    "uvxWith": (
+        'uvx --from sqlrooms --with "snowflake-connector-python>=4.3.0" '
+        "sqlrooms --db-path :memory:"
+    ),
+}
+
+
+@dataclass(frozen=True)
+class SnowflakeConnectorSettings:
+    account: str | None = None
+    user: str | None = None
+    password: str | None = None
+    warehouse: str | None = None
+    database: str | None = None
+    schema: str | None = None
+    role: str | None = None
+    authenticator: str | None = None
+    connection_id: str = "snowflake-default"
+    title: str = "Snowflake"
+
+    def is_enabled(self) -> bool:
+        return bool(self.account and self.user)
+
+    def to_connect_kwargs(self) -> dict[str, Any]:
+        kwargs: dict[str, Any] = {}
+        for key in (
+            "account",
+            "user",
+            "password",
+            "warehouse",
+            "database",
+            "schema",
+            "role",
+            "authenticator",
+        ):
+            value = getattr(self, key)
+            if value:
+                kwargs[key] = value
+        return kwargs
+
+
+@dataclass(frozen=True)
+class SnowflakeBridgeConnector(BaseSqlBridgeConnector):
+    settings: SnowflakeConnectorSettings
+    engine_id: str = "snowflake"
+
+    @property
+    def connection_id(self) -> str:
+        return self.settings.connection_id
+
+    @property
+    def title(self) -> str:
+        return self.settings.title
+
+    def config_dict(self) -> dict[str, Any]:
+        return self.settings.to_connect_kwargs()
+
+    def _connect(self):
+        try:
+            import snowflake.connector  # type: ignore
+        except ImportError as exc:
+            raise RuntimeError(
+                "Snowflake bridge requires `snowflake-connector-python`."
+            ) from exc
+
+        kwargs = self.settings.to_connect_kwargs()
+        if not kwargs.get("account") or not kwargs.get("user"):
+            raise RuntimeError(
+                "Snowflake bridge requires both account and user configuration."
+            )
+        return snowflake.connector.connect(**kwargs)
+
+    def dependency_diagnostics(self) -> dict[str, Any]:
+        try:
+            available = importlib.util.find_spec("snowflake.connector") is not None
+        except ModuleNotFoundError as exc:
+            return {
+                "available": False,
+                "reason": (
+                    f"Failed to inspect dependency via "
+                    f"find_spec('snowflake.connector'): {exc}"
+                ),
+                "error": "Missing Python dependency: snowflake-connector-python",
+                "requiredPackages": ["snowflake-connector-python>=4.3.0"],
+                "installCommands": SNOWFLAKE_INSTALL_COMMANDS,
+            }
+        if available:
+            return {"available": True}
+        return {
+            "available": False,
+            "error": "Missing Python dependency: snowflake-connector-python",
+            "requiredPackages": ["snowflake-connector-python>=4.3.0"],
+            "installCommands": SNOWFLAKE_INSTALL_COMMANDS,
+        }
+
+    def list_catalog(self) -> dict[str, list[dict[str, Any]]]:
+        with self._connect() as conn:
+            with conn.cursor() as cur:
+                cur.execute("SHOW DATABASES")
+                db_rows = cur.fetchall()
+                db_columns = cursor_columns(cur)
+                databases: list[dict[str, Any]] = []
+                for row in rows_to_json_rows(db_rows, db_columns):
+                    normalized = {str(k).lower(): v for k, v in row.items()}
+                    name = normalized.get("name") or normalized.get("database_name")
+                    if name:
+                        databases.append({"database": str(name)})
+
+                current_db = self.settings.database
+                if not current_db:
+                    cur.execute("SELECT CURRENT_DATABASE()")
+                    maybe_db = cur.fetchone()
+                    if maybe_db and maybe_db[0]:
+                        current_db = str(maybe_db[0])
+
+                schemas: list[dict[str, Any]] = []
+                tables: list[dict[str, Any]] = []
+
+                if current_db:
+                    db_ref = quoted_ident(current_db)
+                    cur.execute(
+                        f"""
+                        SELECT SCHEMA_NAME
+                        FROM {db_ref}.INFORMATION_SCHEMA.SCHEMATA
+                        ORDER BY SCHEMA_NAME
+                        """
+                    )
+                    schemas = [
+                        {"database": current_db, "schema": row[0]}
+                        for row in cur.fetchall()
+                    ]
+                    cur.execute(
+                        f"""
+                        SELECT TABLE_SCHEMA, TABLE_NAME, TABLE_TYPE
+                        FROM {db_ref}.INFORMATION_SCHEMA.TABLES
+                        ORDER BY TABLE_SCHEMA, TABLE_NAME
+                        """
+                    )
+                    tables = [
+                        {
+                            "database": current_db,
+                            "schema": row[0],
+                            "table": row[1],
+                            "isView": str(row[2]).upper().endswith("VIEW"),
+                        }
+                        for row in cur.fetchall()
+                    ]
+
+        return {"databases": databases, "schemas": schemas, "tables": tables}

sqlrooms/web/db_bridge/factory.py

@@ -0,0 +1,91 @@
+from __future__ import annotations
+
+from typing import Any
+
+from .connectors import (
+    PostgresBridgeConnector,
+    PostgresConnectorSettings,
+    SnowflakeBridgeConnector,
+    SnowflakeConnectorSettings,
+)
+from .connectors.base import BaseSqlBridgeConnector
+from .registry import DbBridgeRegistry
+
+SUPPORTED_ENGINES: list[str] = ["postgres", "snowflake"]
+
+ENGINE_CONFIG_FIELDS: dict[str, list[dict[str, Any]]] = {
+    "postgres": [
+        {"key": "host", "label": "Host", "placeholder": "localhost", "required": True},
+        {"key": "port", "label": "Port", "placeholder": "5432"},
+        {
+            "key": "database",
+            "label": "Database",
+            "placeholder": "my_db",
+            "required": True,
+        },
+        {"key": "user", "label": "User", "placeholder": "postgres", "required": True},
+        {"key": "password", "label": "Password", "placeholder": "", "secret": True},
+    ],
+    "snowflake": [
+        {
+            "key": "account",
+            "label": "Account",
+            "placeholder": "xy12345.us-east-1",
+            "required": True,
+        },
+        {"key": "user", "label": "User", "placeholder": "my_user", "required": True},
+        {"key": "password", "label": "Password", "placeholder": "", "secret": True},
+        {"key": "warehouse", "label": "Warehouse", "placeholder": "COMPUTE_WH"},
+        {"key": "database", "label": "Database", "placeholder": "MY_DB"},
+        {"key": "schema", "label": "Schema", "placeholder": "PUBLIC"},
+        {"key": "role", "label": "Role", "placeholder": "ACCOUNTADMIN"},
+        {"key": "authenticator", "label": "Authenticator", "placeholder": "snowflake"},
+    ],
+}
+
+
+def build_cli_db_bridge_registry(
+    *,
+    bridge_id: str,
+    connector_settings: list[PostgresConnectorSettings | SnowflakeConnectorSettings]
+    | None = None,
+) -> DbBridgeRegistry:
+    registry = DbBridgeRegistry(bridge_id=bridge_id)
+    for settings in connector_settings or []:
+        if isinstance(settings, PostgresConnectorSettings):
+            registry.register(PostgresBridgeConnector(settings=settings))
+            continue
+        if isinstance(settings, SnowflakeConnectorSettings) and settings.is_enabled():
+            registry.register(SnowflakeBridgeConnector(settings=settings))
+    return registry
+
+
+def build_ephemeral_connector(
+    engine: str,
+    config: dict[str, str],
+) -> BaseSqlBridgeConnector:
+    """Create a throwaway connector from raw config for ad-hoc connection tests."""
+    if engine == "postgres":
+        return PostgresBridgeConnector(
+            settings=PostgresConnectorSettings(
+                host=config.get("host", "localhost"),
+                port=config.get("port", "5432"),
+                database=config.get("database", ""),
+                user=config.get("user", ""),
+                password=config.get("password"),
+            )
+        )
+    if engine == "snowflake":
+        return SnowflakeBridgeConnector(
+            settings=SnowflakeConnectorSettings(
+                account=config.get("account"),
+                user=config.get("user"),
+                password=config.get("password"),
+                warehouse=config.get("warehouse"),
+                database=config.get("database"),
+                schema=config.get("schema"),
+                role=config.get("role"),
+                authenticator=config.get("authenticator"),
+            )
+        )
+    raise ValueError(f"Unsupported engine: {engine!r}")

sqlrooms/web/db_bridge/registry.py

@@ -0,0 +1,113 @@
+from __future__ import annotations
+
+from typing import Any
+
+from .types import DbBridgeConnector
+
+_SECRET_KEYS: frozenset[str] | None = None
+
+
+def _get_secret_keys() -> frozenset[str]:
+    global _SECRET_KEYS
+    if _SECRET_KEYS is None:
+        from .factory import ENGINE_CONFIG_FIELDS
+
+        keys: set[str] = set()
+        for fields in ENGINE_CONFIG_FIELDS.values():
+            for f in fields:
+                if f.get("secret"):
+                    keys.add(f["key"])
+        _SECRET_KEYS = frozenset(keys)
+    return _SECRET_KEYS
+
+
+class UnknownBridgeConnectionError(KeyError):
+    """Raised when a bridge request references an unknown connection id."""
+
+
+class DbBridgeRegistry:
+    def __init__(self, *, bridge_id: str):
+        self.bridge_id = bridge_id
+        self._connectors: dict[str, DbBridgeConnector] = {}
+
+    def register(self, connector: DbBridgeConnector) -> None:
+        if connector.connection_id in self._connectors:
+            raise ValueError(
+                f"Duplicate DB bridge connection id: {connector.connection_id}"
+            )
+        self._connectors[connector.connection_id] = connector
+
+    def has_connections(self) -> bool:
+        return bool(self._connectors)
+
+    def has_engine(self, engine_id: str) -> bool:
+        return any(conn.engine_id == engine_id for conn in self._connectors.values())
+
+    def runtime_connections(self) -> list[dict[str, Any]]:
+        result = []
+        for conn in self._connectors.values():
+            entry: dict[str, Any] = {
+                "id": conn.connection_id,
+                "engineId": conn.engine_id,
+                "title": conn.title,
+                "runtimeSupport": "server",
+                "requiresBridge": True,
+                "bridgeId": self.bridge_id,
+                "isCore": False,
+            }
+            cfg = conn.config_dict()
+            if cfg:
+                secret_keys = _get_secret_keys()
+                entry["config"] = {k: v for k, v in cfg.items() if k not in secret_keys}
+            result.append(entry)
+        return result
+
+    def runtime_diagnostics(self) -> list[dict[str, Any]]:
+        diagnostics: list[dict[str, Any]] = []
+        for conn in self._connectors.values():
+            payload = conn.dependency_diagnostics()
+            payload["id"] = conn.connection_id
+            payload["engineId"] = conn.engine_id
+            payload["title"] = conn.title
+            diagnostics.append(payload)
+        return diagnostics
+
+    def test_connection(self, connection_id: str) -> bool:
+        return self._get_connector(connection_id).test_connection()
+
+    def list_catalog(self, connection_id: str) -> dict[str, list[dict[str, Any]]]:
+        return self._get_connector(connection_id).list_catalog()
+
+    def execute_query(
+        self,
+        connection_id: str,
+        sql: str,
+        query_type: str,
+    ) -> dict[str, Any]:
+        return self._get_connector(connection_id).execute_query(sql, query_type)
+
+    def fetch_arrow_bytes(self, connection_id: str, sql: str) -> bytes:
+        return self._get_connector(connection_id).fetch_arrow_bytes(sql)
+
+    def stream_arrow_batches(
+        self,
+        connection_id: str,
+        sql: str,
+        *,
+        chunk_rows: int = 5000,
+        query_id: str | None = None,
+    ):
+        return self._get_connector(connection_id).stream_arrow_batches(
+            sql, chunk_rows=chunk_rows, query_id=query_id
+        )
+
+    def cancel_query(self, connection_id: str, query_id: str) -> bool:
+        return self._get_connector(connection_id).cancel_query(query_id)
+
+    def _get_connector(self, connection_id: str) -> DbBridgeConnector:
+        connector = self._connectors.get(connection_id)
+        if connector is None:
+            raise UnknownBridgeConnectionError(
+                f"Unknown DB bridge connection: {connection_id}"
+            )
+        return connector

sqlrooms/web/db_bridge/types.py

@@ -0,0 +1,29 @@
+from __future__ import annotations
+
+from typing import Any, Iterable, Protocol
+
+
+class DbBridgeConnector(Protocol):
+    connection_id: str
+    engine_id: str
+    title: str
+
+    def config_dict(self) -> dict[str, Any]:
+        """Return engine-specific configuration as a flat dict."""
+        ...
+
+    def test_connection(self) -> bool: ...
+
+    def list_catalog(self) -> dict[str, list[dict[str, Any]]]: ...
+
+    def execute_query(self, sql: str, query_type: str) -> dict[str, Any]: ...
+
+    def fetch_arrow_bytes(self, sql: str) -> bytes: ...
+
+    def stream_arrow_batches(
+        self, sql: str, chunk_rows: int = 5000, query_id: str | None = None
+    ) -> Iterable[bytes]: ...
+
+    def cancel_query(self, query_id: str) -> bool: ...
+
+    def dependency_diagnostics(self) -> dict[str, Any]: ...

sqlrooms/web/db_bridge/utils.py

@@ -0,0 +1,45 @@
+from __future__ import annotations
+
+from typing import Any, Iterable
+
+import pyarrow as pa
+
+
+def cursor_columns(cursor: Any) -> list[str]:
+    description = getattr(cursor, "description", None) or []
+    columns: list[str] = []
+    for item in description:
+        if hasattr(item, "name"):
+            columns.append(str(item.name))
+        elif isinstance(item, (list, tuple)) and item:
+            columns.append(str(item[0]))
+    return columns
+
+
+def rows_to_json_rows(
+    rows: Iterable[Iterable[Any]],
+    columns: list[str],
+) -> list[dict[str, Any]]:
+    return [dict(zip(columns, row)) for row in rows]
+
+
+def rows_to_arrow_bytes(
+    rows: Iterable[Iterable[Any]],
+    columns: list[str],
+) -> bytes:
+    json_rows = rows_to_json_rows(rows, columns)
+    if json_rows:
+        table = pa.Table.from_pylist(json_rows)
+    elif columns:
+        table = pa.table({column: pa.array([], type=pa.null()) for column in columns})
+    else:
+        table = pa.table({})
+
+    sink = pa.BufferOutputStream()
+    with pa.ipc.new_stream(sink, table.schema) as writer:
+        writer.write_table(table)
+    return sink.getvalue().to_pybytes()
+
+
+def quoted_ident(ident: str) -> str:
+    return '"' + ident.replace('"', '""') + '"'