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

wherobots/db/_transport.py

@@ -0,0 +1,57 @@
+"""Transport termination independent of the WebSocket protocol's send lock."""
+
+import errno
+import select
+import selectors
+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()
+
+
+def socket_writable(ws: ClientConnection) -> bool:
+    """Whether a small write to the transport's socket won't block right now.
+
+    For a send that must not stall the reader thread: a full send buffer
+    (a peer that keeps sending but stops reading) would park sendall(), which
+    also holds the library's protocol mutex, so even keepalive pings stall.
+
+    Never raises. Returns ``True`` (proceed, as without this check) on any
+    failure or for a transport without a real socket, such as a test fake:
+    ``False`` would silently disable the caller's sends. Uses ``poll`` (or a
+    selector), not ``select.select``, which fails for fds above FD_SETSIZE.
+    """
+    try:
+        sock = getattr(ws, "socket", None)
+        if not isinstance(sock, socket.socket):
+            return True
+        fd = sock.fileno()
+        if fd < 0:
+            return True  # Closed: let the send itself report it.
+        if hasattr(select, "poll"):
+            poller = select.poll()
+            poller.register(fd, select.POLLOUT)
+            # POLLERR/POLLHUP count as writable: the send then fails fast.
+            return bool(poller.poll(0))
+        with selectors.DefaultSelector() as selector:
+            selector.register(fd, selectors.EVENT_WRITE)
+            return bool(selector.select(0))
+    except Exception:
+        return True

wherobots/db/connection.py

@@ -1,9 +1,11 @@
 import json
 import logging
+import math
 import textwrap
 import threading
+import time
 import uuid
-from dataclasses import dataclass
+from dataclasses import dataclass, field
 from typing import Any, Callable, Dict
 
 from wherobots.db.redaction import get_statement_type, redact_sql
@@ -12,10 +14,10 @@ import pandas
 import pyarrow
 import cbor2
 import websockets.exceptions
-import websockets.protocol
 import websockets.sync.client
 
-from .constants import DEFAULT_READ_TIMEOUT_SECONDS
+from .constants import DEFAULT_READ_TIMEOUT_SECONDS, DEFAULT_STALE_QUERY_PROBE_SECONDS
+from ._transport import abort_connection, socket_writable
 from .cursor import Cursor
 from .errors import NotSupportedError, OperationalError
 from .models import ExecutionResult, ProgressInfo, Store, StoreResult
@@ -33,6 +35,151 @@ 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
+
+
+_STALE_CHECK_MAX_INTERVAL_SECONDS = 1.0
+"""Upper bound on how often the reader thread scans for stale queries."""
+
+_MAX_PROBE_BACKOFF_DOUBLINGS = 3
+"""Probe spacing doubles from N up to 8N (2**3 * N), then stays there."""
+
+_RESULTS_PROBE_AFTER_FACTOR = 4
+"""A query whose results were already requested gets one probe, after 4N.
+
+Its terminal event has arrived; only the retrieve_results reply is missing.
+The session re-serializes and re-sends the whole result (up to 100 MiB) for
+every retrieve, synchronously on its event loop, so duplicates are costly. A
+lost reply is also the rarer case: retrieve replies are sent inline on the
+session's main loop, not pushed from a callback like state_updated.
+"""
+
+_EXECUTION_NOT_FOUND = "Execution not found"
+"""sql-session's error for a retrieve_results it has no cache entry for.
+
+Its execution cache is a small LRU that can evict a query that is still
+running, whose terminal state_updated still arrives (it's pushed from the
+query's own done-callback). So for a probe, this means "can't answer", not
+"failed".
+"""
+
+_PROBED_STATES = frozenset(
+    {
+        ExecutionState.EXECUTION_REQUESTED,
+        ExecutionState.PENDING,
+        ExecutionState.RUNNING,
+        ExecutionState.RESULTS_REQUESTED,
+    }
+)
+"""States in which the driver is waiting on a session event it can re-ask for."""
+
+
+def _probe_interval(value: Any) -> float | None:
+    """The validated stale-query probe interval, or ``None`` to disable it.
+
+    Like ``_handshake_bound``, returns ``None`` for a non-numeric (including
+    ``bool``), non-finite or non-positive value. Never raises ``Exception``:
+    ``connect()`` callers pass this through, and a bad value must not cost
+    them their connection.
+    """
+    try:
+        if isinstance(value, bool) or not isinstance(value, (int, float)):
+            return None
+        value = float(value)
+        if not math.isfinite(value) or value <= 0:
+            return None
+        return value
+    except Exception:
+        return None
+
+
+_INVALID_VALUE_REPR_CHARS = 80
+"""Longest repr of an invalid setting that a warning quotes."""
+
+
+def _describe_invalid(value: Any) -> str:
+    """A bounded ``repr`` of ``value`` for a warning. Never raises ``Exception``.
+
+    Truncated so a huge value (``10**400``) can't flood the log; falls back to
+    the type name if ``repr`` fails, since this runs in the constructor.
+    """
+    try:
+        text = repr(value)
+    except Exception:
+        try:
+            return f"<unrepresentable {type(value).__name__}>"
+        except Exception:
+            return "<unrepresentable>"
+    if len(text) > _INVALID_VALUE_REPR_CHARS:
+        text = text[: _INVALID_VALUE_REPR_CHARS - 3] + "..."
+    return text
+
+
+@dataclass
+class _Watch:
+    """Stale-query watchdog bookkeeping for one query.
+
+    Only the reader thread touches it once the query is registered; ``__send``
+    stamps ``last_activity`` and ``last_event`` under ``__lock`` just before
+    registration.
+    """
+
+    last_activity: float = field(default_factory=time.monotonic)
+    """``time.monotonic()`` of the last inbound event (or probe) for the query.
+
+    Spaces the probes; see ``last_event`` for the true silence."""
+    last_event: float = field(default_factory=time.monotonic)
+    """``time.monotonic()`` of the last inbound event (or registration).
+
+    Unlike ``last_activity``, probes don't update it, so the probe log reports
+    how long the query has really been silent."""
+    probes_sent: int = 0
+    """Probes sent since the last event that wasn't a reply to a retrieve."""
+    probe_outstanding: bool = False
+    """A probe was sent and no reply to a retrieve has arrived since.
+
+    Other events, a non-terminal ``state_updated`` included, leave it set: they
+    don't answer the probe, whose reply may still be on its way."""
+    probeable: bool = True
+    """False once the session has said it can't answer probes for the query."""
+    results_requested: bool = False
+    """The normal (non-probe) retrieve_results was sent."""
+    results_probed: bool = False
+    """The one probe allowed after ``results_requested`` was sent."""
+
+
 @dataclass
 class Query:
     sql: str
@@ -40,6 +187,7 @@ class Query:
     state: ExecutionState
     handler: Callable[[Any], None]
     store: Store | None = None
+    watch: _Watch = field(default_factory=_Watch)
 
 
 class Connection:
@@ -64,14 +212,40 @@ class Connection:
         results_format: ResultsFormat | None = None,
         data_compression: DataCompression | None = None,
         geometry_representation: GeometryRepresentation | None = None,
+        session_id: str | None = None,
+        stale_query_probe_seconds: float | None = DEFAULT_STALE_QUERY_PROBE_SECONDS,
     ):
         self.__ws = ws
-        self.__read_timeout = read_timeout
         self.__results_format = results_format
         self.__data_compression = data_compression
         self.__geometry_representation = geometry_representation
         self.__progress_handler: ProgressHandler | None = None
-
+        self.__probe_after = _probe_interval(stale_query_probe_seconds)
+        if self.__probe_after is None and stale_query_probe_seconds is not None:
+            logging.warning(
+                "Invalid stale_query_probe_seconds (%s); stale-query probes disabled.",
+                _describe_invalid(stale_query_probe_seconds),
+            )
+        self.__next_stale_check = 0.0
+        # recv() must wake up for the watchdog's checks, also when the caller
+        # asked for no read timeout (websockets accepts None) or a long one.
+        self.__recv_timeout = read_timeout
+        if self.__probe_after is not None:
+            check = min(_STALE_CHECK_MAX_INTERVAL_SECONDS, self.__probe_after)
+            if read_timeout is None or (
+                isinstance(read_timeout, (int, float)) and read_timeout > check
+            ):
+                self.__recv_timeout = check
+
+        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"
@@ -85,7 +259,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
@@ -114,18 +318,174 @@ 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
+            # Checked on every iteration, not only when recv() times out: on a
+            # busy connection, steady traffic for one query would otherwise
+            # starve another query's watchdog.
+            self.__probe_stale_queries()
             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 __probe_stale_queries(self) -> None:
+        """Re-ask the session about queries that have gone quiet.
+
+        Defence in depth for a lost terminal event: the driver otherwise trusts
+        a single pushed ``state_updated`` and would wait forever if it never
+        arrives. A probe is a ``retrieve_results`` request; the session answers
+        with the query's current state (and its results or result location
+        once it has succeeded), or with an error for an unknown execution.
+
+        Runs on the reader thread, which owns every ``Query``'s watchdog fields.
+        Never raises: a watchdog failure must not disturb the receive loop.
+        """
+        probe_after = self.__probe_after
+        if probe_after is None:
+            return
+        try:
+            now = time.monotonic()
+            if now < self.__next_stale_check:
+                return
+            self.__next_stale_check = now + min(
+                _STALE_CHECK_MAX_INTERVAL_SECONDS, probe_after
+            )
+            with self.__lock:
+                if self.__closed:
+                    return
+                candidates = list(self.__queries.values())
+            for query in candidates:
+                watch = query.watch
+                if not watch.probeable or query.state not in _PROBED_STATES:
+                    continue
+                if watch.results_requested:
+                    if watch.results_probed:
+                        continue
+                    wait = probe_after * _RESULTS_PROBE_AFTER_FACTOR
+                else:
+                    doublings = min(watch.probes_sent, _MAX_PROBE_BACKOFF_DOUBLINGS)
+                    wait = probe_after * 2**doublings
+                silence = now - watch.last_activity
+                if silence < wait:
+                    continue
+                if not self.__request_results(query.execution_id, probe=True):
+                    # A user thread is sending, possibly stalled, or the
+                    # socket can't take a write. Don't wait: that would stop
+                    # recv() for every query. Either holds for the rest too;
+                    # retry at the next check.
+                    break
+                # Only a reply to a retrieve clears the flag; other events,
+                # state_updated included, don't. Still set means the last
+                # probe went unanswered.
+                unanswered = watch.probe_outstanding
+                watch.probes_sent += 1
+                watch.probe_outstanding = True
+                watch.results_probed = watch.results_requested
+                # Space the next probe from this one, not from the last event.
+                watch.last_activity = now
+                logging.info(
+                    "No events for query %s in %.1fs while %s; probed the SQL "
+                    "session for its state (probe %d%s).",
+                    query.execution_id,
+                    now - watch.last_event,
+                    query.state,
+                    watch.probes_sent,
+                    ", previous probe unanswered" if unanswered else "",
+                )
+        except Exception:
+            logging.exception("Stale-query probe failed")
+
+    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.
 
@@ -138,6 +498,27 @@ class Connection:
             # Invalid event.
             return
 
+        query = self.__queries.get(execution_id)
+        # Whether this event answers a probe (possibly along with the normal
+        # retrieve). Read before the bookkeeping below resets it.
+        probe_reply = False
+        if query is not None:
+            watch = query.watch
+            # Any inbound event, progress included, shows the query is alive.
+            watch.last_activity = watch.last_event = time.monotonic()
+            if kind == EventKind.EXECUTION_RESULT or kind == EventKind.ERROR:
+                # A reply to a retrieve_results, which may be our own probe: it
+                # resets the timer but not the backoff, so a long-running
+                # silent query isn't probed every N seconds forever.
+                probe_reply = watch.probe_outstanding
+                watch.probe_outstanding = False
+            else:
+                # A genuine event resets the backoff, but leaves an outstanding
+                # probe outstanding: a non-terminal state_updated can arrive
+                # between a probe and its reply, which must still be
+                # recognised as the probe's (see the not-found guard below).
+                watch.probes_sent = 0
+
         # Progress events are independent of the query state machine and don't
         # require a tracked query — the handler is connection-level.
         if kind == EventKind.EXECUTION_PROGRESS:
@@ -157,7 +538,6 @@ class Connection:
                 logging.exception("Progress handler raised an exception")
             return
 
-        query = self.__queries.get(execution_id)
         if not query:
             logging.warning(
                 "Received %s event for unknown execution ID %s", kind, execution_id
@@ -168,16 +548,56 @@ class Connection:
             # 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)
+            with self.__lock:
+                claimed = self.__queries.pop(execution_id, None)
+            if claimed is not None:
+                if probe_reply and result.error is None:
+                    if claimed.watch.results_requested:
+                        # The normal retrieve's reply may have won the race.
+                        logging.info(
+                            "Query %s results arrived after a stale-query probe.",
+                            execution_id,
+                        )
+                    else:
+                        logging.info(
+                            "Query %s recovered from a stale-query probe reply.",
+                            execution_id,
+                        )
+                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)))
+
+        def complete_with_store_result(result_uri: str) -> None:
+            # Results are stored in cloud storage
+            store_result = StoreResult(result_uri=result_uri, size=message.get("size"))
+            logging.info(
+                "Query %s results stored at: %s (size: %s)",
+                execution_id,
+                result_uri,
+                store_result.size,
+            )
+            query.state = ExecutionState.COMPLETED
+            complete_query(ExecutionResult(store_result=store_result))
 
         # 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:
@@ -186,19 +606,7 @@ class Connection:
                 if kind == EventKind.STATE_UPDATED:
                     result_uri = message.get("result_uri")
                     if result_uri:
-                        # Results are stored in cloud storage
-                        store_result = StoreResult(
-                            result_uri=result_uri,
-                            size=message.get("size"),
-                        )
-                        logging.info(
-                            "Query %s results stored at: %s (size: %s)",
-                            execution_id,
-                            result_uri,
-                            store_result.size,
-                        )
-                        query.state = ExecutionState.COMPLETED
-                        complete_query(ExecutionResult(store_result=store_result))
+                        complete_with_store_result(result_uri)
                         return
 
                     if query.store is not None:
@@ -212,21 +620,60 @@ class Connection:
                         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, this execution_result answers a retrieve_results.
+                # For a store-configured query, that can only be a stale-query
+                # probe: the normal path never requests results for one.
+                # Sessions with the probe contract echo the store location.
+                result_uri = message.get("result_uri")
+                if result_uri:
+                    complete_with_store_result(result_uri)
                     return
 
-                # Otherwise, process the results from the execution_result event.
                 results = message.get("results")
-                if not results or not isinstance(results, dict):
+                if query.store is not None and (results is None or results == {}):
+                    # Succeeded without a result location: either an older
+                    # session that doesn't echo it, or an empty store result
+                    # (no store_path). The two are indistinguishable, and
+                    # completing empty would silently drop an old session's
+                    # real result. Keep waiting for the genuine state_updated,
+                    # exactly as without the watchdog, and don't probe again:
+                    # the answer won't change. Known limitation: an empty store
+                    # result whose terminal event was lost can't be recovered.
+                    query.watch.probeable = False
+                    logging.warning(
+                        "Query %s succeeded, but the SQL session's reply has no "
+                        "result location (older SQL session, or an empty store "
+                        "result); cannot recover it, still waiting.",
+                        execution_id,
+                    )
+                    return
+
+                if results is None or results == {}:
                     logging.warning("Got no results back from %s.", execution_id)
                     query.state = ExecutionState.COMPLETED
                     complete_query(ExecutionResult())
                     return
 
-                query.state = ExecutionState.COMPLETED
-                complete_query(
-                    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.",
@@ -238,8 +685,28 @@ class Connection:
                 # details.
                 pass
         elif kind == EventKind.ERROR:
-            query.state = ExecutionState.FAILED
             error = message.get("message")
+            if (
+                probe_reply
+                and error == _EXECUTION_NOT_FOUND
+                and not query.watch.results_requested
+                and query.state != ExecutionState.FAILED
+            ):
+                # The session evicted the query from its cache. If the query is
+                # still running, its terminal event will still arrive; if that
+                # event was already lost, nothing more can be done (still
+                # running and finished-with-a-lost-event look the same from
+                # here), so keep waiting either way. After the normal
+                # retrieve, though, there's nothing left to wait for.
+                query.watch.probeable = False
+                logging.warning(
+                    "SQL session no longer tracks execution %s (likely evicted "
+                    "from its result cache); cannot probe it, waiting for its "
+                    "terminal event.",
+                    execution_id,
+                )
+                return
+            query.state = ExecutionState.FAILED
             complete_query(ExecutionResult(error=OperationalError(error)))
         else:
             logging.warning("Received unknown %s event!", kind)
@@ -264,16 +731,67 @@ 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,
+        blocking: bool = True,
+    ) -> bool:
+        """Send a request; return ``False`` only if non-blocking and busy.
+
+        Busy means another sender holds the send lock, or the socket's send
+        buffer is full.
+        """
+        # Serialization and redaction are local work. Fail before registration,
+        # without poisoning unrelated cursors or misreporting a transport loss.
         request = json.dumps(message)
         # 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)
+        if not self.__send_lock.acquire(blocking=blocking):
+            return False
+        try:
+            with self.__lock:
+                if self.__closed:
+                    if query is not None:
+                        raise self.__connection_error(query.execution_id)
+                    return True
+                if query is not None:
+                    # The watchdog's clock starts at registration.
+                    watch = query.watch
+                    watch.last_activity = watch.last_event = time.monotonic()
+                    self.__queries[query.execution_id] = query
+                elif message.get("execution_id") not in self.__queries:
+                    return True
+            # A non-blocking send runs on the reader thread (a probe). Not
+            # waiting for the lock isn't enough: with the socket's send buffer
+            # full (a peer that keeps sending but stops reading), ws.send()
+            # itself would block, stopping recv() for every query, and would
+            # hold the library's protocol mutex, which keepalive pings need
+            # too. Treat it like a busy lock: skip, and retry at the next
+            # check. Checked outside __lock: it's a syscall on the transport.
+            if not blocking and not socket_writable(self.__ws):
+                return False
+            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 True
+        finally:
+            self.__send_lock.release()
+        self.__fail_pending()
+        return True
 
     @staticmethod
     def __redacted_request(message: Dict[str, Any]) -> str:
@@ -289,7 +807,14 @@ class Connection:
         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.__recv_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):
@@ -318,14 +843,6 @@ 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(
@@ -334,13 +851,27 @@ class Connection:
             get_statement_type(sql),
             textwrap.shorten(redact_sql(sql), width=200),
         )
-        self.__send(request)
+        self.__send(
+            request,
+            Query(
+                sql=sql,
+                execution_id=execution_id,
+                state=ExecutionState.EXECUTION_REQUESTED,
+                handler=handler,
+                store=store,
+            ),
+        )
         return execution_id
 
-    def __request_results(self, execution_id: str) -> None:
+    def __request_results(self, execution_id: str, probe: bool = False) -> bool:
+        """Send retrieve_results; ``False`` if a probe found the sender busy.
+
+        A probe never waits for ``__send_lock`` or a full socket (see
+        ``__probe_stale_queries`` and ``__send``).
+        """
         query = self.__queries.get(execution_id)
         if not query:
-            return
+            return True
 
         request = {
             "kind": RequestKind.RETRIEVE_RESULTS.value,
@@ -353,9 +884,12 @@ class Connection:
         if self.__geometry_representation:
             request["geometry"] = self.__geometry_representation.value
 
-        query.state = ExecutionState.RESULTS_REQUESTED
-        logging.info("Requesting results from %s ...", execution_id)
-        self.__send(request)
+        if not probe:
+            # A probe only asks; the reply carries the query's actual state.
+            query.state = ExecutionState.RESULTS_REQUESTED
+            query.watch.results_requested = True
+            logging.info("Requesting results from %s ...", execution_id)
+        return self.__send(request, blocking=not probe)
 
     def __cancel_query(self, execution_id: str) -> None:
         """Cancels the query with the given execution ID."""

wherobots/db/constants.py

@@ -14,7 +14,11 @@ 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
+# Silence, per query, before the driver re-asks the session for its state.
+DEFAULT_STALE_QUERY_PROBE_SECONDS: float = 30.0
 
 MAX_MESSAGE_SIZE: int = 100 * 2**20  # 100MiB
 PROTOCOL_VERSION: Version = Version("1.0.0")

wherobots/db/cursor.py

@@ -186,6 +186,9 @@ class Cursor:
         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.__queue.put,

wherobots/db/driver.py

@@ -20,10 +20,12 @@ import certifi
 
 from .connection import Connection
 from .constants import (
+    DEFAULT_CLOSE_TIMEOUT_SECONDS,
     DEFAULT_ENDPOINT,
     DEFAULT_READ_TIMEOUT_SECONDS,
     DEFAULT_SESSION_TYPE,
     DEFAULT_SESSION_WAIT_TIMEOUT_SECONDS,
+    DEFAULT_STALE_QUERY_PROBE_SECONDS,
     MAX_MESSAGE_SIZE,
     PARAM_STYLE,
     PROTOCOL_VERSION,
@@ -124,6 +126,7 @@ def connect(
     geometry_representation: Union[GeometryRepresentation, None] = None,
     cancel_event: Union[threading.Event, None] = None,
     extra_headers: Union[Dict[str, str], None] = None,
+    stale_query_probe_seconds: Union[float, None] = DEFAULT_STALE_QUERY_PROBE_SECONDS,
 ) -> Connection:
     """Create a connection to a Wherobots SQL session.
 
@@ -145,6 +148,12 @@ def connect(
         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.
+    :param stale_query_probe_seconds: How long a running query may go without
+        any event from the SQL session before the driver re-asks the session
+        for its state, in case its completion event was lost. Later probes
+        back off to at most 8 times this interval. ``None`` disables probing;
+        an invalid value (non-numeric, non-finite, or not positive) also
+        disables it.
     """
     if not token and not api_key:
         raise ValueError("At least one of `token` or `api_key` is required")
@@ -267,6 +276,10 @@ 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],
+        stale_query_probe_seconds=stale_query_probe_seconds,
     )
 
 
@@ -294,7 +307,13 @@ 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,
+    stale_query_probe_seconds: Union[float, None] = DEFAULT_STALE_QUERY_PROBE_SECONDS,
 ) -> Connection:
+    """Connect to a SQL session's WebSocket endpoint directly.
+
+    :param stale_query_probe_seconds: See :func:`connect`.
+    """
     uri_with_protocol = f"{uri}/{protocol}"
     ssl_context = ssl.create_default_context()
     ssl_context.load_verify_locations(certifi.where())
@@ -321,6 +340,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,
         )
 
@@ -337,4 +357,6 @@ def connect_direct(
         results_format=results_format,
         data_compression=data_compression,
         geometry_representation=geometry_representation,
+        session_id=session_id,
+        stale_query_probe_seconds=stale_query_probe_seconds,
     )

wherobots/db/types.py

@@ -9,6 +9,9 @@ class ExecutionState(LowercaseStrEnum):
     EXECUTION_REQUESTED = auto()
     "Execution of a query has been requested by the driver."
 
+    PENDING = auto()
+    "The SQL session has accepted the query but not started running it yet."
+
     RUNNING = auto()
     "The SQL session has reported the query is running."
 

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

@@ -1,6 +1,6 @@
-Metadata-Version: 2.4
+Metadata-Version: 2.5
 Name: wherobots-python-dbapi
-Version: 0.28.1
+Version: 0.29.0
 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
@@ -302,3 +302,8 @@ users may find useful:
     caller needs to interrupt it (e.g. on client disconnect or timeout).
     The event is checked before each HTTP request, between retry
     attempts, and before the WebSocket handshake.
+* `stale_query_probe_seconds`: how long a running query may go without
+    any message from the SQL session before the driver asks the session
+    for its state again, in case the query's completion event was lost.
+    Repeat probes back off to at most 8 times this interval. Defaults to
+    30 seconds; pass `None` to disable probing.

wherobots_python_dbapi-0.29.0.dist-info/RECORD

@@ -0,0 +1,17 @@
+wherobots/db/__init__.py,sha256=HN_uGOUMor80GB3rZIqnPcYWDAia7QAC-y9HfGKnegs,744
+wherobots/db/_transport.py,sha256=eNe8YN_MqOTga3XNllCLFG69GQbzwEjMXBJsBLM3BEI,2256
+wherobots/db/connection.py,sha256=0_ccXOL9Yq2yalZBbWOoeBaTkympZK30dDs4DWAXHjE,39395
+wherobots/db/constants.py,sha256=m360b8GZ7mVgUcgLUgGu76gwQc1pOcEg7qTRTy9W3eg,984
+wherobots/db/cursor.py,sha256=jF9v4UhfBqZ0iZEAmIYRxWYU4E8UbMFIrwfY_PlNI60,8701
+wherobots/db/driver.py,sha256=1vlwf4RtjlCBeONqkU82zpwcZVT5FAsoL5Vuoxz_xmE,14113
+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=uRobm9S8CERnV0a3uZA5hJ1cLxZGZcIHFZRcRQ6KX10,2615
+wherobots_python_dbapi-0.29.0.dist-info/METADATA,sha256=zyJNqMu87nw1nKnZ-3fUB0znHEJ6OSaxH-oyRT5lPtQ,11509
+wherobots_python_dbapi-0.29.0.dist-info/WHEEL,sha256=W3fkpkm7-wf9vBI5Z-7s0eWkeM-spu78I8Neb98DeEg,87
+wherobots_python_dbapi-0.29.0.dist-info/licenses/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
+wherobots_python_dbapi-0.29.0.dist-info/RECORD,,

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

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

wherobots_python_dbapi-0.28.1.dist-info/RECORD

@@ -1,16 +0,0 @@
-wherobots/db/__init__.py,sha256=HN_uGOUMor80GB3rZIqnPcYWDAia7QAC-y9HfGKnegs,744
-wherobots/db/connection.py,sha256=VqXqFoKeyy9DMXP8-t0GhijWBMYB154HevbDRvyMzF4,14352
-wherobots/db/constants.py,sha256=qV-6CmOya7-f-DoRMzcDzlhcmV5PkLCcqG5TFIsxQ7o,744
-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.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,,