Add decode_continuation() to MessageDecoder for multi-part ROWS responses

When a server streams large result sets with PART markers, continuation
messages have the same ROWS type code but no column_count prefix. Calling
decode() on these would misinterpret row data as a column count. The new
decode_continuation() method reads the next message and delegates to
RowsResponse.decode_rows_continuation() with the caller-provided column
metadata from the initial response.

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

Author: antoineleclair

src/dqlitewire/codec.py

@@ -178,6 +178,36 @@ def is_skipping(self) -> bool:
         """True if still discarding bytes from an oversized message."""
         return self._buffer.is_skipping
 
+    def decode_continuation(
+        self,
+        column_names: list[str],
+        column_count: int,
+        max_rows: int = RowsResponse.DEFAULT_MAX_ROWS,
+    ) -> RowsResponse | None:
+        """Decode a ROWS continuation message from the buffer.
+
+        After receiving a RowsResponse with has_more=True, the server sends
+        additional ROWS messages containing only row data (no column_count or
+        column_names prefix). Use this method instead of decode() for those
+        continuation messages.
+
+        Args:
+            column_names: Column names from the initial RowsResponse.
+            column_count: Number of columns from the initial RowsResponse.
+            max_rows: Maximum rows to decode per message.
+
+        Returns None if no complete message is available.
+        """
+        data = self._buffer.read_message()
+        if data is None:
+            return None
+
+        header = Header.decode(data[:HEADER_SIZE])
+        body = data[HEADER_SIZE : HEADER_SIZE + header.size_words * 8]
+        return RowsResponse.decode_rows_continuation(
+            body, column_names, column_count, max_rows=max_rows
+        )
+
     def decode(self) -> Message | None:
         """Decode the next message from the buffer.
 

tests/test_codec.py

@@ -553,6 +553,68 @@ def test_stmt_response_v0_with_trailing_data_not_detected_as_v1(self) -> None:
         assert decoded.tail_offset is None
 
 
+class TestDecoderContinuation:
+    """Test decode_continuation() for multi-part ROWS responses."""
+
+    def test_decode_continuation_exists(self) -> None:
+        """MessageDecoder should have a decode_continuation method."""
+        decoder = MessageDecoder(is_request=False)
+        assert hasattr(decoder, "decode_continuation")
+
+    def test_decode_continuation_roundtrip(self) -> None:
+        """Multi-part ROWS: initial decode() + continuation decode_continuation()."""
+        from dqlitewire.constants import ROW_DONE_MARKER, ROW_PART_MARKER, ValueType
+        from dqlitewire.messages.base import Header
+        from dqlitewire.messages.responses import RowsResponse
+        from dqlitewire.tuples import encode_row_header, encode_row_values
+        from dqlitewire.types import encode_text, encode_uint64
+
+        types = [ValueType.INTEGER, ValueType.TEXT]
+
+        # Build initial ROWS message with PART marker
+        body1 = encode_uint64(2)  # column_count
+        body1 += encode_text("id") + encode_text("name")
+        body1 += encode_row_header(types)
+        body1 += encode_row_values([1, "alice"], types)
+        body1 += encode_uint64(ROW_PART_MARKER)
+        header1 = Header(size_words=len(body1) // 8, msg_type=7, schema=0)
+        msg1_bytes = header1.encode() + body1
+
+        # Build continuation ROWS message with DONE marker
+        body2 = encode_row_header(types)
+        body2 += encode_row_values([2, "bob"], types)
+        body2 += encode_uint64(ROW_DONE_MARKER)
+        header2 = Header(size_words=len(body2) // 8, msg_type=7, schema=0)
+        msg2_bytes = header2.encode() + body2
+
+        # Feed both messages
+        decoder = MessageDecoder(is_request=False)
+        decoder.feed(msg1_bytes + msg2_bytes)
+
+        # Decode initial response
+        initial = decoder.decode()
+        assert isinstance(initial, RowsResponse)
+        assert initial.has_more is True
+        assert len(initial.rows) == 1
+        assert initial.rows[0] == [1, "alice"]
+
+        # Decode continuation
+        continuation = decoder.decode_continuation(
+            column_names=initial.column_names,
+            column_count=len(initial.column_names),
+        )
+        assert isinstance(continuation, RowsResponse)
+        assert continuation.has_more is False
+        assert len(continuation.rows) == 1
+        assert continuation.rows[0] == [2, "bob"]
+
+    def test_decode_continuation_returns_none_when_no_data(self) -> None:
+        """decode_continuation should return None when no message is available."""
+        decoder = MessageDecoder(is_request=False)
+        result = decoder.decode_continuation(column_names=["x"], column_count=1)
+        assert result is None
+
+
 class TestDecoderSkipMessage:
     """Test skip_message() and is_skipping on MessageDecoder."""