Merge pull request #8 from wherobots/max/params

feat: pass through advanced parameters

Author: mpetazzoni

README.md

@@ -21,6 +21,11 @@ $ pip install git+https://github.com/wherobots/wherobots-python-dbapi-driver
 
 ## Usage
 
+### Basic usage
+
+Basic usage follows the typical pattern of establishing the connection,
+acquiring a cursor, and executing SQL queries through it:
+
 ```python
 from wherobots.db import connect
 from wherobots.db.region import Region
@@ -36,8 +41,8 @@ with connect(
     print(results)
 ```
 
-The `Cursor` also supports the context manager protocol, so you can use
-it within a `with` statement when needed:
+The `Cursor` supports the context manager protocol, so you can use it
+within a `with` statement when needed:
 
 ```python
 with connect(...) as conn:
@@ -49,3 +54,30 @@ with connect(...) as conn:
 It also implements the `close()` method, as suggested by the PEP-2049
 specification, to support situations where the cursor is wrapped in a
 `contextmanager.closing()`.
+
+### Runtime and region selection
+
+You can chose the Wherobots runtime you want to use using the `runtime`
+parameter, passing in one of the `Runtime` enum values. For more
+information on runtime sizing and selection, please consult the
+[Wherobots product documentation](https://docs.wherobots.com).
+
+The only supported Wherobots compute region for now is `aws-us-west-2`,
+in AWS's Oregon (`us-west-2`) region.
+
+### Advanced parameters
+
+The `connect()` method takes some additional parameters that advanced
+users may find useful:
+
+* `results_format`: one of the `ResultsFormat` enum values;
+    Arrow encoding is the default and most efficient format for
+    receiving query results.
+* `data_compression`: one of the `DataCompression` enum values; Brotli
+    compression is the default and the most efficient compression
+    algorithm for receiving query results.
+* `geometry_representation`: one of the `GeometryRepresentation` enum
+    values; selects the encoding of geometry columns returned to the
+    client application. The default is EWKT (string) and the most
+    convenient for human inspection while still being usable by
+    libraries like Shapely.

wherobots/db/connection.py

@@ -1,5 +1,6 @@
 import json
 import logging
+import textwrap
 import threading
 import uuid
 from dataclasses import dataclass
@@ -18,13 +19,11 @@
     ExecutionState,
     ResultsFormat,
     DataCompression,
+    GeometryRepresentation,
 )
 from wherobots.db.cursor import Cursor
 from wherobots.db.errors import NotSupportedError, OperationalError
 
-_DEFAULT_RESULTS_FORMAT = ResultsFormat.ARROW
-_DEFAULT_DATA_COMPRESSION = DataCompression.BROTLI
-
 
 @dataclass
 class Query:
@@ -53,9 +52,16 @@ def __init__(
         self,
         ws: websockets.sync.client.ClientConnection,
         read_timeout: float = DEFAULT_READ_TIMEOUT_SECONDS,
+        results_format: ResultsFormat | None = None,
+        data_compression: DataCompression | None = None,
+        geometry_representation: GeometryRepresentation | None = None,
     ):
         self.__ws = ws
         self.__read_timeout = read_timeout
+        self.__results_format = results_format
+        self.__data_compression = data_compression
+        self.__geometry_representation = geometry_representation
+
         self.__queries: dict[str, Query] = {}
         self.__thread = threading.Thread(
             target=self.__main_loop, daemon=True, name="wherobots-connection"
@@ -155,9 +161,7 @@ def __listen(self):
                     query.handler(reader.read_pandas())
             else:
                 query.handler(
-                    OperationalError(
-                        f"Unsupported results format {result_format}"
-                    )
+                    OperationalError(f"Unsupported results format {result_format}")
                 )
         elif kind == EventKind.ERROR:
             query.state = ExecutionState.FAILED
@@ -167,7 +171,9 @@ def __listen(self):
             logging.warning("Received unknown %s event!", kind)
 
     def __send(self, message: dict[str, Any]) -> None:
-        self.__ws.send(json.dumps(message))
+        request = json.dumps(message)
+        logging.debug("Request: %s", request)
+        self.__ws.send(request)
 
     def __recv(self) -> dict[str, Any]:
         frame = self.__ws.recv(timeout=self.__read_timeout)
@@ -194,6 +200,10 @@ def __execute_sql(self, sql: str, handler: Callable[[Any], None]) -> str:
             state=ExecutionState.EXECUTION_REQUESTED,
             handler=handler,
         )
+
+        logging.info(
+            "Executing SQL query %s: %s", execution_id, textwrap.shorten(sql, width=60)
+        )
         self.__send(request)
         return execution_id
 
@@ -205,9 +215,14 @@ def __request_results(self, execution_id: str) -> None:
         request = {
             "kind": RequestKind.RETRIEVE_RESULTS.value,
             "execution_id": execution_id,
-            "format": _DEFAULT_RESULTS_FORMAT.value,
-            "compression": _DEFAULT_DATA_COMPRESSION.value,
         }
+        if self.__results_format:
+            request["format"] = self.__results_format.value
+        if self.__data_compression:
+            request["compression"] = self.__data_compression.value
+        if self.__geometry_representation:
+            request["geometry"] = self.__geometry_representation.value
+
         query.state = ExecutionState.RESULTS_REQUESTED
         logging.info("Requesting results from %s ...", execution_id)
         self.__send(request)

wherobots/db/constants.py

@@ -58,3 +58,11 @@ class ResultsFormat(LowercaseStrEnum):
 
 class DataCompression(LowercaseStrEnum):
     BROTLI = auto()
+
+
+class GeometryRepresentation(LowercaseStrEnum):
+    WKT = auto()
+    WKB = auto()
+    EWKT = auto()
+    EWKB = auto()
+    GEOJSON = auto()

wherobots/db/driver.py

@@ -18,6 +18,9 @@
     DEFAULT_READ_TIMEOUT_SECONDS,
     DEFAULT_SESSION_WAIT_TIMEOUT_SECONDS,
     MAX_MESSAGE_SIZE,
+    ResultsFormat,
+    DataCompression,
+    GeometryRepresentation,
 )
 from .errors import (
     InterfaceError,
@@ -40,6 +43,9 @@ def connect(
     region: Region = None,
     wait_timeout: float = DEFAULT_SESSION_WAIT_TIMEOUT_SECONDS,
     read_timeout: float = DEFAULT_READ_TIMEOUT_SECONDS,
+    results_format: ResultsFormat | None = None,
+    data_compression: DataCompression | None = None,
+    geometry_representation: GeometryRepresentation | None = None,
 ) -> Connection:
     if not token and not api_key:
         raise ValueError("At least one of `token` or `api_key` is required")
@@ -113,6 +119,9 @@ def get_session_uri() -> str:
         uri=http_to_ws(session_uri),
         headers=headers,
         read_timeout=read_timeout,
+        results_format=results_format,
+        data_compression=data_compression,
+        geometry_representation=geometry_representation,
     )
 
 
@@ -129,6 +138,9 @@ def connect_direct(
     uri: str,
     headers: dict[str, str] = None,
     read_timeout: float = DEFAULT_READ_TIMEOUT_SECONDS,
+    results_format: ResultsFormat | None = None,
+    data_compression: DataCompression | None = None,
+    geometry_representation: GeometryRepresentation | None = None,
 ) -> Connection:
     q = queue.SimpleQueue()
 
@@ -156,4 +168,10 @@ def create_ws_connection():
     if isinstance(result, Exception):
         raise InterfaceError("Failed to connect to SQL session!") from result
 
-    return Connection(result, read_timeout)
+    return Connection(
+        result,
+        read_timeout=read_timeout,
+        results_format=results_format,
+        data_compression=data_compression,
+        geometry_representation=geometry_representation,
+    )