wherobots-python-dbapi 0.22.0__tar.gz → 0.23.0__tar.gz

{wherobots_python_dbapi-0.22.0 → wherobots_python_dbapi-0.23.0}/PKG-INFO

@@ -1,13 +1,13 @@
 Metadata-Version: 2.4
 Name: wherobots-python-dbapi
-Version: 0.22.0
+Version: 0.23.0
 Summary: Python DB-API driver for Wherobots DB
 Project-URL: Homepage, https://github.com/wherobots/wherobots-python-dbapi-driver
 Project-URL: Tracker, https://github.com/wherobots/wherobots-python-dbapi-driver/issues
 Author-email: Maxime Petazzoni <max@wherobots.com>
 License-Expression: Apache-2.0
 License-File: LICENSE
-Requires-Python: <4,>=3.8
+Requires-Python: <4,>=3.10
 Requires-Dist: cbor2>=5.6.3
 Requires-Dist: packaging
 Requires-Dist: pandas
@@ -83,6 +83,46 @@ 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()`.
 
+### Storing results in cloud storage
+
+For large query results, you can store them directly in cloud storage
+instead of retrieving them over the connection. This is useful when
+results are too large to transfer efficiently, or when you want to
+process them later with other tools.
+
+```python
+from wherobots.db import connect, Store, StorageFormat
+from wherobots.db.region import Region
+from wherobots.db.runtime import Runtime
+
+with connect(
+        api_key='...',
+        runtime=Runtime.TINY,
+        region=Region.AWS_US_WEST_2) as conn:
+    curr = conn.cursor()
+
+    # Store results with a presigned URL for easy download
+    curr.execute(
+        "SELECT * FROM wherobots_open_data.overture.places LIMIT 1000",
+        store=Store.for_download()
+    )
+    store_result = curr.get_store_result()
+    print(f"Results stored at: {store_result.result_uri}")
+    print(f"Size: {store_result.size} bytes")
+```
+
+The `Store` class supports the following options:
+
+* `format`: output format - `StorageFormat.PARQUET` (default),
+    `StorageFormat.CSV`, or `StorageFormat.GEOJSON`
+* `single`: if `True`, write results to a single file instead of
+    multiple partitioned files (default: `True`)
+* `generate_presigned_url`: if `True`, generate a presigned URL for
+    downloading results (default: `False`)
+
+Use `Store.for_download()` as a convenient shorthand for storing results
+as a single Parquet file with a presigned URL.
+
 ### Runtime and region selection
 
 You can chose the Wherobots runtime you want to use using the `runtime`

{wherobots_python_dbapi-0.22.0 → wherobots_python_dbapi-0.23.0}/README.md

@@ -59,6 +59,46 @@ 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()`.
 
+### Storing results in cloud storage
+
+For large query results, you can store them directly in cloud storage
+instead of retrieving them over the connection. This is useful when
+results are too large to transfer efficiently, or when you want to
+process them later with other tools.
+
+```python
+from wherobots.db import connect, Store, StorageFormat
+from wherobots.db.region import Region
+from wherobots.db.runtime import Runtime
+
+with connect(
+        api_key='...',
+        runtime=Runtime.TINY,
+        region=Region.AWS_US_WEST_2) as conn:
+    curr = conn.cursor()
+
+    # Store results with a presigned URL for easy download
+    curr.execute(
+        "SELECT * FROM wherobots_open_data.overture.places LIMIT 1000",
+        store=Store.for_download()
+    )
+    store_result = curr.get_store_result()
+    print(f"Results stored at: {store_result.result_uri}")
+    print(f"Size: {store_result.size} bytes")
+```
+
+The `Store` class supports the following options:
+
+* `format`: output format - `StorageFormat.PARQUET` (default),
+    `StorageFormat.CSV`, or `StorageFormat.GEOJSON`
+* `single`: if `True`, write results to a single file instead of
+    multiple partitioned files (default: `True`)
+* `generate_presigned_url`: if `True`, generate a presigned URL for
+    downloading results (default: `False`)
+
+Use `Store.for_download()` as a convenient shorthand for storing results
+as a single Parquet file with a presigned URL.
+
 ### Runtime and region selection
 
 You can chose the Wherobots runtime you want to use using the `runtime`

{wherobots_python_dbapi-0.22.0 → wherobots_python_dbapi-0.23.0}/pyproject.toml

@@ -1,9 +1,9 @@
 [project]
 name = "wherobots-python-dbapi"
-version = "0.22.0"
+version = "0.23.0"
 description = "Python DB-API driver for Wherobots DB"
 authors = [{ name = "Maxime Petazzoni", email = "max@wherobots.com" }]
-requires-python = ">=3.8, <4"
+requires-python = ">=3.10, <4"
 readme = "README.md"
 license = "Apache-2.0"
 dependencies = [

{wherobots_python_dbapi-0.22.0 → wherobots_python_dbapi-0.23.0}/wherobots/db/__init__.py

@@ -10,8 +10,10 @@ from .errors import (
     ProgrammingError,
     NotSupportedError,
 )
+from .models import Store, StoreResult
 from .region import Region
 from .runtime import Runtime
+from .types import StorageFormat
 
 __all__ = [
     "Connection",
@@ -27,4 +29,7 @@ __all__ = [
     "NotSupportedError",
     "Region",
     "Runtime",
+    "Store",
+    "StorageFormat",
+    "StoreResult",
 ]

{wherobots_python_dbapi-0.22.0 → wherobots_python_dbapi-0.23.0}/wherobots/db/connection.py

@@ -4,7 +4,7 @@ import textwrap
 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
@@ -13,8 +13,11 @@ import websockets.exceptions
 import websockets.protocol
 import websockets.sync.client
 
-from wherobots.db.constants import (
-    DEFAULT_READ_TIMEOUT_SECONDS,
+from .constants import DEFAULT_READ_TIMEOUT_SECONDS
+from .cursor import Cursor
+from .errors import NotSupportedError, OperationalError
+from .models import ExecutionResult, Store, StoreResult
+from .types import (
     RequestKind,
     EventKind,
     ExecutionState,
@@ -22,8 +25,6 @@ from wherobots.db.constants import (
     DataCompression,
     GeometryRepresentation,
 )
-from wherobots.db.cursor import Cursor
-from wherobots.db.errors import NotSupportedError, OperationalError
 
 
 @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 @@ class Connection:
         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,26 @@ class Connection:
 
             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
+                        query.handler(ExecutionResult(store_result=store_result))
+                        return
+
+                    # No store configured, request results normally
                     self.__request_results(execution_id)
                     return
 
@@ -144,13 +164,15 @@ class Connection:
                     return
 
                 query.state = ExecutionState.COMPLETED
-                query.handler(self._handle_results(execution_id, results))
+                query.handler(
+                    ExecutionResult(results=self._handle_results(execution_id, results))
+                )
             elif query.state == ExecutionState.CANCELLED:
                 logging.info(
                     "Query %s has been cancelled; returning empty results.",
                     execution_id,
                 )
-                query.handler(pandas.DataFrame())
+                query.handler(ExecutionResult(results=pandas.DataFrame()))
                 self.__queries.pop(execution_id)
             elif query.state == ExecutionState.FAILED:
                 # Don't do anything here; the ERROR event is coming with more
@@ -159,7 +181,7 @@ class Connection:
         elif kind == EventKind.ERROR:
             query.state = ExecutionState.FAILED
             error = message.get("message")
-            query.handler(OperationalError(error))
+            query.handler(ExecutionResult(error=OperationalError(error)))
         else:
             logging.warning("Received unknown %s event!", kind)
 
@@ -200,7 +222,12 @@ class Connection:
             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 +236,19 @@ class Connection:
             "statement": sql,
         }
 
+        if store:
+            request["store"] = {
+                "format": store.format.value,
+                "single": str(store.single).lower(),
+                "generate_presigned_url": str(store.generate_presigned_url).lower(),
+            }
+
         self.__queries[execution_id] = Query(
             sql=sql,
             execution_id=execution_id,
             state=ExecutionState.EXECUTION_REQUESTED,
             handler=handler,
+            store=store,
         )
 
         logging.info(

wherobots_python_dbapi-0.23.0/wherobots/db/constants.py

@@ -0,0 +1,22 @@
+from packaging.version import Version
+
+from .region import Region
+from .runtime import Runtime
+from .session_type import SessionType
+from .types import StorageFormat
+
+
+DEFAULT_ENDPOINT: str = "api.cloud.wherobots.com"  # "api.cloud.wherobots.com"
+STAGING_ENDPOINT: str = "api.staging.wherobots.com"  # "api.staging.wherobots.com"
+
+DEFAULT_RUNTIME: Runtime = Runtime.TINY
+DEFAULT_REGION: Region = Region.AWS_US_WEST_2
+DEFAULT_SESSION_TYPE: SessionType = SessionType.MULTI
+DEFAULT_STORAGE_FORMAT: StorageFormat = StorageFormat.PARQUET
+DEFAULT_READ_TIMEOUT_SECONDS: float = 0.25
+DEFAULT_SESSION_WAIT_TIMEOUT_SECONDS: float = 900
+
+MAX_MESSAGE_SIZE: int = 100 * 2**20  # 100MiB
+PROTOCOL_VERSION: Version = Version("1.0.0")
+
+PARAM_STYLE = "pyformat"

{wherobots_python_dbapi-0.22.0 → wherobots_python_dbapi-0.23.0}/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 .errors import ProgrammingError
+from .models import ExecutionResult, Store, StoreResult
 
 _TYPE_MAP = {
     "object": "STRING",
@@ -20,20 +21,21 @@ class Cursor:
         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,47 +45,78 @@ class Cursor:
     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:
             return self.__results
 
-        result = self.__queue.get()
-        if isinstance(result, DatabaseError):
-            raise result
+        execution_result = self.__queue.get()
+        if not isinstance(execution_result, ExecutionResult):
+            raise ProgrammingError("Unexpected result type")
 
-        self.__rowcount = len(result)
-        self.__results = result
-        if not result.empty:
+        if execution_result.error:
+            raise execution_result.error
+
+        self.__store_result = execution_result.store_result
+        results = execution_result.results
+
+        # Results is None when results are stored in cloud storage
+        if results is None:
+            return None
+
+        self.__rowcount = len(results)
+        self.__results = results
+        if not results.empty:
             self.__description = [
                 (
                     col_name,  # name
-                    _TYPE_MAP.get(str(result[col_name].dtype), "STRING"),  # type_code
+                    _TYPE_MAP.get(str(results[col_name].dtype), "STRING"),  # type_code
                     None,  # display_size
-                    result[col_name].memory_usage(),  # internal_size
+                    results[col_name].memory_usage(),  # internal_size
                     None,  # precision
                     None,  # scale
                     True,  # null_ok; Assuming all columns can accept NULL values
                 )
-                for col_name in result.columns
+                for col_name in results.columns
             ]
 
         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:

{wherobots_python_dbapi-0.22.0 → wherobots_python_dbapi-0.23.0}/wherobots/db/driver.py

@@ -27,11 +27,6 @@ from .constants import (
     MAX_MESSAGE_SIZE,
     PARAM_STYLE,
     PROTOCOL_VERSION,
-    AppStatus,
-    DataCompression,
-    GeometryRepresentation,
-    ResultsFormat,
-    SessionType,
 )
 from .errors import (
     InterfaceError,
@@ -39,6 +34,13 @@ from .errors import (
 )
 from .region import Region
 from .runtime import Runtime
+from .session_type import SessionType
+from .types import (
+    AppStatus,
+    DataCompression,
+    GeometryRepresentation,
+    ResultsFormat,
+)
 
 apilevel = "2.0"
 threadsafety = 1

wherobots_python_dbapi-0.23.0/wherobots/db/models.py

@@ -0,0 +1,80 @@
+from dataclasses import dataclass
+
+import pandas
+
+from .constants import DEFAULT_STORAGE_FORMAT
+from .types import StorageFormat
+
+
+@dataclass(frozen=True)
+class StoreResult:
+    """Result information when a query's results are stored to cloud storage.
+
+    Attributes:
+        result_uri: The URI or presigned URL of the stored result.
+        size: The size of the stored result in bytes, or None if not available.
+    """
+
+    result_uri: str
+    size: int | None = None
+
+
+@dataclass
+class Store:
+    """Configuration for storing query results to cloud storage.
+
+    When passed to cursor.execute(), query results will be written to cloud
+    storage instead of being returned directly over the WebSocket connection.
+
+    Attributes:
+        format: The storage format (parquet, csv, or geojson). Defaults to parquet.
+        single: If True, store as a single file. If False, store as multiple files.
+        generate_presigned_url: If True, generate a presigned URL for the result.
+            Requires single=True.
+    """
+
+    format: StorageFormat
+    single: bool = False
+    generate_presigned_url: bool = False
+
+    def __post_init__(self) -> None:
+        if self.generate_presigned_url and not self.single:
+            raise ValueError("Presigned URL can only be generated when single=True")
+
+    @classmethod
+    def for_download(cls, format: StorageFormat | None = None) -> "Store":
+        """Create a configuration for downloading results via a presigned URL.
+
+        This is a convenience method that creates a configuration with
+        single file mode and presigned URL generation enabled.
+
+        Args:
+            format: The storage format.
+
+        Returns:
+            A Store configured for single-file download with presigned URL.
+        """
+        return cls(
+            format=format or DEFAULT_STORAGE_FORMAT,
+            single=True,
+            generate_presigned_url=True,
+        )
+
+
+@dataclass
+class ExecutionResult:
+    """Result of a query execution.
+
+    This class encapsulates all possible outcomes of a query execution:
+    a DataFrame result, an error, or a store result (when results are
+    written to cloud storage).
+
+    Attributes:
+        results: The query results as a pandas DataFrame, or None if an error occurred.
+        error: The error that occurred during execution, or None if successful.
+        store_result: The store result if results were written to cloud storage.
+    """
+
+    results: pandas.DataFrame | None = None
+    error: Exception | None = None
+    store_result: StoreResult | None = None

wherobots_python_dbapi-0.22.0/wherobots/db/constants.py → wherobots_python_dbapi-0.23.0/wherobots/db/types.py

@@ -1,26 +1,6 @@
 from enum import auto
-from packaging.version import Version
 from strenum import LowercaseStrEnum, StrEnum
 
-from .region import Region
-from .runtime import Runtime
-from .session_type import SessionType
-
-
-DEFAULT_ENDPOINT: str = "api.cloud.wherobots.com"  # "api.cloud.wherobots.com"
-STAGING_ENDPOINT: str = "api.staging.wherobots.com"  # "api.staging.wherobots.com"
-
-DEFAULT_RUNTIME: Runtime = Runtime.TINY
-DEFAULT_REGION: Region = Region.AWS_US_WEST_2
-DEFAULT_SESSION_TYPE: SessionType = SessionType.MULTI
-DEFAULT_READ_TIMEOUT_SECONDS: float = 0.25
-DEFAULT_SESSION_WAIT_TIMEOUT_SECONDS: float = 900
-
-MAX_MESSAGE_SIZE: int = 100 * 2**20  # 100MiB
-PROTOCOL_VERSION: Version = Version("1.0.0")
-
-PARAM_STYLE = "pyformat"
-
 
 class ExecutionState(LowercaseStrEnum):
     IDLE = auto()
@@ -84,6 +64,12 @@ class GeometryRepresentation(LowercaseStrEnum):
     GEOJSON = auto()
 
 
+class StorageFormat(LowercaseStrEnum):
+    PARQUET = auto()
+    CSV = auto()
+    GEOJSON = auto()
+
+
 class AppStatus(StrEnum):
     PENDING = auto()
     PREPARING = auto()