feat: implement query cancellation

Author: mpetazzoni

wherobots/db/connection.py

@@ -131,6 +131,9 @@ def __listen(self):
             # Incoming state transitions are handled here.
             if query.state == ExecutionState.SUCCEEDED:
                 self.__request_results(execution_id)
+            elif query.state == ExecutionState.CANCELLED:
+                logging.info("Query %s has been cancelled.", execution_id)
+                self.__queries.pop(execution_id)
             elif query.state == ExecutionState.FAILED:
                 # Don't do anything here; the ERROR event is coming with more
                 # details.
@@ -230,7 +233,14 @@ def __request_results(self, execution_id: str) -> None:
         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/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_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."
 
@@ -41,12 +45,17 @@ class ExecutionState(LowercaseStrEnum):
     "The driver has completed processing the query results."
 
     def is_terminal_state(self):
-        return self in (ExecutionState.COMPLETED, ExecutionState.FAILED)
+        return self in (
+            ExecutionState.COMPLETED,
+            ExecutionState.CANCELLED,
+            ExecutionState.FAILED,
+        )
 
 
 class RequestKind(LowercaseStrEnum):
     EXECUTE_SQL = auto()
     RETRIEVE_RESULTS = auto()
+    CANCEL = auto()
 
 
 class EventKind(LowercaseStrEnum):

wherobots/db/cursor.py

@@ -89,7 +89,7 @@ def executemany(self, operation: str, seq_of_parameters: list[dict[str, Any]]):
 
     def fetchone(self):
         results = self.__get_results()[self.__current_row :]
-        if not results:
+        if len(results) == 0:
             return None
         self.__current_row += 1
         return results[0]
@@ -105,7 +105,8 @@ def fetchall(self):
 
     def close(self):
         """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

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 .region import Region
 from .runtime import Runtime
-from .connection import Connection
 
 apilevel = "2.0"
 threadsafety = 1
@@ -163,14 +162,15 @@ 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,
     data_compression: Union[DataCompression, None] = None,
     geometry_representation: Union[GeometryRepresentation, None] = None,
 ) -> Connection:
     q = queue.SimpleQueue()
-    uri_with_protocol = f"{uri}/{PROTOCOL_VERSION}"
+    uri_with_protocol = f"{uri}/{protocol}"
 
     try:
         logging.info("Connecting to SQL session at %s ...", uri_with_protocol)