From f4e78884bbbecdb6faa65756750c08542bc379a8 Mon Sep 17 00:00:00 2001 From: Ashhar Hasan Date: Tue, 7 Jul 2026 21:09:33 +0530 Subject: [PATCH] Make result iteration recoverable from errors TrinoResult implemented iteration as a generator. A generator that raises an unhandled exception is finalized so every subsequent next() raises StopIteration which dbapi's fetchone() maps to None. Since spooled segments are downloaded lazily during iteration, a transient error (e.g. an S3 timeout) mid-iteration permanently killed the iterator and silently dropped the remaining rows as if the query had completed normally. Rewrite TrinoResult as a class-based iterator that keeps its state on the instance. An error propagates to the caller while the iterator stays usable, so a persistent error keeps surfacing instead of turning into StopIteration and a retried next() resumes where the failure happened. Also make SegmentIterator retry-safe. Acknowledge and advance past a segment only after it was decoded successfully so retrying after a failed segment download re-decodes the same segment instead of acknowledging and skipping it. --- tests/unit/test_client.py | 92 ++++++++++++++++++++++++++++++ tests/unit/test_client_spooling.py | 53 +++++++++++++++++ trino/client.py | 81 ++++++++++++++++++-------- 3 files changed, 201 insertions(+), 25 deletions(-) diff --git a/tests/unit/test_client.py b/tests/unit/test_client.py index 29f3f388..f7d4fee6 100644 --- a/tests/unit/test_client.py +++ b/tests/unit/test_client.py @@ -1448,3 +1448,95 @@ def test_decoder_factory_raises_with_message_on_missing_lz4(): match=f"lz4 is not installed so json\\+lz4 encoding is not supported: {error_message}" ): factory.create("json+lz4") + + +class _FinishedQuery: + """Query stub that is already finished. All rows come from the initial batch.""" + finished = True + + def fetch(self): + return [] + + +@pytest.mark.parametrize("consecutive_failures", (1, 2, 3)) +def test_trino_result_resumes_after_transient_error_in_rows_iterator(consecutive_failures): + class FlakyIterator: + """Fails a given number of times at the third row, succeeds when retried.""" + def __init__(self, failures): + self._rows = iter([[1], [2], [3]]) + self._served = 0 + self._remaining_failures = failures + + def __iter__(self): + return self + + def __next__(self): + if self._served == 2 and self._remaining_failures > 0: + self._remaining_failures -= 1 + raise IOError("segment download failed") + self._served += 1 + return next(self._rows) + + result = TrinoResult(_FinishedQuery(), FlakyIterator(consecutive_failures)) + it = iter(result) + assert next(it) == [1] + assert next(it) == [2] + # Each retry surfaces the error again until the underlying iterator recovers + for _ in range(consecutive_failures): + with pytest.raises(IOError): + next(it) + # The iterator stays usable and resumes where the failure happened + assert next(it) == [3] + with pytest.raises(StopIteration): + next(it) + assert result.rownumber == 3 + + +def test_trino_result_reraises_persistent_error_instead_of_stopping(): + class FailingIterator: + def __init__(self): + self._count = 0 + + def __iter__(self): + return self + + def __next__(self): + if self._count >= 3: + raise IOError("segment download failed") + self._count += 1 + return [self._count] + + result = TrinoResult(_FinishedQuery(), FailingIterator()) + it = iter(result) + assert [next(it), next(it), next(it)] == [[1], [2], [3]] + # The error keeps surfacing instead of turning into StopIteration + # which dbapi would report as a normally exhausted result set + with pytest.raises(IOError): + next(it) + with pytest.raises(IOError): + next(it) + + +def test_trino_result_resumes_after_transient_fetch_error(): + class FlakyQuery: + def __init__(self): + self.finished = False + self._fetches = 0 + + def fetch(self): + self._fetches += 1 + if self._fetches == 1: + raise IOError("connection reset") + self.finished = True + return [[2]] + + result = TrinoResult(FlakyQuery(), [[1]]) + it = iter(result) + # The next batch is prefetched before the first row is served so the + # fetch error surfaces before any rows + with pytest.raises(IOError): + next(it) + assert next(it) == [1] + assert next(it) == [2] + with pytest.raises(StopIteration): + next(it) diff --git a/tests/unit/test_client_spooling.py b/tests/unit/test_client_spooling.py index e739cb2e..a56d62df 100644 --- a/tests/unit/test_client_spooling.py +++ b/tests/unit/test_client_spooling.py @@ -200,3 +200,56 @@ def test_fetch_passes_request_and_interval_to_segment_iterator(heartbeat_interva assert MockSI.call_args.kwargs["request"] is req assert MockSI.call_args.kwargs["heartbeat_interval"] == heartbeat_interval + + +class _FakeSpooledSegment(SpooledSegment): + """SpooledSegment that records acknowledgments instead of sending requests.""" + def __init__(self, name): + super().__init__( + { + "type": "spooled", + "uri": f"http://storage/{name}", + "ackUri": f"http://storage/{name}/ack", + "metadata": {"uncompressedSize": "10", "segmentSize": "10"}, + }, + request=None, + ) + self.acknowledge_count = 0 + + def acknowledge(self): + self.acknowledge_count += 1 + + +class _FlakyDecoder: + """Decoder that fails the first decode of `failing_segment`, then succeeds.""" + def __init__(self, rows_by_segment, failing_segment): + self._rows_by_segment = rows_by_segment + self._failing_segment = failing_segment + + def decode(self, segment): + if segment is self._failing_segment: + self._failing_segment = None + raise IOError("segment download failed") + return self._rows_by_segment[segment] + + +@pytest.mark.parametrize("failing_segment_index", (0, 1, 2)) +def test_segment_iterator_retries_failed_segment_without_skipping_it(failing_segment_index): + segs = [_FakeSpooledSegment(name) for name in ("s1", "s2", "s3")] + segments = [DecodableSegment("json", None, seg) for seg in segs] + iterator = SegmentIterator(segments, mapper=None) + iterator._decoder = _FlakyDecoder( + {seg: [[index + 1]] for index, seg in enumerate(segs)}, + failing_segment=segs[failing_segment_index], + ) + + rows = [next(iterator) for _ in range(failing_segment_index)] + with pytest.raises(IOError): + next(iterator) + # The failed segment was neither acknowledged nor skipped. Retrying should deliver its rows. + assert segs[failing_segment_index].acknowledge_count == 0 + rows.extend(iterator) + assert rows == [[1], [2], [3]] + with pytest.raises(StopIteration): + next(iterator) + assert [seg.acknowledge_count for seg in segs] == [1, 1, 1] diff --git a/trino/client.py b/trino/client.py index 692d10fa..abf1c2f2 100644 --- a/trino/client.py +++ b/trino/client.py @@ -821,8 +821,13 @@ class TrinoResult: """ Represent the result of a Trino query as an iterator on rows. - This class implements the iterator protocol as a generator type - https://docs.python.org/3/library/stdtypes.html#generator-types + This class implements the iterator protocol on the instance itself instead + of as a generator. A generator that raises an exception is finalized and + every subsequent next() raises StopIteration indistinguishable from normal + exhaustion. Keeping the iteration state on the instance lets a transient + error (e.g. a failed spooled segment download) propagate to the caller + while the iterator stays usable so a retried next() resumes where the + failure happened instead of silently dropping the remaining rows. """ def __init__(self, query, rows: List[Any]): @@ -830,6 +835,10 @@ def __init__(self, query, rows: List[Any]): # Initial rows from the first POST request self._rows = rows self._rownumber = 0 + # Iterator over the batch of rows currently being served + self._current_batch: Optional[Iterator[Any]] = None + # Rows prefetched while the current batch is being served + self._next_rows: Optional[Any] = None @property def rows(self): @@ -844,15 +853,28 @@ def rownumber(self) -> int: return self._rownumber def __iter__(self): - # A query only transitions to a FINISHED state when the results are fully consumed: - # The reception of the data is acknowledged by calling the next_uri before exposing the data through dbapi. - while not self._query.finished or self._rows is not None: - next_rows = self._query.fetch() if not self._query.finished else None - for row in self._rows: - self._rownumber += 1 - yield row + return self - self._rows = next_rows + def __next__(self): + while True: + if self._current_batch is None: + if self._query.finished and self._rows is None: + raise StopIteration + # A query only transitions to a FINISHED state when the results are fully consumed: + # The reception of the data is acknowledged by calling the next_uri before exposing the data through + # dbapi. + self._next_rows = self._query.fetch() if not self._query.finished else None + self._current_batch = iter(self._rows) + + try: + row = next(self._current_batch) + except StopIteration: + self._rows = self._next_rows + self._next_rows = None + self._current_batch = None + continue + self._rownumber += 1 + return row class TrinoQuery: @@ -1362,6 +1384,8 @@ def __init__( self._rows: Iterator[List[List[Any]]] = iter([]) self._finished = False 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 @@ -1381,29 +1405,36 @@ def __next__(self) -> List[Any]: self._load_next_segment() def _load_next_segment(self): - try: + # A segment is acknowledged only after its rows were decoded successfully. If the previous attempt failed + # mid-decode (e.g. the spooled segment download failed) the same segment is retried instead of being skipped. + if self._pending_segment is None: if self._current_segment: segment = self._current_segment.segment if isinstance(segment, SpooledSegment): segment.acknowledge() + self._current_segment = None - self._current_segment = next(self._segments) - if self._decoder is None: - self._decoder = SegmentDecoder(CompressedQueryDataDecoderFactory(self._mapper) - .create(self._current_segment.encoding)) + try: + self._pending_segment = next(self._segments) + except StopIteration: + self._finished = True + return - if isinstance(self._current_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._current_segment.segment) - else: - rows = self._decoder.decode(self._current_segment.segment) + if self._decoder is None: + self._decoder = SegmentDecoder(CompressedQueryDataDecoderFactory(self._mapper) + .create(self._pending_segment.encoding)) - self._rows = iter(rows) + 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) - except StopIteration: - self._finished = True + self._rows = iter(rows) + self._current_segment = self._pending_segment + self._pending_segment = None class SegmentDecoder():