py-flink-sql-gateway 0.1.0.dev0__py3-none-any.whl

flink_gateway/__init__.py

@@ -0,0 +1,104 @@
+"""Flink SQL Gateway — Python client library.
+
+Implements PEP 249 (DB-API 2.0) on top of the Flink SQL Gateway REST API.
+"""
+
+from flink_gateway.client import FlinkSqlGatewayClient
+from flink_gateway.connection import Connection, connect
+from flink_gateway.cursor import Cursor
+from flink_gateway.exceptions import (
+    DatabaseError,
+    Error,
+    FlinkSqlGatewayError,
+    InterfaceError,
+    NotSupportedError,
+    OperationalError,
+    ProgrammingError,
+)
+from flink_gateway.models import (
+    ColumnInfo,
+    CompleteStatementRequest,
+    ConfigureSessionRequest,
+    ExecuteStatementRequest,
+    FetchResultsResponse,
+    InfoResponse,
+    LogicalType,
+    OpenSessionRequest,
+    RefreshMaterializedTableRequest,
+    ResultKind,
+    ResultSet,
+    ResultType,
+    RowData,
+)
+from flink_gateway.types import (
+    BINARY,
+    DATETIME,
+    NUMBER,
+    ROWID,
+    STRING,
+    Binary,
+    Date,
+    DateFromTicks,
+    DBAPITypeObject,
+    FlinkType,
+    Time,
+    TimeFromTicks,
+    Timestamp,
+    TimestampFromTicks,
+)
+
+# PEP 249 module-level globals
+apilevel = "2.0"
+threadsafety = 1  # Threads may share the module but not connections.
+paramstyle = "qmark"
+
+__all__ = [
+    # PEP 249
+    "apilevel",
+    "threadsafety",
+    "paramstyle",
+    "connect",
+    "Connection",
+    "Cursor",
+    # Exceptions
+    "Error",
+    "InterfaceError",
+    "DatabaseError",
+    "OperationalError",
+    "FlinkSqlGatewayError",
+    "ProgrammingError",
+    "NotSupportedError",
+    # Type objects
+    "STRING",
+    "BINARY",
+    "NUMBER",
+    "DATETIME",
+    "ROWID",
+    "DBAPITypeObject",
+    "FlinkType",
+    # PEP 249 constructors
+    "Date",
+    "Time",
+    "Timestamp",
+    "DateFromTicks",
+    "TimeFromTicks",
+    "TimestampFromTicks",
+    "Binary",
+    # Low-level client
+    "FlinkSqlGatewayClient",
+    "FlinkSqlGatewayError",
+    # Models
+    "ColumnInfo",
+    "CompleteStatementRequest",
+    "ConfigureSessionRequest",
+    "ExecuteStatementRequest",
+    "FetchResultsResponse",
+    "InfoResponse",
+    "LogicalType",
+    "OpenSessionRequest",
+    "RefreshMaterializedTableRequest",
+    "ResultKind",
+    "ResultSet",
+    "ResultType",
+    "RowData",
+]

flink_gateway/_version.py

@@ -0,0 +1,34 @@
+# file generated by setuptools-scm
+# don't change, don't track in version control
+
+__all__ = [
+    "__version__",
+    "__version_tuple__",
+    "version",
+    "version_tuple",
+    "__commit_id__",
+    "commit_id",
+]
+
+TYPE_CHECKING = False
+if TYPE_CHECKING:
+    from typing import Tuple
+    from typing import Union
+
+    VERSION_TUPLE = Tuple[Union[int, str], ...]
+    COMMIT_ID = Union[str, None]
+else:
+    VERSION_TUPLE = object
+    COMMIT_ID = object
+
+version: str
+__version__: str
+__version_tuple__: VERSION_TUPLE
+version_tuple: VERSION_TUPLE
+commit_id: COMMIT_ID
+__commit_id__: COMMIT_ID
+
+__version__ = version = '0.1.0.dev0'
+__version_tuple__ = version_tuple = (0, 1, 0, 'dev0')
+
+__commit_id__ = commit_id = None

flink_gateway/client.py

@@ -0,0 +1,302 @@
+"""Flink SQL Gateway REST API client.
+
+This module provides a low-level HTTP client that wraps every endpoint of
+the Apache Flink SQL Gateway REST API.
+"""
+
+from __future__ import annotations
+
+from typing import Any
+
+import httpx
+
+from flink_gateway.exceptions import FlinkSqlGatewayError
+from flink_gateway.models import (
+    CompleteStatementRequest,
+    ConfigureSessionRequest,
+    ExecuteStatementRequest,
+    FetchResultsResponse,
+    InfoResponse,
+    OpenSessionRequest,
+    RefreshMaterializedTableRequest,
+)
+
+
+class FlinkSqlGatewayClient:
+    """Low-level client for the Flink SQL Gateway REST API.
+
+    Args:
+        base_url: Root URL of the SQL Gateway, e.g. ``http://localhost:8083``.
+        http_client: Optional pre-configured ``httpx.Client``.  When *None*, a
+            default client is created with a 30-second timeout.
+        api_version: REST API version prefix (``"v1"``, ``"v2"``, or ``"v3"``).
+            Defaults to ``"v3"``.
+    """
+
+    def __init__(
+        self,
+        base_url: str,
+        http_client: httpx.Client | None = None,
+        api_version: str = "v3",
+    ) -> None:
+        self._base_url = base_url.rstrip("/")
+        self._api_version = api_version or "v3"
+        self._owns_client = http_client is None
+        self._client = http_client or httpx.Client(timeout=30.0)
+
+    # ── Context manager ────────────────────────────────────────────────
+
+    def __enter__(self) -> FlinkSqlGatewayClient:
+        return self
+
+    def __exit__(self, *_: Any) -> None:
+        self.close()
+
+    def close(self) -> None:
+        """Close the underlying HTTP client if we own it."""
+        if self._owns_client:
+            self._client.close()
+
+    # ── URL helpers ────────────────────────────────────────────────────
+
+    def _build_endpoint(self, *segments: str) -> str:
+        parts = [self._base_url, self._api_version, *segments]
+        return "/".join(parts)
+
+    # ── Metadata ───────────────────────────────────────────────────────
+
+    def get_info(self) -> InfoResponse:
+        """GET /info — cluster metadata."""
+        url = self._build_endpoint("info")
+        resp = self._client.get(url)
+        self._check_response(resp, "get info")
+        data = resp.json()
+        return InfoResponse(
+            product_name=data.get("productName", ""),
+            version=data.get("version", ""),
+        )
+
+    def get_api_versions(self) -> list[str]:
+        """GET /api_versions — supported REST API versions."""
+        url = self._build_endpoint("api_versions")
+        resp = self._client.get(url)
+        self._check_response(resp, "get api versions")
+        return resp.json().get("versions", [])
+
+    # ── Session management ─────────────────────────────────────────────
+
+    def open_session(
+        self,
+        request: OpenSessionRequest | None = None,
+    ) -> str:
+        """POST /sessions — open a new session.
+
+        Returns:
+            The session handle string.
+        """
+        url = self._build_endpoint("sessions")
+        body = request.to_dict() if request else {}
+        resp = self._client.post(url, json=body)
+        self._check_response(resp, "open session")
+        return resp.json()["sessionHandle"]
+
+    def close_session(self, session_handle: str) -> None:
+        """DELETE /sessions/{session_handle}."""
+        url = self._build_endpoint("sessions", session_handle)
+        resp = self._client.delete(url)
+        self._check_response(resp, "close session")
+
+    def get_session_config(self, session_handle: str) -> dict[str, str]:
+        """GET /sessions/{session_handle} — current session properties."""
+        url = self._build_endpoint("sessions", session_handle)
+        resp = self._client.get(url)
+        self._check_response(resp, "get session config")
+        return resp.json().get("properties", {})
+
+    def configure_session(
+        self,
+        session_handle: str,
+        request: ConfigureSessionRequest,
+    ) -> None:
+        """POST /sessions/{session_handle}/configure-session."""
+        url = self._build_endpoint("sessions", session_handle, "configure-session")
+        resp = self._client.post(url, json=request.to_dict())
+        self._check_response(resp, "configure session")
+
+    def heartbeat(self, session_handle: str) -> None:
+        """POST /sessions/{session_handle}/heartbeat."""
+        url = self._build_endpoint("sessions", session_handle, "heartbeat")
+        resp = self._client.post(url, json={})
+        self._check_response(resp, "heartbeat")
+
+    def complete_statement(
+        self,
+        session_handle: str,
+        request: CompleteStatementRequest,
+    ) -> list[str]:
+        """GET /sessions/{session_handle}/complete-statement.
+
+        Returns:
+            List of completion candidates.
+        """
+        url = self._build_endpoint("sessions", session_handle, "complete-statement")
+        # The Flink Gateway uses GET with a JSON body for this endpoint.
+        resp = self._client.request("GET", url, json=request.to_dict())
+        self._check_response(resp, "complete statement")
+        return resp.json().get("candidates", [])
+
+    # ── Statement execution ────────────────────────────────────────────
+
+    def execute_statement(
+        self,
+        session_handle: str,
+        request: ExecuteStatementRequest,
+    ) -> str:
+        """POST /sessions/{session_handle}/statements.
+
+        Returns:
+            The operation handle string.
+        """
+        url = self._build_endpoint("sessions", session_handle, "statements")
+        resp = self._client.post(url, json=request.to_dict())
+        self._check_response(resp, "execute statement")
+        return resp.json()["operationHandle"]
+
+    def fetch_results(
+        self,
+        session_handle: str,
+        operation_handle: str,
+        token: str,
+        row_format: str = "",
+    ) -> FetchResultsResponse:
+        """GET /sessions/{sh}/operations/{oh}/result/{token}.
+
+        Args:
+            session_handle: Session handle.
+            operation_handle: Operation handle.
+            token: Pagination token (``"0"`` for the first batch).
+            row_format: Optional row format (e.g. ``"json"``).
+
+        Returns:
+            Parsed :class:`FetchResultsResponse`.
+        """
+        url = self._build_endpoint(
+            "sessions",
+            session_handle,
+            "operations",
+            operation_handle,
+            "result",
+            token,
+        )
+        params: dict[str, str] = {}
+        if row_format:
+            params["rowFormat"] = row_format
+        resp = self._client.get(url, params=params)
+        self._check_response(resp, "fetch results")
+        return FetchResultsResponse.from_dict(resp.json())
+
+    # ── Operation management ───────────────────────────────────────────
+
+    def get_operation_status(
+        self,
+        session_handle: str,
+        operation_handle: str,
+    ) -> str:
+        """GET /sessions/{sh}/operations/{oh}/status.
+
+        Returns:
+            The status string (e.g. ``"RUNNING"``, ``"FINISHED"``).
+        """
+        url = self._build_endpoint(
+            "sessions",
+            session_handle,
+            "operations",
+            operation_handle,
+            "status",
+        )
+        resp = self._client.get(url)
+        self._check_response(resp, "get operation status")
+        return resp.json()["status"]
+
+    def cancel_operation(
+        self,
+        session_handle: str,
+        operation_handle: str,
+    ) -> str:
+        """POST /sessions/{sh}/operations/{oh}/cancel.
+
+        Returns:
+            The operation status after cancellation.
+        """
+        url = self._build_endpoint(
+            "sessions",
+            session_handle,
+            "operations",
+            operation_handle,
+            "cancel",
+        )
+        resp = self._client.post(url, json={})
+        self._check_response(resp, "cancel operation")
+        return resp.json()["status"]
+
+    def close_operation(
+        self,
+        session_handle: str,
+        operation_handle: str,
+    ) -> str:
+        """DELETE /sessions/{sh}/operations/{oh}.
+
+        Returns:
+            The operation status after closure.
+        """
+        url = self._build_endpoint(
+            "sessions",
+            session_handle,
+            "operations",
+            operation_handle,
+        )
+        resp = self._client.delete(url)
+        self._check_response(resp, "close operation")
+        return resp.json()["status"]
+
+    # ── Materialized tables ────────────────────────────────────────────
+
+    def refresh_materialized_table(
+        self,
+        session_handle: str,
+        identifier: str,
+        request: RefreshMaterializedTableRequest | None = None,
+    ) -> str:
+        """POST /sessions/{sh}/materialized-tables/{id}/refresh.
+
+        Returns:
+            The operation handle string.
+        """
+        url = self._build_endpoint(
+            "sessions",
+            session_handle,
+            "materialized-tables",
+            identifier,
+            "refresh",
+        )
+        body = request.to_dict() if request else {}
+        resp = self._client.post(url, json=body)
+        self._check_response(resp, "refresh materialized table")
+        return resp.json()["operationHandle"]
+
+    # ── Internal helpers ───────────────────────────────────────────────
+
+    @staticmethod
+    def _check_response(resp: httpx.Response, action: str) -> None:
+        if resp.status_code == 200:
+            return
+        detail = ""
+        try:
+            errors = resp.json().get("errors", [])
+            if errors:
+                detail = ": " + "; ".join(errors)
+        except (ValueError, KeyError):
+            pass
+        raise FlinkSqlGatewayError(
+            f"{action} failed: {resp.status_code} {resp.reason_phrase}{detail}"
+        )

flink_gateway/connection.py

@@ -0,0 +1,109 @@
+"""PEP 249 Connection implementation for the Flink SQL Gateway."""
+
+from __future__ import annotations
+
+from typing import Any
+
+from flink_gateway.client import FlinkSqlGatewayClient
+from flink_gateway.cursor import Cursor
+from flink_gateway.exceptions import NotSupportedError, ProgrammingError
+from flink_gateway.models import OpenSessionRequest
+
+
+class Connection:
+    """PEP 249 Connection to a Flink SQL Gateway.
+
+    Do not instantiate directly; use :func:`flink_gateway.connect`.
+    """
+
+    def __init__(
+        self,
+        client: FlinkSqlGatewayClient,
+        session_handle: str,
+        *,
+        _owns_client: bool = True,
+    ) -> None:
+        self._client = client
+        self._session_handle = session_handle
+        self._owns_client = _owns_client
+        self._closed = False
+
+    # ── Public read-only properties ────────────────────────────────
+
+    @property
+    def client(self) -> FlinkSqlGatewayClient:
+        """The underlying REST client."""
+        return self._client
+
+    @property
+    def session_handle(self) -> str:
+        """The session handle for this connection."""
+        return self._session_handle
+
+    @property
+    def closed(self) -> bool:
+        """Whether this connection has been closed."""
+        return self._closed
+
+    # ── Context manager ────────────────────────────────────────────
+
+    def __enter__(self) -> Connection:
+        return self
+
+    def __exit__(self, *_: Any) -> None:
+        self.close()
+
+    # ── PEP 249 interface ──────────────────────────────────────────
+
+    def close(self) -> None:
+        """Close the session and release resources."""
+        if not self._closed:
+            self._closed = True
+            try:
+                self._client.close_session(self._session_handle)
+            except Exception:
+                pass
+            if self._owns_client:
+                self._client.close()
+
+    def commit(self) -> None:
+        """No-op — Flink SQL Gateway does not support transactions."""
+
+    def rollback(self) -> None:
+        """Not supported — Flink SQL Gateway does not support transactions."""
+        raise NotSupportedError("Flink SQL Gateway does not support transactions")
+
+    def cursor(self) -> Cursor:
+        """Create a new Cursor bound to this connection."""
+        if self._closed:
+            raise ProgrammingError("connection is closed")
+        return Cursor(self)
+
+
+def connect(
+    url: str,
+    *,
+    properties: dict[str, str] | None = None,
+    api_version: str = "v3",
+) -> Connection:
+    """Open a connection to a Flink SQL Gateway.
+
+    This is the PEP 249 module-level ``connect()`` function.
+
+    Args:
+        url: Gateway URL, e.g. ``"http://localhost:8083"``.
+        properties: Optional session properties.
+        api_version: REST API version (default ``"v3"``).
+
+    Returns:
+        A :class:`Connection` instance.
+    """
+    client = FlinkSqlGatewayClient(url, api_version=api_version)
+    try:
+        session_handle = client.open_session(
+            OpenSessionRequest(properties=properties) if properties else None
+        )
+    except Exception:
+        client.close()
+        raise
+    return Connection(client, session_handle, _owns_client=True)