fix: keep the last streamed result set readable after nextset

src/firebolt/async_db/cursor.py

@@ -489,10 +489,10 @@ async def fetchall(self) -> List[List[ColType]]:
     @check_query_executed
     async def nextset(self) -> bool:
         """
-        Skip to the next available set, discarding any remaining rows
-        from the current set. Result sets returned by one response are
-        visited in order, like those of separate statements; for a streamed
-        result the remaining rows are read and discarded on the way.
+        Skip to the next available set, if there is one, discarding any
+        remaining rows from the current set. If there is no next set, the
+        current set and its remaining rows are kept. Result sets returned by
+        one response are visited in order, like those of separate statements.
 
         Returns:
             bool: True if there is a next result set, False otherwise

src/firebolt/common/row_set/asynchronous/streaming.py

@@ -143,7 +143,8 @@ async def nextset(self) -> bool:
 
         A response may carry several result sets one after another. Those are
         visited first, then the following responses. The remaining rows of the
-        current result set are read and discarded on the way.
+        current result set are read on the way, and kept readable if there is
+        no next result set.
 
         Returns:
             bool: True if there is a next result set, False otherwise
@@ -177,7 +178,8 @@ async def nextset(self) -> bool:
     async def _next_result_set_in_response(self) -> bool:
         """
         Skip the remaining rows of the current result set and start the next
-        result set of the same response, if there is one.
+        result set of the same response, if there is one. Otherwise the
+        skipped rows stay readable.
 
         Returns:
             bool: True if a new result set was started, False if the response
@@ -188,9 +190,10 @@ async def _next_result_set_in_response(self) -> bool:
         """
         if self._response_failed:
             return False
-        self._discard_current_record()
-        while not self._response_consumed:
-            self._count_skipped_record(await self._pop_data_record())
+        if not self._response_consumed:
+            self._count_unread_rows_of_current_record()
+            while not self._response_consumed:
+                self._skip_record(await self._pop_data_record())
         return self._start_result_set_from_record(await self._next_json_lines_record())
 
     @close_on_op_error
@@ -219,13 +222,18 @@ async def __anext__(self) -> List[ColType]:
             StopAsyncIteration: If there are no more rows
             OperationalError: If an error occurs while reading the row
         """
-        if self._current_response is None or self._response_consumed:
+        if self._current_response is None:
             raise StopAsyncIteration
 
         while self._current_record is None or self._current_record_row_idx >= len(
             self._current_record.data
         ):
-            self._current_record = await self._pop_data_record()
+            if self._skipped_records:
+                self._current_record = self._skipped_records.popleft()
+            elif self._response_consumed:
+                raise StopAsyncIteration
+            else:
+                self._current_record = await self._pop_data_record()
             self._current_record_row_idx = 0
             if self._current_record is None:
                 break

src/firebolt/common/row_set/streaming_common.py

@@ -1,5 +1,15 @@
 import json
-from typing import Any, AsyncIterator, Iterator, List, Optional, Type, Union
+from collections import deque
+from typing import (
+    Any,
+    AsyncIterator,
+    Deque,
+    Iterator,
+    List,
+    Optional,
+    Type,
+    Union,
+)
 
 from httpx import Response
 
@@ -77,6 +87,9 @@ def _reset_result_set(self) -> None:
         self._response_consumed = False
         self._current_columns = None
         self._rows_returned = 0
+        # Rows nextset() read past. They stay readable if no next result set
+        # follows, as with a buffered row set.
+        self._skipped_records: Deque[DataRecord] = deque()
 
     @property
     def _current_response(self) -> Optional[Response]:
@@ -127,6 +140,9 @@ def _next_json_lines_record_from_line(
             if self._response_consumed:
                 # The error belongs to a statement that produced no START record
                 self._reset_result_set()
+            # Rows of a failed statement must not be fetched afterwards
+            self._current_record = None
+            self._skipped_records.clear()
             self._response_consumed = True
             self._response_failed = True
             self._current_statistics = record.statistics
@@ -180,23 +196,23 @@ def _fetch_columns_from_record(
             for col in record.result_columns
         ]
 
-    def _discard_current_record(self) -> None:
+    def _count_unread_rows_of_current_record(self) -> None:
         """
-        Drop the current data record, counting its unread rows as read.
+        Count the unread rows of the current data record as read.
         """
         if self._current_record is not None:
             self._rows_returned += (
                 len(self._current_record.data) - self._current_record_row_idx
             )
-            self._current_record = None
-            self._current_record_row_idx = -1
 
-    def _count_skipped_record(self, record: Optional[DataRecord]) -> None:
+    def _skip_record(self, record: Optional[DataRecord]) -> None:
         """
-        Count the rows of a data record that nextset() skipped over.
+        Count the rows of a data record that nextset() read past, and keep
+        the record in case no next result set follows.
         """
         if record is not None:
             self._rows_returned += len(record.data)
+            self._skipped_records.append(record)
 
     def _start_result_set_from_record(self, record: Optional[JSONLinesRecord]) -> bool:
         """

src/firebolt/common/row_set/synchronous/streaming.py

@@ -143,7 +143,8 @@ def nextset(self) -> bool:
 
         A response may carry several result sets one after another. Those are
         visited first, then the following responses. The remaining rows of the
-        current result set are read and discarded on the way.
+        current result set are read on the way, and kept readable if there is
+        no next result set.
 
         Returns:
             bool: True if there is a next result set, False otherwise
@@ -177,7 +178,8 @@ def nextset(self) -> bool:
     def _next_result_set_in_response(self) -> bool:
         """
         Skip the remaining rows of the current result set and start the next
-        result set of the same response, if there is one.
+        result set of the same response, if there is one. Otherwise the
+        skipped rows stay readable.
 
         Returns:
             bool: True if a new result set was started, False if the response
@@ -188,9 +190,10 @@ def _next_result_set_in_response(self) -> bool:
         """
         if self._response_failed:
             return False
-        self._discard_current_record()
-        while not self._response_consumed:
-            self._count_skipped_record(self._pop_data_record())
+        if not self._response_consumed:
+            self._count_unread_rows_of_current_record()
+            while not self._response_consumed:
+                self._skip_record(self._pop_data_record())
         return self._start_result_set_from_record(self._next_json_lines_record())
 
     @close_on_op_error
@@ -219,13 +222,18 @@ def __next__(self) -> List[ColType]:
             StopIteration: If there are no more rows
             OperationalError: If an error occurs while reading the row
         """
-        if self._current_response is None or self._response_consumed:
+        if self._current_response is None:
             raise StopIteration
 
         while self._current_record is None or self._current_record_row_idx >= len(
             self._current_record.data
         ):
-            self._current_record = self._pop_data_record()
+            if self._skipped_records:
+                self._current_record = self._skipped_records.popleft()
+            elif self._response_consumed:
+                raise StopIteration
+            else:
+                self._current_record = self._pop_data_record()
             self._current_record_row_idx = 0
             if self._current_record is None:
                 break

src/firebolt/db/cursor.py

@@ -471,10 +471,10 @@ def fetchall(self) -> List[List[ColType]]:
     @check_query_executed
     def nextset(self) -> bool:
         """
-        Skip to the next available set, discarding any remaining rows
-        from the current set. Result sets returned by one response are
-        visited in order, like those of separate statements; for a streamed
-        result the remaining rows are read and discarded on the way.
+        Skip to the next available set, if there is one, discarding any
+        remaining rows from the current set. If there is no next set, the
+        current set and its remaining rows are kept. Result sets returned by
+        one response are visited in order, like those of separate statements.
 
         Returns:
             bool: True if there is a next result set, False otherwise

tests/unit/async_db/test_cursor.py

@@ -2013,3 +2013,48 @@ async def test_cursor_execute_stream_several_result_sets(
     assert await cursor.fetchall() == [["x"], ["y"]]
     assert cursor.rowcount == 2
     assert await cursor.nextset() is False
+
+
+def _copy_then_select_lines(statistics: Dict[str, Any]) -> List[Dict[str, Any]]:
+    """Two statements without result columns, like COPY and CALL, then a SELECT."""
+    lines: List[Dict[str, Any]] = []
+    for query_id, columns, data in [
+        ("q1", [], []),
+        ("q2", [], []),
+        ("q3", [{"name": "query_id", "type": "text"}], [["q1"], ["q2"]]),
+    ]:
+        lines.append(
+            {
+                "message_type": "START",
+                "result_columns": columns,
+                "query_id": query_id,
+                "query_label": "l",
+                "request_id": "r",
+            }
+        )
+        if data:
+            lines.append({"message_type": "DATA", "data": data})
+        lines.append({"message_type": "FINISH_SUCCESSFULLY", "statistics": statistics})
+    return lines
+
+
+async def test_cursor_execute_stream_walk_to_last_result_set(
+    httpx_mock: HTTPXMock,
+    streaming_query_url: str,
+    query_statistics: Dict[str, Any],
+    cursor: Cursor,
+):
+    """The rows of the last streamed result set survive the nextset that ends the walk."""
+    httpx_mock.add_response(
+        url=streaming_query_url,
+        content="\n".join(
+            json.dumps(line) for line in _copy_then_select_lines(query_statistics)
+        ).encode(),
+    )
+    await cursor.execute_stream("copy ...; call ...; select ...", skip_parsing=True)
+    while await cursor.nextset():
+        pass
+    assert [d.name for d in cursor.description] == ["query_id"]
+    assert cursor.rowcount == 2
+    assert await cursor.fetchall() == [["q1"], ["q2"]]
+    assert await cursor.fetchall() == []

tests/unit/common/row_set/asynchronous/test_streaming.py

@@ -806,8 +806,9 @@ async def test_recovery_after_error(self, streaming_rowset):
                 row = await streaming_rowset.__anext__()
                 assert row == ["success"]
 
-                # Mark as consumed for the test
+                # Mark as consumed for the test, as reading the FINISH record does
                 streaming_rowset._response_consumed = True
+                streaming_rowset._current_record = None
 
                 # Should be able to iterate to the end
                 with pytest.raises(StopAsyncIteration):
@@ -1218,6 +1219,96 @@ async def test_row_count_counts_rows_skipped_by_nextset(rows_fetched: int):
     assert row_set.row_count == 3
 
 
+@pytest.mark.parametrize(
+    "rows_fetched,rows_left",
+    [(0, [[1], [2], [3]]), (1, [[2], [3]]), (2, [[3]]), (3, [])],
+)
+async def test_nextset_without_next_set_keeps_remaining_rows(rows_fetched, rows_left):
+    """When nextset finds no next result set, the unread rows stay readable."""
+    row_set = StreamingAsyncRowSet()
+    await row_set.append_response(
+        _response([_start(FIRST_COLUMNS), _data([[1], [2]]), _data([[3]]), _finish(3)])
+    )
+    for _ in range(rows_fetched):
+        await row_set.__anext__()
+    assert await row_set.nextset() is False
+    assert [c.name for c in row_set.columns] == ["col1"]
+    assert row_set.row_count == 3
+    assert row_set.statistics.rows_read == 3
+    assert await _rows(row_set) == rows_left
+    assert row_set.row_count == 3
+    assert await row_set.nextset() is False
+
+
+async def test_repeated_nextset_without_next_set_keeps_remaining_rows():
+    """Calling nextset again at the end neither loses nor recounts rows."""
+    row_set = StreamingAsyncRowSet()
+    await row_set.append_response(
+        _response([_start(FIRST_COLUMNS), _data([[1], [2]]), _data([[3]]), _finish(3)])
+    )
+    assert await row_set.__anext__() == [1]
+    assert await row_set.nextset() is False
+    assert await row_set.nextset() is False
+    assert row_set.row_count == 3
+    assert await _rows(row_set) == [[2], [3]]
+    assert await row_set.nextset() is False
+    assert await _rows(row_set) == []
+
+
+async def test_walk_to_last_result_set_then_fetch():
+    """Skipping to the last result set with nextset leaves its rows to fetch."""
+    row_set = StreamingAsyncRowSet()
+    await row_set.append_response(
+        _response(
+            [
+                _start([]),
+                _finish(0),
+                _start([]),
+                _finish(0),
+                _start(SECOND_COLUMNS),
+                _data([["a", 3]]),
+                _data([["b", 4]]),
+                _finish(2),
+            ]
+        )
+    )
+    while await row_set.nextset():
+        pass
+    assert [c.name for c in row_set.columns] == ["col2", "col3"]
+    assert row_set.row_count == 2
+    assert await _rows(row_set) == [["a", 3], ["b", 4]]
+
+
+async def test_rows_skipped_before_next_response_are_dropped():
+    """Rows skipped on the way to the next response are not read afterwards."""
+    row_set = StreamingAsyncRowSet()
+    await row_set.append_response(_response(TWO_RESULT_SETS))
+    await row_set.append_response(
+        _response([_start(FIRST_COLUMNS), _data([[5]]), _finish(1)])
+    )
+    assert await row_set.nextset() is True
+    assert await row_set.nextset() is True
+    assert await _rows(row_set) == [[5]]
+    assert await row_set.nextset() is False
+    assert await _rows(row_set) == []
+
+
+async def test_rows_of_statement_failing_during_nextset_are_dropped():
+    """Rows read past before an error are not fetchable after nextset raises."""
+    row_set = StreamingAsyncRowSet()
+    await row_set.append_response(
+        _response(
+            [_start(FIRST_COLUMNS), _data([[1], [2]]), _data([[3]]), FINISH_WITH_ERRORS]
+        )
+    )
+    assert await row_set.__anext__() == [1]
+    with pytest.raises(FireboltStructuredError):
+        await row_set.nextset()
+    assert await _rows(row_set) == []
+    assert await row_set.nextset() is False
+    assert await _rows(row_set) == []
+
+
 async def test_transport_error_after_last_record():
     """A broken connection after the last record surfaces as OperationalError."""
     row_set = StreamingAsyncRowSet()