wherobots-python-dbapi 0.27.0__tar.gz → 0.28.1__tar.gz

{wherobots_python_dbapi-0.27.0 → wherobots_python_dbapi-0.28.1}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: wherobots-python-dbapi
-Version: 0.27.0
+Version: 0.28.1
 Summary: Python DB-API driver for Wherobots DB
 Project-URL: Homepage, https://github.com/wherobots/wherobots-python-dbapi-driver
 Project-URL: Tracker, https://github.com/wherobots/wherobots-python-dbapi-driver/issues
@@ -14,6 +14,7 @@ Requires-Dist: pandas
 Requires-Dist: pandas-stubs>=2.0.3.230814
 Requires-Dist: pyarrow>=14.0.2
 Requires-Dist: requests>=2.31.0
+Requires-Dist: sqlparse>=0.5.0
 Requires-Dist: strenum<0.5,>=0.4.15
 Requires-Dist: tenacity>=8.2.3
 Requires-Dist: types-requests>=2.31.0
@@ -107,6 +108,16 @@ with connect(
     print(results)
 ```
 
+`runtime` and `region` are optional and accept either the provided enums (handy
+for autocomplete) or a plain string — strings are passed to the API as-is, so
+new or BYOC regions (e.g. `region="byoc-acme-us-east-1"`) work without an SDK
+upgrade. When omitted, your organization's configured defaults are used:
+
+```python
+with connect(api_key='...') as conn:  # uses your org's default runtime + region
+    ...
+```
+
 The `Cursor` supports the context manager protocol, so you can use it
 within a `with` statement when needed:
 

{wherobots_python_dbapi-0.27.0 → wherobots_python_dbapi-0.28.1}/README.md

@@ -83,6 +83,16 @@ with connect(
     print(results)
 ```
 
+`runtime` and `region` are optional and accept either the provided enums (handy
+for autocomplete) or a plain string — strings are passed to the API as-is, so
+new or BYOC regions (e.g. `region="byoc-acme-us-east-1"`) work without an SDK
+upgrade. When omitted, your organization's configured defaults are used:
+
+```python
+with connect(api_key='...') as conn:  # uses your org's default runtime + region
+    ...
+```
+
 The `Cursor` supports the context manager protocol, so you can use it
 within a `with` statement when needed:
 

{wherobots_python_dbapi-0.27.0 → wherobots_python_dbapi-0.28.1}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "wherobots-python-dbapi"
-version = "0.27.0"
+version = "0.28.1"
 description = "Python DB-API driver for Wherobots DB"
 authors = [{ name = "Maxime Petazzoni", email = "max@wherobots.com" }]
 requires-python = ">=3.10, <4"
@@ -17,6 +17,7 @@ dependencies = [
     "pandas",
     "types-requests>=2.31.0",
     "pandas-stubs>=2.0.3.230814",
+    "sqlparse>=0.5.0",
 ]
 
 [project.optional-dependencies]

{wherobots_python_dbapi-0.27.0 → wherobots_python_dbapi-0.28.1}/wherobots/db/connection.py

@@ -6,6 +6,8 @@ import uuid
 from dataclasses import dataclass
 from typing import Any, Callable, Dict
 
+from wherobots.db.redaction import get_statement_type, redact_sql
+
 import pandas
 import pyarrow
 import cbor2
@@ -162,6 +164,13 @@ class Connection:
             )
             return
 
+        def complete_query(result: ExecutionResult) -> None:
+            # Terminal delivery: stop tracking the query first. Keeping it in
+            # __queries would retain its handler — and the results the handler
+            # references — for the connection's lifetime (WBC-922).
+            self.__queries.pop(execution_id, None)
+            query.handler(result)
+
         # Incoming state transitions are handled here.
         if kind == EventKind.STATE_UPDATED or kind == EventKind.EXECUTION_RESULT:
             try:
@@ -189,7 +198,7 @@ class Connection:
                             store_result.size,
                         )
                         query.state = ExecutionState.COMPLETED
-                        query.handler(ExecutionResult(store_result=store_result))
+                        complete_query(ExecutionResult(store_result=store_result))
                         return
 
                     if query.store is not None:
@@ -199,7 +208,7 @@ class Connection:
                             execution_id,
                         )
                         query.state = ExecutionState.COMPLETED
-                        query.handler(ExecutionResult())
+                        complete_query(ExecutionResult())
                         return
 
                     # No store configured, request results normally
@@ -211,11 +220,11 @@ class Connection:
                 if not results or not isinstance(results, dict):
                     logging.warning("Got no results back from %s.", execution_id)
                     query.state = ExecutionState.COMPLETED
-                    query.handler(ExecutionResult())
+                    complete_query(ExecutionResult())
                     return
 
                 query.state = ExecutionState.COMPLETED
-                query.handler(
+                complete_query(
                     ExecutionResult(results=self._handle_results(execution_id, results))
                 )
             elif query.state == ExecutionState.CANCELLED:
@@ -223,8 +232,7 @@ class Connection:
                     "Query %s has been cancelled; returning empty results.",
                     execution_id,
                 )
-                query.handler(ExecutionResult(results=pandas.DataFrame()))
-                self.__queries.pop(execution_id)
+                complete_query(ExecutionResult(results=pandas.DataFrame()))
             elif query.state == ExecutionState.FAILED:
                 # Don't do anything here; the ERROR event is coming with more
                 # details.
@@ -232,7 +240,7 @@ class Connection:
         elif kind == EventKind.ERROR:
             query.state = ExecutionState.FAILED
             error = message.get("message")
-            query.handler(ExecutionResult(error=OperationalError(error)))
+            complete_query(ExecutionResult(error=OperationalError(error)))
         else:
             logging.warning("Received unknown %s event!", kind)
 
@@ -260,9 +268,26 @@ class Connection:
 
     def __send(self, message: Dict[str, Any]) -> None:
         request = json.dumps(message)
-        logging.debug("Request: %s", request)
+        # Only compute the redacted request (json.dumps + sqlparse parse) when
+        # DEBUG is actually enabled; the log argument is evaluated eagerly, so an
+        # unguarded call would redact on every request even with DEBUG off.
+        if logging.getLogger().isEnabledFor(logging.DEBUG):
+            logging.debug("Request: %s", self.__redacted_request(message))
         self.__ws.send(request)
 
+    @staticmethod
+    def __redacted_request(message: Dict[str, Any]) -> str:
+        """Serialize a request for logging with any SQL statement redacted.
+
+        The wire payload (sent verbatim by ``__send``) carries the raw
+        ``statement``; this driver is embedded by other services, so logging it
+        -- even at DEBUG -- would leak raw SQL into their log streams (WBC-139).
+        """
+        statement = message.get("statement")
+        if isinstance(statement, str):
+            message = {**message, "statement": redact_sql(statement)}
+        return json.dumps(message)
+
     def __recv(self) -> Dict[str, Any]:
         frame = self.__ws.recv(timeout=self.__read_timeout)
         if isinstance(frame, str):
@@ -301,8 +326,13 @@ class Connection:
             store=store,
         )
 
+        # Redact literal values before logging: this driver is embedded by other
+        # services, so raw SQL here would leak into their log streams (WBC-139).
         logging.info(
-            "Executing SQL query %s: %s", execution_id, textwrap.shorten(sql, width=60)
+            "Executing SQL query %s (%s): %s",
+            execution_id,
+            get_statement_type(sql),
+            textwrap.shorten(redact_sql(sql), width=200),
         )
         self.__send(request)
         return execution_id

{wherobots_python_dbapi-0.27.0 → wherobots_python_dbapi-0.28.1}/wherobots/db/cursor.py

@@ -74,6 +74,8 @@ class Cursor:
         self.__queue: queue.Queue = queue.Queue()
         self.__results: list[Any] | None = None
         self.__store_result: StoreResult | None = None
+        self.__complete: bool = False
+        self.__error: Exception | None = None
         self.__current_execution_id: str | None = None
         self.__current_row: int = 0
 
@@ -93,21 +95,49 @@ class Cursor:
     def rowcount(self) -> int:
         return self.__rowcount
 
-    def __on_execution_result(self, result) -> None:
-        self.__queue.put(result)
+    def __in_flight_execution_id(self) -> str | None:
+        """The current execution's id if its result has not yet arrived.
+
+        Once a terminal result has been fetched (``__complete``) or is
+        waiting in the queue, the execution is finished and must not be
+        cancelled.
+
+        Results are delivered from the connection's reader thread, so one
+        may arrive between the ``empty()`` check and a cancellation that
+        follows it. That race is benign: the cancel targets an execution
+        that already completed — a server-side no-op, same as when a query
+        finishes while a legitimate cancel is in flight — and the delivered
+        result sits in this execution's own queue, so it can never be
+        observed by a later execution's fetches.
+        """
+        if (
+            self.__current_execution_id is not None
+            and not self.__complete
+            and self.__queue.empty()
+        ):
+            return self.__current_execution_id
+        return None
 
     def __get_results(self) -> List[Tuple[Any, ...]] | None:
         if not self.__current_execution_id:
             raise ProgrammingError("No query has been executed yet")
-        if self.__results is not None:
+        if self.__complete:
+            if self.__error:
+                raise self.__error
             return self.__results
 
         execution_result = self.__queue.get()
         if not isinstance(execution_result, ExecutionResult):
             raise ProgrammingError("Unexpected result type")
 
+        # Whatever the outcome — rows, store export, empty result, or error —
+        # the execution has reached a terminal state. ``__results`` alone
+        # cannot signal this: store-backed and empty executions never set it.
+        self.__complete = True
+
         if execution_result.error:
-            raise execution_result.error
+            self.__error = execution_result.error
+            raise self.__error
 
         self.__store_result = execution_result.store_result
         results = execution_result.results
@@ -140,18 +170,25 @@ class Cursor:
         parameters: Dict[str, Any] | None = None,
         store: Store | None = None,
     ) -> None:
-        if self.__current_execution_id:
-            self.__cancel_fn(self.__current_execution_id)
-
+        in_flight = self.__in_flight_execution_id()
+        if in_flight:
+            self.__cancel_fn(in_flight)
+
+        # Each execution gets its own queue, and the handler closes over it:
+        # a late result from a superseded execution lands in the orphaned
+        # queue and can never be observed by fetches of the current one.
+        self.__queue = queue.Queue()
         self.__results = None
         self.__store_result = None
+        self.__complete = False
+        self.__error = None
         self.__current_row = 0
         self.__rowcount = -1
         self.__description = None
 
         self.__current_execution_id = self.__exec_fn(
             _substitute_parameters(operation, parameters),
-            self.__on_execution_result,
+            self.__queue.put,
             store,
         )
 
@@ -193,8 +230,9 @@ class Cursor:
 
     def close(self) -> None:
         """Close the cursor."""
-        if self.__results is None and self.__current_execution_id:
-            self.__cancel_fn(self.__current_execution_id)
+        in_flight = self.__in_flight_execution_id()
+        if in_flight:
+            self.__cancel_fn(in_flight)
 
     def __iter__(self):
         return self

{wherobots_python_dbapi-0.27.0 → wherobots_python_dbapi-0.28.1}/wherobots/db/driver.py

@@ -21,8 +21,6 @@ import certifi
 from .connection import Connection
 from .constants import (
     DEFAULT_ENDPOINT,
-    DEFAULT_REGION,
-    DEFAULT_RUNTIME,
     DEFAULT_READ_TIMEOUT_SECONDS,
     DEFAULT_SESSION_TYPE,
     DEFAULT_SESSION_WAIT_TIMEOUT_SECONDS,
@@ -56,24 +54,65 @@ TRANSIENT_HTTP_STATUS_CODES = {429, 502, 503, 504}
 DEFAULT_HTTP_TIMEOUT = 30
 
 
-def gen_user_agent_header():
+def _resolve_dbapi_version() -> str:
+    """Resolve this package's version, or "unknown" if it can't be found."""
     try:
-        package_version = metadata.version("wherobots-python-dbapi")
+        return metadata.version("wherobots-python-dbapi")
     except PackageNotFoundError:
-        package_version = "unknown"
+        return "unknown"
+
+
+# Resolved once at import: `importlib.metadata.version` scans the installed
+# package database on each call, and the version can't change within a process.
+_DBAPI_VERSION: Final[str] = _resolve_dbapi_version()
+
+
+def gen_user_agent_header():
     python_version = platform.python_version()
     system = platform.system().lower()
     return {
-        "User-Agent": f"wherobots-python-dbapi/{package_version} os/{system} python/{python_version}"
+        "User-Agent": f"wherobots-python-dbapi/{_DBAPI_VERSION} os/{system} python/{python_version}"
     }
 
 
+# Canonical name of the shared, cross-service client-chain header.
+WHEROBOTS_CLIENT_HEADER: Final[str] = "X-Wherobots-Client"
+
+
+def _append_wherobots_client_hop(headers: Dict[str, str]) -> None:
+    """Append this driver's hop to the shared ``X-Wherobots-Client`` header.
+
+    ``X-Wherobots-Client`` is an ordered, append-only, comma-separated list of
+    hops; the leftmost is the origin and each component appends its own hop on
+    the right. DBAPI's hop is ``client=dbapi;ver=<version>``.
+
+    The lookup is case-insensitive because HTTP header names are, and any
+    differently-cased inbound key is collapsed into the canonical
+    ``X-Wherobots-Client`` key so the request carries exactly one such header.
+
+    This header is advisory: it is informational only and must never affect
+    authentication. ``headers`` is mutated in place.
+    """
+    dbapi_hop = f"client=dbapi;ver={_DBAPI_VERSION}"
+
+    # Find any existing hop chain case-insensitively and remove differently
+    # cased duplicates so we don't emit two headers.
+    existing_chain: Union[str, None] = None
+    for key in [k for k in headers if k.lower() == WHEROBOTS_CLIENT_HEADER.lower()]:
+        existing_chain = headers.pop(key)
+
+    if existing_chain:
+        headers[WHEROBOTS_CLIENT_HEADER] = f"{existing_chain}, {dbapi_hop}"
+    else:
+        headers[WHEROBOTS_CLIENT_HEADER] = dbapi_hop
+
+
 def connect(
     host: str = DEFAULT_ENDPOINT,
     token: Union[str, None] = None,
     api_key: Union[str, None] = None,
-    runtime: Union[Runtime, None] = None,
-    region: Union[Region, None] = None,
+    runtime: Union[str, Runtime, None] = None,
+    region: Union[str, Region, None] = None,
     version: Union[str, None] = None,
     wait_timeout: float = DEFAULT_SESSION_WAIT_TIMEOUT_SECONDS,
     read_timeout: float = DEFAULT_READ_TIMEOUT_SECONDS,
@@ -84,29 +123,63 @@ def connect(
     data_compression: Union[DataCompression, None] = None,
     geometry_representation: Union[GeometryRepresentation, None] = None,
     cancel_event: Union[threading.Event, None] = None,
+    extra_headers: Union[Dict[str, str], None] = None,
 ) -> Connection:
+    """Create a connection to a Wherobots SQL session.
+
+    :param runtime: The compute runtime to use. Accepts a ``Runtime`` enum value
+        or a raw string; strings are passed to the API as-is. Override the
+        default runtime set for your organization — only set this if you need a
+        specific runtime instead of the one your administrator has configured.
+        When omitted, your organization's default runtime is used.
+    :param region: The compute region to run in. Accepts a ``Region`` enum value
+        or a raw string (e.g. a BYOC region such as ``byoc-acme-us-east-1``);
+        strings are passed to the API as-is. Override the default region set for
+        your organization — only set this if you intend to use a specific region
+        instead of the one your administrator has configured. When omitted, your
+        organization's default region is used.
+    :param extra_headers: Optional extra HTTP headers to send on the session
+        requests and the WebSocket upgrade. Merged after the driver's own
+        headers, so callers can pass through tracing/correlation headers. If it
+        includes an ``X-Wherobots-Client`` hop chain, this driver appends its
+        own ``client=dbapi;ver=<version>`` hop to the right of it. This header
+        is advisory only and never affects authentication; ``extra_headers``
+        cannot be used to override the ``Authorization``/``X-API-Key`` headers.
+    """
     if not token and not api_key:
         raise ValueError("At least one of `token` or `api_key` is required")
     if token and api_key:
         raise ValueError("`token` and `api_key` can't be both provided")
 
     headers = gen_user_agent_header()
+    # Merge caller-supplied headers first so the driver's own auth headers,
+    # applied below, always win and can never be overridden.
+    if extra_headers:
+        headers.update(extra_headers)
     if token:
         headers["Authorization"] = f"Bearer {token}"
     elif api_key:
         headers["X-API-Key"] = api_key
 
+    # Append this driver's hop to the shared client-chain header. Done after the
+    # auth headers are set so it operates on the final, merged header set.
+    _append_wherobots_client_hop(headers)
+
     host = host or DEFAULT_ENDPOINT
-    runtime = runtime or DEFAULT_RUNTIME
-    region = region or DEFAULT_REGION
     session_type = session_type or DEFAULT_SESSION_TYPE
 
+    # Normalize enum values to their string form and pass raw strings through
+    # untouched. When omitted (None) the field is dropped from the request so
+    # the API applies the organization's configured default.
+    runtime_id = runtime.value if isinstance(runtime, Runtime) else runtime
+    region_name = region.value if isinstance(region, Region) else region
+
     logging.info(
         "Requesting %s%s runtime %sin %s from %s ...",
         "new " if force_new else "",
-        runtime.value,
+        runtime_id or "org-default",
         f"running {version} " if version else "",
-        region.value,
+        region_name or "org-default",
         host,
     )
 
@@ -119,9 +192,11 @@ def connect(
     try:
         resp = requests.post(
             url=f"{host}/sql/session",
-            params={"region": region.value, "force_new": force_new},
+            # `requests` omits query params whose value is None, so an omitted
+            # region is simply not sent and the API applies the org default.
+            params={"region": region_name, "force_new": force_new},
             json={
-                "runtimeId": runtime.value,
+                "runtimeId": runtime_id,
                 "shutdownAfterInactiveSeconds": shutdown_after_inactive_seconds,
                 "version": version,
                 "sessionType": session_type.value,

wherobots_python_dbapi-0.28.1/wherobots/db/redaction.py

@@ -0,0 +1,78 @@
+"""Redaction helpers for SQL statements logged by this driver.
+
+The driver logs SQL statements for observability (e.g. ``Executing SQL query
+<id>: ...``). Raw statements can embed literal PII (for example ``WHERE ssn =
+'123-45-6789'``), and because this library is embedded by other services its log
+output ends up in their log streams. This module replaces literal *values*
+(single-quoted string literals and numeric literals) with a ``?`` placeholder
+while preserving the statement structure: keywords, function names, and
+identifiers (including double-quoted/back-quoted identifiers) are kept intact so
+the redacted form is still useful for debugging and aggregation.
+
+The implementation tokenizes with ``sqlparse``, a lenient, pure-Python tokenizer
+with zero transitive dependencies. It is dialect-agnostic, never raises on
+malformed input, and classifies literals by token type, which lets us redact
+precisely the value-bearing tokens while copying everything else (keywords,
+identifiers, operators, comments, whitespace) through verbatim.
+
+This logic is duplicated from the ``sql-session`` service (PR #197); the two
+repositories do not share a package, so the implementation is intentionally
+replicated here rather than imported.
+
+This is a best-effort redaction for *logging*. It is not a security boundary and
+must not be relied on to sanitize untrusted input for any other purpose.
+"""
+
+import re
+
+import sqlparse
+from sqlparse import tokens as T
+
+REDACTED_PLACEHOLDER = "?"
+
+# Leading SQL keyword -> statement type. Used purely for observability tagging.
+# Intentionally a regex (not sqlparse's ``Statement.get_type()``) because
+# ``get_type()`` returns "UNKNOWN" for SHOW/DESCRIBE/SET, which would regress the
+# observability tagging this module exists to support.
+_STATEMENT_TYPE_RE = re.compile(r"\s*([A-Za-z]+)")
+
+
+def get_statement_type(statement: str) -> str:
+    """Return the upper-cased leading keyword of a statement (e.g. ``SELECT``).
+
+    Returns ``"UNKNOWN"`` when the statement does not begin with a word.
+    """
+    match = _STATEMENT_TYPE_RE.match(statement)
+    if match is None:
+        return "UNKNOWN"
+    return match.group(1).upper()
+
+
+def redact_sql(statement: str) -> str:
+    """Return ``statement`` with string and numeric literals replaced by ``?``.
+
+    Single-quoted string literals (``String.Single``) and numeric literals
+    (``Number.*``) are each collapsed to a single ``?`` placeholder. Everything
+    else -- keywords, function names, plain and quoted identifiers (``"col"`` /
+    `` `col` ``), comments, whitespace, and operators -- is preserved verbatim.
+    """
+    out: list[str] = []
+    for stmt in sqlparse.parse(statement):
+        # sqlparse ships no type stubs, so ``flatten`` is untyped under strict.
+        for token in stmt.flatten():  # type: ignore[no-untyped-call]
+            ttype = token.ttype
+            if ttype in T.String.Single or ttype in T.Number:
+                out.append(REDACTED_PLACEHOLDER)
+            elif ttype in T.Error and token.value in ("'", '"', "`"):
+                # An unterminated quote: sqlparse emits the lone opener as an
+                # Error token and tokenizes the trailing characters as ordinary
+                # text. Redact the opener and fail closed by bailing on the
+                # entire input -- returning here (rather than ``break``, which
+                # only exits the inner loop) ensures that for multi-statement
+                # input no later statement is appended verbatim, which would
+                # leak the literals this function exists to hide.
+                out.append(REDACTED_PLACEHOLDER)
+                return "".join(out)
+            else:
+                out.append(token.value)
+    return "".join(out)