From a1d5e12c01f3ef16e8e1a0993730869202fcd036 Mon Sep 17 00:00:00 2001 From: Ashhar Hasan Date: Sat, 11 Jul 2026 16:46:44 +0530 Subject: [PATCH] Expose query stats through a cursor callback Add an optional stats_callback to Connection.cursor() that is invoked with a copy of the query stats every time the client processes a response from the coordinator: once right after the query is submitted (so the query id and initial state are available while the query is still running) and again on every subsequent poll. This mirrors the progress monitor exposed by the Trino JDBC driver and lets callers log or monitor a query instead of only being able to inspect it after all rows have been fetched. --- README.md | 24 +++++++++++++ tests/integration/test_dbapi_integration.py | 18 ++++++++++ trino/client.py | 12 ++++++- trino/dbapi.py | 37 ++++++++++++++++----- 4 files changed, 81 insertions(+), 10 deletions(-) diff --git a/README.md b/README.md index 68a84706..7e2df2c4 100644 --- a/README.md +++ b/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** diff --git a/tests/integration/test_dbapi_integration.py b/tests/integration/test_dbapi_integration.py index 35ef0b30..e27f384a 100644 --- a/tests/integration/test_dbapi_integration.py +++ b/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") diff --git a/trino/client.py b/trino/client.py index abf1c2f2..5fcdf379 100644 --- a/trino/client.py +++ b/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]: @@ -980,6 +983,7 @@ def execute(self, additional_http_headers=None) -> TrinoResult: self._stats.update({"queryId": self.query_id}) self._update_state(status) self._warnings = getattr(status, "warnings", []) + self._report_stats() if status.next_uri is None: self._finished = True @@ -1019,6 +1023,11 @@ def _update_state(self, status): if status.columns: self._columns = status.columns + def _report_stats(self) -> None: + if self._stats_callback is not None: + # Pass a copy so the callback cannot mutate internal query state. + self._stats_callback(dict(self._stats)) + def fetch(self) -> Union[List[Union[List[Any], Any]], Iterator[List[Any]]]: """Continue fetching data for the current query_id""" try: @@ -1027,6 +1036,7 @@ def fetch(self) -> Union[List[Union[List[Any], Any]], Iterator[List[Any]]]: raise trino.exceptions.TrinoConnectionError("failed to fetch: {}".format(e)) status = self._request.process(response) self._update_state(status) + self._report_stats() if status.next_uri is None: self._finished = True diff --git a/trino/dbapi.py b/trino/dbapi.py index 24506004..69d3047b 100644 --- a/trino/dbapi.py +++ b/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