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

wherobots/db/_transport.py

@@ -0,0 +1,24 @@
+"""Transport termination independent of the WebSocket protocol's send lock."""
+
+import errno
+import socket
+
+from websockets.sync.client import ClientConnection
+
+
+def abort_connection(ws: ClientConnection) -> None:
+    """Disable further writes and wake blocked I/O before publishing failures.
+
+    ClientConnection.close() takes the library's protocol mutex, which a
+    stalled sendall() may hold. Shutdown the owned socket instead; the library's
+    receive thread will observe EOF/error and finish protocol cleanup itself.
+    Already transmitted bytes may still execute remotely.
+    """
+    try:
+        ws.socket.shutdown(socket.SHUT_RDWR)
+    except OSError as exc:
+        if exc.errno not in (errno.ENOTCONN, errno.EBADF, errno.ECONNRESET):
+            raise
+    finally:
+        # Even if shutdown fails, prevent any later send on this socket object.
+        ws.socket.close()

wherobots/db/connection.py

@@ -1,19 +1,22 @@
 import json
 import logging
+import math
 import textwrap
 import threading
 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
 import websockets.exceptions
-import websockets.protocol
 import websockets.sync.client
 
 from .constants import DEFAULT_READ_TIMEOUT_SECONDS
+from ._transport import abort_connection
 from .cursor import Cursor
 from .errors import NotSupportedError, OperationalError
 from .models import ExecutionResult, ProgressInfo, Store, StoreResult
@@ -31,6 +34,41 @@ ProgressHandler = Callable[[ProgressInfo], None]
 """A callable invoked with a :class:`ProgressInfo` on every progress event."""
 
 
+class _TransportError(Exception):
+    """An I/O failure raised while receiving from the WebSocket."""
+
+
+_READER_JOIN_SECONDS = 1.0
+"""How long close() waits for the reader thread once teardown has finished."""
+
+
+def _handshake_bound(ws: Any) -> float | None:
+    """The transport's ``close_timeout`` if it bounds the close handshake.
+
+    websockets documents ``close_timeout=None`` as disabling the timeout, which
+    would make a graceful close unbounded against an unresponsive peer. Return
+    ``None`` for that, and for a missing, non-numeric (including ``bool`` and
+    mock), non-finite, non-positive, or too-large value; callers then abort
+    instead. Too large means the ``bound + 1s`` wait in ``close()`` would
+    exceed ``threading.TIMEOUT_MAX``.
+    Never raises ``Exception``: a failure here would strand pending queries.
+    """
+    try:
+        value = getattr(ws, "close_timeout", None)
+        if isinstance(value, bool) or not isinstance(value, (int, float)):
+            return None
+        # Range-check a plain float, so a subclass can't override the
+        # comparisons. float() raises OverflowError for a huge int.
+        value = float(value)
+        if not math.isfinite(value) or value <= 0:
+            return None
+        if value > threading.TIMEOUT_MAX - 1.0:
+            return None
+        return value
+    except Exception:
+        return None
+
+
 @dataclass
 class Query:
     sql: str
@@ -62,6 +100,7 @@ class Connection:
         results_format: ResultsFormat | None = None,
         data_compression: DataCompression | None = None,
         geometry_representation: GeometryRepresentation | None = None,
+        session_id: str | None = None,
     ):
         self.__ws = ws
         self.__read_timeout = read_timeout
@@ -70,6 +109,15 @@ class Connection:
         self.__geometry_representation = geometry_representation
         self.__progress_handler: ProgressHandler | None = None
 
+        self.__session_id = session_id
+        self.__lock = threading.Lock()
+        self.__send_lock = threading.Lock()
+        self.__shutdown_done = threading.Event()
+        self.__shutdown_owner: int | None = None
+        # Resolved when shutdown starts (websockets reads close_timeout live),
+        # so the graceful/abort decision and a second closer's wait agree.
+        self.__close_bound: float | None = None
+        self.__closed = False
         self.__queries: dict[str, Query] = {}
         self.__thread = threading.Thread(
             target=self.__main_loop, daemon=True, name="wherobots-connection"
@@ -83,7 +131,37 @@ class Connection:
         self.close()
 
     def close(self) -> None:
-        self.__ws.close()
+        """Close the transport, fail pending work, then wait for the reader.
+
+        The WebSocket close handshake is attempted when no send is in flight,
+        so the server can distinguish a client exit from a crash; otherwise the
+        socket is aborted. The handshake is bounded by the ``close_timeout`` of
+        the underlying ``ClientConnection`` (1s via ``connect``/
+        ``connect_direct``; the library default of 10s for a caller-built
+        socket). If ``close_timeout`` doesn't bound it (``None``, missing, or
+        not a finite positive number), the socket is aborted instead.
+
+        Pending work has been failed by the time ``close()`` returns, also when
+        another thread is closing concurrently. After teardown, ``close()``
+        waits up to 1s for the reader thread, so a single call takes at most
+        about ``close_timeout`` + 1s. A second, concurrent ``close()`` can take
+        up to ``close_timeout`` + 2s (about 12s on a caller-built socket using
+        the library's 10s default).
+
+        Closing doesn't imply that server-side writes were rolled back. A
+        decoder or callback can outlive the bounded reader join.
+        """
+        self.__fail_pending(graceful=True)
+        # A handler may close its own connection during terminal delivery.
+        if self.__shutdown_owner == threading.get_ident():
+            return
+        # Another thread may own shutdown and be mid-handshake, which is bounded
+        # by close_timeout; outlast it so close() never returns before pending
+        # work has been failed.
+        self.__shutdown_done.wait((self.__close_bound or 0.0) + 1.0)
+        # Timed from here, not from entry, so a slow handshake can't consume it.
+        if self.__thread is not threading.current_thread():
+            self.__thread.join(timeout=_READER_JOIN_SECONDS)
 
     def commit(self) -> None:
         raise NotSupportedError
@@ -112,18 +190,103 @@ class Connection:
     def __main_loop(self) -> None:
         """Main background loop listening for messages from the SQL session."""
         logging.info("Starting background connection handling loop...")
-        while self.__ws.protocol.state < websockets.protocol.State.CLOSING:
+        try:
+            self.__receive_loop()
+        finally:
+            self.__fail_pending()
+
+    def __receive_loop(self) -> None:
+        # recv drains buffered results before raising ConnectionClosed.
+        while True:
+            if self.__shutdown_done.is_set():
+                return
             try:
                 self.__listen()
             except TimeoutError:
                 # Expected, retry next time
                 continue
-            except websockets.exceptions.ConnectionClosedOK:
+            except websockets.exceptions.ConnectionClosed:
                 logging.info("Connection closed; stopping main loop.")
                 return
+            except _TransportError:
+                logging.exception("SQL session transport failed; stopping main loop")
+                return
             except Exception as e:
                 logging.exception("Error handling message from SQL session", exc_info=e)
 
+    def __connection_error(self, execution_id: str) -> OperationalError:
+        message = (
+            f"SQL connection lost (session={self.__session_id or 'unknown'}, "
+            f"execution={execution_id}). Commit outcome is unknown; "
+            "verify the operation before retrying writes."
+        )
+        return OperationalError(message)
+
+    def __fail_pending(self, graceful: bool = False) -> None:
+        # Stop admission first. Do not wait for __send_lock: its owner may be
+        # blocked in network I/O. __closed means closing until shutdown_done.
+        # Resolve the bound before taking __lock: it reads a transport
+        # attribute, which may be a property, and nothing foreign may run under
+        # __lock or between the latch and the delivery try/finally.
+        bound = _handshake_bound(self.__ws)
+        with self.__lock:
+            if self.__closed:
+                return
+            self.__closed = True
+            self.__shutdown_owner = threading.get_ident()
+            self.__close_bound = bound
+        # __closed is a one-way latch: nothing below may be skipped, or pending
+        # queries are stranded with no path to recovery.
+        try:
+            self.__terminate_transport(graceful)
+            with self.__lock:
+                pending = list(self.__queries.values())
+                self.__queries.clear()
+            for query in pending:
+                try:
+                    query.handler(
+                        ExecutionResult(
+                            error=self.__connection_error(query.execution_id)
+                        )
+                    )
+                except Exception:
+                    logging.exception(
+                        "Could not deliver connection failure to query handler"
+                    )
+        finally:
+            self.__shutdown_owner = None
+            self.__shutdown_done.set()
+
+    def __terminate_transport(self, graceful: bool) -> None:
+        """Disable the transport. Never raises: delivery must not be skipped."""
+        # A non-blocking acquire tells us whether a sender is inside ws.send()
+        # right now, possibly stalled holding the library's protocol mutex,
+        # which close() would deadlock on. Releasing immediately is safe: the
+        # __closed latch is already set, and __send re-checks it under __lock
+        # after taking __send_lock, so no sender can reach ws.send() from here
+        # on. Holding the lock across the handshake would only park unrelated
+        # senders (including the reader's own __request_results) for up to
+        # close_timeout instead of letting them fail fast. On a reader-side
+        # failure the transport is already broken and a close frame is
+        # pointless, so callers abort directly. Without a handshake bound the
+        # close could block forever, so abort then too.
+        if (
+            graceful
+            and self.__close_bound is not None
+            and self.__send_lock.acquire(blocking=False)
+        ):
+            self.__send_lock.release()
+            try:
+                self.__ws.close()
+                return
+            except Exception:
+                logging.debug("Graceful close failed; aborting", exc_info=True)
+        try:
+            abort_connection(self.__ws)
+        except Exception:
+            # The adapter still closes the socket object in its finally block.
+            logging.exception("Could not abort SQL session transport")
+
     def __listen(self) -> None:
         """Waits for the next message from the SQL session and processes it.
 
@@ -162,13 +325,36 @@ 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).
+            with self.__lock:
+                claimed = self.__queries.pop(execution_id, None)
+            if claimed is not None:
+                claimed.handler(result)
+
+        def fail_query(action: str, error: Exception) -> None:
+            # This is a query-local failure, not evidence of transport loss.
+            # Exception text may contain result data; report only its type.
+            query.state = ExecutionState.FAILED
+            message = (
+                f"Could not {action} SQL results "
+                f"(session={self.__session_id or 'unknown'}, "
+                f"execution={execution_id}; {type(error).__name__}). "
+                "The statement may have completed; verify the operation before "
+                "retrying writes."
+            )
+            logging.error("%s", message)
+            complete_query(ExecutionResult(error=OperationalError(message)))
+
         # Incoming state transitions are handled here.
         if kind == EventKind.STATE_UPDATED or kind == EventKind.EXECUTION_RESULT:
             try:
                 query.state = ExecutionState[message["state"].upper()]
                 logging.info("Query %s is now %s.", execution_id, query.state)
-            except KeyError:
-                logging.warning("Invalid state update message for %s", execution_id)
+            except (KeyError, AttributeError, TypeError) as error:
+                fail_query("interpret", error)
                 return
 
             if query.state == ExecutionState.SUCCEEDED:
@@ -189,7 +375,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,32 +385,44 @@ class Connection:
                             execution_id,
                         )
                         query.state = ExecutionState.COMPLETED
-                        query.handler(ExecutionResult())
+                        complete_query(ExecutionResult())
                         return
 
                     # No store configured, request results normally
-                    self.__request_results(execution_id)
+                    try:
+                        self.__request_results(execution_id)
+                    except Exception as error:
+                        # Transport failures are handled by __send; a local
+                        # retrieval failure must not orphan this execution.
+                        fail_query("request", error)
                     return
 
                 # Otherwise, process the results from the execution_result event.
                 results = message.get("results")
-                if not results or not isinstance(results, dict):
+                if results is None or results == {}:
                     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(
-                    ExecutionResult(results=self._handle_results(execution_id, results))
-                )
+                try:
+                    if not isinstance(results, dict):
+                        raise TypeError("Expected a result object")
+                    decoded = self._handle_results(execution_id, results)
+                except Exception as error:
+                    # Even OSError/TimeoutError here belong to decoding, not
+                    # recv. Claim and deliver exactly once, including if close
+                    # concurrently claims this execution.
+                    fail_query("decode", error)
+                else:
+                    query.state = ExecutionState.COMPLETED
+                    complete_query(ExecutionResult(results=decoded))
             elif query.state == ExecutionState.CANCELLED:
                 logging.info(
                     "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 +430,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)
 
@@ -256,15 +454,63 @@ class Connection:
             with pyarrow.ipc.open_stream(stream) as reader:
                 return reader.read_pandas()
         else:
-            return OperationalError(f"Unsupported results format {result_format}")
+            raise NotSupportedError("Unsupported results format")
 
-    def __send(self, message: Dict[str, Any]) -> None:
+    def __send(self, message: Dict[str, Any], query: Query | None = None) -> None:
+        # Serialization and redaction are local work. Fail before registration,
+        # without poisoning unrelated cursors or misreporting a transport loss.
         request = json.dumps(message)
-        logging.debug("Request: %s", request)
-        self.__ws.send(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))
+        with self.__send_lock:
+            with self.__lock:
+                if self.__closed:
+                    if query is not None:
+                        raise self.__connection_error(query.execution_id)
+                    return
+                if query is not None:
+                    self.__queries[query.execution_id] = query
+                elif message.get("execution_id") not in self.__queries:
+                    return
+            try:
+                self.__ws.send(request)
+            except (websockets.exceptions.ConnectionClosed, OSError):
+                pass  # Terminate outside the send gate, before delivering errors.
+            except Exception:
+                # API/programming errors aren't evidence of connection loss.
+                if query is not None:
+                    with self.__lock:
+                        self.__queries.pop(query.execution_id, None)
+                raise
+            else:
+                return
+        self.__fail_pending()
+
+    @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)
+        try:
+            frame = self.__ws.recv(timeout=self.__read_timeout)
+        except TimeoutError:
+            raise  # Idle polls are expected, not terminal I/O failures.
+        except OSError as e:
+            # Distinguish transport I/O failures from OSErrors raised later by
+            # protocol parsing or result decoding; only the former are terminal.
+            raise _TransportError from e
         if isinstance(frame, str):
             message = json.loads(frame)
         elif isinstance(frame, bytes):
@@ -293,18 +539,24 @@ class Connection:
         if store:
             request["store"] = store.to_dict()
 
-        self.__queries[execution_id] = Query(
-            sql=sql,
-            execution_id=execution_id,
-            state=ExecutionState.EXECUTION_REQUESTED,
-            handler=handler,
-            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,
+            Query(
+                sql=sql,
+                execution_id=execution_id,
+                state=ExecutionState.EXECUTION_REQUESTED,
+                handler=handler,
+                store=store,
+            ),
         )
-        self.__send(request)
         return execution_id
 
     def __request_results(self, execution_id: str) -> None:

wherobots/db/constants.py

@@ -14,6 +14,8 @@ DEFAULT_REGION: Region = Region.AWS_US_WEST_2
 DEFAULT_SESSION_TYPE: SessionType = SessionType.MULTI
 DEFAULT_STORAGE_FORMAT: StorageFormat = StorageFormat.PARQUET
 DEFAULT_READ_TIMEOUT_SECONDS: float = 0.25
+# Bound on the WebSocket close handshake performed by a graceful close().
+DEFAULT_CLOSE_TIMEOUT_SECONDS: float = 1.0
 DEFAULT_SESSION_WAIT_TIMEOUT_SECONDS: float = 900
 
 MAX_MESSAGE_SIZE: int = 100 * 2**20  # 100MiB

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,28 @@ 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
 
+        # A rejected submission must not leave the previous execution ID paired
+        # with this new empty queue (which would make a later fetch wait forever).
+        self.__current_execution_id = None
         self.__current_execution_id = self.__exec_fn(
             _substitute_parameters(operation, parameters),
-            self.__on_execution_result,
+            self.__queue.put,
             store,
         )
 
@@ -193,8 +233,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

@@ -20,6 +20,7 @@ import certifi
 
 from .connection import Connection
 from .constants import (
+    DEFAULT_CLOSE_TIMEOUT_SECONDS,
     DEFAULT_ENDPOINT,
     DEFAULT_READ_TIMEOUT_SECONDS,
     DEFAULT_SESSION_TYPE,
@@ -54,18 +55,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 +124,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 +139,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 +153,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
 
@@ -210,6 +268,9 @@ def connect(
         data_compression=data_compression,
         geometry_representation=geometry_representation,
         cancel_event=cancel_event,
+        session_id=urllib.parse.urlparse(session_id_url)
+        .path.rstrip("/")
+        .rsplit("/", 1)[-1],
     )
 
 
@@ -237,6 +298,7 @@ def connect_direct(
     data_compression: Union[DataCompression, None] = None,
     geometry_representation: Union[GeometryRepresentation, None] = None,
     cancel_event: Union[threading.Event, None] = None,
+    session_id: str | None = None,
 ) -> Connection:
     uri_with_protocol = f"{uri}/{protocol}"
     ssl_context = ssl.create_default_context()
@@ -264,6 +326,7 @@ def connect_direct(
             additional_headers=headers,
             max_size=MAX_MESSAGE_SIZE,
             open_timeout=DEFAULT_HTTP_TIMEOUT,
+            close_timeout=DEFAULT_CLOSE_TIMEOUT_SECONDS,
             ssl=ssl_context,
         )
 
@@ -280,4 +343,5 @@ def connect_direct(
         results_format=results_format,
         data_compression=data_compression,
         geometry_representation=geometry_representation,
+        session_id=session_id,
     )

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.2.dist-info}/METADATA

@@ -1,6 +1,6 @@
-Metadata-Version: 2.4
+Metadata-Version: 2.5
 Name: wherobots-python-dbapi
-Version: 0.28.0
+Version: 0.28.2
 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.2.dist-info/RECORD

@@ -0,0 +1,17 @@
+wherobots/db/__init__.py,sha256=HN_uGOUMor80GB3rZIqnPcYWDAia7QAC-y9HfGKnegs,744
+wherobots/db/_transport.py,sha256=b26gp20zreMAjvXxfBlP_Nl6xrxjmG3k1POiDhfVphw,887
+wherobots/db/connection.py,sha256=o3UsjmFcntMavPUg6aQQTDMV1uPBMGmTrKXvSnm7Xt4,24955
+wherobots/db/constants.py,sha256=SbJVxGSuN-2kBGu4WxNVxUlPi-8hhoG6eeuOewaaq_I,861
+wherobots/db/cursor.py,sha256=jF9v4UhfBqZ0iZEAmIYRxWYU4E8UbMFIrwfY_PlNI60,8701
+wherobots/db/driver.py,sha256=oxztjXrdsjqglOYEHn5JDMgmKh8_VCWj-YDGEb4OZsg,13245
+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.2.dist-info/METADATA,sha256=NrBYfI2FBGEDhuejl5zhV0f-HMUjpFEU2yBomo8u4Ho,11173
+wherobots_python_dbapi-0.28.2.dist-info/WHEEL,sha256=W3fkpkm7-wf9vBI5Z-7s0eWkeM-spu78I8Neb98DeEg,87
+wherobots_python_dbapi-0.28.2.dist-info/licenses/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
+wherobots_python_dbapi-0.28.2.dist-info/RECORD,,

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

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

wherobots_python_dbapi-0.28.0.dist-info/RECORD

@@ -1,16 +0,0 @@
-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/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/errors.py,sha256=G8a-DB6pgB0pUld59Wyq8PuywL-JWxGJfvrr-g8BPuU,310
-wherobots/db/models.py,sha256=uk0EqZMH_RFkQusSw2AvyFivGvGJBu_WYP8FHa6JJGM,4178
-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,,