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