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..6468f941ec0 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,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() diff --git a/tests/unit/common/row_set/synchronous/test_streaming.py b/tests/unit/common/row_set/synchronous/test_streaming.py index 5fb54d95cde..4296c88076e 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,96 @@ 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() == []