Preserve decoded schema on Exec/Query request round-trip

The upstream C client sends schema=1 for ExecRequest / QueryRequest /
ExecSqlRequest / QuerySqlRequest regardless of param count, but the
Python encoder's count-based heuristic downgraded the schema byte to
0 on re-encode whenever params numbered <= 255. A proxy or mock
server that decoded then re-encoded a real client's bytes therefore
emitted a different wire shape than it received. Carry the decoded
schema on a private dataclass field and honour it on encode so the
round-trip stays byte-identical; default construction still uses the
count heuristic for callers that do not hold a hint.

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

Author: antoineleclair

src/dqlitewire/messages/requests.py

@@ -198,12 +198,20 @@ class ExecRequest(Message):
     db_id: int
     stmt_id: int
     params: Sequence[Any] = field(default_factory=list)
+    # Preserves the header schema byte seen on decode so a decode →
+    # re-encode round-trip emits byte-identical output even when the
+    # upstream C client used schema=1 with ≤255 params (which the count
+    # heuristic alone would otherwise downgrade to schema=0). Excluded
+    # from repr/compare so it stays an internal round-trip hint.
+    _decoded_schema: int | None = field(default=None, repr=False, compare=False)
 
     def __post_init__(self) -> None:
         _check_uint32("db_id", self.db_id)
         _check_uint32("stmt_id", self.stmt_id)
 
     def _get_schema(self) -> int:
+        if self._decoded_schema is not None:
+            return self._decoded_schema
         return 1 if len(self.params) > 255 else 0
 
     def encode_body(self) -> bytes:
@@ -217,7 +225,7 @@ def decode_body(cls, data: bytes, schema: int = 0) -> "ExecRequest":
         db_id = decode_uint32(data)
         stmt_id = decode_uint32(data[4:])
         params, _ = decode_params_tuple(data[8:], schema=schema, buffer_offset=8)
-        return cls(db_id, stmt_id, params)
+        return cls(db_id, stmt_id, params, _decoded_schema=schema)
 
 
 @dataclass
@@ -233,12 +241,15 @@ class QueryRequest(Message):
     db_id: int
     stmt_id: int
     params: Sequence[Any] = field(default_factory=list)
+    _decoded_schema: int | None = field(default=None, repr=False, compare=False)
 
     def __post_init__(self) -> None:
         _check_uint32("db_id", self.db_id)
         _check_uint32("stmt_id", self.stmt_id)
 
     def _get_schema(self) -> int:
+        if self._decoded_schema is not None:
+            return self._decoded_schema
         return 1 if len(self.params) > 255 else 0
 
     def encode_body(self) -> bytes:
@@ -252,7 +263,7 @@ def decode_body(cls, data: bytes, schema: int = 0) -> "QueryRequest":
         db_id = decode_uint32(data)
         stmt_id = decode_uint32(data[4:])
         params, _ = decode_params_tuple(data[8:], schema=schema, buffer_offset=8)
-        return cls(db_id, stmt_id, params)
+        return cls(db_id, stmt_id, params, _decoded_schema=schema)
 
 
 @dataclass
@@ -294,11 +305,14 @@ class ExecSqlRequest(Message):
     db_id: int
     sql: str
     params: Sequence[Any] = field(default_factory=list)
+    _decoded_schema: int | None = field(default=None, repr=False, compare=False)
 
     def __post_init__(self) -> None:
         _check_uint64("db_id", self.db_id)
 
     def _get_schema(self) -> int:
+        if self._decoded_schema is not None:
+            return self._decoded_schema
         return 1 if len(self.params) > 255 else 0
 
     def encode_body(self) -> bytes:
@@ -314,7 +328,7 @@ def decode_body(cls, data: bytes, schema: int = 0) -> "ExecSqlRequest":
         sql, offset = decode_text(data[8:])
         offset += 8
         params, _ = decode_params_tuple(data[offset:], schema=schema, buffer_offset=offset)
-        return cls(db_id, sql, params)
+        return cls(db_id, sql, params, _decoded_schema=schema)
 
 
 @dataclass
@@ -330,11 +344,14 @@ class QuerySqlRequest(Message):
     db_id: int
     sql: str
     params: Sequence[Any] = field(default_factory=list)
+    _decoded_schema: int | None = field(default=None, repr=False, compare=False)
 
     def __post_init__(self) -> None:
         _check_uint64("db_id", self.db_id)
 
     def _get_schema(self) -> int:
+        if self._decoded_schema is not None:
+            return self._decoded_schema
         return 1 if len(self.params) > 255 else 0
 
     def encode_body(self) -> bytes:
@@ -350,7 +367,7 @@ def decode_body(cls, data: bytes, schema: int = 0) -> "QuerySqlRequest":
         sql, offset = decode_text(data[8:])
         offset += 8
         params, _ = decode_params_tuple(data[offset:], schema=schema, buffer_offset=offset)
-        return cls(db_id, sql, params)
+        return cls(db_id, sql, params, _decoded_schema=schema)
 
 
 @dataclass

tests/test_messages_requests.py

@@ -625,3 +625,84 @@ def test_bool_rejected_for_uint64_client_id(self) -> None:
 
         with pytest.raises(TypeError, match="client_id must be int"):
             ClientRequest(client_id=True)
+
+
+class TestParamsBodySchemaRoundtrip:
+    """Upstream C clients emit schema=1 for Exec/Query*Request
+    unconditionally, regardless of param count. Decoding a small-param
+    schema=1 body and re-encoding must be byte-identical, so a proxy
+    or mock server that round-trips the bytes faithfully does not
+    downgrade the schema bit seen over the wire.
+    """
+
+    @pytest.mark.parametrize(
+        "cls_name",
+        ["ExecRequest", "QueryRequest"],
+    )
+    def test_prepared_schema_1_small_params_roundtrip(self, cls_name: str) -> None:
+        """ExecRequest / QueryRequest: 4-byte db_id + 4-byte stmt_id +
+        PARAMS32 body with 3 params. Construct with schema=1 explicitly,
+        then verify decode → re-encode == original bytes.
+        """
+        from dqlitewire.codec import encode_message
+        from dqlitewire.messages import ExecRequest, QueryRequest
+
+        classes = {"ExecRequest": ExecRequest, "QueryRequest": QueryRequest}
+        cls = classes[cls_name]
+
+        from dqlitewire.codec import decode_message
+
+        # Manually force schema=1 by setting _decoded_schema on the
+        # source message.
+        original = cls(db_id=1, stmt_id=2, params=[42, "hello", None], _decoded_schema=1)
+        original_bytes = encode_message(original)
+
+        # Header schema byte is the 6th byte (after 4-byte size_words
+        # + 1-byte msg_type). Confirm the wire reflects schema=1.
+        assert original_bytes[5] == 1, "expected schema=1 in the header"
+
+        decoded = decode_message(original_bytes, is_request=True)
+        assert isinstance(decoded, cls)
+        assert list(decoded.params) == list(original.params)
+
+        re_encoded = encode_message(decoded)
+        assert re_encoded == original_bytes
+
+    @pytest.mark.parametrize(
+        "cls_name",
+        ["ExecSqlRequest", "QuerySqlRequest"],
+    )
+    def test_sql_schema_1_small_params_roundtrip(self, cls_name: str) -> None:
+        from dqlitewire.codec import decode_message, encode_message
+        from dqlitewire.messages import ExecSqlRequest, QuerySqlRequest
+
+        classes = {"ExecSqlRequest": ExecSqlRequest, "QuerySqlRequest": QuerySqlRequest}
+        cls = classes[cls_name]
+
+        original = cls(db_id=1, sql="SELECT 1", params=[1, 2, 3], _decoded_schema=1)
+        original_bytes = encode_message(original)
+        assert original_bytes[5] == 1
+
+        decoded = decode_message(original_bytes, is_request=True)
+        assert isinstance(decoded, cls)
+        assert list(decoded.params) == list(original.params)
+        assert decoded.sql == "SELECT 1"
+
+        re_encoded = encode_message(decoded)
+        assert re_encoded == original_bytes
+
+    def test_default_construction_uses_heuristic(self) -> None:
+        """When a caller constructs a fresh ExecRequest without a
+        decoded schema hint, the count heuristic still applies: ≤255
+        params → schema=0.
+        """
+        from dqlitewire.messages import ExecRequest
+
+        msg = ExecRequest(db_id=1, stmt_id=2, params=[1, 2, 3])
+        assert msg._get_schema() == 0
+
+    def test_large_params_force_schema_1_without_hint(self) -> None:
+        from dqlitewire.messages import ExecRequest
+
+        msg = ExecRequest(db_id=1, stmt_id=2, params=list(range(300)))
+        assert msg._get_schema() == 1