feat: implement new results handling logic

Add support for the new encoding of results as a CBOR2-encoded binary
WebSocket data frame of a results envelope, containing a JSON or Arrow
encoding of a Pandas DataFrame result.

Plus some nicer output in the smoke test script.

Author: mpetazzoni

poetry.lock

@@ -16,6 +16,9 @@ python = "^3.11"
 requests = "^2.31.0"
 websockets = "^12.0"
 tenacity = "^8.2.3"
+pyarrow = "^15.0.2"
+cbor2 = "^5.6.3"
+pandas = "^2.2.2"
 
 [tool.poetry.group.dev.dependencies]
 mypy = "^1.8.0"
@@ -24,7 +27,7 @@ black = "^24.2.0"
 pre-commit = "^3.6.2"
 conventional-pre-commit = "^3.1.0"
 shapely = "^2.0.3"
-tabulate = "^0.9.0"
+rich = "^13.7.1"
 
 [build-system]
 requires = ["poetry-core"]

pyproject.toml

@@ -5,8 +5,9 @@
 import logging
 import sys
 
-import shapely
-import tabulate
+import pandas
+from rich.console import Console
+from rich.table import Table
 
 from wherobots.db import connect, connect_direct
 from wherobots.db.region import Region
@@ -24,6 +25,9 @@
         default=logging.INFO,
     )
     parser.add_argument("--ws-url", help="Direct URL to connect to")
+    parser.add_argument(
+        "--wide", help="Enable wide output", action="store_const", const=80, default=30
+    )
     parser.add_argument("sql", help="SQL query to execute")
     args = parser.parse_args()
 
@@ -60,11 +64,13 @@
     with conn_func() as conn:
         cursor = conn.cursor()
         cursor.execute(args.sql)
-        results = cursor.fetchall()
-
-    for row in results:
-        for key, value in row.items():
-            if "geometry" in key:
-                row[key] = shapely.to_geojson(shapely.from_wkt(value))
+        results: pandas.DataFrame = cursor.fetchall()
 
-    print(tabulate.tabulate(results, headers="keys", tablefmt="rounded_outline"))
+    table = Table()
+    table.add_column("#")
+    for column in results.columns:
+        table.add_column(column, max_width=args.wide, no_wrap=True)
+    for row in results.itertuples(name=None):
+        r = [str(x) for x in row]
+        table.add_row(*r)
+    Console().print(table)

tests/smoke.py

@@ -5,17 +5,30 @@
 from dataclasses import dataclass
 from typing import Callable, Any
 
-from wherobots.db.constants import RequestKind, EventKind, ExecutionState
+import cbor2
+import pyarrow
+
+from wherobots.db.constants import (
+    RequestKind,
+    EventKind,
+    ExecutionState,
+    ResultsFormat,
+    DataCompression,
+)
 from wherobots.db.cursor import Cursor
 from wherobots.db.errors import NotSupportedError, DatabaseError, OperationalError
 
 
+_DEFAULT_RESULTS_FORMAT = ResultsFormat.ARROW
+_DEFAULT_DATA_COMPRESSION = DataCompression.BROTLI
+
+
 @dataclass
 class Query:
     sql: str
     execution_id: str
     state: ExecutionState
-    handler: Callable[[list[Any] | DatabaseError], None]
+    handler: Callable[[Any], None]
 
 
 class Connection:
@@ -101,23 +114,35 @@ def __listen(self):
                             query.handler(OperationalError("Query execution failed"))
 
                 case EventKind.EXECUTION_RESULT:
-                    results_format = message.get("results_format")
+                    results = message.get("results")
+                    if not results or not isinstance(results, dict):
+                        logging.warning("Got no results back from %s.", execution_id)
+                        continue
+
+                    result_bytes = results.get("result_bytes")
+                    result_format = results.get("format")
+                    result_compression = results.get("compression")
                     logging.info(
-                        "Received %s results from %s.",
-                        results_format,
+                        "Received %d bytes of %s-compressed %s results from %s.",
+                        len(result_bytes),
+                        result_compression,
+                        result_format,
                         execution_id,
                     )
 
                     query.state = ExecutionState.COMPLETED
-                    match results_format:
-                        case "json":
-                            query.handler(json.loads(message.get("results")))
-                        case "arrow":
-                            pass
+                    match result_format:
+                        case ResultsFormat.JSON:
+                            query.handler(json.loads(result_bytes.decode("utf-8")))
+                        case 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())
                         case _:
                             query.handler(
                                 OperationalError(
-                                    f"Unsupported results format {results_format}"
+                                    f"Unsupported results format {result_format}"
                                 )
                             )
                 case _:
@@ -129,15 +154,16 @@ def __send(self, message: dict[str, Any]) -> None:
 
     def __recv(self) -> dict[str, Any]:
         frame = self.__ws.recv()
-        if isinstance(frame, bytes):
-            frame = frame.decode("utf-8")
-        message = json.loads(frame)
+        if isinstance(frame, str):
+            message = json.loads(frame)
+        elif isinstance(frame, bytes):
+            message = cbor2.loads(frame)
+        else:
+            raise ValueError("Unexpected frame type received")
         logging.debug("Received message: %s", message)
         return message
 
-    def __execute_sql(
-        self, sql: str, handler: Callable[[list[Any] | DatabaseError], None]
-    ) -> str:
+    def __execute_sql(self, sql: str, handler: Callable[[Any], None]) -> str:
         """Triggers the execution of the given SQL query."""
         execution_id = str(uuid.uuid4())
         request = {
@@ -164,7 +190,8 @@ def __request_results(self, execution_id: str) -> None:
         request = {
             "kind": RequestKind.RETRIEVE_RESULTS.value,
             "execution_id": execution_id,
-            "results_format": "json",
+            "format": _DEFAULT_RESULTS_FORMAT.value,
+            "compression": _DEFAULT_DATA_COMPRESSION.value,
         }
         query.state = ExecutionState.RESULTS_REQUESTED
         logging.info("Requesting results from %s ...", execution_id)

wherobots/db/connection.py

@@ -45,3 +45,12 @@ class RequestKind(StrEnum):
 class EventKind(StrEnum):
     STATE_UPDATED = auto()
     EXECUTION_RESULT = auto()
+
+
+class ResultsFormat(StrEnum):
+    JSON = auto()
+    ARROW = auto()
+
+
+class DataCompression(StrEnum):
+    BROTLI = auto()

wherobots/db/constants.py

@@ -31,7 +31,7 @@ def description(self) -> str | None:
     def rowcount(self) -> int:
         return self.__rowcount
 
-    def __on_execution_result(self, result: list[Any] | DatabaseError) -> None:
+    def __on_execution_result(self, result) -> None:
         self.__queue.put(result)
 
     def __get_results(self) -> list[Any] | None: