thunderduck-sqlalchemy 0.1.0__py3-none-any.whl

thunderduck_sqlalchemy/__init__.py

@@ -0,0 +1,13 @@
+"""SQLAlchemy dialect and DBAPI for thunderduck.io.
+
+Read-only by design: every statement is submitted to console-api, which runs
+it as a Kubernetes Job, and results are paged back over HTTP. See README.md
+for the connection URL format and the caveats that follow from that model.
+"""
+
+from .dbapi import connect
+from .dialect import ThunderduckDialect
+
+__version__ = "0.1.0"
+
+__all__ = ["__version__", "connect", "ThunderduckDialect"]

thunderduck_sqlalchemy/catalog.py

@@ -0,0 +1,76 @@
+"""Flatten console-api's `/catalog/schema` tree into `(schema, table)` pairs.
+
+## Two catalog shapes, one rule
+
+thunderduck names are not uniformly three-part. Nessie catalogs are 2-level
+(`catalog.table`) and put their tables directly under the catalog; flat-file
+catalogs are 3-level (`catalog.schema.table`) and put them under named
+namespaces. SQLAlchemy and Superset both model exactly two levels.
+
+The rule that reconciles them: **the composite `schema` is the whole dotted
+prefix and the `table` is the leaf.** So `iris` in the nessie catalog
+`thunder_duck_demo` has schema `"thunder_duck_demo"`, while `lego_sets` has
+schema `"lego.public"`. The dialect's identifier preparer later splits that
+prefix back apart and quotes each segment, so both compile correctly.
+
+A catalog carrying an `error` (console-api could not introspect it) is skipped
+rather than raised on: one unreachable catalog must not break reflection of
+every other one.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from typing import Any
+
+
+@dataclass(frozen=True)
+class TableRef:
+    schema: str
+    table: str
+    kind: str = "table"
+    columns: list[dict[str, str]] = field(default_factory=list)
+
+
+def flatten(tree: dict[str, Any] | None) -> list[TableRef]:
+    refs: list[TableRef] = []
+    for catalog in (tree or {}).get("catalogs") or []:
+        if catalog.get("error"):
+            continue
+        catalog_name = catalog.get("name")
+        if not catalog_name:
+            continue
+        for table in catalog.get("tables") or []:
+            ref = _make_ref(catalog_name, table)
+            if ref:
+                refs.append(ref)
+        for namespace in catalog.get("namespaces") or []:
+            namespace_name = namespace.get("name")
+            if not namespace_name:
+                continue
+            schema = f"{catalog_name}.{namespace_name}"
+            for table in namespace.get("tables") or []:
+                ref = _make_ref(schema, table)
+                if ref:
+                    refs.append(ref)
+    return refs
+
+
+def _make_ref(schema: str, table: dict[str, Any]) -> TableRef | None:
+    name = table.get("name")
+    if not name:
+        return None
+    kind = str(table.get("kind") or "table").lower()
+    columns = [
+        {"name": str(c.get("name")), "type": str(c.get("type") or "")}
+        for c in (table.get("columns") or [])
+        if c.get("name")
+    ]
+    return TableRef(schema=schema, table=str(name), kind=kind, columns=columns)
+
+
+def split_schema(schema: str) -> list[str]:
+    """Split a composite schema into its segments (`"lego.public"` -> 2)."""
+    if not schema:
+        return []
+    return schema.split(".")

thunderduck_sqlalchemy/client.py

@@ -0,0 +1,216 @@
+"""HTTP transport for console-api.
+
+This is the only module that speaks HTTP. It deliberately knows nothing about
+cursors or SQLAlchemy: it submits SQL, reports execution state, pages results
+and cancels runs. The Java JDBC driver's `ApiClient.java` is the reference
+implementation of the same protocol.
+
+Two design notes:
+
+* **Polling, not streaming.** Submitting returns immediately with an execution
+  id; the run happens in a Kubernetes Job. `await_execution` polls
+  `GET /executions/{id}` until a terminal status. console-api also offers an
+  SSE endpoint (`POST /queries/stream`) which would cut first-row latency;
+  that is a deliberate follow-up, not used here.
+* **The token never appears in a URL or an error message.** It travels only in
+  the `Authorization` header, and error construction never interpolates it.
+"""
+
+from __future__ import annotations
+
+import time
+from dataclasses import dataclass
+from typing import Any
+
+import httpx
+
+from . import exceptions as exc
+
+
+@dataclass(frozen=True)
+class Execution:
+    """State of one query run (a row of `query_executions`)."""
+
+    uuid: str
+    status: str
+    error_message: str | None = None
+    result_location: str | None = None
+    result_row_count: int | None = None
+
+
+@dataclass(frozen=True)
+class ResultPage:
+    """One page of a stored result set."""
+
+    columns: list[str]
+    column_types: list[str]
+    rows: list[list[Any]]
+    total: int
+    offset: int
+
+
+class NoStoredResult(Exception):
+    """The execution stored no result set (a non-SELECT statement).
+
+    Deliberately *not* a PEP 249 error: `dbapi.Cursor` catches it and presents
+    an empty result set, which is what a DBAPI caller expects after running
+    something that returns no rows.
+    """
+
+
+#: Statuses that mean the run is over, one way or another.
+TERMINAL_STATUSES = frozenset({"success", "error", "cancelled"})
+
+
+class ThunderduckClient:
+    def __init__(
+        self,
+        base_url: str,
+        token: str,
+        *,
+        poll_interval: float = 1.0,
+        timeout: float = 300.0,
+        transport: httpx.BaseTransport | None = None,
+    ) -> None:
+        self._token = token
+        self._poll_interval = poll_interval
+        self._timeout = timeout
+        self._http = httpx.Client(
+            base_url=base_url,
+            transport=transport,
+            timeout=httpx.Timeout(30.0),
+            headers={
+                "Authorization": f"Bearer {token}",
+                "Accept": "application/json",
+            },
+        )
+
+    # ---- public API ----
+
+    def submit(self, sql: str) -> str:
+        payload = self._request("POST", "/queries", json={"sql": sql})
+        exec_id = payload.get("id")
+        if not exec_id:
+            raise exc.InternalError("console-api did not return an execution id")
+        return str(exec_id)
+
+    def get_execution(self, exec_id: str) -> Execution:
+        payload = self._request("GET", f"/executions/{exec_id}")
+        return Execution(
+            uuid=str(payload.get("uuid", exec_id)),
+            status=str(payload.get("status", "")),
+            error_message=payload.get("error_message"),
+            result_location=payload.get("result_location"),
+            result_row_count=payload.get("result_row_count"),
+        )
+
+    def await_execution(self, exec_id: str) -> Execution:
+        """Poll until the run reaches a terminal status, then return it.
+
+        Raises on `error` and `cancelled` so callers only ever see a
+        successful `Execution` come back.
+        """
+        deadline = time.monotonic() + self._timeout
+        while True:
+            execution = self.get_execution(exec_id)
+            if execution.status in TERMINAL_STATUSES:
+                if execution.status == "error":
+                    raise exc.ProgrammingError(
+                        execution.error_message or "query failed without a message"
+                    )
+                if execution.status == "cancelled":
+                    raise exc.OperationalError("query was cancelled")
+                return execution
+            if time.monotonic() >= deadline:
+                raise exc.OperationalError(
+                    f"query timed out after {self._timeout}s (still {execution.status})"
+                )
+            if self._poll_interval:
+                time.sleep(self._poll_interval)
+
+    def fetch_page(self, exec_id: str, limit: int, offset: int) -> ResultPage:
+        payload = self._request(
+            "GET",
+            f"/executions/{exec_id}/result",
+            params={"limit": limit, "offset": offset},
+            no_result_is_signal=True,
+        )
+        return ResultPage(
+            columns=list(payload.get("columns") or []),
+            column_types=list(payload.get("column_types") or []),
+            rows=[list(r) for r in (payload.get("rows") or [])],
+            total=int(payload.get("total") or 0),
+            offset=int(payload.get("offset") or offset),
+        )
+
+    def cancel(self, exec_id: str) -> None:
+        """Best-effort cancel; a finished or unknown run is not an error.
+
+        "Best effort" covers the run's *state* (already finished, never
+        existed) — not the credentials. A 401/403 means the token itself was
+        rejected, which must not be swallowed: doing so would make a stop
+        button appear to work while quietly hiding a rejected token.
+        """
+        try:
+            self._request("POST", f"/executions/{exec_id}/cancel")
+        except exc.InterfaceError:
+            raise
+        except exc.Error:
+            pass
+
+    def catalog_schema(self) -> dict[str, Any]:
+        return self._request("GET", "/catalog/schema")
+
+    def close(self) -> None:
+        self._http.close()
+
+    # ---- internals ----
+
+    def _request(
+        self,
+        method: str,
+        path: str,
+        *,
+        json: Any = None,
+        params: Any = None,
+        no_result_is_signal: bool = False,
+    ) -> dict[str, Any]:
+        try:
+            response = self._http.request(method, path, json=json, params=params)
+        except httpx.HTTPError as err:
+            # Never chain the original: httpx puts the full URL in its message,
+            # and a caller may have embedded a token in a query string.
+            raise exc.OperationalError(
+                f"could not reach thunderduck: {type(err).__name__}"
+            ) from None
+
+        if response.is_success:
+            if not response.content:
+                return {}
+            try:
+                return response.json()
+            except ValueError:
+                raise exc.InternalError("console-api returned a non-JSON success body") from None
+
+        detail = self._detail(response)
+        status = response.status_code
+        if status == 401 or status == 403:
+            raise exc.InterfaceError(f"authentication failed (check the API token): {detail}")
+        if status == 404 and no_result_is_signal:
+            raise NoStoredResult(detail)
+        if status == 409:
+            raise exc.OperationalError(detail)
+        if status == 400 or status == 422:
+            raise exc.ProgrammingError(detail)
+        raise exc.OperationalError(f"console-api returned {status}: {detail}")
+
+    @staticmethod
+    def _detail(response: httpx.Response) -> str:
+        """FastAPI puts the human-readable message in a JSON `detail` field."""
+        try:
+            body = response.json()
+        except ValueError:
+            return response.text.strip()[:200] or f"HTTP {response.status_code}"
+        if isinstance(body, dict) and "detail" in body:
+            return str(body["detail"])
+        return str(body)[:200]

thunderduck_sqlalchemy/dbapi.py

@@ -0,0 +1,306 @@
+"""PEP 249 interface over the thunderduck REST API.
+
+## What a "connection" is here
+
+Nothing on the server. console-api is stateless per request, so a
+`Connection` is just a configured HTTP client. That has three consequences
+worth knowing:
+
+* There is no transaction. `commit()` and `rollback()` are no-ops and the
+  dialect reports read-only.
+* Opening a connection costs nothing, so `NullPool` is the recommended
+  SQLAlchemy pool and pre-ping is pointless (a `SELECT 1` would start a
+  Kubernetes Job -- see `dialect.do_ping`).
+* A `Cursor` owns one execution. `cancel()` reaches the run through its
+  execution id, which is what the Superset engine spec uses to make the stop
+  button work.
+
+## Row paging
+
+`execute()` waits for the run to finish, then fetches the first page. The
+server clamps `limit` to its own maximum (`config.result_max_rows()`), so the
+cursor must keep requesting pages by `offset` until it has `total` rows --
+never assume the requested `limit` was honoured.
+"""
+
+from __future__ import annotations
+
+import logging
+from typing import Any
+
+import httpx
+
+from . import exceptions as exc
+from .client import NoStoredResult, ThunderduckClient
+from .exceptions import (
+    DatabaseError,
+    DataError,
+    Error,
+    IntegrityError,
+    InterfaceError,
+    InternalError,
+    NotSupportedError,
+    OperationalError,
+    ProgrammingError,
+)
+from .types import render_sql, sqlalchemy_type
+
+_log = logging.getLogger(__name__)
+
+apilevel = "2.0"
+#: 1 == "threads may share the module, but not connections".
+threadsafety = 1
+paramstyle = "pyformat"
+
+__all__ = [
+    "apilevel",
+    "threadsafety",
+    "paramstyle",
+    "connect",
+    "Connection",
+    "Cursor",
+    "Error",
+    "InterfaceError",
+    "DatabaseError",
+    "DataError",
+    "OperationalError",
+    "IntegrityError",
+    "InternalError",
+    "ProgrammingError",
+    "NotSupportedError",
+]
+
+DEFAULT_ARRAYSIZE = 1000
+
+
+def connect(
+    *,
+    host: str,
+    token: str,
+    port: int | None = None,
+    use_ssl: bool = True,
+    poll_interval: float = 1.0,
+    timeout: float = 300.0,
+    arraysize: int = DEFAULT_ARRAYSIZE,
+    transport: httpx.BaseTransport | None = None,
+) -> Connection:
+    if not host:
+        raise exc.InterfaceError("a thunderduck host is required")
+    if not token:
+        raise exc.InterfaceError(
+            "an API token is required -- pass it as the URL password or ?token=…"
+        )
+    scheme = "https" if use_ssl else "http"
+    netloc = f"{host}:{port}" if port else host
+    client = ThunderduckClient(
+        base_url=f"{scheme}://{netloc}",
+        token=token,
+        poll_interval=poll_interval,
+        timeout=timeout,
+        transport=transport,
+    )
+    return Connection(client, arraysize=arraysize)
+
+
+class Connection:
+    def __init__(self, client: ThunderduckClient, *, arraysize: int = DEFAULT_ARRAYSIZE) -> None:
+        self.client = client
+        self._arraysize = arraysize
+        self._closed = False
+
+    def cursor(self) -> Cursor:
+        if self._closed:
+            raise exc.ProgrammingError("connection is closed")
+        return Cursor(self.client, arraysize=self._arraysize)
+
+    def close(self) -> None:
+        if not self._closed:
+            self._closed = True
+            self.client.close()
+
+    def commit(self) -> None:
+        """No-op: thunderduck has no transactions (read-only driver)."""
+
+    def rollback(self) -> None:
+        """No-op: see `commit`."""
+
+    def __enter__(self) -> Connection:
+        return self
+
+    def __exit__(self, *_exc: object) -> None:
+        self.close()
+
+
+class Cursor:
+    def __init__(self, client: ThunderduckClient, *, arraysize: int = DEFAULT_ARRAYSIZE) -> None:
+        self._client = client
+        self.arraysize = arraysize
+        self._closed = False
+        self._exec_id: str | None = None
+        self._description: list[tuple] | None = None
+        self._rowcount = -1
+        self._buffer: list[tuple] = []
+        self._buffer_pos = 0
+        self._next_offset = 0
+        self._total = 0
+        self._exhausted = True
+
+    # ---- introspection ----
+
+    @property
+    def description(self) -> list[tuple] | None:
+        return self._description
+
+    @property
+    def rowcount(self) -> int:
+        return self._rowcount
+
+    @property
+    def execution_id(self) -> str | None:
+        """The thunderduck execution uuid of the last `execute`, if any."""
+        return self._exec_id
+
+    # ---- execution ----
+
+    def execute(self, operation: str, parameters: Any = None) -> Cursor:
+        self._check_open()
+        self._reset()
+        sql = render_sql(operation, parameters)
+        self._exec_id = self._client.submit(sql)
+        self._client.await_execution(self._exec_id)
+        self._load_first_page()
+        return self
+
+    def executemany(self, operation: str, seq_of_parameters: Any) -> None:
+        raise exc.NotSupportedError(
+            "executemany is not supported: thunderduck is read-only and every "
+            "statement is a separate Kubernetes Job"
+        )
+
+    def cancel(self) -> None:
+        """Best-effort cancel of the in-flight (or last) execution."""
+        if self._exec_id:
+            self._client.cancel(self._exec_id)
+
+    # ---- fetching ----
+
+    def fetchone(self) -> tuple | None:
+        self._check_executed()
+        if not self._ensure_buffered():
+            return None
+        row = self._buffer[self._buffer_pos]
+        self._buffer_pos += 1
+        return row
+
+    def fetchmany(self, size: int | None = None) -> list[tuple]:
+        self._check_executed()
+        want = self.arraysize if size is None else size
+        out: list[tuple] = []
+        while len(out) < want:
+            row = self.fetchone()
+            if row is None:
+                break
+            out.append(row)
+        return out
+
+    def fetchall(self) -> list[tuple]:
+        self._check_executed()
+        out: list[tuple] = []
+        while True:
+            row = self.fetchone()
+            if row is None:
+                return out
+            out.append(row)
+
+    def close(self) -> None:
+        self._closed = True
+
+    def __iter__(self):
+        while True:
+            row = self.fetchone()
+            if row is None:
+                return
+            yield row
+
+    # ---- internals ----
+
+    def _reset(self) -> None:
+        self._description = None
+        self._rowcount = -1
+        self._buffer = []
+        self._buffer_pos = 0
+        self._next_offset = 0
+        self._total = 0
+        self._exhausted = True
+
+    def _load_first_page(self) -> None:
+        assert self._exec_id is not None
+        try:
+            page = self._client.fetch_page(self._exec_id, limit=self.arraysize, offset=0)
+        except NoStoredResult:
+            # A non-SELECT statement: no result set, and that is fine.
+            self._description = None
+            self._rowcount = -1
+            self._buffer = []
+            self._exhausted = True
+            return
+        self._description = [
+            (
+                name,
+                sqlalchemy_type(type_name),
+                None,  # display_size
+                None,  # internal_size
+                None,  # precision
+                None,  # scale
+                True,  # null_ok -- /catalog/schema carries no nullability
+            )
+            for name, type_name in zip(
+                page.columns,
+                list(page.column_types) + [""] * (len(page.columns) - len(page.column_types)),
+                strict=False,
+            )
+        ]
+        self._total = page.total
+        self._rowcount = page.total
+        self._buffer = [tuple(r) for r in page.rows]
+        self._buffer_pos = 0
+        self._next_offset = len(page.rows)
+        self._exhausted = self._next_offset >= self._total or not page.rows
+
+    def _ensure_buffered(self) -> bool:
+        """True if a row is available at `_buffer_pos`, paging if needed."""
+        if self._buffer_pos < len(self._buffer):
+            return True
+        if self._exhausted:
+            return False
+        assert self._exec_id is not None
+        page = self._client.fetch_page(
+            self._exec_id, limit=self.arraysize, offset=self._next_offset
+        )
+        if not page.rows:
+            # The result set claims more rows than it can serve. Never silently
+            # return a short result -- a truncated dashboard looks correct.
+            _log.warning(
+                "thunderduck result set is short: served %d of %d claimed rows "
+                "for execution %s; the stored result may have been truncated",
+                self._next_offset,
+                self._total,
+                self._exec_id,
+            )
+            self._exhausted = True
+            return False
+        self._buffer = [tuple(r) for r in page.rows]
+        self._buffer_pos = 0
+        self._next_offset += len(page.rows)
+        self._total = page.total or self._total
+        self._exhausted = self._next_offset >= self._total
+        return True
+
+    def _check_open(self) -> None:
+        if self._closed:
+            raise exc.ProgrammingError("cursor is closed")
+
+    def _check_executed(self) -> None:
+        self._check_open()
+        if self._exec_id is None:
+            raise exc.ProgrammingError("no query has been executed on this cursor")