fix(FB-5242): read every result set a response carries

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()

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)
 

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 @@
         """
         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 @@
             # 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 @@
         """
         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 @@
             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 @@
         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 @@
             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

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: