feat: refactor to support multiple concurrent queries

mpetazzoni · mpetazzoni · commit 2d246d849300 · 2024-04-10T19:01:22.000-07:00
Major refactoring of the protocol handling logic, moving it to the
Connection. The Connection object is now responsible for all
interactions with the SQL session, query tracking, and the corresponding
state machine.

The Cursor becomes a much simpler affair, simply requesting a query
execution from the Connection, providing a callback handler for results.
Internally, a simple queue is used to allow cursor.fetch*() methods to
block until results become available.
diff --git a/README.md b/README.md
@@ -22,7 +22,6 @@ $ pip install wherobots-python-dbapi-driver
 ## Usage
 
 ```python
-from contextlib import closing
 import tabulate
 
 from wherobots.db import connect
@@ -33,8 +32,8 @@ with connect(
         api_key='...',
         runtime=Runtime.SEDONA,
         region=Region.AWS_US_WEST_2) as conn:
-    with closing(conn.cursor()) as curr:
-        curr.execute("SHOW SCHEMAS IN wherobots_open_data")
-        results = curr.fetchall()
-        print(tabulate.tabulate(results, headers="keys", tablefmt="pretty"))
+    curr = conn.cursor()
+    curr.execute("SHOW SCHEMAS IN wherobots_open_data")
+    results = curr.fetchall()
+    print(tabulate.tabulate(results, headers="keys", tablefmt="pretty"))
 ```
diff --git a/tests/smoke.py b/tests/smoke.py
@@ -3,16 +3,14 @@
 import argparse
 import functools
 import logging
-from contextlib import closing
+import sys
 
 import shapely
-import sys
 import tabulate
 
 from wherobots.db import connect, connect_direct
-from wherobots.db.runtime import Runtime
 from wherobots.db.region import Region
-
+from wherobots.db.runtime import Runtime
 
 if __name__ == "__main__":
     parser = argparse.ArgumentParser()
@@ -60,9 +58,9 @@
         )
 
     with conn_func() as conn:
-        with closing(conn.cursor()) as cursor:
-            cursor.execute(args.sql)
-            results = cursor.fetchall()
+        cursor = conn.cursor()
+        cursor.execute(args.sql)
+        results = cursor.fetchall()
 
     for row in results:
         for key, value in row.items():
diff --git a/wherobots/db/connection.py b/wherobots/db/connection.py
@@ -1,5 +1,21 @@
+import json
+import logging
+import uuid
+from concurrent.futures import ThreadPoolExecutor
+from dataclasses import dataclass
+from typing import Callable, Any
+
+from wherobots.db.constants import RequestKind, EventKind, ExecutionState
 from wherobots.db.cursor import Cursor
-from wherobots.db.errors import NotSupportedError
+from wherobots.db.errors import NotSupportedError, DatabaseError, OperationalError
+
+
+@dataclass
+class Query:
+    sql: str
+    execution_id: str
+    state: ExecutionState
+    handler: Callable[[list[Any] | DatabaseError], None]
 
 
 class Connection:
@@ -8,19 +24,34 @@ class Connection:
 
     The connection is backed by the WebSocket connected to the Wherobots SQL session instance.
     Transactions are not supported, so commit() and rollback() raise NotSupportedError.
+
+    This class handles all the interactions with the remote SQL session, and the details of the
+    Wherobots Spatial SQL API protocol. It supports multiple concurrent cursors, each one executing
+    a single query at a time.
+
+    A background thread listens for events from the SQL session, and handles update to the
+    corresponding query state. Queries are tracked by their unique execution ID.
+
+    Note: the Connection object MUST be used as a context manager.
     """
 
     def __init__(self, ws):
         self.__ws = ws
+        self.__queries: dict[str, Query] = {}
 
     def __enter__(self):
+        self.__executor = ThreadPoolExecutor(
+            max_workers=1, thread_name_prefix="wherobots-sql-connection"
+        )
+        self.__executor.submit(self.__listen)
         return self
 
     def __exit__(self, exc_type, exc_val, exc_tb):
         self.close()
 
     def close(self):
         self.__ws.close()
+        self.__executor.shutdown(wait=True)
 
     def commit(self):
         raise NotSupportedError
@@ -29,4 +60,110 @@ def rollback(self):
         raise NotSupportedError
 
     def cursor(self) -> Cursor:
-        return Cursor(self.__ws.send, self.__ws.recv)
+        return Cursor(self.__execute_sql, self.__cancel_query)
+
+    def __listen(self):
+        """Main background loop listening for messages from the SQL session.
+
+        The code in this method is purposefully defensive to avoid unexpected situations killing the thread.
+        """
+        while True:
+            message = json.loads(self.__ws.recv())
+            logging.debug("Received message: %s", message)
+
+            execution_id = message.get("execution_id")
+            if not execution_id:
+                continue
+
+            query = self.__queries.get(execution_id)
+            if not query:
+                logging.warning(
+                    "Received %s event for unknown execution ID %s", kind, execution_id
+                )
+                continue
+
+            kind = message.get("kind")
+            match kind:
+                case EventKind.STATE_UPDATED:
+                    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
+                        )
+                        continue
+
+                    # Incoming state transitions are handled here.
+                    match query.state:
+                        case ExecutionState.SUCCEEDED:
+                            self.__request_results(execution_id)
+                        case ExecutionState.FAILED:
+                            query.handler(OperationalError("Query execution failed"))
+
+                case EventKind.EXECUTION_RESULT:
+                    results_format = message.get("results_format")
+
+                    # TODO: Support other results formats.
+                    if results_format != "json":
+                        query.handler(
+                            OperationalError(
+                                f"Unsupported results format {results_format}"
+                            )
+                        )
+                        continue
+
+                    logging.info(
+                        "Received %s results from %s.",
+                        results_format,
+                        execution_id,
+                    )
+                    query.state = ExecutionState.COMPLETED
+                    query.handler([json.loads(message.get("results"))])
+                case _:
+                    logging.warning("Received unknown %s event!", kind)
+
+    def __send(self, message: dict[str, Any]) -> None:
+        logging.debug("Sending %s", message)
+        self.__ws.send(json.dumps(message))
+
+    def __execute_sql(
+        self, sql: str, handler: Callable[[list[Any] | DatabaseError], None]
+    ) -> str:
+        """Triggers the execution of the given SQL query."""
+        execution_id = str(uuid.uuid4())
+        request = {
+            "kind": RequestKind.EXECUTE_SQL.value,
+            "execution_id": execution_id,
+            "statement": sql,
+        }
+
+        self.__queries[execution_id] = Query(
+            sql=sql,
+            execution_id=execution_id,
+            state=ExecutionState.EXECUTION_REQUESTED,
+            handler=handler,
+        )
+        self.__send(request)
+        return execution_id
+
+    def __request_results(self, execution_id: str) -> None:
+        query = self.__queries.get(execution_id)
+        if not query:
+            return
+
+        # TODO: Switch to Arrow encoding of results when supported.
+        request = {
+            "kind": RequestKind.RETRIEVE_RESULTS.value,
+            "execution_id": execution_id,
+            "results_format": "json",
+        }
+        query.state = ExecutionState.RESULTS_REQUESTED
+        logging.info("Requesting results from %s ...", execution_id)
+        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.
diff --git a/wherobots/db/cursor.py b/wherobots/db/cursor.py
@@ -1,28 +1,20 @@
-from concurrent.futures import CancelledError, Future, ThreadPoolExecutor
-import json
-import logging
-import uuid
-from typing import Any, Callable
+import queue
+from typing import Any
 
-from .constants import ExecutionState, RequestKind, EventKind
-from .errors import ProgrammingError, OperationalError
+from .errors import ProgrammingError, DatabaseError
 
 
 class Cursor:
 
-    def __init__(self, send_func: Callable[[str], None], recv_func: Callable[[], str]):
-        self.__send_func = send_func
-        self.__recv_func = recv_func
+    def __init__(self, exec_fn, cancel_fn):
+        self.__exec_fn = exec_fn
+        self.__cancel_fn = cancel_fn
 
+        self.__queue: queue.Queue = queue.Queue()
+        self.__results: list[Any] | None = None
         self.__current_execution_id: str | None = None
-        self.__current_execution_state: ExecutionState = ExecutionState.IDLE
-        self.__current_execution_results: Future[list[Any]] | None = None
         self.__current_row: int = 0
 
-        self.__executor = ThreadPoolExecutor(
-            max_workers=1, thread_name_prefix="wherobots-sql-cursor"
-        )
-
         # Description and row count are set by the last executed operation.
         # Their default values are defined by PEP-0249.
         self.__description: str | None = None
@@ -39,104 +31,36 @@ def description(self) -> str | None:
     def rowcount(self) -> int:
         return self.__rowcount
 
-    def close(self):
-        self.__executor.shutdown()
+    def __on_execution_result(self, result: list[Any] | DatabaseError) -> None:
+        self.__queue.put(result)
 
-    def __send_request(self, request):
-        return self.__send_func(json.dumps(request))
+    def __get_results(self) -> list[Any] | None:
+        if not self.__current_execution_id:
+            raise ProgrammingError("No query has been executed yet")
+        if self.__results is not None:
+            return self.__results
+
+        result = self.__queue.get()
+        if isinstance(result, DatabaseError):
+            raise result
+        self.__rowcount = len(result)
+        self.__results = result
+        return self.__results
 
     def execute(self, operation: str, parameters: dict[str, Any] = None):
-        self.__current_execution_id = str(uuid.uuid4())
-        self.__current_execution_state = ExecutionState.EXECUTION_REQUESTED
+        if self.__current_execution_id:
+            self.__cancel_fn(self.__current_execution_id)
+
+        self.__results = None
         self.__current_row = 0
         self.__rowcount = -1
 
-        def _execute(request):
-            """This function is executed in a separate thread to send the request, wait for, and gather the results."""
-            logging.info("Executing SQL: %s", request["statement"])
-            self.__send_request(request)
-            return self.__gather_results()
-
         sql = operation.format(**(parameters or {}))
-        exec_request = {
-            "kind": RequestKind.EXECUTE_SQL.value,
-            "execution_id": self.__current_execution_id,
-            "statement": sql,
-        }
-
-        if self.__current_execution_results:
-            self.__current_execution_results.cancel()
-        self.__current_execution_results = self.__executor.submit(
-            _execute, exec_request
-        )
+        self.__current_execution_id = self.__exec_fn(sql, self.__on_execution_result)
 
     def executemany(self, operation: str, seq_of_parameters: list[dict[str, Any]]):
         raise NotImplementedError
 
-    def __get_results(self) -> list[Any] | None:
-        if not self.__current_execution_results:
-            raise ProgrammingError("No query has been executed yet")
-        try:
-            return self.__current_execution_results.result()
-        except CancelledError:
-            raise ProgrammingError(
-                "Query execution was cancelled while waiting for results"
-            )
-
-    def __gather_results(self):
-        """
-        Reads from the SQL session for updates on the current execution state. Once the query has completed
-        successfully, requests and processes the results.
-        """
-        while not self.__current_execution_state.is_terminal_state():
-            response = json.loads(self.__recv_func())
-            logging.debug("Received response: %s", response)
-
-            kind = EventKind[response["kind"].upper()]
-            match kind:
-                case EventKind.STATE_UPDATED:
-                    self.__current_execution_state = ExecutionState[
-                        response["state"].upper()
-                    ]
-                    logging.info(
-                        "Query %s is %s.",
-                        self.__current_execution_id,
-                        self.__current_execution_state,
-                    )
-
-                    match self.__current_execution_state:
-                        case ExecutionState.SUCCEEDED:
-                            results_request = {
-                                "kind": RequestKind.RETRIEVE_RESULTS.value,
-                                "execution_id": self.__current_execution_id,
-                                "results_format": "json",
-                            }
-                            logging.info(
-                                "Requesting results from %s ...",
-                                self.__current_execution_id,
-                            )
-                            self.__send_request(results_request)
-                            self.__current_execution_state = (
-                                ExecutionState.RESULTS_REQUESTED
-                            )
-                        case ExecutionState.FAILED:
-                            raise OperationalError("Execution failed")
-                case EventKind.EXECUTION_RESULT:
-                    self.__current_execution_state = ExecutionState.COMPLETED
-                    results_format = response["results_format"]
-                    logging.info(
-                        "Received %s results from %s.",
-                        results_format,
-                        self.__current_execution_id,
-                    )
-                    if results_format != "json":
-                        raise OperationalError(
-                            f"Unsupported results format {results_format}"
-                        )
-                    # TODO: When full results are sent, we won't need to wrap them in a list to simulate a row.
-                    self.__rowcount = 1
-                    return [json.loads(response["results"])]
-
     def fetchone(self):
         results = self.__get_results()[self.__current_row :]
         if not results:
