wherobots-python-dbapi 0.28.1__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,5 +1,6 @@
 import json
 import logging
+import math
 import textwrap
 import threading
 import uuid
@@ -12,10 +13,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 ._transport import abort_connection
 from .cursor import Cursor
 from .errors import NotSupportedError, OperationalError
 from .models import ExecutionResult, ProgressInfo, Store, StoreResult
@@ -33,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
@@ -64,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
@@ -72,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"
@@ -85,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
@@ -114,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.
 
@@ -168,16 +329,32 @@ 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:
+                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:
@@ -212,21 +389,34 @@ 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, 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
                     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.",
@@ -264,16 +454,40 @@ 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)
         # 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)
+        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:
@@ -289,7 +503,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.__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):
@@ -318,14 +539,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,7 +547,16 @@ 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:

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

@@ -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,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,
@@ -267,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],
     )
 
 
@@ -294,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()
@@ -321,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,
         )
 
@@ -337,4 +343,5 @@ def connect_direct(
         results_format=results_format,
         data_compression=data_compression,
         geometry_representation=geometry_representation,
+        session_id=session_id,
     )

{wherobots_python_dbapi-0.28.1.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.1
+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

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.1.dist-info → wherobots_python_dbapi-0.28.2.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,,