wherobots-python-dbapi 0.28.0__py3-none-any.whl → 0.28.1__py3-none-any.whl

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/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/db/driver.py

@@ -54,18 +54,59 @@ 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,
@@ -82,6 +123,7 @@ 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.
 
@@ -96,6 +138,13 @@ def connect(
         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")
@@ -103,11 +152,19 @@ def connect(
         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
     session_type = session_type or DEFAULT_SESSION_TYPE
 

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)

{wherobots_python_dbapi-0.28.0.dist-info → wherobots_python_dbapi-0.28.1.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: wherobots-python-dbapi
-Version: 0.28.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

{wherobots_python_dbapi-0.28.0.dist-info → wherobots_python_dbapi-0.28.1.dist-info}/RECORD

@@ -1,16 +1,16 @@
-wherobots/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 wherobots/db/__init__.py,sha256=HN_uGOUMor80GB3rZIqnPcYWDAia7QAC-y9HfGKnegs,744
-wherobots/db/connection.py,sha256=BWQwVzFgujEwPKvhDePM9q45M2opCk6Jq998L0ekMyY,12795
+wherobots/db/connection.py,sha256=VqXqFoKeyy9DMXP8-t0GhijWBMYB154HevbDRvyMzF4,14352
 wherobots/db/constants.py,sha256=qV-6CmOya7-f-DoRMzcDzlhcmV5PkLCcqG5TFIsxQ7o,744
-wherobots/db/cursor.py,sha256=Z0rNdhJ985IP9XnZpjrCzvM0OqDFVgCNfJn8lrbijAo,6809
-wherobots/db/driver.py,sha256=fHfYtujvafmAnJh9uf_0oZUNVtyf8XLsgs9sqfiDMk4,10304
+wherobots/db/cursor.py,sha256=vaQY1WcTSehx-XmCkn-vRsv5YcTddxj0HR1ZaGAtEis,8495
+wherobots/db/driver.py,sha256=5iZXEsR2TM0xVPdcPLwlOifIkJHsYHLM0zKkwMk4UTo,12975
 wherobots/db/errors.py,sha256=G8a-DB6pgB0pUld59Wyq8PuywL-JWxGJfvrr-g8BPuU,310
 wherobots/db/models.py,sha256=uk0EqZMH_RFkQusSw2AvyFivGvGJBu_WYP8FHa6JJGM,4178
+wherobots/db/redaction.py,sha256=issl5e5QvZQbs4d5UGUJIW7SIZIYfQ04MIyuT0Eo5_g,3744
 wherobots/db/region.py,sha256=GoFzazvgOjE1kVVT1QKAsRK88c_1XSAebR3ME4foSlQ,400
 wherobots/db/runtime.py,sha256=eAR1oJEtssgkp-IRPTLOrlPI-_5SM_ASw6T6hQ_qF4I,985
 wherobots/db/session_type.py,sha256=jZNwCl0nxvY-0reIn8RR_-fIYSSvIcaGx_4pRc71Lww,137
 wherobots/db/types.py,sha256=WXKT_gbdQCH2_yXD_FECk0BpB7yfNWOLewCfax0U_Lk,2516
-wherobots_python_dbapi-0.28.0.dist-info/METADATA,sha256=qUqbExsrE8Q6GwBE8f_mWL8-yNnGAp4vY1qgc_2OU0c,11142
-wherobots_python_dbapi-0.28.0.dist-info/WHEEL,sha256=QccIxa26bgl1E6uMy58deGWi-0aeIkkangHcxk2kWfw,87
-wherobots_python_dbapi-0.28.0.dist-info/licenses/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
-wherobots_python_dbapi-0.28.0.dist-info/RECORD,,
+wherobots_python_dbapi-0.28.1.dist-info/METADATA,sha256=EuIBcH4k9NCG9eIQOkyJvowab-iD0LVDXbWFQ1jg28c,11173
+wherobots_python_dbapi-0.28.1.dist-info/WHEEL,sha256=lCkmxWfQsSc9CfIClYeavTdQeEX2toPqufh9gI35EQA,87
+wherobots_python_dbapi-0.28.1.dist-info/licenses/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
+wherobots_python_dbapi-0.28.1.dist-info/RECORD,,

{wherobots_python_dbapi-0.28.0.dist-info → wherobots_python_dbapi-0.28.1.dist-info}/WHEEL

@@ -1,4 +1,4 @@
 Wheel-Version: 1.0
-Generator: hatchling 1.29.0
+Generator: hatchling 1.31.0
 Root-Is-Purelib: true
 Tag: py3-none-any