Skip to content

Commit b3d0926

Browse files
Validate column_names vs column_count in decode_rows_continuation
A mismatch between len(column_names) and column_count would produce a RowsResponse with inconsistent data. Now raises DecodeError early if they disagree. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
1 parent 7f18910 commit b3d0926

2 files changed

Lines changed: 26 additions & 0 deletions

File tree

src/dqlitewire/messages/responses.py

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -378,6 +378,11 @@ def decode_rows_continuation(
378378
method to decode those continuation messages, passing the column_names
379379
and column_count from the initial response.
380380
"""
381+
if len(column_names) != column_count:
382+
raise DecodeError(
383+
f"column_names length ({len(column_names)}) does not match "
384+
f"column_count ({column_count})"
385+
)
381386
offset = 0
382387
rows: list[list[Any]] = []
383388
all_row_types: list[list[ValueType]] = []

tests/test_messages_responses.py

Lines changed: 21 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -512,6 +512,27 @@ def build_body(n_rows: int) -> bytes:
512512
decoded = RowsResponse.decode_body(build_body(2), max_rows=3)
513513
assert len(decoded.rows) == 2
514514

515+
def test_continuation_column_count_mismatch_raises(self) -> None:
516+
"""decode_rows_continuation should reject mismatched column_names/column_count."""
517+
import pytest
518+
519+
from dqlitewire.exceptions import DecodeError
520+
from dqlitewire.tuples import encode_row_header, encode_row_values
521+
from dqlitewire.types import encode_uint64
522+
523+
types = [ValueType.INTEGER, ValueType.TEXT]
524+
body = encode_row_header(types)
525+
body += encode_row_values([1, "hello"], types)
526+
body += encode_uint64(0xFFFFFFFFFFFFFFFF) # DONE
527+
528+
# column_names has 3 elements but column_count is 2 — mismatch
529+
with pytest.raises(DecodeError, match="column_names.*does not match.*column_count"):
530+
RowsResponse.decode_rows_continuation(
531+
body,
532+
column_names=["a", "b", "c"],
533+
column_count=2,
534+
)
535+
515536

516537
class TestEmptyResponse:
517538
def test_encode(self) -> None:

0 commit comments

Comments
 (0)