Expose query stats through a cursor callback

README.md

@@ -54,6 +54,30 @@ rows for example `Cursor.fetchone()` or `Cursor.fetchmany()`. By default
 `Cursor.fetchmany()` fetches one row. Please set
 `trino.dbapi.Cursor.arraysize` accordingly.
 
+**Monitoring query progress**
+
+Since queries can take a while to complete, it is useful to be able to monitor
+their progress (e.g. to log the query ID and its state) while they are still
+running, rather than only after all rows have been fetched. Pass a
+`stats_callback` to `Connection.cursor()` to have it invoked with the query
+stats every time the client polls the coordinator:
+
+```python
+def log_progress(stats):
+    print(stats["queryId"], stats.get("state"))
+
+cur = conn.cursor(stats_callback=log_progress)
+cur.execute("SELECT * FROM system.runtime.nodes")
+rows = cur.fetchall()
+```
+
+The callback fires once as soon as the query is submitted, so the query ID and
+initial state are available while the query is still running, and again after
+every subsequent poll of the coordinator. Each invocation receives a copy of
+the query's current stats dictionary (the same dictionary returned by
+`Cursor.stats`), so mutating it has no effect on the client. Any exception
+raised by the callback propagates to the caller of `execute()`/`fetch()`.
+
 ### SQLAlchemy
 
 **Prerequisite**

tests/integration/test_dbapi_integration.py

@@ -102,6 +102,24 @@ def test_select_query(trino_connection):
     assert cur.stats is not None
 
 
+def test_stats_callback_invoked_while_query_is_running(trino_connection):
+    captured_stats = []
+    cur = trino_connection.cursor(stats_callback=captured_stats.append)
+    # A larger result set spans multiple pages, so the client polls the
+    # coordinator several times and the callback is invoked more than once.
+    cur.execute("SELECT * FROM tpch.sf1.customer LIMIT 5000")
+    rows = cur.fetchall()
+
+    assert len(rows) == 5000
+    # The callback fires once as soon as the query is submitted and again on
+    # every subsequent poll, so it is invoked while the query is still running.
+    assert len(captured_stats) >= 2
+    # The query id is exposed through the callback for every invocation.
+    assert all(stats["queryId"] == cur.query_id for stats in captured_stats)
+    # The final invocation reflects the terminal state of the query.
+    assert captured_stats[-1]["state"] == "FINISHED"
+
+
 def test_select_query_result_iteration(trino_connection):
     cur0 = trino_connection.cursor()
     cur0.execute("SELECT custkey FROM tpch.sf1.customer LIMIT 10")

tests/unit/test_client.py

@@ -1145,6 +1145,44 @@ def text(self):
         assert isinstance(result, TrinoResult)
 
 
+def test_stats_callback_cannot_mutate_query_stats():
+    received = []
+
+    def stats_callback(stats):
+        received.append(stats)
+        stats["state"] = "MUTATED"
+        stats["rootStage"]["subStages"][0]["stageId"] = "999"
+
+    query = TrinoQuery(
+        request=TrinoRequest(
+            host="coordinator",
+            port=8080,
+            client_session=ClientSession(user="test"),
+            http_scheme="http",
+        ),
+        query="SELECT 1",
+        stats_callback=stats_callback,
+    )
+    query._stats = {
+        "queryId": "q1",
+        "state": "RUNNING",
+        "rootStage": {"stageId": "0", "subStages": [{"stageId": "1"}]},
+    }
+
+    query._report_stats()
+
+    assert received == [{
+        "queryId": "q1",
+        "state": "MUTATED",
+        "rootStage": {"stageId": "0", "subStages": [{"stageId": "999"}]},
+    }]
+    assert query.stats == {
+        "queryId": "q1",
+        "state": "RUNNING",
+        "rootStage": {"stageId": "0", "subStages": [{"stageId": "1"}]},
+    }
+
+
 def test_delay_exponential_without_jitter():
     max_delay = 1200.0
     get_delay = _DelayExponential(base=5, jitter=False, max_delay=max_delay)

trino/client.py

@@ -55,6 +55,7 @@
 from enum import Enum
 from time import sleep
 from typing import Any
+from typing import Callable
 from typing import cast
 from typing import Dict
 from typing import List
@@ -885,7 +886,8 @@ def __init__(
             request: TrinoRequest,
             query: str,
             legacy_primitive_types: bool = False,
-            fetch_mode: Literal["mapped", "segments"] = "mapped"
+            fetch_mode: Literal["mapped", "segments"] = "mapped",
+            stats_callback: Optional[Callable[[Dict[str, Any]], None]] = None
     ) -> None:
         self._query_id: Optional[str] = None
         self._stats: Dict[Any, Any] = {}
@@ -903,6 +905,7 @@ def __init__(
         self._legacy_primitive_types = legacy_primitive_types
         self._row_mapper: Optional[RowMapper] = None
         self._fetch_mode = fetch_mode
+        self._stats_callback = stats_callback
 
     @property
     def query_id(self) -> Optional[str]:
@@ -1018,6 +1021,12 @@ def _update_state(self, status):
                                                          legacy_primitive_types=self._legacy_primitive_types)
         if status.columns:
             self._columns = status.columns
+        self._report_stats()
+
+    def _report_stats(self) -> None:
+        if self._stats_callback is not None:
+            # Pass a deep copy so the callback cannot mutate internal query state.
+            self._stats_callback(copy.deepcopy(self._stats))
 
     def fetch(self) -> Union[List[Union[List[Any], Any]], Iterator[List[Any]]]:
         """Continue fetching data for the current query_id"""

trino/dbapi.py

@@ -26,6 +26,7 @@
 from threading import Lock
 from time import time
 from typing import Any
+from typing import Callable
 from typing import Dict
 from typing import List
 from typing import NamedTuple
@@ -297,7 +298,11 @@ def _create_request(self):
             self.request_timeout,
         )
 
-    def cursor(self, cursor_style: str = "row", legacy_primitive_types: bool = None):
+    def cursor(
+            self,
+            cursor_style: str = "row",
+            legacy_primitive_types: bool = None,
+            stats_callback: Optional[Callable[[Dict[str, Any]], None]] = None):
         """Return a new :py:class:`Cursor` object using the connection."""
         if self.isolation_level != IsolationLevel.AUTOCOMMIT:
             if self.transaction is None:
@@ -320,7 +325,8 @@ def cursor(self, cursor_style: str = "row", legacy_primitive_types: bool = None)
                 legacy_primitive_types
                 if legacy_primitive_types is not None
                 else self.legacy_primitive_types
-            )
+            ),
+            stats_callback=stats_callback
         )
 
     def _use_legacy_prepared_statements(self):
@@ -395,7 +401,8 @@ def __init__(
             self,
             connection,
             request,
-            legacy_primitive_types: bool = False):
+            legacy_primitive_types: bool = False,
+            stats_callback: Optional[Callable[[Dict[str, Any]], None]] = None):
         if not isinstance(connection, Connection):
             raise ValueError(
                 "connection must be a Connection object: {}".format(type(connection))
@@ -407,6 +414,7 @@ def __init__(
         self._iterator = None
         self._query = None
         self._legacy_primitive_types = legacy_primitive_types
+        self._stats_callback = stats_callback
 
     def __iter__(self):
         return self._iterator
@@ -510,7 +518,11 @@ def _execute_prepared_statement(
         params
     ):
         sql = 'EXECUTE ' + statement_name + ' USING ' + ','.join(map(self._format_prepared_param, params))
-        return trino.client.TrinoQuery(self._request, query=sql, legacy_primitive_types=self._legacy_primitive_types)
+        return trino.client.TrinoQuery(
+            self._request,
+            query=sql,
+            legacy_primitive_types=self._legacy_primitive_types,
+            stats_callback=self._stats_callback)
 
     def _execute_immediate_statement(self, statement: str, params):
         """
@@ -522,7 +534,10 @@ def _execute_immediate_statement(self, statement: str, params):
         sql = "EXECUTE IMMEDIATE '" + statement.replace("'", "''") + \
               "' USING " + ",".join(map(self._format_prepared_param, params))
         return trino.client.TrinoQuery(
-            self.connection._create_request(), query=sql, legacy_primitive_types=self._legacy_primitive_types)
+            self.connection._create_request(),
+            query=sql,
+            legacy_primitive_types=self._legacy_primitive_types,
+            stats_callback=self._stats_callback)
 
     def _format_prepared_param(self, param):
         """
@@ -645,7 +660,8 @@ def execute(self, operation, params=None):
 
         else:
             self._query = trino.client.TrinoQuery(self._request, query=operation,
-                                                  legacy_primitive_types=self._legacy_primitive_types)
+                                                  legacy_primitive_types=self._legacy_primitive_types,
+                                                  stats_callback=self._stats_callback)
             self._iterator = iter(self._query.execute())
         return self
 
@@ -764,8 +780,10 @@ def __init__(
             self,
             connection,
             request,
-            legacy_primitive_types: bool = False):
-        super().__init__(connection, request, legacy_primitive_types=legacy_primitive_types)
+            legacy_primitive_types: bool = False,
+            stats_callback: Optional[Callable[[Dict[str, Any]], None]] = None):
+        super().__init__(
+            connection, request, legacy_primitive_types=legacy_primitive_types, stats_callback=stats_callback)
         if self.connection._client_session.encoding is None:
             raise ValueError("SegmentCursor can only be used if encoding is set on the connection")
 
@@ -776,7 +794,8 @@ def execute(self, operation, params=None):
 
         self._query = trino.client.TrinoQuery(self._request, query=operation,
                                               legacy_primitive_types=self._legacy_primitive_types,
-                                              fetch_mode="segments")
+                                              fetch_mode="segments",
+                                              stats_callback=self._stats_callback)
         self._iterator = iter(self._query.execute())
         return self