wherobots-python-dbapi 0.8.0__tar.gz → 0.9.1__tar.gz

{wherobots_python_dbapi-0.8.0 → wherobots_python_dbapi-0.9.1}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: wherobots-python-dbapi
-Version: 0.8.0
+Version: 0.9.1
 Summary: Python DB-API driver for Wherobots DB
 License: Apache 2.0
 Author: Maxime Petazzoni
@@ -12,6 +12,7 @@ Classifier: Programming Language :: Python :: 3.9
 Classifier: Programming Language :: Python :: 3.10
 Classifier: Programming Language :: Python :: 3.11
 Classifier: Programming Language :: Python :: 3.12
+Classifier: Programming Language :: Python :: 3.13
 Requires-Dist: StrEnum (>=0.4.15,<0.5.0)
 Requires-Dist: cbor2 (>=5.6.3,<6.0.0)
 Requires-Dist: numpy (<2)
@@ -19,7 +20,7 @@ Requires-Dist: pandas (>=2.1.0,<3.0.0)
 Requires-Dist: pyarrow (>=14.0.2,<15.0.0)
 Requires-Dist: requests (>=2.31.0,<3.0.0)
 Requires-Dist: tenacity (>=8.2.3,<9.0.0)
-Requires-Dist: websockets (>=12.0,<13.0)
+Requires-Dist: websockets (==13.0)
 Description-Content-Type: text/markdown
 
 # wherobots-python-dbapi

{wherobots_python_dbapi-0.8.0 → wherobots_python_dbapi-0.9.1}/pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "wherobots-python-dbapi"
-version = "0.8.0"
+version = "0.9.1"
 description = "Python DB-API driver for Wherobots DB"
 authors = ["Maxime Petazzoni <max@wherobots.com>"]
 license = "Apache 2.0"
@@ -14,7 +14,7 @@ Tracker = "https://github.com/wherobots/wherobots-python-dbapi-driver/issues"
 [tool.poetry.dependencies]
 python = "^3.9"
 requests = "^2.31.0"
-websockets = "^12.0"
+websockets = "13.0"
 tenacity = "^8.2.3"
 pyarrow = "^14.0.2"
 cbor2 = "^5.6.3"
@@ -29,6 +29,7 @@ pytest = "^8.0.2"
 black = "^24.2.0"
 pre-commit = "^3.6.2"
 conventional-pre-commit = "^3.1.0"
+types-requests = "^2.32.0.20241016"
 rich = "^13.7.1"
 
 [build-system]

{wherobots_python_dbapi-0.8.0 → wherobots_python_dbapi-0.9.1}/wherobots/db/connection.py

@@ -7,6 +7,7 @@ from dataclasses import dataclass
 from typing import Any, Callable, Union
 
 import cbor2
+import pandas
 import pyarrow
 import websockets.exceptions
 import websockets.protocol
@@ -74,19 +75,19 @@ class Connection:
     def __exit__(self, exc_type, exc_val, exc_tb):
         self.close()
 
-    def close(self):
+    def close(self) -> None:
         self.__ws.close()
 
-    def commit(self):
+    def commit(self) -> None:
         raise NotSupportedError
 
-    def rollback(self):
+    def rollback(self) -> None:
         raise NotSupportedError
 
     def cursor(self) -> Cursor:
         return Cursor(self.__execute_sql, self.__cancel_query)
 
-    def __main_loop(self):
+    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:
@@ -101,7 +102,7 @@ class Connection:
             except Exception as e:
                 logging.exception("Error handling message from SQL session", exc_info=e)
 
-    def __listen(self):
+    def __listen(self) -> None:
         """Waits for the next message from the SQL session and processes it.
 
         The code in this method is purposefully defensive to avoid unexpected situations killing the thread.
@@ -120,7 +121,8 @@ class Connection:
             )
             return
 
-        if kind == EventKind.STATE_UPDATED:
+        # 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)
@@ -128,43 +130,32 @@ class Connection:
                 logging.warning("Invalid state update message for %s", execution_id)
                 return
 
-            # Incoming state transitions are handled here.
             if query.state == ExecutionState.SUCCEEDED:
-                self.__request_results(execution_id)
+                # On a state_updated event telling us the query succeeded,
+                # ask for results.
+                if kind == EventKind.STATE_UPDATED:
+                    self.__request_results(execution_id)
+                    return
+
+                # Otherwise, process the results from the execution_result event.
+                results = message.get("results")
+                if not results or not isinstance(results, dict):
+                    logging.warning("Got no results back from %s.", execution_id)
+                    return
+
+                query.state = ExecutionState.COMPLETED
+                query.handler(self._handle_results(execution_id, results))
+            elif query.state == ExecutionState.CANCELLED:
+                logging.info(
+                    "Query %s has been cancelled; returning empty results.",
+                    execution_id,
+                )
+                query.handler(pandas.DataFrame())
+                self.__queries.pop(execution_id)
             elif query.state == ExecutionState.FAILED:
                 # Don't do anything here; the ERROR event is coming with more
                 # details.
                 pass
-
-        elif kind == EventKind.EXECUTION_RESULT:
-            results = message.get("results")
-            if not results or not isinstance(results, dict):
-                logging.warning("Got no results back from %s.", execution_id)
-                return
-
-            result_bytes = results.get("result_bytes")
-            result_format = results.get("format")
-            result_compression = results.get("compression")
-            logging.info(
-                "Received %d bytes of %s-compressed %s results from %s.",
-                len(result_bytes),
-                result_compression,
-                result_format,
-                execution_id,
-            )
-
-            query.state = ExecutionState.COMPLETED
-            if result_format == ResultsFormat.JSON:
-                query.handler(json.loads(result_bytes.decode("utf-8")))
-            elif result_format == ResultsFormat.ARROW:
-                buffer = pyarrow.py_buffer(result_bytes)
-                stream = pyarrow.input_stream(buffer, result_compression)
-                with pyarrow.ipc.open_stream(stream) as reader:
-                    query.handler(reader.read_pandas())
-            else:
-                query.handler(
-                    OperationalError(f"Unsupported results format {result_format}")
-                )
         elif kind == EventKind.ERROR:
             query.state = ExecutionState.FAILED
             error = message.get("message")
@@ -172,6 +163,28 @@ class Connection:
         else:
             logging.warning("Received unknown %s event!", kind)
 
+    def _handle_results(self, execution_id: str, results: dict[str, Any]) -> Any:
+        result_bytes = results.get("result_bytes")
+        result_format = results.get("format")
+        result_compression = results.get("compression")
+        logging.info(
+            "Received %d bytes of %s-compressed %s results from %s.",
+            len(result_bytes),
+            result_compression,
+            result_format,
+            execution_id,
+        )
+
+        if result_format == ResultsFormat.JSON:
+            return json.loads(result_bytes.decode("utf-8"))
+        elif result_format == ResultsFormat.ARROW:
+            buffer = pyarrow.py_buffer(result_bytes)
+            stream = pyarrow.input_stream(buffer, result_compression)
+            with pyarrow.ipc.open_stream(stream) as reader:
+                return reader.read_pandas()
+        else:
+            return OperationalError(f"Unsupported results format {result_format}")
+
     def __send(self, message: dict[str, Any]) -> None:
         request = json.dumps(message)
         logging.debug("Request: %s", request)
@@ -230,7 +243,14 @@ class Connection:
         self.__send(request)
 
     def __cancel_query(self, execution_id: str) -> None:
-        query = self.__queries.pop(execution_id)
-        if query:
-            logging.info("Cancelled query %s.", execution_id)
-            # TODO: when protocol supports it, send cancellation request.
+        """Cancels the query with the given execution ID."""
+        query = self.__queries.get(execution_id)
+        if not query:
+            return
+
+        request = {
+            "kind": RequestKind.CANCEL.value,
+            "execution_id": execution_id,
+        }
+        logging.info("Cancelling query %s...", execution_id)
+        self.__send(request)

{wherobots_python_dbapi-0.8.0 → wherobots_python_dbapi-0.9.1}/wherobots/db/constants.py

@@ -1,4 +1,5 @@
 from enum import auto
+from packaging.version import Version
 from strenum import LowercaseStrEnum, StrEnum
 
 from .region import Region
@@ -15,7 +16,7 @@ DEFAULT_SESSION_WAIT_TIMEOUT_SECONDS: float = 900
 DEFAULT_REUSE_SESSION: bool = True
 
 MAX_MESSAGE_SIZE: int = 100 * 2**20  # 100MiB
-PROTOCOL_VERSION: str = "1.0.0"
+PROTOCOL_VERSION: Version = Version("1.0.0")
 
 
 class ExecutionState(LowercaseStrEnum):
@@ -31,6 +32,9 @@ class ExecutionState(LowercaseStrEnum):
     SUCCEEDED = auto()
     "The SQL session has reported the query has completed successfully."
 
+    CANCELLED = auto()
+    "The SQL session has reported the query has been cancelled."
+
     FAILED = auto()
     "The SQL session has reported the query has failed."
 
@@ -40,13 +44,18 @@ class ExecutionState(LowercaseStrEnum):
     COMPLETED = auto()
     "The driver has completed processing the query results."
 
-    def is_terminal_state(self):
-        return self in (ExecutionState.COMPLETED, ExecutionState.FAILED)
+    def is_terminal_state(self) -> bool:
+        return self in (
+            ExecutionState.COMPLETED,
+            ExecutionState.CANCELLED,
+            ExecutionState.FAILED,
+        )
 
 
 class RequestKind(LowercaseStrEnum):
     EXECUTE_SQL = auto()
     RETRIEVE_RESULTS = auto()
+    CANCEL = auto()
 
 
 class EventKind(LowercaseStrEnum):
@@ -88,7 +97,7 @@ class AppStatus(StrEnum):
     DESTROY_FAILED = auto()
     DESTROYED = auto()
 
-    def is_starting(self):
+    def is_starting(self) -> bool:
         return self in (
             AppStatus.PENDING,
             AppStatus.PREPARING,
@@ -98,7 +107,7 @@ class AppStatus(StrEnum):
             AppStatus.INITIALIZING,
         )
 
-    def is_terminal_state(self):
+    def is_terminal_state(self) -> bool:
         return self in (
             AppStatus.PREPARE_FAILED,
             AppStatus.DEPLOY_FAILED,

{wherobots_python_dbapi-0.8.0 → wherobots_python_dbapi-0.9.1}/wherobots/db/cursor.py

@@ -1,7 +1,7 @@
 import queue
 from typing import Any, Optional, List, Tuple
 
-from .errors import ProgrammingError, DatabaseError
+from .errors import DatabaseError, ProgrammingError
 
 _TYPE_MAP = {
     "object": "STRING",
@@ -16,7 +16,7 @@ _TYPE_MAP = {
 
 class Cursor:
 
-    def __init__(self, exec_fn, cancel_fn):
+    def __init__(self, exec_fn, cancel_fn) -> None:
         self.__exec_fn = exec_fn
         self.__cancel_fn = cancel_fn
 
@@ -72,7 +72,7 @@ class Cursor:
 
         return self.__results
 
-    def execute(self, operation: str, parameters: dict[str, Any] = None):
+    def execute(self, operation: str, parameters: dict[str, Any] = None) -> None:
         if self.__current_execution_id:
             self.__cancel_fn(self.__current_execution_id)
 
@@ -84,37 +84,40 @@ class Cursor:
         sql = operation.format(**(parameters or {}))
         self.__current_execution_id = self.__exec_fn(sql, self.__on_execution_result)
 
-    def executemany(self, operation: str, seq_of_parameters: list[dict[str, Any]]):
+    def executemany(
+        self, operation: str, seq_of_parameters: list[dict[str, Any]]
+    ) -> None:
         raise NotImplementedError
 
-    def fetchone(self):
+    def fetchone(self) -> Any:
         results = self.__get_results()[self.__current_row :]
-        if not results:
+        if len(results) == 0:
             return None
         self.__current_row += 1
         return results[0]
 
-    def fetchmany(self, size: int = None):
+    def fetchmany(self, size: int = None) -> list[Any]:
         size = size or self.arraysize
         results = self.__get_results()[self.__current_row : self.__current_row + size]
         self.__current_row += size
         return results
 
-    def fetchall(self):
+    def fetchall(self) -> list[Any]:
         return self.__get_results()[self.__current_row :]
 
-    def close(self):
+    def close(self) -> None:
         """Close the cursor."""
-        pass
+        if self.__results is None and self.__current_execution_id:
+            self.__cancel_fn(self.__current_execution_id)
 
     def __iter__(self):
         return self
 
-    def __next__(self):
+    def __next__(self) -> None:
         raise StopIteration
 
     def __enter__(self):
         return self
 
-    def __exit__(self, exc_type, exc_val, exc_tb):
+    def __exit__(self, exc_type, exc_val, exc_tb) -> None:
         self.close()

{wherobots_python_dbapi-0.8.0 → wherobots_python_dbapi-0.9.1}/wherobots/db/driver.py

@@ -3,19 +3,19 @@
 A PEP-0249 compatible driver for interfacing with Wherobots DB.
 """
 
+from importlib import metadata
+from importlib.metadata import PackageNotFoundError
 import logging
+from packaging.version import Version
 import platform
-import urllib.parse
 import queue
-from importlib import metadata
-from importlib.metadata import PackageNotFoundError
-
 import requests
 import tenacity
-import threading
 from typing import Union
+import urllib.parse
 import websockets.sync.client
 
+from .connection import Connection
 from .constants import (
     DEFAULT_ENDPOINT,
     DEFAULT_REGION,
@@ -36,7 +36,6 @@ from .errors import (
 )
 from .region import Region
 from .runtime import Runtime
-from .connection import Connection
 
 apilevel = "2.0"
 threadsafety = 1
@@ -163,6 +162,7 @@ def http_to_ws(uri: str) -> str:
 
 def connect_direct(
     uri: str,
+    protocol: Version = PROTOCOL_VERSION,
     headers: dict[str, str] = None,
     read_timeout: float = DEFAULT_READ_TIMEOUT_SECONDS,
     results_format: Union[ResultsFormat, None] = None,
@@ -170,34 +170,20 @@ def connect_direct(
     geometry_representation: Union[GeometryRepresentation, None] = None,
 ) -> Connection:
     q = queue.SimpleQueue()
-    uri_with_protocol = f"{uri}/{PROTOCOL_VERSION}"
-
-    def create_ws_connection():
-        try:
-            logging.info("Connecting to SQL session at %s ...", uri_with_protocol)
-            ws = websockets.sync.client.connect(
-                uri=uri_with_protocol,
-                additional_headers=headers,
-                max_size=MAX_MESSAGE_SIZE,
-            )
-            q.put(ws)
-        except Exception as e:
-            q.put(e)
-
-    dt = threading.Thread(
-        name="wherobots-ws-connector",
-        target=create_ws_connection,
-        daemon=True,
-    )
-    dt.start()
-    dt.join()
+    uri_with_protocol = f"{uri}/{protocol}"
 
-    result = q.get()
-    if isinstance(result, Exception):
-        raise InterfaceError("Failed to connect to SQL session!") from result
+    try:
+        logging.info("Connecting to SQL session at %s ...", uri_with_protocol)
+        ws = websockets.sync.client.connect(
+            uri=uri_with_protocol,
+            additional_headers=headers,
+            max_size=MAX_MESSAGE_SIZE,
+        )
+    except Exception as e:
+        raise InterfaceError("Failed to connect to SQL session!") from e
 
     return Connection(
-        result,
+        ws,
         read_timeout=read_timeout,
         results_format=results_format,
         data_compression=data_compression,