Add execution progress handler following sqlite3 set_progress_handler pattern

Implement connection-level progress reporting as a PEP 249 vendor extension:

- ProgressInfo NamedTuple and ProgressHandler type alias in connection.py
- set_progress_handler(handler) method on Connection
- Automatically sends enable_progress_events: true when handler is set
- Handles execution_progress WebSocket events independently of the
  query state machine, with exception isolation
- Export ProgressInfo from wherobots.db
- Add --progress flag to tests/smoke.py for manual testing with rich output

Author: peterfoldes

tests/smoke.py

@@ -11,8 +11,8 @@
 from rich.table import Table
 
 from wherobots.db import connect, connect_direct, errors
+from wherobots.db.connection import Connection, ProgressInfo
 from wherobots.db.constants import DEFAULT_ENDPOINT, DEFAULT_SESSION_TYPE
-from wherobots.db.connection import Connection
 from wherobots.db.region import Region
 from wherobots.db.runtime import Runtime
 from wherobots.db.session_type import SessionType
@@ -54,6 +54,11 @@
     parser.add_argument(
         "--wide", help="Enable wide output", action="store_const", const=80, default=30
     )
+    parser.add_argument(
+        "--progress",
+        help="Enable execution progress reporting",
+        action="store_true",
+    )
     parser.add_argument("sql", nargs="+", help="SQL query to execute")
     args = parser.parse_args()
 
@@ -134,6 +139,26 @@ def execute(conn: Connection, sql: str) -> pandas.DataFrame | StoreResult:
 
     try:
         with conn_func() as conn:
+            if args.progress:
+                console = Console(stderr=True)
+
+                def _on_progress(info: ProgressInfo) -> None:
+                    pct = (
+                        f"{info.tasks_completed / info.tasks_total * 100:.0f}%"
+                        if info.tasks_total
+                        else "?"
+                    )
+                    console.print(
+                        f"  [dim]\\[progress][/dim] "
+                        f"[bold]{pct}[/bold]  "
+                        f"{info.tasks_completed}/{info.tasks_total} tasks "
+                        f"[dim]({info.tasks_active} active)[/dim]  "
+                        f"[dim]{info.execution_id[:8]}[/dim]",
+                        highlight=False,
+                    )
+
+                conn.set_progress_handler(_on_progress)
+
             with concurrent.futures.ThreadPoolExecutor() as pool:
                 futures = [pool.submit(execute, conn, s) for s in args.sql]
                 for future in concurrent.futures.as_completed(futures):

wherobots/db/__init__.py

@@ -1,4 +1,4 @@
-from .connection import Connection
+from .connection import Connection, ProgressInfo
 from .cursor import Cursor
 from .driver import connect, connect_direct
 from .errors import (
@@ -18,6 +18,7 @@
 __all__ = [
     "Connection",
     "Cursor",
+    "ProgressInfo",
     "connect",
     "connect_direct",
     "Error",

wherobots/db/connection.py

@@ -4,7 +4,7 @@
 import threading
 import uuid
 from dataclasses import dataclass
-from typing import Any, Callable, Dict
+from typing import Any, Callable, Dict, NamedTuple
 
 import pandas
 import pyarrow
@@ -27,6 +27,22 @@
 )
 
 
+class ProgressInfo(NamedTuple):
+    """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
+
+
+ProgressHandler = Callable[[ProgressInfo], None]
+"""A callable invoked with a :class:`ProgressInfo` on every progress event."""
+
+
 @dataclass
 class Query:
     sql: str
@@ -64,6 +80,7 @@ def __init__(
         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 +106,21 @@ def rollback(self) -> None:
     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 +148,24 @@ def __listen(self) -> None:
             # 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 not None:
+                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,6 +286,9 @@ def __execute_sql(
             "statement": sql,
         }
 
+        if self.__progress_handler is not None:
+            request["enable_progress_events"] = True
+
         if store:
             request["store"] = {
                 "format": store.format.value,

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):