From e0947450faa81878dcbbfdeb8a5d2ac143f5a2eb Mon Sep 17 00:00:00 2001 From: Ashhar Hasan Date: Tue, 25 Aug 2026 18:07:05 +0530 Subject: [PATCH 1/3] Count heartbeat requests client-side in the integration test Before the change the test grepped the coordinator's HTTP request log through docker exec, which needed request logging enabled and a log-flush retry loop. Record HEAD requests on the client's own `requests.Session` instead. This removes the only caller of get_trino_container(), so delete it too. --- etc/config.properties | 4 -- tests/development_server.py | 25 -------- tests/integration/test_dbapi_integration.py | 63 +++++++++------------ 3 files changed, 28 insertions(+), 64 deletions(-) diff --git a/etc/config.properties b/etc/config.properties index bd517b61..3710913e 100644 --- a/etc/config.properties +++ b/etc/config.properties @@ -15,7 +15,3 @@ protocol.spooling.retrieval-mode=coordinator_proxy # Enable dynamic catalog management catalog.management=dynamic - -# Enable HTTP request log as it's grepped in integration tests for the heartbeat mechanism -http-server.log.enabled=true -http-server.log.immediate-flush=true diff --git a/tests/development_server.py b/tests/development_server.py index e404b015..9a0d9ffc 100644 --- a/tests/development_server.py +++ b/tests/development_server.py @@ -3,7 +3,6 @@ from contextlib import contextmanager from pathlib import Path -import docker.errors from testcontainers.core.container import DockerContainer from testcontainers.core.network import Network from testcontainers.core.waiting_utils import wait_for_logs @@ -20,30 +19,6 @@ TRINO_HOST = "localhost" -def get_trino_container(port: int): - """Find and return a running trino container. - Returns None if no matching container is found. - """ - client = docker.from_env() - try: - container = client.containers.get(TRINO_CONTAINER_NAME) - except docker.errors.NotFound: - return None - - if not any(tag.startswith(f"{TRINO_IMAGE_REPO}:") for tag in (container.image.tags or [])): - return None - - host_ports = [ - binding["HostPort"] - for bindings in container.ports.values() if bindings - for binding in bindings - ] - if str(port) not in host_ports: - return None - - return container - - def create_bucket(s3_client): bucket_name = "spooling" try: diff --git a/tests/integration/test_dbapi_integration.py b/tests/integration/test_dbapi_integration.py index 49556077..0206003b 100644 --- a/tests/integration/test_dbapi_integration.py +++ b/tests/integration/test_dbapi_integration.py @@ -27,7 +27,6 @@ from tzlocal import get_localzone_name # type: ignore import trino -from tests.development_server import get_trino_container from tests.integration.conftest import trino_version from trino import constants from trino.client import InlineSegment @@ -2055,29 +2054,44 @@ def test_spooled_segments_lazy_description(trino_connection): assert len(cur.fetchall()) == 60175 +class _HeadCountingSession(requests.Session): + """Records every HEAD request the client sends.""" + + def __init__(self): + super().__init__() + self.sent_head_urls = [] + + def head(self, url, **kwargs): + self.sent_head_urls.append(url) + return super().head(url, **kwargs) + + +def _heartbeat_head_urls(session, query_id): + heartbeat_path = f"/v1/statement/executing/{query_id}/" + return [url for url in session.sent_head_urls if heartbeat_path in url] + + +def _assert_heartbeats_sent(session, query_id): + heartbeat_urls = _heartbeat_head_urls(session, query_id) + assert len(heartbeat_urls) > 0, ( + f"Expected at least one heartbeat HEAD request for query {query_id}, " + f"but the client sent HEAD only to {session.sent_head_urls}" + ) + + @pytest.mark.skipif( trino_version() <= 466, reason="spooling protocol was introduced in version 466" ) def test_heartbeat_head_requests_during_spooled_download(run_trino): - """Verify that heartbeat HEAD requests are sent to the coordinator while - downloading spooled segments from external storage.""" host, port = run_trino - container = get_trino_container(port) - assert container, "Cannot find a running Trino container" - + session = _HeadCountingSession() conn = trino.dbapi.Connection( host=host, port=port, user="test", source="test", max_attempts=1, encoding="json", heartbeat_interval=0.1, + http_session=session, ) - log_path = "/data/trino/var/log/http-request.log" - - # Capture the current size of the HTTP request log - exit_code, output = container.exec_run(["wc", "-l", log_path]) - assert exit_code == 0, f"Cannot read Trino HTTP request log, is the log path `{log_path}` correct?" - logfile_lines = int(output.decode().split()[0]) - cur = conn.cursor() cur.execute("""SELECT l.* FROM tpch.tiny.lineitem l, TABLE(sequence( @@ -2088,28 +2102,7 @@ def test_heartbeat_head_requests_during_spooled_download(run_trino): cur.fetchall() cur.close() - head_request_found = False - # Sometimes trino needs time to flush the logs so we make few attempts - # to check the log with sleep inbetween. - for attempt in range(10): - if attempt: - t.sleep(1.0) - - _, output = container.exec_run(["tail", "-n", f"+{logfile_lines}", log_path]) - loglines = output.decode().splitlines() - - pattern = f"/v1/statement/executing/{query_id}/" - for line in loglines: - if "HEAD" in line and pattern in line: - head_request_found = True - break - - if head_request_found: - break - - assert head_request_found, ( - f"Expected heartbeat HEAD requests in http-request.log not found. Log tail:\n{''.join(loglines)}" - ) + _assert_heartbeats_sent(session, query_id) def get_cursor(legacy_prepared_statements, run_trino): From 8bb0a268c64b4d693426e5fc8527dd5aab6f2244 Mon Sep 17 00:00:00 2001 From: Ashhar Hasan Date: Tue, 25 Aug 2026 21:14:58 +0530 Subject: [PATCH 2/3] Send heartbeats while the caller consumes fetched rows The client made no requests between fetch() calls. A caller that held a batch of rows for longer than `query.client.timeout` got the query abandoned by the coordinator. The existing heartbeat ran only during spooled segment downloads, on a thread that shared the not-thread-safe `requests.Session` with regular traffic. Mirror the Java client instead. Check the clock for every served row. Once a full `heartbeat_interval` passes without a request, send one HEAD to the current `nextUri` on the caller's own thread. This renews the coordinator's abandonment timer without a background thread. An actively fetching query sends no extra traffic. Servers without HEAD support keep the old abandonment behavior. --- tests/integration/test_dbapi_integration.py | 29 ++ tests/unit/test_client.py | 284 ++++++++++++++++++++ tests/unit/test_client_spooling.py | 151 +---------- trino/client.py | 186 +++++++------ 4 files changed, 424 insertions(+), 226 deletions(-) diff --git a/tests/integration/test_dbapi_integration.py b/tests/integration/test_dbapi_integration.py index 0206003b..6549e639 100644 --- a/tests/integration/test_dbapi_integration.py +++ b/tests/integration/test_dbapi_integration.py @@ -2105,6 +2105,35 @@ def test_heartbeat_head_requests_during_spooled_download(run_trino): _assert_heartbeats_sent(session, query_id) +@pytest.mark.parametrize("heartbeat_interval, expect_heartbeats", [(0.2, True), (None, False)]) +def test_heartbeat_head_requests_while_caller_is_slow(run_trino, heartbeat_interval, expect_heartbeats): + host, port = run_trino + session = _HeadCountingSession() + conn = trino.dbapi.Connection( + host=host, port=port, user="test", source="test", + max_attempts=1, encoding=None, heartbeat_interval=heartbeat_interval, + http_session=session, + ) + + cur = conn.cursor() + cur.execute("SELECT * FROM tpch.tiny.lineitem") + query_id = cur.query_id + # Each fetched row checks whether a full interval passed since the last + # processed response and sends a HEAD if so. Model a slow consumer that + # asks for one row at a time with a pause between requests longer than + # the interval. + for _ in range(3): + assert cur.fetchone() is not None + t.sleep(0.3) + cur.fetchall() + cur.close() + + if expect_heartbeats: + _assert_heartbeats_sent(session, query_id) + else: + assert _heartbeat_head_urls(session, query_id) == [] + + def get_cursor(legacy_prepared_statements, run_trino): host, port = run_trino diff --git a/tests/unit/test_client.py b/tests/unit/test_client.py index c5d33dbd..2ae52ac3 100644 --- a/tests/unit/test_client.py +++ b/tests/unit/test_client.py @@ -10,6 +10,7 @@ # See the License for the specific language governing permissions and # limitations under the License. import base64 +import math import threading import time import urllib @@ -48,6 +49,7 @@ from trino.auth import _OAuth2KeyRingTokenCache from trino.auth import _OAuth2TokenBearer from trino.client import _DelayExponential +from trino.client import _Heartbeat from trino.client import _retry_with from trino.client import _RetryWithExponentialBackoff from trino.client import ClientSession @@ -1282,6 +1284,282 @@ def stats_callback(stats): } +_HEARTBEAT_NEXT_URI = "http://coordinator/v1/statement/q/1" + + +class _FakeHeartbeatRequest(TrinoRequest): + """Base for fake TrinoRequest subclasses used in heartbeat tests.""" + + def __init__(self, client_session): + super().__init__( + host="coordinator", + port=8080, + client_session=client_session, + http_scheme="http", + ) + self.head_calls = [] + # cancel() issues a DELETE. Serve 204 so it succeeds without real HTTP. + self._delete = self._serve_no_content + + def _serve_no_content(self, url, **kwargs): + return self._canned_response(None, status_code=204) + + def _canned_response(self, payload, status_code=200): + response = requests.Response() + response.status_code = status_code + if payload is not None: + body = json.dumps(payload) + # orjson.dumps returns bytes, stdlib json.dumps returns str + response._content = body if isinstance(body, bytes) else body.encode("utf-8") + return response + + +class _ScriptedHeadRequest(_FakeHeartbeatRequest): + """Serves scripted HEAD responses (status codes or exceptions) and records the calls.""" + + def __init__(self, responses): + super().__init__(ClientSession(user="test")) + self._responses = iter(responses) + self.head_timeouts = [] + + def _head(self, url, timeout): + self.head_calls.append(url) + self.head_timeouts.append(timeout) + item = next(self._responses) + if isinstance(item, Exception): + raise item + return self._canned_response(None, status_code=item) + + +def _scripted_query(responses, next_uri=_HEARTBEAT_NEXT_URI): + """A query whose request serves scripted HEAD responses.""" + req = _ScriptedHeadRequest(responses) + query = TrinoQuery(req, query="SELECT 1") + query._next_uri = next_uri + return query, req + + +def test_heartbeat_sends_head_to_next_uri_and_reschedules(): + query, req = _scripted_query([200]) + query._heartbeat._deadline = 0.0 + + query._maybe_heartbeat() + + assert req.head_calls == [_HEARTBEAT_NEXT_URI] + assert req.head_timeouts == [(_Heartbeat.HEAD_TIMEOUT_CAP, _Heartbeat.HEAD_TIMEOUT_CAP)] + # A successful beat pushes the deadline forward instead of disabling it + assert time.monotonic() < query._heartbeat._deadline < math.inf + + +def test_heartbeat_before_the_deadline_sends_nothing(): + # Default construction sets a deadline far in the future + query, req = _scripted_query([]) + + query._maybe_heartbeat() + + assert req.head_calls == [] + + +# 404 means this query is gone on the server, the next fetch() will show why anyway. +# 405 means the server does not support HEAD requests at all. +@pytest.mark.parametrize("status_code", (404, 405)) +def test_gone_or_unsupported_response_disables_heartbeats(status_code): + query, req = _scripted_query([status_code]) + query._heartbeat._deadline = 0.0 + + query._maybe_heartbeat() + + assert req.head_calls == [_HEARTBEAT_NEXT_URI] + assert query._heartbeat._deadline == math.inf + + # A disabled deadline is at math.inf so a later call is never due again on its own + query._maybe_heartbeat() + + assert req.head_calls == [_HEARTBEAT_NEXT_URI] + + +def test_heartbeat_disables_after_max_consecutive_failures(): + max_failures = _Heartbeat.MAX_FAILURES + query, req = _scripted_query([Exception("boom")] * max_failures) + + for _ in range(max_failures): + query._heartbeat._deadline = 0.0 + query._maybe_heartbeat() + + assert len(req.head_calls) == max_failures + assert query._heartbeat._deadline == math.inf + + query._maybe_heartbeat() + + assert len(req.head_calls) == max_failures + + +def test_heartbeat_success_resets_the_failure_counter(): + # Only three CONSECUTIVE failures disable heartbeats. The 200 in the middle resets the count. + responses = [Exception("boom")] * 2 + [200] + [Exception("boom")] * 3 + query, req = _scripted_query(responses) + + for _ in range(len(responses)): + query._heartbeat._deadline = 0.0 + query._maybe_heartbeat() + + assert len(req.head_calls) == len(responses) + assert query._heartbeat._deadline == math.inf + + +def test_heartbeat_error_response_resets_the_failure_counter(): + # A response of any status proves the server is reachable. Only network + # failures count toward disabling heartbeats. + responses = [Exception("boom")] * 2 + [503] + [Exception("boom")] * 2 + query, req = _scripted_query(responses) + + for _ in range(len(responses)): + query._heartbeat._deadline = 0.0 + query._maybe_heartbeat() + + assert len(req.head_calls) == len(responses) + assert query._heartbeat._deadline < math.inf + + +def test_heartbeat_skips_when_next_uri_is_none(): + query, req = _scripted_query([], next_uri=None) + query._heartbeat._deadline = 0.0 + + query._maybe_heartbeat() + + assert req.head_calls == [] + + +def test_heartbeat_skips_when_the_query_is_finished(): + query, req = _scripted_query([]) + query._finished = True + query._heartbeat._deadline = 0.0 + + query._maybe_heartbeat() + + assert req.head_calls == [] + + +class _HeartbeatRecordingRequest(_FakeHeartbeatRequest): + """Serves canned response pages and records heartbeat HEAD calls.""" + + def __init__(self, pages, heartbeat_interval): + super().__init__(ClientSession(user="test", heartbeat_interval=heartbeat_interval)) + self._pages = iter(pages) + # Replace the transports behind get/post so the public methods still run + # their normal response processing. + self._get = self._serve_page + self._post = self._serve_page + + def _serve_page(self, url, **kwargs): + return self._canned_response(next(self._pages)) + + def _head(self, url, timeout): + self.head_calls.append(url) + return self._canned_response(None) + + +def _heartbeat_page(next_uri=None, data=None): + page = { + "id": "q1", + "infoUri": "http://coordinator/query.html?q1", + "stats": {"state": "RUNNING"}, + "columns": [{"name": "x", "type": "integer", "typeSignature": {"rawType": "integer", "arguments": []}}], + } + if next_uri is not None: + page["nextUri"] = next_uri + if data is not None: + page["data"] = data + return page + + +_HEARTBEAT_URI_1 = "http://coordinator/v1/statement/executing/q1/1" +_HEARTBEAT_URI_2 = "http://coordinator/v1/statement/executing/q1/2" + + +@httprettified +def test_heartbeat_head_bypasses_the_retry_wrapper(): + # The retry wrapper resends body-less 200 responses and a HEAD response never has a + # body. Heartbeat HEADs go straight to the session, so max_attempts must not apply. + url = "http://coordinator:8080/v1/statement/executing/q1/1" + httpretty.register_uri(httpretty.HEAD, url, status=200) + request = TrinoRequest( + host="coordinator", + port=8080, + client_session=ClientSession(user="test"), + http_scheme="http", + max_attempts=3, + ) + + request._head(url, timeout=5) + + assert len([r for r in httpretty.latest_requests() if r.method == "HEAD"]) == 1 + + +@httprettified +def test_heartbeat_head_follows_redirects(): + # A gateway may answer with a redirect to the coordinator. The heartbeat only + # renews the server-side deadline if the HEAD follows it. + gateway_url = "http://gateway:8080/v1/statement/executing/q1/1" + coordinator_url = "http://coordinator:8080/v1/statement/executing/q1/1" + httpretty.register_uri(httpretty.HEAD, gateway_url, status=302, adding_headers={"Location": coordinator_url}) + httpretty.register_uri(httpretty.HEAD, coordinator_url, status=200) + request = TrinoRequest( + host="gateway", + port=8080, + client_session=ClientSession(user="test"), + http_scheme="http", + ) + + response = request._head(gateway_url, timeout=5) + + assert response.status_code == 200 + + head_requests = [request for request in httpretty.latest_requests() if request.method == "HEAD"] + requested_hosts = [request.headers.get("Host") for request in head_requests] + assert requested_hosts == ["gateway:8080", "coordinator:8080"] + + +def test_heartbeat_disabled_when_interval_is_zero(): + request = _FakeHeartbeatRequest(ClientSession(user="test", heartbeat_interval=0.0)) + query = TrinoQuery(request, query="SELECT 1") + + assert query._heartbeat._deadline == math.inf + + +def test_heartbeat_sends_nothing_after_cancel(): + query, req = _scripted_query([]) + + query.cancel() + query._heartbeat._deadline = 0.0 + query._maybe_heartbeat() + + assert req.head_calls == [] + + +def test_fetch_defers_the_heartbeat(): + # A processed fetch() response proves the client is alive, so it pushes the deadline + # forward. A beat that was due before the fetch is no longer due after it. + request = _HeartbeatRecordingRequest( + pages=[ + _heartbeat_page(next_uri=_HEARTBEAT_URI_2, data=[[1]]), + _heartbeat_page(data=[[2]]), + ], + heartbeat_interval=30, + ) + request._next_uri = _HEARTBEAT_URI_1 + query = TrinoQuery(request, query="SELECT 1") + query._heartbeat._deadline = 0.0 + + rows = query.fetch() + + assert query._heartbeat._deadline > time.monotonic() + + list(TrinoResult(query, rows)) + + assert request.head_calls == [] + + def test_delay_exponential_without_jitter(): max_delay = 1200.0 get_delay = _DelayExponential(base=5, jitter=False, max_delay=max_delay) @@ -1594,6 +1872,9 @@ class _FinishedQuery: def fetch(self): return [] + def _maybe_heartbeat(self): + pass + @pytest.mark.parametrize("consecutive_failures", (1, 2, 3)) def test_trino_result_resumes_after_transient_error_in_rows_iterator(consecutive_failures): @@ -1667,6 +1948,9 @@ def fetch(self): self.finished = True return [[2]] + def _maybe_heartbeat(self): + pass + result = TrinoResult(FlakyQuery(), [[1]]) it = iter(result) # The next batch is prefetched before the first row is served so the diff --git a/tests/unit/test_client_spooling.py b/tests/unit/test_client_spooling.py index 21912d37..22678f0e 100644 --- a/tests/unit/test_client_spooling.py +++ b/tests/unit/test_client_spooling.py @@ -10,157 +10,18 @@ # See the License for the specific language governing permissions and # limitations under the License. import json -import time from unittest import mock import pytest -from trino.client import _RequestHeartbeat from trino.client import ClientSession from trino.client import DecodableSegment -from trino.client import InlineSegment from trino.client import SegmentIterator from trino.client import SpooledSegment from trino.client import TrinoQuery from trino.client import TrinoRequest -def _mock_trino_request(): - req = TrinoRequest( - host="coordinator", - port=8080, - client_session=ClientSession(user="test"), - http_scheme="http", - ) - req._next_uri = "http://coordinator/v1/statement/q/1" - return req - - -def _head_response(status_code): - return mock.Mock(status_code=status_code, ok=(200 <= status_code < 300)) - - -@pytest.fixture -def ensure_max_failures_3(): - # Some tests assume _RequestHeartbeart.MAX_FAILURES is set to 3 - with mock.patch.object(_RequestHeartbeat, "MAX_FAILURES", 3): - yield - - -def test_heartbeat_sends_head_to_next_uri(): - req = _mock_trino_request() - with mock.patch.object(req, "head", return_value=_head_response(200)) as mock_head: - with _RequestHeartbeat(req, interval=0.01): - time.sleep(0.1) - assert mock_head.call_count >= 2 - mock_head.assert_called_with(req.next_uri) - - -@pytest.mark.parametrize("status_code", (404, 405)) -def test_heartbeat_stops_on_404_405(status_code): - req = _mock_trino_request() - with mock.patch.object(req, "head", return_value=_head_response(status_code)) as mock_head: - with _RequestHeartbeat(req, interval=0.01): - time.sleep(0.1) - # 404/405 means the server does not support heartbeat requests; they should stop after the first one - assert mock_head.call_count == 1 - - -def test_heartbeat_stops_after_max_failures_non_2xx(ensure_max_failures_3): - req = _mock_trino_request() - with mock.patch.object(req, "head", return_value=_head_response(500)) as mock_head: - with _RequestHeartbeat(req, interval=0.01): - time.sleep(0.1) - assert mock_head.call_count == _RequestHeartbeat.MAX_FAILURES - - -def test_heartbeat_stops_after_max_failures_on_exception(ensure_max_failures_3): - req = _mock_trino_request() - with mock.patch.object(req, "head", side_effect=Exception("network error")) as mock_head: - with _RequestHeartbeat(req, interval=0.01): - time.sleep(0.1) - assert mock_head.call_count == _RequestHeartbeat.MAX_FAILURES - - -def test_heartbeat_resets_failure_count_on_success(ensure_max_failures_3): - req = _mock_trino_request() - # Failure counter resets on 200 so the heartbeat keeps running past initial failures - responses = [_head_response(500), _head_response(500)] + [_head_response(200)] * 20 - with mock.patch.object(req, "head", side_effect=responses) as mock_head: - with _RequestHeartbeat(req, interval=0.01): - time.sleep(0.1) - assert mock_head.call_count > _RequestHeartbeat.MAX_FAILURES - - -def test_heartbeat_skips_when_next_uri_is_none(): - req = _mock_trino_request() - req._next_uri = None - with mock.patch.object(req, "head") as mock_head: - with _RequestHeartbeat(req, interval=0.01): - time.sleep(0.1) - mock_head.assert_not_called() - - -def test_heartbeat_stop_is_immediate(): - req = _mock_trino_request() - with mock.patch.object(req, "head", return_value=_head_response(200)): - hb = _RequestHeartbeat(req, interval=30) - start = time.monotonic() - with hb: - pass - elapsed = time.monotonic() - start - assert elapsed < 1.0 - - -def _spooled_iterator(request, heartbeat_interval, rows=None): - """SegmentIterator with one SpooledSegment and a pre-set mock decoder.""" - segment = DecodableSegment("json", None, mock.Mock(spec=SpooledSegment)) - mapper = mock.Mock() - it = SegmentIterator([segment], mapper, request=request, heartbeat_interval=heartbeat_interval) - it._decoder = mock.Mock() - it._decoder.decode.return_value = rows if rows is not None else [[1, 2]] - return it - - -@pytest.mark.parametrize( - "trino_request, interval", - [(None, 1.0), (_mock_trino_request(), None), (_mock_trino_request(), 0.0)] -) -def test_iterator_value_error_when_only_request_or_heartbeat_interval_specified(trino_request, interval): - with pytest.raises(ValueError): - _ = _spooled_iterator(trino_request, interval) - - -def test_heartbeat_starts_during_spooled_segment_download(): - req = _mock_trino_request() - iterator = _spooled_iterator(req, heartbeat_interval=30.0) - with mock.patch("trino.client._RequestHeartbeat") as MockHB: - next(iterator) - MockHB.assert_called_once_with(req, 30.0) - # Make sure MockHB instance is used as a context manager - MockHB.return_value.__enter__.assert_called_once() - MockHB.return_value.__exit__.assert_called_once() - - -def test_no_heartbeat_for_inline_segment(): - segment = DecodableSegment("json", None, mock.Mock(spec=InlineSegment)) - mapper = mock.Mock() - iterator = SegmentIterator([segment], mapper, request=_mock_trino_request(), heartbeat_interval=30.0) - iterator._decoder = mock.Mock() - iterator._decoder.decode.return_value = [[1, 2]] - with mock.patch("trino.client._RequestHeartbeat") as MockHB: - next(iterator) - MockHB.assert_not_called() - - -@pytest.mark.parametrize("interval", (None, 0.0)) -def test_no_heartbeat_when_interval_none_or_zero(interval): - iterator = _spooled_iterator(request=None, heartbeat_interval=interval) - with mock.patch("trino.client._RequestHeartbeat") as MockHB: - next(iterator) - MockHB.assert_not_called() - - def _spooled_fetch_response(): """Minimal spooled protocol GET response JSON.""" resp = mock.Mock() @@ -185,21 +46,17 @@ def _spooled_fetch_response(): return resp -@pytest.mark.parametrize("heartbeat_interval", (30.0, None)) -def test_fetch_passes_request_and_interval_to_segment_iterator(heartbeat_interval): - session = ClientSession(user="test", encoding="json", heartbeat_interval=heartbeat_interval) +def test_fetch_returns_segment_iterator(): + session = ClientSession(user="test", encoding="json") req = TrinoRequest(host="coordinator", port=8080, client_session=session, http_scheme="http") req._next_uri = "http://coordinator/v1/statement/q1/1" query = TrinoQuery(req, query="SELECT 1") query._row_mapper = mock.Mock() with mock.patch.object(req, "get", return_value=_spooled_fetch_response()): - with mock.patch("trino.client.SegmentIterator") as MockSI: - MockSI.return_value = iter([]) - query.fetch() + result = query.fetch() - assert MockSI.call_args.kwargs["request"] is req - assert MockSI.call_args.kwargs["heartbeat_interval"] == heartbeat_interval + assert isinstance(result, SegmentIterator) class _FakeSpooledSegment(SpooledSegment): diff --git a/trino/client.py b/trino/client.py index 9d4956bf..541ab78a 100644 --- a/trino/client.py +++ b/trino/client.py @@ -40,6 +40,7 @@ import copy import functools import itertools +import math import os import random import re @@ -53,6 +54,7 @@ from datetime import datetime from email.utils import parsedate_to_datetime from enum import Enum +from time import monotonic from time import sleep from typing import Any from typing import Callable @@ -639,7 +641,6 @@ def max_attempts(self, value: int) -> None: self._get = self._http_session.get self._post = self._http_session.post self._delete = self._http_session.delete - self._head = self._http_session.head return with_retry = _retry_with( @@ -658,7 +659,6 @@ def max_attempts(self, value: int) -> None: self._get = with_retry(self._http_session.get) self._post = with_retry(self._http_session.post) self._delete = with_retry(self._http_session.delete) - self._head = with_retry(self._http_session.head) def get_url(self, path: str) -> str: return "{protocol}://{host}:{port}{path}".format( @@ -706,12 +706,16 @@ def get(self, url: str) -> Response: def delete(self, url: str) -> Response: return self._delete(url, timeout=self._request_timeout, proxies=PROXIES) - def head(self, url: str) -> Response: - return self._head( + def _head(self, url: str, timeout: Union[float, Tuple[float, float]]) -> Response: + # requests disables redirect-following for HEAD by default. Follow them so a + # heartbeat behind a redirecting gateway still reaches the coordinator (like + # the Java client where OkHttp follows redirects). + return self._http_session.head( url, headers=self.http_headers, - timeout=self._request_timeout, + timeout=timeout, proxies=PROXIES, + allow_redirects=True, ) @staticmethod @@ -886,9 +890,87 @@ def __next__(self): self._current_batch = None continue self._rownumber += 1 + self._query._maybe_heartbeat() return row +class _Heartbeat: + """ + Heartbeat schedule for one query. + Trino cancels a query which doesn't poll within `query.client.timeout`. + A HEAD to next_uri renews the server-side deadline without consuming a page. + """ + + MAX_FAILURES = 3 + # Upper bound on a single heartbeat HEAD since it blocks the caller's row loop. + HEAD_TIMEOUT_CAP = 5.0 + + def __init__(self, request: TrinoRequest, interval: Optional[float]) -> None: + self._request = request + # An interval of 0 disables heartbeats, same as None. + self._interval = interval + # Deadline for the next heartbeat. math.inf means heartbeats are disabled. + self._deadline = monotonic() + self._interval if self._interval else math.inf + self._failures = 0 + self._head_timeout = self._head_timeout_for(request._request_timeout) + + def due(self) -> bool: + # This check runs once per served row. The disabled case skips the clock read. + if self._deadline == math.inf: + return False + return monotonic() >= self._deadline + + def defer(self) -> None: + # A processed response proves the client is alive. Push the deadline forward. + if self._deadline != math.inf: + self._deadline = monotonic() + self._interval + + def beat(self, next_uri: str) -> None: + # Send one heartbeat HEAD to next_uri. The caller checks that a beat is due. + self.defer() + try: + response = self._request._head(next_uri, self._head_timeout) + except Exception: + self._record_failure() + return + + # A response of any status proves the server is reachable. Only network + # failures count toward disabling heartbeats, matching the Java client. + self._failures = 0 + if response.status_code == 405: + logger.warning("The server does not support heartbeat calls") + self._disable() + elif response.status_code == 404: + # This query is gone on the server. + self._disable() + + def _disable(self) -> None: + self._deadline = math.inf + + def _record_failure(self) -> None: + self._failures += 1 + if self._failures >= self.MAX_FAILURES: + logger.warning(f"Disabling heartbeats after {self.MAX_FAILURES} consecutive errors") + self._disable() + + @classmethod + def _head_timeout_for(cls, request_timeout: Union[float, Tuple[float, float], None]) -> Tuple[float, float]: + # requests takes one timeout or a (connect, read) pair. Either part may be None. + if isinstance(request_timeout, tuple): + connect_timeout, read_timeout = request_timeout + else: + connect_timeout = request_timeout + read_timeout = request_timeout + + # Cap each timeout on its own. A short connect timeout must not shorten the read timeout. + if connect_timeout is None or connect_timeout > cls.HEAD_TIMEOUT_CAP: + connect_timeout = cls.HEAD_TIMEOUT_CAP + if read_timeout is None or read_timeout > cls.HEAD_TIMEOUT_CAP: + read_timeout = cls.HEAD_TIMEOUT_CAP + + return connect_timeout, read_timeout + + class TrinoQuery: """Represent the execution of a SQL statement by Trino.""" @@ -917,6 +999,7 @@ def __init__( self._row_mapper: Optional[RowMapper] = None self._fetch_mode = fetch_mode self._stats_callback = stats_callback + self._heartbeat = _Heartbeat(request, request._client_session.heartbeat_interval) @property def query_id(self) -> Optional[str]: @@ -1041,6 +1124,7 @@ def execute(self, additional_http_headers=None) -> TrinoResult: return self._result def _update_state(self, status): + self._heartbeat.defer() self._stats.update(status.stats) self._update_type = status.update_type self._update_count = status.update_count @@ -1057,6 +1141,18 @@ def _report_stats(self) -> None: # Pass a deep copy so the callback cannot mutate internal query state. self._stats_callback(copy.deepcopy(self._stats)) + def _maybe_heartbeat(self) -> None: + """ + Send a heartbeat after a full interval without a request. + TrinoResult calls this for every row it serves so the due check is done first. + """ + if not self._heartbeat.due(): + return + if self.finished or self.cancelled or self._next_uri is None: + return + + self._heartbeat.beat(self._next_uri) + def fetch(self) -> Union[List[Union[List[Any], Any]], Iterator[List[Any]]]: """Continue fetching data for the current query_id""" try: @@ -1078,13 +1174,10 @@ def fetch(self) -> Union[List[Union[List[Any], Any]], Iterator[List[Any]]]: spooled = self._to_segments(rows) if self._fetch_mode == "segments": return spooled - # Return iterator directly, do NOT materialize with list() - return SegmentIterator( - spooled, - self._row_mapper, - request=self._request, - heartbeat_interval=self._request._client_session.heartbeat_interval, - ) + # Return iterator directly, do NOT materialize with list(). + # Rows stream through TrinoResult whose per-row heartbeat check covers + # slow segment downloads and slow consumers. + return SegmentIterator(spooled, self._row_mapper) elif isinstance(status.rows, list): return self._row_mapper.map(rows) else: @@ -1114,7 +1207,7 @@ def _to_segments(self, rows: _SpooledProtocolResponseTO) -> List[DecodableSegmen def cancel(self) -> None: """Cancel the current query""" - if self._next_uri is None: + if self.cancelled or self._next_uri is None: return logger.debug("cancelling query: %s", self.query_id) @@ -1370,66 +1463,11 @@ def __repr__(self): return (f"DecodableSegment(encoding={self._encoding}, metadata={self._metadata}, segment={self._segment})") -class _RequestHeartbeat: - """ - Heartbeat loop for a trino request. Periodically sends HEAD requests to the request's next URI. - This prevents the coordinator from abandoning a query if the client is silent for a longer - period of time, for example when downloading a spooled segment from an external storage. - """ - MAX_FAILURES = 3 - - def __init__(self, request: TrinoRequest, interval: float) -> None: - self._request = request - self._interval = interval - # The event for telling the heartbeat thread to exit - self._stop_event = threading.Event() - - def __enter__(self) -> _RequestHeartbeat: - threading.Thread(target=self._run, daemon=True).start() - return self - - def __exit__(self, *_) -> None: - self._stop_event.set() - - def _run(self) -> None: - """ - Run the heartbeat loop. - - Exit when the self._stop_event is set, the query completed - or if the error count exceeds _MAX_FAILURES. - """ - failures = 0 - - while not self._stop_event.wait(timeout=self._interval): - uri = self._request.next_uri - if uri is None: - return - - try: - response = self._request.head(uri) - if response.status_code in (404, 405): - logger.warning("The server does not support heartbeat calls") - return - if not response.ok: - failures += 1 - else: - failures = 0 - except Exception: - failures += 1 - - if failures >= self.MAX_FAILURES: - logger.warning(f"Stopping the heartbeat after {self.MAX_FAILURES} consecutive errors") - return - - class SegmentIterator: def __init__( self, segments: Union[DecodableSegment, List[DecodableSegment]], mapper: RowMapper, - *, - request: Optional[TrinoRequest] = None, - heartbeat_interval: Optional[float] = None, ) -> None: self._segments = iter(segments if isinstance(segments, List) else [segments]) self._mapper = mapper @@ -1439,10 +1477,6 @@ def __init__( self._current_segment: Optional[DecodableSegment] = None # Segment whose decoding failed. Retried on the next call instead of being acknowledged and skipped. self._pending_segment: Optional[DecodableSegment] = None - if (request is not None) != bool(heartbeat_interval): - raise ValueError("request and heartbeat_interval must be both provided or both omitted") - self._request = request - self._heartbeat_interval = heartbeat_interval def __iter__(self) -> Iterator[List[Any]]: return self @@ -1477,13 +1511,7 @@ def _load_next_segment(self): self._decoder = SegmentDecoder(CompressedQueryDataDecoderFactory(self._mapper) .create(self._pending_segment.encoding)) - if isinstance(self._pending_segment.segment, SpooledSegment) and self._request and self._heartbeat_interval: - # Downloading a spooled segment may take some time. In the meantime, send heartbeat - # requests so the coordinator doesn't think we lost interest and close the query. - with _RequestHeartbeat(self._request, self._heartbeat_interval): - rows = self._decoder.decode(self._pending_segment.segment) - else: - rows = self._decoder.decode(self._pending_segment.segment) + rows = self._decoder.decode(self._pending_segment.segment) self._rows = iter(rows) self._current_segment = self._pending_segment From 99c40ee3d41454de8b02ce65da85d6560e4a8629 Mon Sep 17 00:00:00 2001 From: Ashhar Hasan Date: Tue, 25 Aug 2026 17:37:12 +0530 Subject: [PATCH 3/3] Cancel outstanding queries when the connection closes Before the change Connection.close() closed the HTTP session and left running queries to hit `query.client.timeout` on the coordinator. Close the connection's cursors first; closing a cursor cancels its running query. --- tests/unit/test_dbapi.py | 22 ++++++++++++++++++++++ trino/dbapi.py | 14 ++++++++++++-- 2 files changed, 34 insertions(+), 2 deletions(-) diff --git a/tests/unit/test_dbapi.py b/tests/unit/test_dbapi.py index 47689336..dce6bfc6 100644 --- a/tests/unit/test_dbapi.py +++ b/tests/unit/test_dbapi.py @@ -28,6 +28,9 @@ from tests.unit.oauth_test_utils import RedirectHandler from tests.unit.oauth_test_utils import SERVER_ADDRESS from tests.unit.oauth_test_utils import TOKEN_RESOURCE +from tests.unit.test_client import _heartbeat_page +from tests.unit.test_client import _HEARTBEAT_URI_1 +from tests.unit.test_client import _HeartbeatRecordingRequest from trino import constants from trino.auth import BasicAuthentication from trino.auth import OAuth2Authentication @@ -545,3 +548,22 @@ def test_format_prepared_param_binary(value, expected): assert cursor._format_prepared_param(value) == expected # Round trip through Binary(), as SQLAlchemy's _Binary.bind_processor does. assert cursor._format_prepared_param(Binary(value)) == expected + + +def _canned_cursor(pages): + """A cursor whose request serves canned pages and records heartbeat HEAD calls.""" + conn = Connection(host="coordinator", user="test") + cur = conn.cursor() + cur._request = _HeartbeatRecordingRequest(pages=pages, heartbeat_interval=0.02) + return conn, cur + + +def test_connection_close_closes_its_cursors(): + conn, cur = _canned_cursor( + pages=[_heartbeat_page(next_uri=_HEARTBEAT_URI_1, data=[[1]])], + ) + + cur.execute("SELECT 1") + conn.close() + + assert cur._query.cancelled diff --git a/trino/dbapi.py b/trino/dbapi.py index 35d576e3..e8d47b7f 100644 --- a/trino/dbapi.py +++ b/trino/dbapi.py @@ -20,6 +20,7 @@ import datetime import math import uuid +import weakref from collections import OrderedDict from decimal import Decimal from itertools import islice @@ -250,6 +251,7 @@ def __init__( self._transaction = None self.legacy_primitive_types = legacy_primitive_types self.legacy_prepared_statements = legacy_prepared_statements + self._cursors: "weakref.WeakSet[Cursor]" = weakref.WeakSet() @property def isolation_level(self): @@ -271,7 +273,13 @@ def __exit__(self, exc_type, exc_value, traceback): self.close() def close(self): - # TODO cancel outstanding queries? + # Closing a cursor cancels its running query. Best effort: the coordinator + # reaps whatever this misses after query.client.timeout. + for cursor in list(self._cursors): + try: + cursor.close() + except Exception as e: + logger.warning("Failed to close a cursor while closing the connection: %s", e) self._http_session.close() def start_transaction(self): @@ -323,7 +331,7 @@ def cursor( "row": Cursor }.get(cursor_style.lower(), Cursor) - return cursor_class( + cursor = cursor_class( self, request, legacy_primitive_types=( @@ -333,6 +341,8 @@ def cursor( ), stats_callback=stats_callback ) + self._cursors.add(cursor) + return cursor def _use_legacy_prepared_statements(self): if self.legacy_prepared_statements is not None: