feat: add store support to cursor.execute() and connection

Author: mpetazzoni

wherobots/db/connection.py

@@ -4,7 +4,7 @@
 import threading
 import uuid
 from dataclasses import dataclass
-from typing import Any, Callable, Union, Dict
+from typing import Any, Callable, Dict
 
 import pandas
 import pyarrow
@@ -24,6 +24,7 @@
 )
 from wherobots.db.cursor import Cursor
 from wherobots.db.errors import NotSupportedError, OperationalError
+from wherobots.db.store import Store, StoreResult
 
 
 @dataclass
@@ -32,6 +33,7 @@ class Query:
     execution_id: str
     state: ExecutionState
     handler: Callable[[Any], None]
+    store: Store | None = None
 
 
 class Connection:
@@ -53,9 +55,9 @@ def __init__(
         self,
         ws: websockets.sync.client.ClientConnection,
         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,
+        results_format: ResultsFormat | None = None,
+        data_compression: DataCompression | None = None,
+        geometry_representation: GeometryRepresentation | None = None,
     ):
         self.__ws = ws
         self.__read_timeout = read_timeout
@@ -132,8 +134,27 @@ def __listen(self) -> None:
 
             if query.state == ExecutionState.SUCCEEDED:
                 # On a state_updated event telling us the query succeeded,
-                # ask for results.
+                # check if results are stored in cloud storage or need to be fetched.
                 if kind == EventKind.STATE_UPDATED:
+                    result_uri = message.get("result_uri")
+                    if result_uri:
+                        # Results are stored in cloud storage
+                        store_result = StoreResult(
+                            result_uri=result_uri,
+                            size=message.get("size"),
+                        )
+                        logging.info(
+                            "Query %s results stored at: %s (size: %s)",
+                            execution_id,
+                            result_uri,
+                            store_result.size,
+                        )
+                        query.state = ExecutionState.COMPLETED
+                        # Return empty DataFrame with the StoreResult
+                        query.handler((pandas.DataFrame(), store_result))
+                        return
+
+                    # No store configured, request results normally
                     self.__request_results(execution_id)
                     return
 
@@ -200,7 +221,12 @@ def __recv(self) -> Dict[str, Any]:
             raise ValueError("Unexpected frame type received")
         return message
 
-    def __execute_sql(self, sql: str, handler: Callable[[Any], None]) -> str:
+    def __execute_sql(
+        self,
+        sql: str,
+        handler: Callable[[Any], None],
+        store: Store | None = None,
+    ) -> str:
         """Triggers the execution of the given SQL query."""
         execution_id = str(uuid.uuid4())
         request = {
@@ -209,11 +235,19 @@ def __execute_sql(self, sql: str, handler: Callable[[Any], None]) -> str:
             "statement": sql,
         }
 
+        if store:
+            request["store"] = {
+                "format": store.format.value if store.format else None,
+                "single": store.single,
+                "generate_presigned_url": store.generate_presigned_url,
+            }
+
         self.__queries[execution_id] = Query(
             sql=sql,
             execution_id=execution_id,
             state=ExecutionState.EXECUTION_REQUESTED,
             handler=handler,
+            store=store,
         )
 
         logging.info(

wherobots/db/cursor.py

@@ -1,7 +1,8 @@
 import queue
-from typing import Any, Optional, List, Tuple, Dict
+from typing import Any, List, Tuple, Dict
 
 from .errors import DatabaseError, ProgrammingError
+from .store import Store, StoreResult
 
 _TYPE_MAP = {
     "object": "STRING",
@@ -20,20 +21,21 @@ def __init__(self, exec_fn, cancel_fn) -> None:
         self.__cancel_fn = cancel_fn
 
         self.__queue: queue.Queue = queue.Queue()
-        self.__results: Optional[list[Any]] = None
-        self.__current_execution_id: Optional[str] = None
+        self.__results: list[Any] | None = None
+        self.__store_result: StoreResult | None = None
+        self.__current_execution_id: str | None = None
         self.__current_row: int = 0
 
         # Description and row count are set by the last executed operation.
         # Their default values are defined by PEP-0249.
-        self.__description: Optional[List[Tuple]] = None
+        self.__description: List[Tuple] | None = None
         self.__rowcount: int = -1
 
         # Array-size is also defined by PEP-0249 and is expected to be read/writable.
         self.arraysize: int = 1
 
     @property
-    def description(self) -> Optional[List[Tuple]]:
+    def description(self) -> List[Tuple] | None:
         return self.__description
 
     @property
@@ -43,7 +45,7 @@ def rowcount(self) -> int:
     def __on_execution_result(self, result) -> None:
         self.__queue.put(result)
 
-    def __get_results(self) -> Optional[List[Tuple[Any, ...]]]:
+    def __get_results(self) -> List[Tuple[Any, ...]] | None:
         if not self.__current_execution_id:
             raise ProgrammingError("No query has been executed yet")
         if self.__results is not None:
@@ -53,6 +55,10 @@ def __get_results(self) -> Optional[List[Tuple[Any, ...]]]:
         if isinstance(result, DatabaseError):
             raise result
 
+        # Unpack store result if present (result is a tuple of (DataFrame, StoreResult))
+        if isinstance(result, tuple):
+            result, self.__store_result = result
+
         self.__rowcount = len(result)
         self.__results = result
         if not result.empty:
@@ -71,19 +77,40 @@ def __get_results(self) -> Optional[List[Tuple[Any, ...]]]:
 
         return self.__results
 
-    def execute(self, operation: str, parameters: Dict[str, Any] = None) -> None:
+    def execute(
+        self,
+        operation: str,
+        parameters: Dict[str, Any] | None = None,
+        store: Store | None = None,
+    ) -> None:
         if self.__current_execution_id:
             self.__cancel_fn(self.__current_execution_id)
 
         self.__results = None
+        self.__store_result = None
         self.__current_row = 0
         self.__rowcount = -1
         self.__description = None
 
         self.__current_execution_id = self.__exec_fn(
-            operation % (parameters or {}), self.__on_execution_result
+            operation % (parameters or {}), self.__on_execution_result, store
         )
 
+    def get_store_result(self) -> StoreResult | None:
+        """Get the store result for the last executed query.
+
+        Returns the StoreResult containing the URI and size of the stored
+        results, or None if the query was not configured to store results.
+
+        This method blocks until the query completes.
+        """
+        if not self.__current_execution_id:
+            raise ProgrammingError("No query has been executed yet")
+
+        # Ensure we've waited for the result
+        self.__get_results()
+        return self.__store_result
+
     def executemany(
         self, operation: str, seq_of_parameters: List[Dict[str, Any]]
     ) -> None: