From 43b75af911e5b959e7151a2d4363f3f1a847271d Mon Sep 17 00:00:00 2001 From: Giuseppe Mazzotta <16498973+gm42@users.noreply.github.com> Date: Mon, 26 Oct 2026 15:31:42 +0100 Subject: [PATCH] fix(FB-5242): read every result set a response carries A response holds one result document per statement it ran. The in-memory row set parsed the body as a single JSON document and failed with "Extra data" on the second one, and the streaming row set stopped at the first FINISH record, so later result sets were never reachable through nextset(). Both row sets now walk the result sets inside a response before moving on to the next response. On the streaming path nextset() reads and discards the rest of the current result set, counting the skipped rows into rowcount, skips blank lines, turns a transport error into an OperationalError, and reports an error record met where the next START was expected against a reset result set. The in-memory path only keeps result sets from a body that parsed completely, accepts JSON whitespace only between documents, and reports a non-object document or trailing text after a complete document as a DataError that carries that text. Along the way, _reset() assigned a misspelled attribute, so the line iterator of a finished response was reused for the next one; the error message for an unexpected record crashed on the raw message_type string; and a DATA record without rows raised IndexError on fetch. --- src/firebolt/async_db/cursor.py | 9 +- .../common/row_set/asynchronous/streaming.py | 53 +++- .../common/row_set/streaming_common.py | 77 ++++- .../common/row_set/synchronous/in_memory.py | 86 ++++-- .../common/row_set/synchronous/streaming.py | 53 +++- src/firebolt/db/cursor.py | 9 +- tests/unit/async_db/test_cursor.py | 75 +++++ .../row_set/asynchronous/test_in_memory.py | 74 ++++- .../row_set/asynchronous/test_streaming.py | 252 ++++++++++++++++- .../row_set/synchronous/test_in_memory.py | 134 ++++++++- .../row_set/synchronous/test_streaming.py | 267 +++++++++++++++++- tests/unit/db/test_cursor.py | 75 +++++ 12 files changed, 1113 insertions(+), 51 deletions(-) diff --git a/src/firebolt/async_db/cursor.py b/src/firebolt/async_db/cursor.py index dbf617a703d..cd261b73752 100644 --- a/src/firebolt/async_db/cursor.py +++ b/src/firebolt/async_db/cursor.py @@ -490,10 +490,17 @@ async def fetchall(self) -> List[List[ColType]]: async def nextset(self) -> bool: """ Skip to the next available set, discarding any remaining rows - from the current set. + 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. Returns: bool: True if there is a next result set, False otherwise + + Raises: + OperationalError: If the response cannot be read + FireboltStructuredError: If the response reports an error for the + current or the next result set """ assert self._row_set is not None return await self._row_set.nextset() diff --git a/src/firebolt/common/row_set/asynchronous/streaming.py b/src/firebolt/common/row_set/asynchronous/streaming.py index 63be6cbe4af..55b0cfff5c5 100644 --- a/src/firebolt/common/row_set/asynchronous/streaming.py +++ b/src/firebolt/common/row_set/asynchronous/streaming.py @@ -78,13 +78,18 @@ async def _next_json_lines_record(self) -> Optional[JSONLinesRecord]: """ if self._current_response is None: return None - if self._lines_iter is None: - try: + try: + if self._lines_iter is None: self._lines_iter = self._current_response.aiter_lines() - except HTTPError as err: - raise OperationalError("Failed to read response stream.") from err + while True: + next_line = await anext(self._lines_iter, None) + # A blank line separates nothing + if next_line is None or next_line.strip(): + break + except HTTPError as err: + self._response_failed = True + raise OperationalError("Failed to read response stream.") from err - next_line = await anext(self._lines_iter, None) return self._next_json_lines_record_from_line(next_line) @property @@ -136,13 +141,25 @@ async def nextset(self) -> bool: """ Move to the next result set. + 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. + Returns: bool: True if there is a next result set, False otherwise Raises: OperationalError: If the response stream cannot be closed or if an error - occurs while fetching new columns + occurs while reading the response + FireboltStructuredError: If the response reports an error for the + current or the next result set """ + if ( + self._responses + and self._current_response is not None + and await self._next_result_set_in_response() + ): + return True if self._current_row_set_idx + 1 < len(self._responses): if self._current_response is not None: try: @@ -156,6 +173,26 @@ async def nextset(self) -> bool: return True return False + @close_on_op_error + 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. + + Returns: + bool: True if a new result set was started, False if the response + carries no more result sets + + Raises: + OperationalError: If an error occurs while reading the response + """ + if self._response_failed: + return False + self._discard_current_record() + while not self._response_consumed: + self._count_skipped_record(await self._pop_data_record()) + return self._start_result_set_from_record(await self._next_json_lines_record()) + @close_on_op_error async def _pop_data_record(self) -> Optional[DataRecord]: """ @@ -185,11 +222,13 @@ async def __anext__(self) -> List[ColType]: if self._current_response is None or self._response_consumed: raise StopAsyncIteration - if self._current_record is None or self._current_record_row_idx >= len( + 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() self._current_record_row_idx = 0 + if self._current_record is None: + break return self._get_next_data_row_from_current_record(StopAsyncIteration) diff --git a/src/firebolt/common/row_set/streaming_common.py b/src/firebolt/common/row_set/streaming_common.py index 86416d537cb..98a194b3dbf 100644 --- a/src/firebolt/common/row_set/streaming_common.py +++ b/src/firebolt/common/row_set/streaming_common.py @@ -8,6 +8,7 @@ DataRecord, ErrorRecord, JSONLinesRecord, + MessageType, StartRecord, SuccessRecord, parse_json_lines_record, @@ -20,6 +21,14 @@ ) +def _message_type_name(record: JSONLinesRecord) -> str: + """Name of a record's message type, whether it holds the enum or its raw value.""" + message_type = record.message_type + if isinstance(message_type, MessageType): + return message_type.value + return str(message_type) + + class StreamingRowSetCommonBase: """ A mixin class that provides common functionality for streaming row sets. @@ -46,12 +55,23 @@ def _reset(self) -> None: """ Reset the state of the streaming row set. - Resets internal counters, iterators, and cached data for the next row set. + Resets internal counters, iterators, and cached data for the next response. Note: Does not reset _current_row_set_idx to allow for multiple row sets. """ + self._lines_iter: Optional[Union[AsyncIterator[str], Iterator[str]]] = None + # Set once reading the response failed; no further result set is read from it. + self._response_failed = False + self._reset_result_set() + + def _reset_result_set(self) -> None: + """ + Reset the state of the current result set. + + Keeps the response and its line iterator, so that the next result set + can be read from the same response. + """ self._current_row_count = -1 self._current_statistics = None - self._lines_ite: Optional[Union[AsyncIterator[str], Iterator[str]]] = None self._current_record = None self._current_record_row_idx = -1 self._response_consumed = False @@ -97,13 +117,18 @@ def _next_json_lines_record_from_line( # Skip parsing floats to properly parse them later record = json.loads(next_line, parse_float=str) except json.JSONDecodeError as err: + self._response_failed = True raise OperationalError( f"Invalid JSON line response format: {next_line}" ) from err record = parse_json_lines_record(record) if isinstance(record, ErrorRecord): + if self._response_consumed: + # The error belongs to a statement that produced no START record + self._reset_result_set() self._response_consumed = True + self._response_failed = True self._current_statistics = record.statistics self._handle_error_record(record) return record @@ -138,13 +163,15 @@ def _fetch_columns_from_record( """ if record is None: self._response_consumed = True + self._response_failed = True raise OperationalError( "Unexpected end of response stream while reading columns." ) if not isinstance(record, StartRecord): self._response_consumed = True + self._response_failed = True raise OperationalError( - f"Unexpected json line message type {record.message_type.value}, " + f"Unexpected json line message type {_message_type_name(record)}, " "expected START" ) @@ -153,6 +180,46 @@ def _fetch_columns_from_record( for col in record.result_columns ] + def _discard_current_record(self) -> None: + """ + Drop the current data record, counting its unread rows 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: + """ + Count the rows of a data record that nextset() skipped over. + """ + if record is not None: + self._rows_returned += len(record.data) + + def _start_result_set_from_record(self, record: Optional[JSONLinesRecord]) -> bool: + """ + Start the next result set of the current response. + + Args: + record: The record that follows the end of the previous result set, + or None if the response stream is exhausted. + + Returns: + bool: True if a new result set was started, False if the response + carries no more result sets. + + Raises: + OperationalError: If the record is not a START record. + """ + if record is None: + return False + columns = self._fetch_columns_from_record(record) + self._reset_result_set() + self._current_columns = columns + return True + def _pop_data_record_from_record( self, record: Optional[JSONLinesRecord] ) -> Optional[DataRecord]: @@ -172,6 +239,7 @@ def _pop_data_record_from_record( if record is None: if not self._response_consumed: self._response_consumed = True + self._response_failed = True raise OperationalError( "Unexpected end of response stream while reading data." ) @@ -184,8 +252,9 @@ def _pop_data_record_from_record( self._response_consumed = True return None if not isinstance(record, DataRecord): + self._response_failed = True raise OperationalError( - f"Unexpected json line message type {record.message_type.value}, " + f"Unexpected json line message type {_message_type_name(record)}, " "expected DATA" ) return record diff --git a/src/firebolt/common/row_set/synchronous/in_memory.py b/src/firebolt/common/row_set/synchronous/in_memory.py index ddf0a4f78c9..5b84a824e2a 100644 --- a/src/firebolt/common/row_set/synchronous/in_memory.py +++ b/src/firebolt/common/row_set/synchronous/in_memory.py @@ -1,5 +1,6 @@ import json -from typing import Iterator, List, Optional +import re +from typing import Any, Dict, Iterator, List, Optional from httpx import Response @@ -8,6 +9,9 @@ from firebolt.common.row_set.types import Column, RowsResponse, Statistics from firebolt.utils.exception import DataError, FireboltStructuredError +# Whitespace as JSON defines it (RFC 8259), which is narrower than str.isspace() +_JSON_WHITESPACE = re.compile(r"[ \t\n\r]*") + class InMemoryRowSet(BaseSyncRowSet): """ @@ -41,29 +45,65 @@ def append_response_stream(self, stream: Iterator[bytes]) -> None: content = b"".join(stream).decode("utf-8") if len(content) == 0: self.append_empty_response() - else: + return + try: + row_sets = [ + self._parse_query_data(query_data) + for query_data in self._iter_documents(content) + ] + if not row_sets: + raise ValueError("no result set found in response") + except (KeyError, TypeError, ValueError) as err: + raise DataError(f"Invalid query data format: {str(err)}") + # Only a fully parsed response contributes result sets + self._row_sets.extend(row_sets) + + @staticmethod + def _iter_documents(content: str) -> Iterator[Dict[str, Any]]: + """ + Yield every JSON document of a response, one per result set. + """ + # Skip parsing floats to properly parse them later + decoder = json.JSONDecoder(parse_float=str) + position = 0 + while True: + whitespace = _JSON_WHITESPACE.match(content, position) + if whitespace is not None: + position = whitespace.end() + if position == len(content): + return try: - # Skip parsing floats to properly parse them later - query_data = json.loads(content, parse_float=str) - - if "errors" in query_data and len(query_data["errors"]) > 0: - raise FireboltStructuredError(query_data) - - columns = [ - Column( - d["name"], parse_type(d["type"]), None, None, None, None, None - ) - for d in query_data["meta"] - ] - # Extract rows - rows = query_data["data"] - row_count = len(rows) - statistics = Statistics(**query_data.get("statistics", {})) - self._row_sets.append( - RowsResponse(row_count, columns, statistics, rows) - ) - except (KeyError, ValueError) as err: - raise DataError(f"Invalid query data format: {str(err)}") + document, position = decoder.raw_decode(content, position) + except json.JSONDecodeError as err: + if position == 0: + raise + # Keep the text the engine wrote after a complete document, which + # is where it reports an error of a later statement + raise ValueError( + f"{err}; unexpected content after a result set: " + f"{content[position:position + 500]!r}" + ) from err + yield document + + @staticmethod + def _parse_query_data(query_data: Any) -> RowsResponse: + """ + Build a result set from a parsed JSON document. + """ + if not isinstance(query_data, dict): + raise ValueError(f"expected a JSON object, got {type(query_data).__name__}") + if "errors" in query_data and len(query_data["errors"]) > 0: + raise FireboltStructuredError(query_data) + + columns = [ + Column(d["name"], parse_type(d["type"]), None, None, None, None, None) + for d in query_data["meta"] + ] + # Extract rows + rows = query_data["data"] + row_count = len(rows) + statistics = Statistics(**query_data.get("statistics", {})) + return RowsResponse(row_count, columns, statistics, rows) @property def _row_set(self) -> RowsResponse: diff --git a/src/firebolt/common/row_set/synchronous/streaming.py b/src/firebolt/common/row_set/synchronous/streaming.py index b356720d669..6ec9d58003f 100644 --- a/src/firebolt/common/row_set/synchronous/streaming.py +++ b/src/firebolt/common/row_set/synchronous/streaming.py @@ -77,13 +77,18 @@ def _next_json_lines_record(self) -> Optional[JSONLinesRecord]: """ if self._current_response is None: return None - if self._lines_iter is None: - try: + try: + if self._lines_iter is None: self._lines_iter = self._current_response.iter_lines() - except HTTPError as err: - raise OperationalError("Failed to read response stream.") from err + while True: + next_line = next(self._lines_iter, None) + # A blank line separates nothing + if next_line is None or next_line.strip(): + break + except HTTPError as err: + self._response_failed = True + raise OperationalError("Failed to read response stream.") from err - next_line = next(self._lines_iter, None) return self._next_json_lines_record_from_line(next_line) @property @@ -136,13 +141,25 @@ def nextset(self) -> bool: """ Move to the next result set. + 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. + Returns: bool: True if there is a next result set, False otherwise Raises: OperationalError: If the response stream cannot be closed or if an error - occurs while fetching new columns + occurs while reading the response + FireboltStructuredError: If the response reports an error for the + current or the next result set """ + if ( + self._responses + and self._current_response is not None + and self._next_result_set_in_response() + ): + return True if self._current_row_set_idx + 1 < len(self._responses): if self._current_response is not None: try: @@ -156,6 +173,26 @@ def nextset(self) -> bool: return True return False + @close_on_op_error + 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. + + Returns: + bool: True if a new result set was started, False if the response + carries no more result sets + + Raises: + OperationalError: If an error occurs while reading the response + """ + if self._response_failed: + return False + self._discard_current_record() + while not self._response_consumed: + self._count_skipped_record(self._pop_data_record()) + return self._start_result_set_from_record(self._next_json_lines_record()) + @close_on_op_error def _pop_data_record(self) -> Optional[DataRecord]: """ @@ -185,11 +222,13 @@ def __next__(self) -> List[ColType]: if self._current_response is None or self._response_consumed: raise StopIteration - if self._current_record is None or self._current_record_row_idx >= len( + while self._current_record is None or self._current_record_row_idx >= len( self._current_record.data ): self._current_record = self._pop_data_record() self._current_record_row_idx = 0 + if self._current_record is None: + break return self._get_next_data_row_from_current_record(StopIteration) diff --git a/src/firebolt/db/cursor.py b/src/firebolt/db/cursor.py index f862adb8f4c..cdb74912584 100644 --- a/src/firebolt/db/cursor.py +++ b/src/firebolt/db/cursor.py @@ -472,10 +472,17 @@ def fetchall(self) -> List[List[ColType]]: def nextset(self) -> bool: """ Skip to the next available set, discarding any remaining rows - from the current set. + 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. Returns: bool: True if there is a next result set, False otherwise + + Raises: + OperationalError: If the response cannot be read + FireboltStructuredError: If the response reports an error for the + current or the next result set """ assert self._row_set is not None return self._row_set.nextset() diff --git a/tests/unit/async_db/test_cursor.py b/tests/unit/async_db/test_cursor.py index e59a2685027..36856d50a87 100644 --- a/tests/unit/async_db/test_cursor.py +++ b/tests/unit/async_db/test_cursor.py @@ -1938,3 +1938,78 @@ async def test_cursor_plaintext_error( assert "Plaintext error message" in str( excinfo.value ), "Invalid error message for plaintext error response" + + +async def test_cursor_execute_several_result_sets( + httpx_mock: HTTPXMock, + query_url: str, + query_statistics: Dict[str, Any], + cursor: Cursor, +): + """A response carrying several result sets is walked with nextset.""" + first = { + "meta": [{"name": "a", "type": "int"}], + "data": [[1]], + "rows": 1, + "statistics": query_statistics, + } + second = { + "meta": [{"name": "b", "type": "text"}], + "data": [["x"], ["y"]], + "rows": 2, + "statistics": query_statistics, + } + httpx_mock.add_response( + url=query_url, content=(json.dumps(first) + "\n" + json.dumps(second)).encode() + ) + await cursor.execute("select 1; select 'x'", skip_parsing=True) + assert cursor.rowcount == 1 + assert [d.name for d in cursor.description] == ["a"] + assert await cursor.fetchall() == [[1]] + assert await cursor.nextset() is True + assert cursor.rowcount == 2 + assert [d.name for d in cursor.description] == ["b"] + assert await cursor.fetchall() == [["x"], ["y"]] + assert await cursor.nextset() is False + + +async def test_cursor_execute_stream_several_result_sets( + httpx_mock: HTTPXMock, + streaming_query_url: str, + query_statistics: Dict[str, Any], + cursor: Cursor, +): + """A streamed response carrying several result sets is walked with nextset.""" + lines = [ + { + "message_type": "START", + "result_columns": [{"name": "a", "type": "integer"}], + "query_id": "q1", + "query_label": "l", + "request_id": "r", + }, + {"message_type": "DATA", "data": [[1]]}, + {"message_type": "FINISH_SUCCESSFULLY", "statistics": query_statistics}, + { + "message_type": "START", + "result_columns": [{"name": "b", "type": "text"}], + "query_id": "q2", + "query_label": "l", + "request_id": "r", + }, + {"message_type": "DATA", "data": [["x"], ["y"]]}, + {"message_type": "FINISH_SUCCESSFULLY", "statistics": query_statistics}, + ] + httpx_mock.add_response( + url=streaming_query_url, + content="\n".join(json.dumps(line) for line in lines).encode(), + ) + await cursor.execute_stream("select 1; select 'x'", skip_parsing=True) + assert [d.name for d in cursor.description] == ["a"] + assert await cursor.fetchall() == [[1]] + assert cursor.rowcount == 1 + assert await cursor.nextset() is True + assert [d.name for d in cursor.description] == ["b"] + assert await cursor.fetchall() == [["x"], ["y"]] + assert cursor.rowcount == 2 + assert await cursor.nextset() is False diff --git a/tests/unit/common/row_set/asynchronous/test_in_memory.py b/tests/unit/common/row_set/asynchronous/test_in_memory.py index 93fe55b31e7..8010766fb36 100644 --- a/tests/unit/common/row_set/asynchronous/test_in_memory.py +++ b/tests/unit/common/row_set/asynchronous/test_in_memory.py @@ -6,7 +6,7 @@ from httpx import Response from firebolt.common.row_set.asynchronous.in_memory import InMemoryAsyncRowSet -from firebolt.utils.exception import DataError +from firebolt.utils.exception import DataError, FireboltStructuredError class TestInMemoryAsyncRowSet: @@ -454,3 +454,75 @@ async def test_append_response_with_decimals( assert ( str(row[2]) == "1231232.123459999990457054844258706536" ), "Decimal value mismatch" + + +def _document(columns, rows) -> str: + statistics = { + "elapsed": 0.1, + "rows_read": len(rows), + "bytes_read": 100, + "time_before_execution": 0.01, + "time_to_execute": 0.09, + } + return json.dumps({"meta": columns, "data": rows, "statistics": statistics}) + + +FIRST_DOCUMENT = _document([{"name": "col1", "type": "int"}], [[1], [2]]) +SECOND_DOCUMENT = _document( + [{"name": "col2", "type": "string"}, {"name": "col3", "type": "int"}], [["a", 3]] +) +ERROR_DOCUMENT = json.dumps({"errors": [{"message": "error"}]}) + + +def _response(content: str) -> Response: + mock = MagicMock(spec=Response) + + async def aiter_bytes(): + yield content.encode("utf-8") + + mock.aiter_bytes.return_value = aiter_bytes() + return mock + + +async def _rows(row_set: InMemoryAsyncRowSet): + return [row async for row in row_set] + + +async def test_several_result_sets_in_one_response(): + """All result sets carried by one response are available through nextset.""" + row_set = InMemoryAsyncRowSet() + await row_set.append_response(_response(FIRST_DOCUMENT + "\n" + SECOND_DOCUMENT)) + + assert [c.name for c in row_set.columns] == ["col1"] + assert row_set.row_count == 2 + assert await _rows(row_set) == [[1], [2]] + assert row_set.statistics.rows_read == 2 + + assert await row_set.nextset() is True + assert [c.name for c in row_set.columns] == ["col2", "col3"] + assert row_set.row_count == 1 + assert await _rows(row_set) == [["a", 3]] + assert row_set.statistics.rows_read == 1 + + assert await row_set.nextset() is False + + +async def test_error_after_result_set(): + """An error document after a result set is reported as an error.""" + row_set = InMemoryAsyncRowSet() + with pytest.raises(FireboltStructuredError): + await row_set.append_response(_response(FIRST_DOCUMENT + "\n" + ERROR_DOCUMENT)) + + +async def test_invalid_data_after_result_set(): + """Anything but a JSON document after a result set is a data error.""" + row_set = InMemoryAsyncRowSet() + with pytest.raises(DataError): + await row_set.append_response(_response(FIRST_DOCUMENT + "\nnot json")) + + +async def test_non_object_document(): + """A document that is not a JSON object is a data error, not a TypeError.""" + row_set = InMemoryAsyncRowSet() + with pytest.raises(DataError, match="expected a JSON object"): + await row_set.append_response(_response(FIRST_DOCUMENT + "\n[1]")) diff --git a/tests/unit/common/row_set/asynchronous/test_streaming.py b/tests/unit/common/row_set/asynchronous/test_streaming.py index 754ce219913..77a4d977248 100644 --- a/tests/unit/common/row_set/asynchronous/test_streaming.py +++ b/tests/unit/common/row_set/asynchronous/test_streaming.py @@ -2,6 +2,7 @@ from decimal import Decimal from unittest.mock import MagicMock, PropertyMock, patch +import httpx import pytest from httpx import HTTPError, Response @@ -249,14 +250,16 @@ async def test_nextset_no_more_sets( start_record, ): """Test nextset when there are no more result sets.""" - # Setup mocks - mock_next_json_lines_record.return_value = start_record + # Setup mocks: the current result set is fully read and the response + # stream is exhausted after it + mock_next_json_lines_record.return_value = None mock_fetch_columns_from_record.return_value = [ Column("col1", int, None, None, None, None, None) ] streaming_rowset._responses = [mock_response] streaming_rowset._current_columns = mock_fetch_columns_from_record.return_value + streaming_rowset._response_consumed = True assert await streaming_rowset.nextset() is False @@ -276,11 +279,14 @@ async def test_nextset_with_more_sets(self, mock_fetch_columns, streaming_rowset # Setup mocks mock_fetch_columns.side_effect = [columns1, columns2] - # Setup two responses + # Setup two responses, each carrying a single, fully read result set response1 = MagicMock(spec=Response) + response1.aiter_lines.return_value = self._async_iter([]) response2 = MagicMock(spec=Response) + response2.aiter_lines.return_value = self._async_iter([]) streaming_rowset._responses = [response1, response2] streaming_rowset._current_columns = columns1 + streaming_rowset._response_consumed = True # Verify first result set assert streaming_rowset.columns[0].name == "col1" @@ -297,6 +303,7 @@ async def test_nextset_with_more_sets(self, mock_fetch_columns, streaming_rowset # Update current columns to match what mock_fetch_columns returned streaming_rowset._current_columns = columns2 + streaming_rowset._response_consumed = True # Verify second result set assert streaming_rowset.columns[0].name == "col3" @@ -1020,3 +1027,242 @@ async def test_streaming_with_decimals( assert ( str(row[2]) == "1231232.123459999990457054844258706536" ), "Decimal value mismatch" + + +def _start(columns) -> str: + return json.dumps( + { + "message_type": "START", + "result_columns": columns, + "query_id": "q", + "query_label": "l", + "request_id": "r", + } + ) + + +def _data(rows) -> str: + return json.dumps({"message_type": "DATA", "data": rows}) + + +def _finish(rows_read: int) -> str: + return json.dumps( + { + "message_type": "FINISH_SUCCESSFULLY", + "statistics": { + "elapsed": 0.1, + "rows_read": rows_read, + "bytes_read": 100, + "time_before_execution": 0.01, + "time_to_execute": 0.09, + }, + } + ) + + +FINISH_WITH_ERRORS = json.dumps( + { + "message_type": "FINISH_WITH_ERRORS", + "errors": [{"message": "error"}], + "query_id": "q", + "query_label": "l", + "request_id": "r", + "statistics": { + "elapsed": 0.1, + "rows_read": 0, + "bytes_read": 0, + "time_before_execution": 0.01, + "time_to_execute": 0.09, + }, + } +) + +FIRST_COLUMNS = [{"name": "col1", "type": "int"}] +SECOND_COLUMNS = [{"name": "col2", "type": "string"}, {"name": "col3", "type": "int"}] +TWO_RESULT_SETS = [ + _start(FIRST_COLUMNS), + _data([[1], [2]]), + _finish(2), + _start(SECOND_COLUMNS), + _data([["a", 3]]), + _data([["b", 4]]), + _finish(2), +] + + +def _response(lines) -> Response: + mock = MagicMock(spec=Response) + + async def aiter_lines(): + for line in lines: + yield line + + mock.aiter_lines.return_value = aiter_lines() + mock.is_closed = False + return mock + + +async def _rows(row_set: StreamingAsyncRowSet): + return [row async for row in row_set] + + +async def test_several_result_sets_in_one_response(): + """All result sets carried by one response are available through nextset.""" + row_set = StreamingAsyncRowSet() + await row_set.append_response(_response(TWO_RESULT_SETS)) + + assert [c.name for c in row_set.columns] == ["col1"] + assert row_set.row_count == -1 + assert await _rows(row_set) == [[1], [2]] + assert row_set.row_count == 2 + assert row_set.statistics.rows_read == 2 + + assert await row_set.nextset() is True + assert [c.name for c in row_set.columns] == ["col2", "col3"] + assert row_set.row_count == -1 + assert await _rows(row_set) == [["a", 3], ["b", 4]] + assert row_set.row_count == 2 + + assert await row_set.nextset() is False + + +async def test_nextset_skips_remaining_rows_of_result_set(): + """nextset discards whatever is left of the current result set.""" + row_set = StreamingAsyncRowSet() + await row_set.append_response(_response(TWO_RESULT_SETS)) + assert await row_set.__anext__() == [1] + assert await row_set.nextset() is True + assert await _rows(row_set) == [["a", 3], ["b", 4]] + assert await row_set.nextset() is False + + +async def test_result_sets_across_several_responses(): + """Result sets of one response come before those of the next response.""" + row_set = StreamingAsyncRowSet() + first = _response(TWO_RESULT_SETS) + second = _response([_start(FIRST_COLUMNS), _data([[5]]), _finish(1)]) + await row_set.append_response(first) + await row_set.append_response(second) + + assert await _rows(row_set) == [[1], [2]] + assert await row_set.nextset() is True + assert await _rows(row_set) == [["a", 3], ["b", 4]] + first.aclose.assert_not_called() + + assert await row_set.nextset() is True + first.aclose.assert_called_once() + assert await _rows(row_set) == [[5]] + assert await row_set.nextset() is False + + +async def test_error_in_later_result_set(): + """An error of a later result set is raised when its rows are read.""" + row_set = StreamingAsyncRowSet() + await row_set.append_response( + _response( + [ + _start(FIRST_COLUMNS), + _data([[1]]), + _finish(1), + _start(SECOND_COLUMNS), + FINISH_WITH_ERRORS, + ] + ) + ) + assert await _rows(row_set) == [[1]] + assert await row_set.nextset() is True + with pytest.raises(FireboltStructuredError): + await row_set.__anext__() + assert await row_set.nextset() is False + + +async def test_truncated_result_set_on_nextset(): + """A result set cut before its end is an error on nextset.""" + row_set = StreamingAsyncRowSet() + await row_set.append_response(_response([_start(FIRST_COLUMNS), _data([[1]])])) + with pytest.raises(OperationalError, match="Unexpected end of response stream"): + await row_set.nextset() + + +async def test_unexpected_record_after_result_set(): + """A record other than START after a result set is an error on nextset.""" + row_set = StreamingAsyncRowSet() + await row_set.append_response( + _response([_start(FIRST_COLUMNS), _finish(0), _data([[1]])]) + ) + with pytest.raises(OperationalError, match="expected START"): + await row_set.nextset() + + +class _FailingStream(httpx.AsyncByteStream): + """A body that breaks after its payload, like a connection cut mid-transfer.""" + + def __init__(self, payload: bytes) -> None: + self._payload = payload + + async def __aiter__(self): + yield self._payload + raise httpx.RemoteProtocolError("peer closed connection") + + +@pytest.mark.parametrize("rows_fetched", [0, 1]) +async def test_row_count_counts_rows_skipped_by_nextset(rows_fetched: int): + """row_count covers the rows nextset discarded, not only the fetched ones.""" + 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 row_set.row_count == 3 + + +async def test_transport_error_after_last_record(): + """A broken connection after the last record surfaces as OperationalError.""" + row_set = StreamingAsyncRowSet() + body = "\n".join([_start(FIRST_COLUMNS), _data([[1]]), _finish(1)]) + "\n" + response = httpx.Response(200, stream=_FailingStream(body.encode("utf-8"))) + await row_set.append_response(response) + assert await _rows(row_set) == [[1]] + with pytest.raises(OperationalError, match="Failed to read response stream"): + await row_set.nextset() + assert response.is_closed + + +async def test_blank_lines_between_and_after_result_sets(): + """Blank lines in the stream separate nothing and break nothing.""" + row_set = StreamingAsyncRowSet() + body = "\n".join([_start(FIRST_COLUMNS), _data([[1]]), _finish(1), "", " "]) + body += "\n" + "\n".join([_start(SECOND_COLUMNS), _data([["a", 3]]), _finish(1)]) + body += "\n\n" + await row_set.append_response(httpx.Response(200, content=body.encode("utf-8"))) + assert await _rows(row_set) == [[1]] + assert await row_set.nextset() is True + assert await _rows(row_set) == [["a", 3]] + assert await row_set.nextset() is False + + +async def test_error_record_without_start_after_result_set(): + """An error of a statement that produced no START is reported by nextset.""" + row_set = StreamingAsyncRowSet() + await row_set.append_response( + _response([_start(FIRST_COLUMNS), _data([[1]]), _finish(1), FINISH_WITH_ERRORS]) + ) + assert await _rows(row_set) == [[1]] + with pytest.raises(FireboltStructuredError): + await row_set.nextset() + assert row_set.row_count == -1 + assert row_set.columns is None + + +async def test_empty_data_records_are_skipped(): + """A DATA record without rows does not end or break the iteration.""" + row_set = StreamingAsyncRowSet() + await row_set.append_response( + _response( + [_start(FIRST_COLUMNS), _data([]), _data([[1]]), _data([]), _finish(1)] + ) + ) + assert await _rows(row_set) == [[1]] + assert row_set.row_count == 1 diff --git a/tests/unit/common/row_set/synchronous/test_in_memory.py b/tests/unit/common/row_set/synchronous/test_in_memory.py index 09ad02a06be..bcd6068e99a 100644 --- a/tests/unit/common/row_set/synchronous/test_in_memory.py +++ b/tests/unit/common/row_set/synchronous/test_in_memory.py @@ -6,7 +6,7 @@ from httpx import Response from firebolt.common.row_set.synchronous.in_memory import InMemoryRowSet -from firebolt.utils.exception import DataError +from firebolt.utils.exception import DataError, FireboltStructuredError class TestInMemoryRowSet: @@ -450,3 +450,135 @@ def test_parse_row(self, in_memory_rowset, mock_response): assert isinstance(parsed_row[1], str) assert parsed_row[0] == 1 assert parsed_row[1] == "one" + + +def _document(columns, rows) -> str: + statistics = { + "elapsed": 0.1, + "rows_read": len(rows), + "bytes_read": 100, + "time_before_execution": 0.01, + "time_to_execute": 0.09, + } + return json.dumps({"meta": columns, "data": rows, "statistics": statistics}) + + +FIRST_DOCUMENT = _document([{"name": "col1", "type": "int"}], [[1], [2]]) +SECOND_DOCUMENT = _document( + [{"name": "col2", "type": "string"}, {"name": "col3", "type": "int"}], [["a", 3]] +) +EMPTY_DOCUMENT = _document([], []) +ERROR_DOCUMENT = json.dumps({"errors": [{"message": "error"}]}) + + +def _response(content: str) -> Response: + mock = MagicMock(spec=Response) + mock.iter_bytes.return_value = iter([content.encode("utf-8")]) + return mock + + +def _assert_two_result_sets(row_set: InMemoryRowSet) -> None: + assert [c.name for c in row_set.columns] == ["col1"] + assert row_set.row_count == 2 + assert list(row_set) == [[1], [2]] + assert row_set.statistics.rows_read == 2 + + assert row_set.nextset() is True + assert [c.name for c in row_set.columns] == ["col2", "col3"] + assert row_set.row_count == 1 + assert list(row_set) == [["a", 3]] + assert row_set.statistics.rows_read == 1 + + assert row_set.nextset() is False + + +@pytest.mark.parametrize("separator", ["", "\n", "\n\n "]) +def test_several_result_sets_in_one_response(separator: str): + """All result sets carried by one response are available through nextset.""" + row_set = InMemoryRowSet() + row_set.append_response(_response(FIRST_DOCUMENT + separator + SECOND_DOCUMENT)) + _assert_two_result_sets(row_set) + + +def test_result_sets_in_one_response_then_another_response(): + """Result sets of one response come before those of the next response.""" + row_set = InMemoryRowSet() + row_set.append_response(_response(FIRST_DOCUMENT + "\n" + SECOND_DOCUMENT)) + row_set.append_response(_response(EMPTY_DOCUMENT)) + assert list(row_set) == [[1], [2]] + assert row_set.nextset() is True + assert list(row_set) == [["a", 3]] + assert row_set.nextset() is True + assert row_set.columns == [] + assert row_set.row_count == 0 + assert row_set.nextset() is False + + +def test_empty_result_set_between_others(): + """A result set without columns does not break the following ones.""" + row_set = InMemoryRowSet() + row_set.append_response( + _response(FIRST_DOCUMENT + "\n" + EMPTY_DOCUMENT + "\n" + SECOND_DOCUMENT) + ) + assert list(row_set) == [[1], [2]] + assert row_set.nextset() is True + assert row_set.columns == [] + assert row_set.row_count == 0 + assert row_set.nextset() is True + assert list(row_set) == [["a", 3]] + assert row_set.nextset() is False + + +def test_error_after_result_set(): + """An error document after a result set is reported as an error.""" + row_set = InMemoryRowSet() + with pytest.raises(FireboltStructuredError): + row_set.append_response(_response(FIRST_DOCUMENT + "\n" + ERROR_DOCUMENT)) + + +def test_invalid_data_after_result_set(): + """Anything but a JSON document after a result set is a data error.""" + row_set = InMemoryRowSet() + with pytest.raises(DataError): + row_set.append_response(_response(FIRST_DOCUMENT + "\nnot json")) + + +def test_whitespace_only_response(): + """A response without any document is a data error.""" + row_set = InMemoryRowSet() + with pytest.raises(DataError): + row_set.append_response(_response(" \n")) + + +def test_failed_parse_leaves_no_result_set(): + """A response that fails to parse contributes no result set at all.""" + row_set = InMemoryRowSet() + with pytest.raises(DataError): + row_set.append_response(_response(FIRST_DOCUMENT + "\nnot json")) + assert row_set.nextset() is False + with pytest.raises(DataError, match="No results available"): + row_set.columns + + +@pytest.mark.parametrize("document", ["[1]", "1", "null", '"text"']) +def test_non_object_document(document: str): + """A document that is not a JSON object is a data error, not a TypeError.""" + row_set = InMemoryRowSet() + with pytest.raises(DataError, match="expected a JSON object"): + row_set.append_response(_response(FIRST_DOCUMENT + "\n" + document)) + + +def test_only_json_whitespace_separates_documents(): + """Whitespace JSON does not know, such as a no-break space, is a data error.""" + row_set = InMemoryRowSet() + with pytest.raises(DataError): + row_set.append_response(_response(FIRST_DOCUMENT + "\u00a0" + SECOND_DOCUMENT)) + + +def test_invalid_content_after_result_set_is_reported(): + """The engine's text after a complete document is kept in the error.""" + row_set = InMemoryRowSet() + with pytest.raises(DataError, match="unexpected content after a result set"): + row_set.append_response( + _response(FIRST_DOCUMENT + '\n\t],\n "errors": [{"description": "x"}]}') + ) diff --git a/tests/unit/common/row_set/synchronous/test_streaming.py b/tests/unit/common/row_set/synchronous/test_streaming.py index deae15cd412..5fb54d95cde 100644 --- a/tests/unit/common/row_set/synchronous/test_streaming.py +++ b/tests/unit/common/row_set/synchronous/test_streaming.py @@ -2,6 +2,7 @@ from decimal import Decimal from unittest.mock import MagicMock, PropertyMock, patch +import httpx import pytest from httpx import HTTPError, Response @@ -253,14 +254,16 @@ def test_nextset_no_more_sets( start_record, ): """Test nextset when there are no more result sets.""" - # Setup mocks - mock_next_json_lines_record.return_value = start_record + # Setup mocks: the current result set is fully read and the response + # stream is exhausted after it + mock_next_json_lines_record.return_value = None mock_fetch_columns_from_record.return_value = [ Column("col1", int, None, None, None, None, None) ] streaming_rowset._responses = [mock_response] streaming_rowset._current_columns = mock_fetch_columns_from_record.return_value + streaming_rowset._response_consumed = True assert streaming_rowset.nextset() is False @@ -280,11 +283,14 @@ def test_nextset_with_more_sets(self, mock_fetch_columns, streaming_rowset): # Setup mocks mock_fetch_columns.side_effect = [columns1, columns2] - # Setup two responses + # Setup two responses, each carrying a single, fully read result set response1 = MagicMock(spec=Response) + response1.iter_lines.return_value = iter([]) response2 = MagicMock(spec=Response) + response2.iter_lines.return_value = iter([]) streaming_rowset._responses = [response1, response2] streaming_rowset._current_columns = columns1 + streaming_rowset._response_consumed = True # Verify first result set assert streaming_rowset.columns[0].name == "col1" @@ -301,6 +307,7 @@ def test_nextset_with_more_sets(self, mock_fetch_columns, streaming_rowset): # Update current columns to match what mock_fetch_columns returned streaming_rowset._current_columns = columns2 + streaming_rowset._response_consumed = True # Verify second result set assert streaming_rowset.columns[0].name == "col3" @@ -1016,3 +1023,257 @@ def test_streaming_with_decimals( assert ( str(row[2]) == "1231232.123459999990457054844258706536" ), "Decimal value mismatch" + + +def _start(columns) -> str: + return json.dumps( + { + "message_type": "START", + "result_columns": columns, + "query_id": "q", + "query_label": "l", + "request_id": "r", + } + ) + + +def _data(rows) -> str: + return json.dumps({"message_type": "DATA", "data": rows}) + + +def _finish(rows_read: int) -> str: + return json.dumps( + { + "message_type": "FINISH_SUCCESSFULLY", + "statistics": { + "elapsed": 0.1, + "rows_read": rows_read, + "bytes_read": 100, + "time_before_execution": 0.01, + "time_to_execute": 0.09, + }, + } + ) + + +FINISH_WITH_ERRORS = json.dumps( + { + "message_type": "FINISH_WITH_ERRORS", + "errors": [{"message": "error"}], + "query_id": "q", + "query_label": "l", + "request_id": "r", + "statistics": { + "elapsed": 0.1, + "rows_read": 0, + "bytes_read": 0, + "time_before_execution": 0.01, + "time_to_execute": 0.09, + }, + } +) + +FIRST_COLUMNS = [{"name": "col1", "type": "int"}] +SECOND_COLUMNS = [{"name": "col2", "type": "string"}, {"name": "col3", "type": "int"}] +TWO_RESULT_SETS = [ + _start(FIRST_COLUMNS), + _data([[1], [2]]), + _finish(2), + _start(SECOND_COLUMNS), + _data([["a", 3]]), + _data([["b", 4]]), + _finish(2), +] + + +def _response(lines) -> Response: + mock = MagicMock(spec=Response) + mock.iter_lines.return_value = iter(lines) + mock.is_closed = False + return mock + + +def _assert_two_result_sets(row_set: StreamingRowSet) -> None: + assert [c.name for c in row_set.columns] == ["col1"] + assert row_set.row_count == -1 + assert list(row_set) == [[1], [2]] + assert row_set.row_count == 2 + assert row_set.statistics.rows_read == 2 + + assert row_set.nextset() is True + assert [c.name for c in row_set.columns] == ["col2", "col3"] + assert row_set.row_count == -1 + assert list(row_set) == [["a", 3], ["b", 4]] + assert row_set.row_count == 2 + assert row_set.statistics.rows_read == 2 + + assert row_set.nextset() is False + + +def test_several_result_sets_in_one_response(): + """All result sets carried by one response are available through nextset.""" + row_set = StreamingRowSet() + row_set.append_response(_response(TWO_RESULT_SETS)) + _assert_two_result_sets(row_set) + + +def test_nextset_skips_remaining_rows_of_result_set(): + """nextset discards whatever is left of the current result set.""" + row_set = StreamingRowSet() + row_set.append_response(_response(TWO_RESULT_SETS)) + assert next(row_set) == [1] + assert row_set.nextset() is True + assert list(row_set) == [["a", 3], ["b", 4]] + assert row_set.nextset() is False + + +def test_empty_result_set_before_another(): + """A result set without columns does not break the following one.""" + row_set = StreamingRowSet() + row_set.append_response( + _response( + [_start([]), _finish(0), _start(FIRST_COLUMNS), _data([[1]]), _finish(1)] + ) + ) + assert row_set.columns == [] + assert list(row_set) == [] + assert row_set.row_count == 0 + assert row_set.nextset() is True + assert list(row_set) == [[1]] + assert row_set.nextset() is False + + +def test_result_sets_across_several_responses(): + """Result sets of one response come before those of the next response.""" + row_set = StreamingRowSet() + first = _response(TWO_RESULT_SETS) + second = _response([_start(FIRST_COLUMNS), _data([[5]]), _finish(1)]) + row_set.append_response(first) + row_set.append_response(second) + + assert list(row_set) == [[1], [2]] + assert row_set.nextset() is True + assert list(row_set) == [["a", 3], ["b", 4]] + first.close.assert_not_called() + + assert row_set.nextset() is True + first.close.assert_called_once() + assert [c.name for c in row_set.columns] == ["col1"] + assert list(row_set) == [[5]] + assert row_set.nextset() is False + + +def test_error_in_later_result_set(): + """An error of a later result set is raised when its rows are read.""" + row_set = StreamingRowSet() + row_set.append_response( + _response( + [ + _start(FIRST_COLUMNS), + _data([[1]]), + _finish(1), + _start(SECOND_COLUMNS), + FINISH_WITH_ERRORS, + ] + ) + ) + assert list(row_set) == [[1]] + assert row_set.nextset() is True + with pytest.raises(FireboltStructuredError): + next(row_set) + assert row_set.nextset() is False + + +def test_truncated_result_set_on_nextset(): + """A result set cut before its end is an error on nextset.""" + row_set = StreamingRowSet() + row_set.append_response(_response([_start(FIRST_COLUMNS), _data([[1]])])) + with pytest.raises(OperationalError, match="Unexpected end of response stream"): + row_set.nextset() + + +def test_unexpected_record_after_result_set(): + """A record other than START after a result set is an error on nextset.""" + row_set = StreamingRowSet() + row_set.append_response( + _response([_start(FIRST_COLUMNS), _finish(0), _data([[1]])]) + ) + with pytest.raises(OperationalError, match="expected START"): + row_set.nextset() + + +class _FailingStream(httpx.SyncByteStream): + """A body that breaks after its payload, like a connection cut mid-transfer.""" + + def __init__(self, payload: bytes) -> None: + self._payload = payload + + def __iter__(self): + yield self._payload + raise httpx.RemoteProtocolError("peer closed connection") + + +@pytest.mark.parametrize("rows_fetched", [0, 1]) +def test_row_count_counts_rows_skipped_by_nextset(rows_fetched: int): + """row_count covers the rows nextset discarded, not only the fetched ones.""" + 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 row_set.row_count == 3 + + +def test_transport_error_after_last_record(): + """A broken connection after the last record surfaces as OperationalError.""" + row_set = StreamingRowSet() + body = "\n".join([_start(FIRST_COLUMNS), _data([[1]]), _finish(1)]) + "\n" + response = httpx.Response(200, stream=_FailingStream(body.encode("utf-8"))) + row_set.append_response(response) + assert list(row_set) == [[1]] + with pytest.raises(OperationalError, match="Failed to read response stream"): + row_set.nextset() + assert response.is_closed + + +def test_blank_lines_between_and_after_result_sets(): + """Blank lines in the stream separate nothing and break nothing.""" + row_set = StreamingRowSet() + body = "\n".join([_start(FIRST_COLUMNS), _data([[1]]), _finish(1), "", " "]) + body += "\n" + "\n".join([_start(SECOND_COLUMNS), _data([["a", 3]]), _finish(1)]) + body += "\n\n" + row_set.append_response(httpx.Response(200, content=body.encode("utf-8"))) + assert list(row_set) == [[1]] + assert row_set.nextset() is True + assert list(row_set) == [["a", 3]] + assert row_set.nextset() is False + + +def test_error_record_without_start_after_result_set(): + """An error of a statement that produced no START is reported by nextset.""" + row_set = StreamingRowSet() + row_set.append_response( + _response([_start(FIRST_COLUMNS), _data([[1]]), _finish(1), FINISH_WITH_ERRORS]) + ) + assert list(row_set) == [[1]] + assert row_set.row_count == 1 + with pytest.raises(FireboltStructuredError): + row_set.nextset() + # The row set reflects the failed statement, not the finished one + assert row_set.row_count == -1 + assert row_set.columns is None + assert row_set.statistics.rows_read == 0 + + +def test_empty_data_records_are_skipped(): + """A DATA record without rows does not end or break the iteration.""" + row_set = StreamingRowSet() + row_set.append_response( + _response( + [_start(FIRST_COLUMNS), _data([]), _data([[1]]), _data([]), _finish(1)] + ) + ) + assert list(row_set) == [[1]] + assert row_set.row_count == 1 diff --git a/tests/unit/db/test_cursor.py b/tests/unit/db/test_cursor.py index 8d7ddf88aab..a6f84e7c4f0 100644 --- a/tests/unit/db/test_cursor.py +++ b/tests/unit/db/test_cursor.py @@ -1826,3 +1826,78 @@ def test_cursor_plaintext_error( assert "Plaintext error message" in str( excinfo.value ), "Invalid error message for plaintext error response" + + +def test_cursor_execute_several_result_sets( + httpx_mock: HTTPXMock, + query_url: str, + query_statistics: Dict[str, Any], + cursor: Cursor, +): + """A response carrying several result sets is walked with nextset.""" + first = { + "meta": [{"name": "a", "type": "int"}], + "data": [[1]], + "rows": 1, + "statistics": query_statistics, + } + second = { + "meta": [{"name": "b", "type": "text"}], + "data": [["x"], ["y"]], + "rows": 2, + "statistics": query_statistics, + } + httpx_mock.add_response( + url=query_url, content=(json.dumps(first) + "\n" + json.dumps(second)).encode() + ) + cursor.execute("select 1; select 'x'", skip_parsing=True) + assert cursor.rowcount == 1 + assert [d.name for d in cursor.description] == ["a"] + assert cursor.fetchall() == [[1]] + assert cursor.nextset() is True + assert cursor.rowcount == 2 + assert [d.name for d in cursor.description] == ["b"] + assert cursor.fetchall() == [["x"], ["y"]] + assert cursor.nextset() is False + + +def test_cursor_execute_stream_several_result_sets( + httpx_mock: HTTPXMock, + streaming_query_url: str, + query_statistics: Dict[str, Any], + cursor: Cursor, +): + """A streamed response carrying several result sets is walked with nextset.""" + lines = [ + { + "message_type": "START", + "result_columns": [{"name": "a", "type": "integer"}], + "query_id": "q1", + "query_label": "l", + "request_id": "r", + }, + {"message_type": "DATA", "data": [[1]]}, + {"message_type": "FINISH_SUCCESSFULLY", "statistics": query_statistics}, + { + "message_type": "START", + "result_columns": [{"name": "b", "type": "text"}], + "query_id": "q2", + "query_label": "l", + "request_id": "r", + }, + {"message_type": "DATA", "data": [["x"], ["y"]]}, + {"message_type": "FINISH_SUCCESSFULLY", "statistics": query_statistics}, + ] + httpx_mock.add_response( + url=streaming_query_url, + content="\n".join(json.dumps(line) for line in lines).encode(), + ) + cursor.execute_stream("select 1; select 'x'", skip_parsing=True) + assert [d.name for d in cursor.description] == ["a"] + assert cursor.fetchall() == [[1]] + assert cursor.rowcount == 1 + assert cursor.nextset() is True + assert [d.name for d in cursor.description] == ["b"] + assert cursor.fetchall() == [["x"], ["y"]] + assert cursor.rowcount == 2 + assert cursor.nextset() is False