protocol: pass negotiated ProtocolFeatures to message serialization

Connection.send_msg now passes the connection's negotiated
ProtocolFeatures to the encoder; _ProtocolHandler.encode_message accepts
it as a new required protocol_features argument (passed by keyword from
send_msg) and forwards it to every message's send_body, which gains the
same parameter. Messages carry connection-independent request data;
send_body decides the wire format from (protocol_version,
protocol_features), so fields belonging to a negotiated protocol
extension are emitted exactly on the connections that negotiated it — on
every send path, including the control-connection fallback, and without
mutating shared message objects per attempt.

This is pure plumbing: no message consumes the parameter yet, so no
bytes on the wire change. It is groundwork for the
SCYLLA_USE_METADATA_ID (#770) and TABLETS_ROUTING_V2 (#913) extensions,
which must serialize extension fields based on what the serving
connection negotiated. The encode side becomes symmetric with
decode_message, which already receives protocol_features.

This changes the contracted signature of encode_message: custom protocol
handlers overriding it must accept the protocol_features keyword
argument. The argument is deliberately required, with no default and no
fallback for old-style encoders: extensions are negotiated per
connection at STARTUP before the per-request handler is known, so an
encoder unaware of protocol_features could silently omit fields a
negotiated extension requires; omitting it fails fast with TypeError
instead.

Tests: send_msg hands the connection's features to the encoder;
encode_message forwards them into send_body (plain and compressed
paths) and raises TypeError when the argument is omitted; a
byte-identity suite pins frames for representative messages (v3/v4/v5)
to the exact bytes produced before this change, both without features
and with all-default features.

Co-authored-by: Dawid Mędrek <dawid.medrek@scylladb.com>

Author: 2 people

cassandra/connection.py

@@ -1222,7 +1222,8 @@ def send_msg(self, msg, request_id, cb, encoder=ProtocolHandler.encode_message,
         # this allows us to inject custom functions per request to encode, decode messages
         self._requests[request_id] = (cb, decoder, result_metadata)
         msg = encoder(msg, request_id, self.protocol_version, compressor=self.compressor,
-                      allow_beta_protocol_version=self.allow_beta_protocol_version)
+                      allow_beta_protocol_version=self.allow_beta_protocol_version,
+                      protocol_features=self.features)
 
         if self._is_checksumming_enabled:
             buffer = io.BytesIO()

cassandra/protocol.py

@@ -424,7 +424,7 @@ def __init__(self, cqlversion, options):
         self.cqlversion = cqlversion
         self.options = options
 
-    def send_body(self, f, protocol_version):
+    def send_body(self, f, protocol_version, protocol_features=None):
         optmap = self.options.copy()
         optmap['CQL_VERSION'] = self.cqlversion
         write_stringmap(f, optmap)
@@ -459,7 +459,7 @@ class CredentialsMessage(_MessageType):
     def __init__(self, creds):
         self.creds = creds
 
-    def send_body(self, f, protocol_version):
+    def send_body(self, f, protocol_version, protocol_features=None):
         if protocol_version > 1:
             raise UnsupportedOperation(
                 "Credentials-based authentication is not supported with "
@@ -490,7 +490,7 @@ class AuthResponseMessage(_MessageType):
     def __init__(self, response):
         self.response = response
 
-    def send_body(self, f, protocol_version):
+    def send_body(self, f, protocol_version, protocol_features=None):
         write_longstring(f, self.response)
 
 
@@ -510,7 +510,7 @@ class OptionsMessage(_MessageType):
     opcode = 0x05
     name = 'OPTIONS'
 
-    def send_body(self, f, protocol_version):
+    def send_body(self, f, protocol_version, protocol_features=None):
         pass
 
 
@@ -558,7 +558,7 @@ def __init__(self, query_params, consistency_level,
         self.skip_meta = skip_meta
         self.keyspace = keyspace
 
-    def _write_query_params(self, f, protocol_version):
+    def _write_query_params(self, f, protocol_version, protocol_features=None):
         write_consistency_level(f, self.consistency_level)
         flags = 0x00
         if self.query_params is not None:
@@ -620,9 +620,9 @@ def __init__(self, query, consistency_level, serial_consistency_level=None,
         super(QueryMessage, self).__init__(query_params, consistency_level, serial_consistency_level, fetch_size,
                                            paging_state, timestamp, False, continuous_paging_options, keyspace)
 
-    def send_body(self, f, protocol_version):
+    def send_body(self, f, protocol_version, protocol_features=None):
         write_longstring(f, self.query)
-        self._write_query_params(f, protocol_version)
+        self._write_query_params(f, protocol_version, protocol_features)
 
 
 class ExecuteMessage(_QueryMessage):
@@ -638,14 +638,14 @@ def __init__(self, query_id, query_params, consistency_level,
         super(ExecuteMessage, self).__init__(query_params, consistency_level, serial_consistency_level, fetch_size,
                                              paging_state, timestamp, skip_meta, continuous_paging_options)
 
-    def _write_query_params(self, f, protocol_version):
-        super(ExecuteMessage, self)._write_query_params(f, protocol_version)
+    def _write_query_params(self, f, protocol_version, protocol_features=None):
+        super(ExecuteMessage, self)._write_query_params(f, protocol_version, protocol_features)
 
-    def send_body(self, f, protocol_version):
+    def send_body(self, f, protocol_version, protocol_features=None):
         write_string(f, self.query_id)
         if ProtocolVersion.uses_prepared_metadata(protocol_version):
             write_string(f, self.result_metadata_id)
-        self._write_query_params(f, protocol_version)
+        self._write_query_params(f, protocol_version, protocol_features)
 
 
 CUSTOM_TYPE = object()
@@ -870,7 +870,7 @@ def __init__(self, query, keyspace=None):
         self.query = query
         self.keyspace = keyspace
 
-    def send_body(self, f, protocol_version):
+    def send_body(self, f, protocol_version, protocol_features=None):
         write_longstring(f, self.query)
 
         flags = 0x00
@@ -914,7 +914,7 @@ def __init__(self, batch_type, queries, consistency_level,
         self.timestamp = timestamp
         self.keyspace = keyspace
 
-    def send_body(self, f, protocol_version):
+    def send_body(self, f, protocol_version, protocol_features=None):
         write_byte(f, self.batch_type.value)
         write_short(f, len(self.queries))
         for prepared, string_or_query_id, params in self.queries:
@@ -972,7 +972,7 @@ class RegisterMessage(_MessageType):
     def __init__(self, event_list):
         self.event_list = event_list
 
-    def send_body(self, f, protocol_version):
+    def send_body(self, f, protocol_version, protocol_features=None):
         write_stringlist(f, self.event_list)
 
 
@@ -1046,7 +1046,7 @@ def __init__(self, op_type, op_id, next_pages=0):
         self.op_id = op_id
         self.next_pages = next_pages
 
-    def send_body(self, f, protocol_version):
+    def send_body(self, f, protocol_version, protocol_features=None):
         write_int(f, self.op_type)
         write_int(f, self.op_id)
         if self.op_type == ReviseRequestMessage.RevisionType.PAGING_BACKPRESSURE:
@@ -1079,14 +1079,20 @@ class _ProtocolHandler(object):
     """Instance of :class:`cassandra.policies.ColumnEncryptionPolicy` in use by this handler"""
 
     @classmethod
-    def encode_message(cls, msg, stream_id, protocol_version, compressor, allow_beta_protocol_version):
+    def encode_message(cls, msg, stream_id, protocol_version, compressor, allow_beta_protocol_version,
+                       protocol_features):
         """
         Encodes a message using the specified frame parameters, and compressor
 
         :param msg: the message, typically of cassandra.protocol._MessageType, generated by the driver
         :param stream_id: protocol stream id for the frame header
         :param protocol_version: version for the frame header, and used encoding contents
         :param compressor: optional compression function to be used on the body
+        :param protocol_features: :class:`~cassandra.protocol_features.ProtocolFeatures` negotiated on the connection
+            this message is sent over, forwarded to ``send_body``. Messages carry
+            connection-independent request data; ``send_body`` decides the wire format from
+            ``(protocol_version, protocol_features)``, so fields belonging to a negotiated
+            protocol extension are emitted exactly on the connections that negotiated it.
         """
         flags = 0
         if msg.custom_payload:
@@ -1108,7 +1114,7 @@ def encode_message(cls, msg, stream_id, protocol_version, compressor, allow_beta
             body = io.BytesIO()
             if msg.custom_payload:
                 write_bytesmap(body, msg.custom_payload)
-            msg.send_body(body, protocol_version)
+            msg.send_body(body, protocol_version, protocol_features)
             body = body.getvalue()
 
             if len(body) > 0:
@@ -1120,7 +1126,7 @@ def encode_message(cls, msg, stream_id, protocol_version, compressor, allow_beta
         else:
             if msg.custom_payload:
                 write_bytesmap(buff, msg.custom_payload)
-            msg.send_body(buff, protocol_version)
+            msg.send_body(buff, protocol_version, protocol_features)
 
             length = buff.tell() - 9
 

tests/unit/test_connection.py

@@ -291,6 +291,26 @@ def test_set_keyspace_async_escapes_quotes(self):
         assert query_msg.query == 'USE "my""ks"', (
             "Double quotes in keyspace name must be escaped as double-double quotes")
 
+    def test_send_msg_passes_negotiated_features_to_encoder(self):
+        """
+        send_msg must hand the connection's negotiated ProtocolFeatures to the
+        encoder, so message serialization can emit fields belonging to protocol
+        extensions exactly on the connections that negotiated them.
+        """
+        c = self.make_connection()
+        c.push = Mock()
+        captured = {}
+
+        def encoder(msg, stream_id, protocol_version, compressor, allow_beta_protocol_version,
+                    protocol_features=None):
+            captured['protocol_features'] = protocol_features
+            return b'encoded-frame'
+
+        c.send_msg(Mock(), 1, cb=Mock(), encoder=encoder, decoder=Mock())
+
+        assert captured['protocol_features'] is c.features
+        c.push.assert_called_once_with(b'encoded-frame')
+
     def test_set_connection_class(self):
         cluster = Cluster(connection_class='test')
         assert 'test' == cluster.connection_class

tests/unit/test_protocol.py

@@ -16,11 +16,13 @@
 
 from unittest.mock import Mock
 
-from cassandra import ProtocolVersion, UnsupportedOperation
+from cassandra import ConsistencyLevel, ProtocolVersion, UnsupportedOperation
 from cassandra.protocol import (
     PrepareMessage, QueryMessage, ExecuteMessage, UnsupportedOperation,
-    BatchMessage
+    BatchMessage, StartupMessage, OptionsMessage, RegisterMessage,
+    AuthResponseMessage, ProtocolHandler, _MessageType
 )
+from cassandra.protocol_features import ProtocolFeatures
 from cassandra.query import BatchType
 import pytest
 
@@ -185,3 +187,120 @@ def test_batch_message_with_keyspace(self):
              (b'\x00\x03',),
              (b'\x00\x00\x00\x80',), (b'\x00\x02',), (b'ks',))
         )
+
+
+class ProtocolFeaturesPlumbingTest(unittest.TestCase):
+    """
+    The negotiated ProtocolFeatures must flow from encode_message into each
+    message's send_body, so serialization can emit fields belonging to
+    protocol extensions exactly on the connections that negotiated them.
+    """
+
+    class CapturingMessage(_MessageType):
+        opcode = 0x00
+        name = 'CAPTURE'
+
+        def __init__(self):
+            self.seen_features = []
+
+        def send_body(self, f, protocol_version, protocol_features=None):
+            self.seen_features.append(protocol_features)
+
+    def test_encode_message_forwards_protocol_features_to_send_body(self):
+        features = ProtocolFeatures()
+        msg = self.CapturingMessage()
+        ProtocolHandler.encode_message(msg, stream_id=0, protocol_version=4, compressor=None,
+                                       allow_beta_protocol_version=False, protocol_features=features)
+        assert msg.seen_features == [features]
+        assert msg.seen_features[0] is features
+
+    def test_encode_message_forwards_protocol_features_when_compressing(self):
+        features = ProtocolFeatures()
+        msg = self.CapturingMessage()
+        ProtocolHandler.encode_message(msg, stream_id=0, protocol_version=4, compressor=lambda body: body,
+                                       allow_beta_protocol_version=False, protocol_features=features)
+        assert msg.seen_features[0] is features
+
+    def test_encode_message_fails_without_protocol_features(self):
+        msg = self.CapturingMessage()
+
+        with pytest.raises(TypeError, match='positional argument'):
+            ProtocolHandler.encode_message(msg, stream_id=0, protocol_version=4, compressor=None,
+                                       allow_beta_protocol_version=False)
+
+
+class FrameByteIdentityTest(unittest.TestCase):
+    """
+    Threading ProtocolFeatures into serialization is pure plumbing: with no
+    extension consuming it (and for all-default features), every frame must be
+    byte-identical to what the driver produced before the parameter existed.
+    The expected frames below were captured from the pre-change encoder.
+    """
+
+    EXPECTED_FRAMES = {
+        'startup_v4': '0400000701000000160001000b43514c5f56455253494f4e0005332e342e35',
+        'options_v4': '040000070500000000',
+        'register_v4': '040000070b000000220002000f544f504f4c4f47595f4348414e4745000d5354415455535f4348414e4745',
+        'auth_response_v4': '040000070f0000000e0000000a00757365720070617373',
+        'prepare_v4': '0400000709000000220000001e53454c454354202a2046524f4d206b732e74205748455245206b203d203f',
+        'prepare_v5_keyspace': '0500000709000000270000001b53454c454354202a2046524f4d2074205748455245206b203d203f0000000100026b73',
+        'query_v3': '0300000707000000270000001253454c454354202a2046524f4d206b732e74000434000013880008000462d53c8abac0',
+        'execute_v3': '030000070a00000033000412345678000a2d000200000002000100000003616263000000640000000b504147494e475354415445000000003ade68b1',
+        'batch_v3': '030000070d00000043000002000000001f494e5345525420494e544f206b732e7420286b292056414c5545532028312900000100041234567800010000000200020001200000000006a11e3d',
+        'query_v4': '0400000707000000270000001253454c454354202a2046524f4d206b732e74000434000013880008000462d53c8abac0',
+        'execute_v4': '040000070a00000033000412345678000a2d000200000002000100000003616263000000640000000b504147494e475354415445000000003ade68b1',
+        'batch_v4': '040000070d00000043000002000000001f494e5345525420494e544f206b732e7420286b292056414c5545532028312900000100041234567800010000000200020001200000000006a11e3d',
+        'query_v5': '05000007070000002a0000001253454c454354202a2046524f4d206b732e74000400000034000013880008000462d53c8abac0',
+        'execute_v5': '050000070a0000003c0004123456780004aabbccdd000a0000002d000200000002000100000003616263000000640000000b504147494e475354415445000000003ade68b1',
+        'batch_v5': '050000070d00000046000002000000001f494e5345525420494e544f206b732e7420286b292056414c5545532028312900000100041234567800010000000200020001000000200000000006a11e3d',
+    }
+
+    @staticmethod
+    def _make_cases():
+        cases = [
+            ('startup_v4', StartupMessage(cqlversion="3.4.5", options={}), 4),
+            ('options_v4', OptionsMessage(), 4),
+            ('register_v4', RegisterMessage(["TOPOLOGY_CHANGE", "STATUS_CHANGE"]), 4),
+            ('auth_response_v4', AuthResponseMessage(b"\x00user\x00pass"), 4),
+            ('prepare_v4', PrepareMessage("SELECT * FROM ks.t WHERE k = ?"), 4),
+            ('prepare_v5_keyspace', PrepareMessage("SELECT * FROM t WHERE k = ?", keyspace="ks"), 5),
+        ]
+        for pv in (3, 4, 5):
+            cases.append((
+                'query_v%d' % pv,
+                QueryMessage("SELECT * FROM ks.t", ConsistencyLevel.QUORUM,
+                             serial_consistency_level=ConsistencyLevel.SERIAL,
+                             fetch_size=5000, timestamp=1234567890123456),
+                pv,
+            ))
+            cases.append((
+                'execute_v%d' % pv,
+                ExecuteMessage(b"\x12\x34\x56\x78", [b"\x00\x01", b"abc"],
+                               ConsistencyLevel.LOCAL_ONE, fetch_size=100,
+                               paging_state=b"PAGINGSTATE",
+                               result_metadata_id=b"\xaa\xbb\xcc\xdd" if pv >= 5 else None,
+                               timestamp=987654321),
+                pv,
+            ))
+            cases.append((
+                'batch_v%d' % pv,
+                BatchMessage(BatchType.LOGGED,
+                             [(False, "INSERT INTO ks.t (k) VALUES (1)", []),
+                              (True, b"\x12\x34\x56\x78", [b"\x00\x02"])],
+                             ConsistencyLevel.ONE, timestamp=111222333),
+                pv,
+            ))
+        return cases
+
+    def _assert_frames(self, protocol_features):
+        for name, msg, pv in self._make_cases():
+            frame = ProtocolHandler.encode_message(
+                msg, stream_id=7, protocol_version=pv, compressor=None,
+                allow_beta_protocol_version=False, protocol_features=protocol_features)
+            assert frame.hex() == self.EXPECTED_FRAMES[name], name
+
+    def test_frames_without_features(self):
+        self._assert_frames(None)
+
+    def test_frames_with_default_features(self):
+        self._assert_frames(ProtocolFeatures())