From 7d9f75dacf63bd2199b4d464c6ffa50374431665 Mon Sep 17 00:00:00 2001 From: Pascal Schulze Date: Wed, 28 Oct 2026 14:06:51 +0100 Subject: [PATCH 1/2] fix: keep the last streamed result set readable after nextset When a streamed response carried several result sets, nextset() read and dropped the remaining rows of the current set before it found out whether another set followed. On the last set it returned False with those rows already gone, so walking to the last set with `while cursor.nextset()` and then fetching returned nothing. Buffered results kept them. nextset() now keeps the rows it reads past and serves them if no next set follows; they are dropped once the next set starts. Rows of a statement that fails while they are skipped are dropped too. Co-Authored-By: Claude Opus 5.5 --- src/firebolt/async_db/cursor.py | 8 +- .../common/row_set/asynchronous/streaming.py | 22 +++-- .../common/row_set/streaming_common.py | 30 ++++-- .../common/row_set/synchronous/streaming.py | 22 +++-- src/firebolt/db/cursor.py | 8 +- tests/unit/async_db/test_cursor.py | 45 +++++++++ .../row_set/asynchronous/test_streaming.py | 92 ++++++++++++++++++- .../row_set/synchronous/test_streaming.py | 92 ++++++++++++++++++- tests/unit/db/test_cursor.py | 45 +++++++++ 9 files changed, 333 insertions(+), 31 deletions(-) diff --git a/src/firebolt/async_db/cursor.py b/src/firebolt/async_db/cursor.py index cd261b73752..8106727b147 100644 --- a/src/firebolt/async_db/cursor.py +++ b/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 diff --git a/src/firebolt/common/row_set/asynchronous/streaming.py b/src/firebolt/common/row_set/asynchronous/streaming.py index 55b0cfff5c5..b07a72142b9 100644 --- a/src/firebolt/common/row_set/asynchronous/streaming.py +++ b/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 diff --git a/src/firebolt/common/row_set/streaming_common.py b/src/firebolt/common/row_set/streaming_common.py index 98a194b3dbf..cedf3c635d5 100644 --- a/src/firebolt/common/row_set/streaming_common.py +++ b/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: """ diff --git a/src/firebolt/common/row_set/synchronous/streaming.py b/src/firebolt/common/row_set/synchronous/streaming.py index 6ec9d58003f..e98722f7669 100644 --- a/src/firebolt/common/row_set/synchronous/streaming.py +++ b/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 diff --git a/src/firebolt/db/cursor.py b/src/firebolt/db/cursor.py index cdb74912584..11dfcd6bbc5 100644 --- a/src/firebolt/db/cursor.py +++ b/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 diff --git a/tests/unit/async_db/test_cursor.py b/tests/unit/async_db/test_cursor.py index 36856d50a87..aee23f46732 100644 --- a/tests/unit/async_db/test_cursor.py +++ b/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() == [] diff --git a/tests/unit/common/row_set/asynchronous/test_streaming.py b/tests/unit/common/row_set/asynchronous/test_streaming.py index 77a4d977248..8ca1a4168be 100644 --- a/tests/unit/common/row_set/asynchronous/test_streaming.py +++ b/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,95 @@ 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() diff --git a/tests/unit/common/row_set/synchronous/test_streaming.py b/tests/unit/common/row_set/synchronous/test_streaming.py index 5fb54d95cde..307c3ded2cc 100644 --- a/tests/unit/common/row_set/synchronous/test_streaming.py +++ b/tests/unit/common/row_set/synchronous/test_streaming.py @@ -802,8 +802,9 @@ def test_recovery_after_error(self, streaming_rowset): row = next(streaming_rowset) 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(StopIteration): @@ -1226,6 +1227,95 @@ 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, [])], +) +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 = StreamingRowSet() + row_set.append_response( + _response([_start(FIRST_COLUMNS), _data([[1], [2]]), _data([[3]]), _finish(3)]) + ) + for _ in range(rows_fetched): + next(row_set) + assert 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 list(row_set) == rows_left + assert row_set.row_count == 3 + assert row_set.nextset() is False + + +def test_repeated_nextset_without_next_set_keeps_remaining_rows(): + """Calling nextset again at the end neither loses nor recounts rows.""" + row_set = StreamingRowSet() + row_set.append_response( + _response([_start(FIRST_COLUMNS), _data([[1], [2]]), _data([[3]]), _finish(3)]) + ) + assert next(row_set) == [1] + assert row_set.nextset() is False + assert row_set.nextset() is False + assert row_set.row_count == 3 + assert list(row_set) == [[2], [3]] + assert row_set.nextset() is False + assert list(row_set) == [] + + +def test_walk_to_last_result_set_then_fetch(): + """Skipping to the last result set with nextset leaves its rows to fetch.""" + row_set = StreamingRowSet() + row_set.append_response( + _response( + [ + _start([]), + _finish(0), + _start([]), + _finish(0), + _start(SECOND_COLUMNS), + _data([["a", 3]]), + _data([["b", 4]]), + _finish(2), + ] + ) + ) + while row_set.nextset(): + pass + assert [c.name for c in row_set.columns] == ["col2", "col3"] + assert row_set.row_count == 2 + assert list(row_set) == [["a", 3], ["b", 4]] + + +def test_rows_skipped_before_next_response_are_dropped(): + """Rows skipped on the way to the next response are not read afterwards.""" + row_set = StreamingRowSet() + row_set.append_response(_response(TWO_RESULT_SETS)) + row_set.append_response( + _response([_start(FIRST_COLUMNS), _data([[5]]), _finish(1)]) + ) + assert row_set.nextset() is True + assert row_set.nextset() is True + assert list(row_set) == [[5]] + assert row_set.nextset() is False + assert list(row_set) == [] + +def test_rows_of_statement_failing_during_nextset_are_dropped(): + """Rows read past before an error are not fetchable after nextset raises.""" + row_set = StreamingRowSet() + row_set.append_response( + _response( + [_start(FIRST_COLUMNS), _data([[1], [2]]), _data([[3]]), FINISH_WITH_ERRORS] + ) + ) + assert next(row_set) == [1] + with pytest.raises(FireboltStructuredError): + row_set.nextset() + assert list(row_set) == [] + assert row_set.nextset() is False + assert list(row_set) == [] + + def test_transport_error_after_last_record(): """A broken connection after the last record surfaces as OperationalError.""" row_set = StreamingRowSet() diff --git a/tests/unit/db/test_cursor.py b/tests/unit/db/test_cursor.py index a6f84e7c4f0..dcb92250a57 100644 --- a/tests/unit/db/test_cursor.py +++ b/tests/unit/db/test_cursor.py @@ -1901,3 +1901,48 @@ def test_cursor_execute_stream_several_result_sets( assert cursor.fetchall() == [["x"], ["y"]] assert cursor.rowcount == 2 assert 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 + + +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(), + ) + cursor.execute_stream("copy ...; call ...; select ...", skip_parsing=True) + while cursor.nextset(): + pass + assert [d.name for d in cursor.description] == ["query_id"] + assert cursor.rowcount == 2 + assert cursor.fetchall() == [["q1"], ["q2"]] + assert cursor.fetchall() == [] From f3317106924c6c1659dc0a13b267ae9c3af7a6a8 Mon Sep 17 00:00:00 2001 From: Pascal Schulze Date: Wed, 28 Oct 2026 14:28:26 +0100 Subject: [PATCH 2/2] style: apply black formatting to streaming tests Co-Authored-By: Claude Opus 5.5 --- tests/unit/common/row_set/asynchronous/test_streaming.py | 1 + tests/unit/common/row_set/synchronous/test_streaming.py | 1 + 2 files changed, 2 insertions(+) diff --git a/tests/unit/common/row_set/asynchronous/test_streaming.py b/tests/unit/common/row_set/asynchronous/test_streaming.py index 8ca1a4168be..6468f941ec0 100644 --- a/tests/unit/common/row_set/asynchronous/test_streaming.py +++ b/tests/unit/common/row_set/asynchronous/test_streaming.py @@ -1292,6 +1292,7 @@ async def test_rows_skipped_before_next_response_are_dropped(): 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() diff --git a/tests/unit/common/row_set/synchronous/test_streaming.py b/tests/unit/common/row_set/synchronous/test_streaming.py index 307c3ded2cc..4296c88076e 100644 --- a/tests/unit/common/row_set/synchronous/test_streaming.py +++ b/tests/unit/common/row_set/synchronous/test_streaming.py @@ -1300,6 +1300,7 @@ def test_rows_skipped_before_next_response_are_dropped(): assert row_set.nextset() is False assert list(row_set) == [] + def test_rows_of_statement_failing_during_nextset_are_dropped(): """Rows read past before an error are not fetchable after nextset raises.""" row_set = StreamingRowSet()