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