duckgresql 1.4.4.0__py3-none-any.whl

duckgresql/__init__.py

@@ -0,0 +1,123 @@
+"""DuckGresQL Python SDK — DuckDB-compatible client for remote databases.
+
+Quick start::
+
+    import duckgresql
+
+    conn = duckgresql.connect(token="dkgql_...", database="mydb")
+    result = conn.execute("SELECT * FROM users LIMIT 10")
+    print(result.fetchall())
+    conn.close()
+"""
+
+from __future__ import annotations
+
+from duckgresql._config import (
+    DEFAULT_FLIGHT_PORT,
+    DEFAULT_HOST,
+    DEFAULT_REST_PORT,
+    DEFAULT_REST_SCHEME,
+    DEFAULT_USE_TLS,
+)
+from duckgresql._types import JobStatus
+from duckgresql._version import __version__
+from duckgresql.async_connection import DuckgresqlAsync
+from duckgresql.async_job import AsyncJob, AsyncJobAsync
+from duckgresql.connection import Duckgresql
+from duckgresql.exceptions import (
+    AuthenticationError,
+    ConnectionError,
+    DuckgresqlError,
+    JobError,
+    QueryError,
+    TimeoutError,
+)
+from duckgresql.result import DuckgresqlResult
+
+
+def connect(
+    host: str = DEFAULT_HOST,
+    *,
+    token: str,
+    database: str,
+    port: int = DEFAULT_FLIGHT_PORT,
+    use_tls: bool = DEFAULT_USE_TLS,
+    rest_port: int = DEFAULT_REST_PORT,
+    rest_scheme: str = DEFAULT_REST_SCHEME,
+) -> Duckgresql:
+    """Create a synchronous connection to a DuckGresQL server.
+
+    Parameters
+    ----------
+    host : str
+        Server hostname or IP.  Defaults to ``DUCKGRESQL_HOST`` env var, or
+        the value baked in at release time.
+    token : str
+        API token (``dkgql_…`` prefix).
+    database : str
+        Database name or UUID.
+    port : int
+        Arrow Flight SQL (gRPC) port.  Defaults to ``DUCKGRESQL_FLIGHT_PORT``
+        env var, or the value baked in at release time.
+    use_tls : bool
+        Use TLS for Flight SQL.  Defaults to ``DUCKGRESQL_USE_TLS`` env var.
+    rest_port : int
+        REST API port.  Defaults to ``DUCKGRESQL_REST_PORT`` env var, or the
+        value baked in at release time.
+    rest_scheme : str
+        ``"http"`` or ``"https"``.  Defaults to ``DUCKGRESQL_REST_SCHEME``
+        env var.
+    """
+    return Duckgresql(
+        host,
+        token=token,
+        database=database,
+        port=port,
+        use_tls=use_tls,
+        rest_port=rest_port,
+        rest_scheme=rest_scheme,
+    )
+
+
+async def connect_async(
+    host: str = DEFAULT_HOST,
+    *,
+    token: str,
+    database: str,
+    port: int = DEFAULT_FLIGHT_PORT,
+    use_tls: bool = DEFAULT_USE_TLS,
+    rest_port: int = DEFAULT_REST_PORT,
+    rest_scheme: str = DEFAULT_REST_SCHEME,
+) -> DuckgresqlAsync:
+    """Create an asynchronous connection to a DuckGresQL server.
+
+    Same parameters as :func:`connect`.  Must be ``await``-ed.
+    """
+    return await DuckgresqlAsync.create(
+        host,
+        token=token,
+        database=database,
+        port=port,
+        use_tls=use_tls,
+        rest_port=rest_port,
+        rest_scheme=rest_scheme,
+    )
+
+
+__all__ = [
+    "__version__",
+    "connect",
+    "connect_async",
+    "Duckgresql",
+    "DuckgresqlAsync",
+    "DuckgresqlResult",
+    "AsyncJob",
+    "AsyncJobAsync",
+    "JobStatus",
+    "DuckgresqlError",
+    "ConnectionError",
+    "AuthenticationError",
+    "QueryError",
+    "JobError",
+    "TimeoutError",
+]

duckgresql/_config.py

@@ -0,0 +1,13 @@
+"""Connection defaults — hardcoded at release time.
+
+This file is auto-generated by ``scripts/inject_release_defaults.py``.
+Do not edit manually.
+"""
+
+from __future__ import annotations
+
+DEFAULT_HOST: str = 'api.duckgresql.io'
+DEFAULT_FLIGHT_PORT: int = 47470
+DEFAULT_REST_PORT: int = 443
+DEFAULT_USE_TLS: bool = False
+DEFAULT_REST_SCHEME: str = 'https'

duckgresql/_flight.py

@@ -0,0 +1,152 @@
+"""Low-level Arrow Flight SQL client for DuckGresQL."""
+
+from __future__ import annotations
+
+from typing import cast
+
+import pyarrow as pa
+import pyarrow.flight as flight
+
+from duckgresql.exceptions import AuthenticationError, ConnectionError, QueryError
+
+# ---------------------------------------------------------------------------
+# Minimal protobuf encoder for Flight SQL command descriptors
+# ---------------------------------------------------------------------------
+# The Arrow Flight SQL protocol requires GetFlightInfo/DoGet to receive a
+# FlightDescriptor whose `cmd` bytes are a serialised google.protobuf.Any
+# wrapping the appropriate command message (e.g. CommandStatementQuery).
+# Passing raw SQL bytes causes "proto: cannot parse invalid wire-format data".
+
+_COMMAND_TYPE_URL = (
+    "type.googleapis.com/arrow.flight.protocol.sql.CommandStatementQuery"
+)
+
+
+def _varint(n: int) -> bytes:
+    result = bytearray()
+    while n > 0x7F:
+        result.append((n & 0x7F) | 0x80)
+        n >>= 7
+    result.append(n)
+    return bytes(result)
+
+
+def _pb_string(field: int, value: str) -> bytes:
+    data = value.encode("utf-8")
+    return _varint((field << 3) | 2) + _varint(len(data)) + data
+
+
+def _pb_bytes_field(field: int, value: bytes) -> bytes:
+    return _varint((field << 3) | 2) + _varint(len(value)) + value
+
+
+def _flight_sql_command(query: str) -> bytes:
+    """Return bytes for google.protobuf.Any(CommandStatementQuery{query})."""
+    # CommandStatementQuery { string query = 1; }
+    cmd = _pb_string(1, query)
+    # google.protobuf.Any { string type_url = 1; bytes value = 2; }
+    return _pb_string(1, _COMMAND_TYPE_URL) + _pb_bytes_field(2, cmd)
+
+
+class FlightSQLClient:
+    """Thin wrapper around :class:`pyarrow.flight.FlightClient` for Flight SQL.
+
+    Authentication uses BasicAuth where *username* is the API token and
+    *password* is the database name.  The server returns a ``conn_`` bearer
+    token that is sent on every subsequent RPC.
+    """
+
+    def __init__(
+        self,
+        host: str,
+        port: int,
+        token: str,
+        database: str,
+        *,
+        use_tls: bool = False,
+    ) -> None:
+        scheme = "grpc+tls" if use_tls else "grpc"
+        location = f"{scheme}://{host}:{port}"
+        try:
+            self._client = flight.FlightClient(location)
+        except Exception as exc:
+            raise ConnectionError(f"Failed to connect to Flight SQL at {location}: {exc}") from exc
+
+        # BasicAuth handshake — username=token, password=database
+        try:
+            header_pair = self._client.authenticate_basic_token(token, database)
+            self._auth_header: tuple[bytes, bytes] = header_pair
+        except flight.FlightUnauthenticatedError as exc:
+            raise AuthenticationError(f"Flight SQL authentication failed: {exc}") from exc
+        except Exception as exc:
+            raise ConnectionError(f"Flight SQL handshake failed: {exc}") from exc
+
+        self._closed = False
+
+    # ------------------------------------------------------------------
+    # Internal helpers
+    # ------------------------------------------------------------------
+
+    def _call_options(self) -> flight.FlightCallOptions:
+        """Build call options with the bearer token from the handshake."""
+        return flight.FlightCallOptions(headers=[self._auth_header])
+
+    # ------------------------------------------------------------------
+    # Public API
+    # ------------------------------------------------------------------
+
+    def execute_query(self, query: str) -> pa.Table:
+        """Execute a read query and return the full result as a :class:`pyarrow.Table`."""
+        try:
+            descriptor = flight.FlightDescriptor.for_command(_flight_sql_command(query))
+            opts = self._call_options()
+            info = self._client.get_flight_info(descriptor, opts)
+
+            # Flight SQL returns one or more endpoints; read the first.
+            if not info.endpoints:
+                return pa.table({})
+
+            ticket = info.endpoints[0].ticket
+            reader = self._client.do_get(ticket, opts)
+            return reader.read_all()
+        except flight.FlightUnauthenticatedError as exc:
+            raise AuthenticationError(str(exc)) from exc
+        except Exception as exc:
+            raise QueryError(f"Query execution failed: {exc}") from exc
+
+    def execute_update(self, query: str) -> int:
+        """Execute a DML statement and return the number of affected rows."""
+        try:
+            opts = self._call_options()
+            # For DML we use the same descriptor path; the server decides
+            # based on the SQL whether to return rows or an update count.
+            descriptor = flight.FlightDescriptor.for_command(_flight_sql_command(query))
+            info = self._client.get_flight_info(descriptor, opts)
+
+            if not info.endpoints:
+                return 0
+
+            ticket = info.endpoints[0].ticket
+            reader = self._client.do_get(ticket, opts)
+            table = reader.read_all()
+
+            # The server returns a single-row table with the affected count
+            # when the statement is DML. If it returns regular rows, return
+            # the row count instead.
+            if table.num_columns == 1 and table.column_names[0] == "affected_rows":
+                return int(table.column(0)[0].as_py())
+            return cast(int, table.num_rows)
+        except flight.FlightUnauthenticatedError as exc:
+            raise AuthenticationError(str(exc)) from exc
+        except Exception as exc:
+            raise QueryError(f"Update execution failed: {exc}") from exc
+
+    def close(self) -> None:
+        """Close the underlying Flight client."""
+        if not self._closed:
+            self._client.close()
+            self._closed = True
+
+    @property
+    def closed(self) -> bool:
+        return self._closed

duckgresql/_rest.py

@@ -0,0 +1,125 @@
+"""Synchronous REST client for DuckGresQL async query endpoints."""
+
+from __future__ import annotations
+
+from typing import Any, cast
+
+import httpx
+
+from duckgresql.exceptions import AuthenticationError, ConnectionError, JobError
+
+
+class RestClient:
+    """Thin synchronous wrapper around :mod:`httpx` for DuckGresQL REST API.
+
+    Handles the ``/connect`` handshake and async-job lifecycle endpoints.
+    """
+
+    def __init__(self, base_url: str, *, timeout: float = 30.0) -> None:
+        self._base_url = base_url.rstrip("/")
+        self._http = httpx.Client(base_url=self._base_url, timeout=timeout)
+        self._closed = False
+
+    # ------------------------------------------------------------------
+    # Connection
+    # ------------------------------------------------------------------
+
+    def connect(self, token: str, database: str) -> str:
+        """Exchange an API token + database for a connection token.
+
+        Returns the ``conn_`` prefixed token string.
+        """
+        try:
+            resp = self._http.post(
+                "/connect",
+                json={"database": database},
+                headers={"Authorization": f"Bearer {token}"},
+            )
+        except httpx.HTTPError as exc:
+            raise ConnectionError(f"REST connect failed: {exc}") from exc
+
+        if resp.status_code == 401:
+            body = resp.json()
+            msg = body.get("error", {}).get("message", "Authentication failed")
+            raise AuthenticationError(msg)
+        if resp.status_code >= 400:
+            raise ConnectionError(f"REST connect returned {resp.status_code}: {resp.text}")
+
+        return cast(str, resp.json()["connection_token"])
+
+    # ------------------------------------------------------------------
+    # Async job endpoints
+    # ------------------------------------------------------------------
+
+    def _headers(self, conn_token: str) -> dict[str, str]:
+        return {"Authorization": f"Bearer {conn_token}"}
+
+    def submit_async(
+        self,
+        conn_token: str,
+        query: str,
+        bindings: Any | None = None,
+    ) -> str:
+        """Submit an async query and return the ``job_id``."""
+        payload: dict[str, Any] = {"query": query}
+        if bindings is not None:
+            payload["bindings"] = bindings
+
+        resp = self._http.post(
+            "/query/async",
+            json=payload,
+            headers=self._headers(conn_token),
+        )
+        self._check_response(resp)
+        return cast(str, resp.json()["job_id"])
+
+    def get_job(self, conn_token: str, job_id: str) -> dict[str, Any]:
+        """Get status/metadata for a specific job."""
+        resp = self._http.get(
+            f"/query/jobs/{job_id}",
+            headers=self._headers(conn_token),
+        )
+        self._check_response(resp)
+        return cast(dict[str, Any], resp.json())
+
+    def get_job_result(self, conn_token: str, job_id: str) -> dict[str, Any]:
+        """Get the result rows for a completed job."""
+        resp = self._http.get(
+            f"/query/jobs/{job_id}/result",
+            headers=self._headers(conn_token),
+        )
+        self._check_response(resp)
+        return cast(dict[str, Any], resp.json())
+
+    def cancel_job(self, conn_token: str, job_id: str) -> None:
+        """Request cancellation of a pending/running job."""
+        resp = self._http.post(
+            f"/query/jobs/{job_id}/cancel",
+            headers=self._headers(conn_token),
+        )
+        self._check_response(resp)
+
+    # ------------------------------------------------------------------
+    # Helpers
+    # ------------------------------------------------------------------
+
+    @staticmethod
+    def _check_response(resp: httpx.Response) -> None:
+        if resp.status_code == 401:
+            raise AuthenticationError("Connection token invalid or expired")
+        if resp.status_code >= 400:
+            try:
+                body = resp.json()
+                msg = body.get("error", {}).get("message", resp.text)
+            except Exception:
+                msg = resp.text
+            raise JobError(f"REST request failed ({resp.status_code}): {msg}")
+
+    def close(self) -> None:
+        if not self._closed:
+            self._http.close()
+            self._closed = True
+
+    @property
+    def closed(self) -> bool:
+        return self._closed

duckgresql/_rest_async.py

@@ -0,0 +1,119 @@
+"""Asynchronous REST client for DuckGresQL async query endpoints."""
+
+from __future__ import annotations
+
+from typing import Any, cast
+
+import httpx
+
+from duckgresql.exceptions import AuthenticationError, ConnectionError, JobError
+
+
+class AsyncRestClient:
+    """Thin asynchronous wrapper around :mod:`httpx` for DuckGresQL REST API."""
+
+    def __init__(self, base_url: str, *, timeout: float = 30.0) -> None:
+        self._base_url = base_url.rstrip("/")
+        self._http = httpx.AsyncClient(base_url=self._base_url, timeout=timeout)
+        self._closed = False
+
+    # ------------------------------------------------------------------
+    # Connection
+    # ------------------------------------------------------------------
+
+    async def connect(self, token: str, database: str) -> str:
+        """Exchange an API token + database for a connection token."""
+        try:
+            resp = await self._http.post(
+                "/connect",
+                json={"database": database},
+                headers={"Authorization": f"Bearer {token}"},
+            )
+        except httpx.HTTPError as exc:
+            raise ConnectionError(f"REST connect failed: {exc}") from exc
+
+        if resp.status_code == 401:
+            body = resp.json()
+            msg = body.get("error", {}).get("message", "Authentication failed")
+            raise AuthenticationError(msg)
+        if resp.status_code >= 400:
+            raise ConnectionError(f"REST connect returned {resp.status_code}: {resp.text}")
+
+        return cast(str, resp.json()["connection_token"])
+
+    # ------------------------------------------------------------------
+    # Async job endpoints
+    # ------------------------------------------------------------------
+
+    def _headers(self, conn_token: str) -> dict[str, str]:
+        return {"Authorization": f"Bearer {conn_token}"}
+
+    async def submit_async(
+        self,
+        conn_token: str,
+        query: str,
+        bindings: Any | None = None,
+    ) -> str:
+        """Submit an async query and return the ``job_id``."""
+        payload: dict[str, Any] = {"query": query}
+        if bindings is not None:
+            payload["bindings"] = bindings
+
+        resp = await self._http.post(
+            "/query/async",
+            json=payload,
+            headers=self._headers(conn_token),
+        )
+        self._check_response(resp)
+        return cast(str, resp.json()["job_id"])
+
+    async def get_job(self, conn_token: str, job_id: str) -> dict[str, Any]:
+        """Get status/metadata for a specific job."""
+        resp = await self._http.get(
+            f"/query/jobs/{job_id}",
+            headers=self._headers(conn_token),
+        )
+        self._check_response(resp)
+        return cast(dict[str, Any], resp.json())
+
+    async def get_job_result(self, conn_token: str, job_id: str) -> dict[str, Any]:
+        """Get the result rows for a completed job."""
+        resp = await self._http.get(
+            f"/query/jobs/{job_id}/result",
+            headers=self._headers(conn_token),
+        )
+        self._check_response(resp)
+        return cast(dict[str, Any], resp.json())
+
+    async def cancel_job(self, conn_token: str, job_id: str) -> None:
+        """Request cancellation of a pending/running job."""
+        resp = await self._http.post(
+            f"/query/jobs/{job_id}/cancel",
+            headers=self._headers(conn_token),
+        )
+        self._check_response(resp)
+
+    # ------------------------------------------------------------------
+    # Helpers
+    # ------------------------------------------------------------------
+
+    @staticmethod
+    def _check_response(resp: httpx.Response) -> None:
+        if resp.status_code == 401:
+            raise AuthenticationError("Connection token invalid or expired")
+        if resp.status_code >= 400:
+            try:
+                body = resp.json()
+                msg = body.get("error", {}).get("message", resp.text)
+            except Exception:
+                msg = resp.text
+            raise JobError(f"REST request failed ({resp.status_code}): {msg}")
+
+    async def close(self) -> None:
+        if not self._closed:
+            await self._http.aclose()
+            self._closed = True
+
+    @property
+    def closed(self) -> bool:
+        return self._closed

duckgresql/_types.py

@@ -0,0 +1,35 @@
+"""Internal types and constants for the DuckGresQL Python SDK."""
+
+from __future__ import annotations
+
+import enum
+
+
+class JobStatus(enum.Enum):
+    """Status of an async query job."""
+
+    PENDING = "pending"
+    RUNNING = "running"
+    COMPLETED = "completed"
+    FAILED = "failed"
+    CANCELLED = "cancelled"
+
+
+# SQL prefixes that indicate a read query (returns rows).
+_READ_PREFIXES = frozenset({
+    "SELECT",
+    "WITH",
+    "EXPLAIN",
+    "SHOW",
+    "DESCRIBE",
+    "PRAGMA",
+    "TABLE",
+    "FROM",
+    "VALUES",
+})
+
+
+def _is_read_query(sql: str) -> bool:
+    """Return True if *sql* looks like a read query (SELECT, etc.)."""
+    first_word = sql.lstrip().split(None, 1)[0].upper() if sql.strip() else ""
+    return first_word in _READ_PREFIXES

duckgresql/_version.py

@@ -0,0 +1 @@
+__version__ = "1.4.4.0"