Implement heartbeat when downloading spooled segments

Author: azawlocki-sbdt

trino/client.py

@@ -37,6 +37,7 @@
 import abc
 import atexit
 import base64
+import contextlib
 import copy
 import functools
 import itertools
@@ -853,10 +854,10 @@ def __iter__(self):
             self._rows = next_rows
 
 
-class QueryHeartbeat:
-    """Periodically sends HEAD requests to the current nextUri to prevent the
-    coordinator from abandoning a query while the client is downloading spooled
-    result segments from external storage."""
+class _QueryHeartbeat:
+    """Periodically sends HEAD requests to the current nextUri to prevent the coordinator
+    from abandoning a query if the client is silent for a longer period of time, for example
+    while the client is downloading spooled result segments from external storage."""
 
     _MAX_FAILURES = 3
 
@@ -881,6 +882,7 @@ def _run(self) -> None:
             try:
                 response = self._request.head(uri)
                 if response.status_code in (404, 405):
+                    # 404/405 means the server does not support heartbeat calls
                     return
                 if not response.ok:
                     self._failures += 1
@@ -894,6 +896,16 @@ def _run(self) -> None:
                     return
 
 
+@contextlib.contextmanager
+def query_heartbeat(request: TrinoRequest, interval: float) -> Iterator[None]:
+    heartbeat = _QueryHeartbeat(request, interval)
+    heartbeat.start()
+    try:
+        yield
+    finally:
+        heartbeat.stop()
+
+
 class TrinoQuery:
     """Represent the execution of a SQL statement by Trino."""
 
@@ -920,7 +932,6 @@ def __init__(
         self._legacy_primitive_types = legacy_primitive_types
         self._row_mapper: Optional[RowMapper] = None
         self._fetch_mode = fetch_mode
-        self._heartbeat: Optional[QueryHeartbeat] = None
 
     @property
     def query_id(self) -> Optional[str]:
@@ -998,13 +1009,8 @@ def execute(self, additional_http_headers=None) -> TrinoResult:
         self._stats.update({"queryId": self.query_id})
         self._update_state(status)
         self._warnings = getattr(status, "warnings", [])
-        interval = self._request._client_session.heartbeat_interval
-        if interval is not None:
-            self._heartbeat = QueryHeartbeat(self._request, interval)
-            self._heartbeat.start()
         if status.next_uri is None:
             self._finished = True
-            self._stop_heartbeat()
 
         rows = self._row_mapper.map(status.rows) if self._row_mapper else status.rows
         self._result = TrinoResult(self, rows)
@@ -1052,7 +1058,6 @@ def fetch(self) -> Union[List[Union[List[Any], Any]], Iterator[List[Any]]]:
         self._update_state(status)
         if status.next_uri is None:
             self._finished = True
-            self._stop_heartbeat()
 
         if not self._row_mapper:
             return []
@@ -1065,7 +1070,9 @@ def fetch(self) -> Union[List[Union[List[Any], Any]], Iterator[List[Any]]]:
             if self._fetch_mode == "segments":
                 return spooled
             # Return iterator directly, do NOT materialize with list()
-            return SegmentIterator(spooled, self._row_mapper)
+            return SegmentIterator(
+                spooled, self._row_mapper, self._request, self._request._client_session.heartbeat_interval
+            )
         elif isinstance(status.rows, list):
             return self._row_mapper.map(rows)
         else:
@@ -1100,16 +1107,11 @@ def cancel(self) -> None:
             raise trino.exceptions.TrinoConnectionError("failed to cancel query: {}".format(e))
         if response.status_code == requests.codes.no_content:
             self._cancelled = True
-            self._stop_heartbeat()
             logger.debug("query cancelled: %s", self.query_id)
             return
 
         self._request.raise_response_error(response)
 
-    def _stop_heartbeat(self) -> None:
-        if self._heartbeat is not None:
-            self._heartbeat.stop()
-
     def is_finished(self) -> bool:
         import warnings
         warnings.warn("is_finished is deprecated, use finished instead", DeprecationWarning)
@@ -1342,13 +1344,21 @@ def __repr__(self):
 
 
 class SegmentIterator:
-    def __init__(self, segments: Union[DecodableSegment, List[DecodableSegment]], mapper: RowMapper) -> None:
+    def __init__(
+        self,
+        segments: Union[DecodableSegment, List[DecodableSegment]],
+        mapper: RowMapper,
+        request: TrinoRequest,
+        heartbeat_interval: Optional[float] = None,
+    ) -> None:
         self._segments = iter(segments if isinstance(segments, List) else [segments])
         self._mapper = mapper
         self._decoder = None
         self._rows: Iterator[List[List[Any]]] = iter([])
         self._finished = False
         self._current_segment: Optional[DecodableSegment] = None
+        self._request = request
+        self._heartbeat_interval = heartbeat_interval
 
     def __iter__(self) -> Iterator[List[Any]]:
         return self
@@ -1374,7 +1384,13 @@ def _load_next_segment(self):
             if self._decoder is None:
                 self._decoder = SegmentDecoder(CompressedQueryDataDecoderFactory(self._mapper)
                                                .create(self._current_segment.encoding))
-            self._rows = iter(self._decoder.decode(self._current_segment.segment))
+            if isinstance(self._current_segment.segment, SpooledSegment) and self._heartbeat_interval:
+                # Downloading a spooled segment may take some time. In the meantime, we send heartbeat
+                # requests so the server doesn't think we lost interest and close the connection.
+                with query_heartbeat(self._request, self._heartbeat_interval):
+                    self._rows = iter(self._decoder.decode(self._current_segment.segment))
+            else:
+                self._rows = iter(self._decoder.decode(self._current_segment.segment))
         except StopIteration:
             self._finished = True