wherobots-python-dbapi 0.23.2__tar.gz → 0.25.0__tar.gz

{wherobots_python_dbapi-0.23.2 → wherobots_python_dbapi-0.25.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: wherobots-python-dbapi
-Version: 0.23.2
+Version: 0.25.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
@@ -123,6 +123,59 @@ The `Store` class supports the following options:
 Use `Store.for_download()` as a convenient shorthand for storing results
 as a single Parquet file with a presigned URL.
 
+#### Store options
+
+You can pass format-specific Spark write options through the `options`
+parameter. These correspond to the options available in Spark's
+`DataFrameWriter` and are applied after the server's default options,
+allowing you to override them.
+
+```python
+# CSV without headers and a custom delimiter
+store = Store.for_download(
+    format=StorageFormat.CSV,
+    options={"header": "false", "delimiter": "|"},
+)
+
+# GeoJSON preserving null fields
+store = Store.for_download(
+    format=StorageFormat.GEOJSON,
+    options={"ignoreNullFields": "false"},
+)
+```
+
+### Execution progress
+
+You can monitor the progress of running queries by registering a
+progress handler on the connection.
+
+```python
+from wherobots.db import connect, ProgressInfo
+from wherobots.db.region import Region
+from wherobots.db.runtime import Runtime
+
+def on_progress(info: ProgressInfo) -> None:
+    print(f"{info.tasks_completed}/{info.tasks_total} tasks "
+          f"({info.tasks_active} active)")
+
+with connect(
+        api_key='...',
+        runtime=Runtime.TINY,
+        region=Region.AWS_US_WEST_2) as conn:
+    conn.set_progress_handler(on_progress)
+    curr = conn.cursor()
+    curr.execute("SELECT ...")
+    results = curr.fetchall()
+```
+
+The handler receives a `ProgressInfo` object with `execution_id`,
+`tasks_total`, `tasks_completed`, and `tasks_active` fields. Pass
+`None` to `set_progress_handler()` to disable progress reporting.
+
+Progress events are best-effort and may not be available for all query
+types or server versions. The handler is simply not invoked when no
+progress information is available.
+
 ### Runtime and region selection
 
 You can chose the Wherobots runtime you want to use using the `runtime`

{wherobots_python_dbapi-0.23.2 → wherobots_python_dbapi-0.25.0}/README.md

@@ -99,6 +99,59 @@ The `Store` class supports the following options:
 Use `Store.for_download()` as a convenient shorthand for storing results
 as a single Parquet file with a presigned URL.
 
+#### Store options
+
+You can pass format-specific Spark write options through the `options`
+parameter. These correspond to the options available in Spark's
+`DataFrameWriter` and are applied after the server's default options,
+allowing you to override them.
+
+```python
+# CSV without headers and a custom delimiter
+store = Store.for_download(
+    format=StorageFormat.CSV,
+    options={"header": "false", "delimiter": "|"},
+)
+
+# GeoJSON preserving null fields
+store = Store.for_download(
+    format=StorageFormat.GEOJSON,
+    options={"ignoreNullFields": "false"},
+)
+```
+
+### Execution progress
+
+You can monitor the progress of running queries by registering a
+progress handler on the connection.
+
+```python
+from wherobots.db import connect, ProgressInfo
+from wherobots.db.region import Region
+from wherobots.db.runtime import Runtime
+
+def on_progress(info: ProgressInfo) -> None:
+    print(f"{info.tasks_completed}/{info.tasks_total} tasks "
+          f"({info.tasks_active} active)")
+
+with connect(
+        api_key='...',
+        runtime=Runtime.TINY,
+        region=Region.AWS_US_WEST_2) as conn:
+    conn.set_progress_handler(on_progress)
+    curr = conn.cursor()
+    curr.execute("SELECT ...")
+    results = curr.fetchall()
+```
+
+The handler receives a `ProgressInfo` object with `execution_id`,
+`tasks_total`, `tasks_completed`, and `tasks_active` fields. Pass
+`None` to `set_progress_handler()` to disable progress reporting.
+
+Progress events are best-effort and may not be available for all query
+types or server versions. The handler is simply not invoked when no
+progress information is available.
+
 ### Runtime and region selection
 
 You can chose the Wherobots runtime you want to use using the `runtime`

{wherobots_python_dbapi-0.23.2 → wherobots_python_dbapi-0.25.0}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "wherobots-python-dbapi"
-version = "0.23.2"
+version = "0.25.0"
 description = "Python DB-API driver for Wherobots DB"
 authors = [{ name = "Maxime Petazzoni", email = "max@wherobots.com" }]
 requires-python = ">=3.10, <4"

{wherobots_python_dbapi-0.23.2 → wherobots_python_dbapi-0.25.0}/wherobots/db/__init__.py

@@ -10,7 +10,7 @@ from .errors import (
     ProgrammingError,
     NotSupportedError,
 )
-from .models import Store, StoreResult
+from .models import ProgressInfo, Store, StoreResult
 from .region import Region
 from .runtime import Runtime
 from .types import StorageFormat
@@ -18,6 +18,7 @@ from .types import StorageFormat
 __all__ = [
     "Connection",
     "Cursor",
+    "ProgressInfo",
     "connect",
     "connect_direct",
     "Error",

{wherobots_python_dbapi-0.23.2 → wherobots_python_dbapi-0.25.0}/wherobots/db/connection.py

@@ -16,7 +16,7 @@ import websockets.sync.client
 from .constants import DEFAULT_READ_TIMEOUT_SECONDS
 from .cursor import Cursor
 from .errors import NotSupportedError, OperationalError
-from .models import ExecutionResult, Store, StoreResult
+from .models import ExecutionResult, ProgressInfo, Store, StoreResult
 from .types import (
     RequestKind,
     EventKind,
@@ -27,6 +27,10 @@ from .types import (
 )
 
 
+ProgressHandler = Callable[[ProgressInfo], None]
+"""A callable invoked with a :class:`ProgressInfo` on every progress event."""
+
+
 @dataclass
 class Query:
     sql: str
@@ -64,6 +68,7 @@ class Connection:
         self.__results_format = results_format
         self.__data_compression = data_compression
         self.__geometry_representation = geometry_representation
+        self.__progress_handler: ProgressHandler | None = None
 
         self.__queries: dict[str, Query] = {}
         self.__thread = threading.Thread(
@@ -89,6 +94,21 @@ class Connection:
     def cursor(self) -> Cursor:
         return Cursor(self.__execute_sql, self.__cancel_query)
 
+    def set_progress_handler(self, handler: ProgressHandler | None) -> None:
+        """Register a callback invoked for execution progress events.
+
+        When a handler is set, every ``execute_sql`` request automatically
+        includes ``enable_progress_events: true`` so the SQL session streams
+        progress updates for running queries.
+
+        Pass ``None`` to disable progress reporting.
+
+        This follows the `sqlite3 Connection.set_progress_handler()
+        <https://docs.python.org/3/library/sqlite3.html#sqlite3.Connection.set_progress_handler>`_
+        pattern (PEP 249 vendor extension).
+        """
+        self.__progress_handler = handler
+
     def __main_loop(self) -> None:
         """Main background loop listening for messages from the SQL session."""
         logging.info("Starting background connection handling loop...")
@@ -116,6 +136,25 @@ class Connection:
             # Invalid event.
             return
 
+        # Progress events are independent of the query state machine and don't
+        # require a tracked query — the handler is connection-level.
+        if kind == EventKind.EXECUTION_PROGRESS:
+            handler = self.__progress_handler
+            if handler is None:
+                return
+            try:
+                handler(
+                    ProgressInfo(
+                        execution_id=execution_id,
+                        tasks_total=message.get("tasks_total", 0),
+                        tasks_completed=message.get("tasks_completed", 0),
+                        tasks_active=message.get("tasks_active", 0),
+                    )
+                )
+            except Exception:
+                logging.exception("Progress handler raised an exception")
+            return
+
         query = self.__queries.get(execution_id)
         if not query:
             logging.warning(
@@ -236,12 +275,11 @@ class Connection:
             "statement": sql,
         }
 
+        if self.__progress_handler is not None:
+            request["enable_progress_events"] = True
+
         if store:
-            request["store"] = {
-                "format": store.format.value,
-                "single": str(store.single).lower(),
-                "generate_presigned_url": str(store.generate_presigned_url).lower(),
-            }
+            request["store"] = store.to_dict()
 
         self.__queries[execution_id] = Query(
             sql=sql,

{wherobots_python_dbapi-0.23.2 → wherobots_python_dbapi-0.25.0}/wherobots/db/models.py

@@ -1,4 +1,5 @@
 from dataclasses import dataclass
+from typing import Any, Dict
 
 import pandas
 
@@ -31,18 +32,32 @@ class Store:
         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.
+        options: Optional dict of format-specific Spark DataFrameWriter options
+            (e.g. ``{"header": "false", "delimiter": "|"}`` for CSV). These are
+            applied after the server's default options, so they can override them.
+            An empty dict is normalized to None.
     """
 
     format: StorageFormat
     single: bool = False
     generate_presigned_url: bool = False
+    options: dict[str, str] | None = None
 
     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")
+        # Normalize empty options to None and defensively copy.
+        if self.options:
+            self.options = dict(self.options)
+        else:
+            self.options = None
 
     @classmethod
-    def for_download(cls, format: StorageFormat | None = None) -> "Store":
+    def for_download(
+        cls,
+        format: StorageFormat | None = None,
+        options: dict[str, str] | None = None,
+    ) -> "Store":
         """Create a configuration for downloading results via a presigned URL.
 
         This is a convenience method that creates a configuration with
@@ -50,6 +65,7 @@ class Store:
 
         Args:
             format: The storage format.
+            options: Optional format-specific Spark DataFrameWriter options.
 
         Returns:
             A Store configured for single-file download with presigned URL.
@@ -58,8 +74,25 @@ class Store:
             format=format or DEFAULT_STORAGE_FORMAT,
             single=True,
             generate_presigned_url=True,
+            options=options,
         )
 
+    def to_dict(self) -> Dict[str, Any]:
+        """Serialize this Store to a dict for the WebSocket request.
+
+        Returns a dict suitable for inclusion as the ``"store"`` field in an
+        ``execute_sql`` request.  The ``options`` key is omitted when there
+        are no user-supplied options (backward compatible).
+        """
+        d: Dict[str, Any] = {
+            "format": self.format.value,
+            "single": str(self.single).lower(),
+            "generate_presigned_url": str(self.generate_presigned_url).lower(),
+        }
+        if self.options:
+            d["options"] = self.options
+        return d
+
 
 @dataclass
 class ExecutionResult:
@@ -78,3 +111,16 @@ class ExecutionResult:
     results: pandas.DataFrame | None = None
     error: Exception | None = None
     store_result: StoreResult | None = None
+
+
+@dataclass(frozen=True)
+class ProgressInfo:
+    """Progress information for a running query.
+
+    Mirrors the ``execution_progress`` event sent by the SQL session.
+    """
+
+    execution_id: str
+    tasks_total: int
+    tasks_completed: int
+    tasks_active: int

{wherobots_python_dbapi-0.23.2 → wherobots_python_dbapi-0.25.0}/wherobots/db/types.py

@@ -45,6 +45,7 @@ class EventKind(LowercaseStrEnum):
     STATE_UPDATED = auto()
     EXECUTION_RESULT = auto()
     ERROR = auto()
+    EXECUTION_PROGRESS = auto()
 
 
 class ResultsFormat(LowercaseStrEnum):