Make result iteration recoverable from errors

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)

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]

trino/client.py

@@ -821,15 +821,24 @@ 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]):
         self._query = query
         # 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():