Add support for SCYLLA_USE_METADATA_ID

SCYLLA_USE_METADATA_ID extension allows using METADATA_ID (which was
introduced in CQLv5) in CQLv4.
This commit:
- Introduce support for SCYLLA_USE_METADATA_ID in protocol extension
negotation
- Reuse CQLv5 metadata id implemnetation if use_metadata_id feature is
enabled
- Modify existing unit tests for introduced changes

Author: andrzej-jackowski-scylladb

cassandra/connection.py

@@ -1110,7 +1110,7 @@ def send_msg(self, msg, request_id, cb, encoder=ProtocolHandler.encode_message,
 # queue the decoder function with the request
 # 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,
+msg=encoder(msg, request_id, self.protocol_version, self.features, compressor=self.compressor,
 allow_beta_protocol_version=self.allow_beta_protocol_version)
 
 ifself._is_checksumming_enabled:

cassandra/protocol.py

@@ -421,7 +421,7 @@ def __init__(self, cqlversion, options):
 self.cqlversion=cqlversion
 self.options=options
 
-defsend_body(self, f, protocol_version):
+defsend_body(self, f, protocol_version, protocol_features):
 optmap=self.options.copy()
 optmap['CQL_VERSION'] =self.cqlversion
 write_stringmap(f, optmap)
@@ -456,7 +456,7 @@ class CredentialsMessage(_MessageType):
 def__init__(self, creds):
 self.creds=creds
 
-defsend_body(self, f, protocol_version):
+defsend_body(self, f, protocol_version, protocol_features):
 ifprotocol_version>1:
 raiseUnsupportedOperation(
 "Credentials-based authentication is not supported with "
@@ -487,7 +487,7 @@ class AuthResponseMessage(_MessageType):
 def__init__(self, response):
 self.response=response
 
-defsend_body(self, f, protocol_version):
+defsend_body(self, f, protocol_version, protocol_features):
 write_longstring(f, self.response)
 
 
@@ -507,7 +507,7 @@ class OptionsMessage(_MessageType):
 opcode=0x05
 name='OPTIONS'
 
-defsend_body(self, f, protocol_version):
+defsend_body(self, f, protocol_version, protocol_features):
 pass
 
 
@@ -645,7 +645,7 @@ def __init__(self, query, consistency_level, serial_consistency_level=None,
 super(QueryMessage, self).__init__(None, consistency_level, serial_consistency_level, fetch_size,
 paging_state, timestamp, False, continuous_paging_options, keyspace)
 
-defsend_body(self, f, protocol_version):
+defsend_body(self, f, protocol_version, protocol_features):
 write_longstring(f, self.query)
 self._write_query_params(f, protocol_version)
 
@@ -681,9 +681,9 @@ def _write_query_params(self, f, protocol_version):
 else:
 super(ExecuteMessage, self)._write_query_params(f, protocol_version)
 
-defsend_body(self, f, protocol_version):
+defsend_body(self, f, protocol_version, protocol_features):
 write_string(f, self.query_id)
-ifProtocolVersion.uses_prepared_metadata(protocol_version):
+ifProtocolVersion.uses_prepared_metadata(protocol_version)orprotocol_features.use_metadata_id:
 write_string(f, self.result_metadata_id)
 self._write_query_params(f, protocol_version)
 
@@ -734,15 +734,15 @@ class ResultMessage(_MessageType):
 def__init__(self, kind):
 self.kind=kind
 
-defrecv(self, f, protocol_version, user_type_map, result_metadata, column_encryption_policy):
+defrecv(self, f, protocol_version, protocol_features, user_type_map, result_metadata, column_encryption_policy):
 ifself.kind==RESULT_KIND_VOID:
 return
 elifself.kind==RESULT_KIND_ROWS:
 self.recv_results_rows(f, protocol_version, user_type_map, result_metadata, column_encryption_policy)
 elifself.kind==RESULT_KIND_SET_KEYSPACE:
 self.new_keyspace=read_string(f)
 elifself.kind==RESULT_KIND_PREPARED:
-self.recv_results_prepared(f, protocol_version, user_type_map)
+self.recv_results_prepared(f, protocol_version, protocol_features, user_type_map)
 elifself.kind==RESULT_KIND_SCHEMA_CHANGE:
 self.recv_results_schema_change(f, protocol_version)
 else:
@@ -752,7 +752,7 @@ def recv(self, f, protocol_version, user_type_map, result_metadata, column_encry
 defrecv_body(cls, f, protocol_version, protocol_features, user_type_map, result_metadata, column_encryption_policy):
 kind=read_int(f)
 msg=cls(kind)
-msg.recv(f, protocol_version, user_type_map, result_metadata, column_encryption_policy)
+msg.recv(f, protocol_version, protocol_features, user_type_map, result_metadata, column_encryption_policy)
 returnmsg
 
 defrecv_results_rows(self, f, protocol_version, user_type_map, result_metadata, column_encryption_policy):
@@ -785,9 +785,9 @@ def decode_row(row):
 col_md[3].cql_parameterized_type(),
 str(e)))
 
-defrecv_results_prepared(self, f, protocol_version, user_type_map):
+defrecv_results_prepared(self, f, protocol_version, protocol_features, user_type_map):
 self.query_id=read_binary_string(f)
-ifProtocolVersion.uses_prepared_metadata(protocol_version):
+ifProtocolVersion.uses_prepared_metadata(protocol_version)orprotocol_features.use_metadata_id:
 self.result_metadata_id=read_binary_string(f)
 else:
 self.result_metadata_id=None
@@ -909,7 +909,7 @@ def __init__(self, query, keyspace=None):
 self.query=query
 self.keyspace=keyspace
 
-defsend_body(self, f, protocol_version):
+defsend_body(self, f, protocol_version, protocol_features):
 write_longstring(f, self.query)
 
 flags=0x00
@@ -953,7 +953,7 @@ def __init__(self, batch_type, queries, consistency_level,
 self.timestamp=timestamp
 self.keyspace=keyspace
 
-defsend_body(self, f, protocol_version):
+defsend_body(self, f, protocol_version, protocol_features):
 write_byte(f, self.batch_type.value)
 write_short(f, len(self.queries))
 forprepared, string_or_query_id, paramsinself.queries:
@@ -1012,7 +1012,7 @@ class RegisterMessage(_MessageType):
 def__init__(self, event_list):
 self.event_list=event_list
 
-defsend_body(self, f, protocol_version):
+defsend_body(self, f, protocol_version, protocol_features):
 write_stringlist(f, self.event_list)
 
 
@@ -1086,7 +1086,7 @@ def __init__(self, op_type, op_id, next_pages=0):
 self.op_id=op_id
 self.next_pages=next_pages
 
-defsend_body(self, f, protocol_version):
+defsend_body(self, f, protocol_version, protocol_features):
 write_int(f, self.op_type)
 write_int(f, self.op_id)
 ifself.op_type==ReviseRequestMessage.RevisionType.PAGING_BACKPRESSURE:
@@ -1122,7 +1122,7 @@ class _ProtocolHandler(object):
 """Instance of :class:`cassandra.policies.ColumnEncryptionPolicy` in use by this handler"""
 
 @classmethod
-defencode_message(cls, msg, stream_id, protocol_version, compressor, allow_beta_protocol_version):
+defencode_message(cls, msg, stream_id, protocol_version, protocol_features, compressor, allow_beta_protocol_version):
 """
  Encodes a message using the specified frame parameters, and compressor
 
@@ -1138,7 +1138,7 @@ def encode_message(cls, msg, stream_id, protocol_version, compressor, allow_beta
 raiseUnsupportedOperation("Custom key/value payloads can only be used with protocol version 4 or higher")
 flags|=CUSTOM_PAYLOAD_FLAG
 write_bytesmap(body, msg.custom_payload)
-msg.send_body(body, protocol_version)
+msg.send_body(body, protocol_version, protocol_features)
 body=body.getvalue()
 
 # With checksumming, the compression is done at the segment frame encoding

cassandra/protocol_features.py

@@ -7,25 +7,29 @@
 
 RATE_LIMIT_ERROR_EXTENSION="SCYLLA_RATE_LIMIT_ERROR"
 TABLETS_ROUTING_V1="TABLETS_ROUTING_V1"
+USE_METADATA_ID="SCYLLA_USE_METADATA_ID"
 
 classProtocolFeatures(object):
 rate_limit_error=None
 shard_id=0
 sharding_info=None
 tablets_routing_v1=False
+use_metadata_id=False
 
-def__init__(self, rate_limit_error=None, shard_id=0, sharding_info=None, tablets_routing_v1=False):
+def__init__(self, rate_limit_error=None, shard_id=0, sharding_info=None, tablets_routing_v1=False, use_metadata_id=False):
 self.rate_limit_error=rate_limit_error
 self.shard_id=shard_id
 self.sharding_info=sharding_info
 self.tablets_routing_v1=tablets_routing_v1
+self.use_metadata_id=use_metadata_id
 
 @staticmethod
 defparse_from_supported(supported):
 rate_limit_error=ProtocolFeatures.maybe_parse_rate_limit_error(supported)
 shard_id, sharding_info=ProtocolFeatures.parse_sharding_info(supported)
 tablets_routing_v1=ProtocolFeatures.parse_tablets_info(supported)
-returnProtocolFeatures(rate_limit_error, shard_id, sharding_info, tablets_routing_v1)
+use_metadata_id=ProtocolFeatures.parse_metadata_id_info(supported)
+returnProtocolFeatures(rate_limit_error, shard_id, sharding_info, tablets_routing_v1, use_metadata_id)
 
 @staticmethod
 defmaybe_parse_rate_limit_error(supported):
@@ -49,6 +53,8 @@ def add_startup_options(self, options):
 options[RATE_LIMIT_ERROR_EXTENSION] =""
 ifself.tablets_routing_v1:
 options[TABLETS_ROUTING_V1] =""
+ifself.use_metadata_id:
+options[USE_METADATA_ID] =""
 
 @staticmethod
 defparse_sharding_info(options):
@@ -72,3 +78,7 @@ def parse_sharding_info(options):
 @staticmethod
 defparse_tablets_info(options):
 returnTABLETS_ROUTING_V1inoptions
+
+@staticmethod
+defparse_metadata_id_info(options):
+returnUSE_METADATA_IDinoptions

tests/unit/test_protocol.py

@@ -23,6 +23,7 @@
 _PAGE_SIZE_FLAG, _WITH_PAGING_STATE_FLAG,
 BatchMessage
 )
+fromcassandra.protocol_featuresimportProtocolFeatures
 fromcassandra.queryimportBatchType
 fromcassandra.marshalimportuint32_unpack
 fromcassandra.clusterimportContinuousPagingOptions
@@ -43,24 +44,24 @@ def test_prepare_message(self):
 message=PrepareMessage("a")
 io=Mock()
 
-message.send_body(io, 4)
+message.send_body(io, 4, ProtocolFeatures())
 self._check_calls(io, [(b'\x00\x00\x00\x01',), (b'a',)])
 
 io.reset_mock()
-message.send_body(io, 5)
+message.send_body(io, 5, ProtocolFeatures())
 
 self._check_calls(io, [(b'\x00\x00\x00\x01',), (b'a',), (b'\x00\x00\x00\x00',)])
 
 deftest_execute_message(self):
 message=ExecuteMessage('1', [], 4)
 io=Mock()
 
-message.send_body(io, 4)
+message.send_body(io, 4, ProtocolFeatures())
 self._check_calls(io, [(b'\x00\x01',), (b'1',), (b'\x00\x04',), (b'\x01',), (b'\x00\x00',)])
 
 io.reset_mock()
 message.result_metadata_id='foo'
-message.send_body(io, 5)
+message.send_body(io, 5, ProtocolFeatures())
 
 self._check_calls(io, [(b'\x00\x01',), (b'1',),
  (b'\x00\x03',), (b'foo',),
@@ -80,11 +81,11 @@ def test_query_message(self):
 message=QueryMessage("a", 3)
 io=Mock()
 
-message.send_body(io, 4)
+message.send_body(io, 4, ProtocolFeatures())
 self._check_calls(io, [(b'\x00\x00\x00\x01',), (b'a',), (b'\x00\x03',), (b'\x00',)])
 
 io.reset_mock()
-message.send_body(io, 5)
+message.send_body(io, 5, ProtocolFeatures())
 self._check_calls(io, [(b'\x00\x00\x00\x01',), (b'a',), (b'\x00\x03',), (b'\x00\x00\x00\x00',)])
 
 def_check_calls(self, io, expected):
@@ -112,10 +113,10 @@ def test_continuous_paging(self):
 io=Mock()
 forversionin [versionforversioninProtocolVersion.SUPPORTED_VERSIONS
 ifnotProtocolVersion.has_continuous_paging_support(version)]:
-self.assertRaises(UnsupportedOperation, message.send_body, io, version)
+self.assertRaises(UnsupportedOperation, message.send_body, io, version, ProtocolFeatures())
 
 io.reset_mock()
-message.send_body(io, ProtocolVersion.DSE_V1)
+message.send_body(io, ProtocolVersion.DSE_V1, ProtocolFeatures())
 
 # continuous paging adds two write calls to the buffer
 self.assertEqual(len(io.write.mock_calls), 6)
@@ -142,7 +143,7 @@ def test_prepare_flag(self):
 message=PrepareMessage("a")
 io=Mock()
 forversioninProtocolVersion.SUPPORTED_VERSIONS:
-message.send_body(io, version)
+message.send_body(io, version, ProtocolFeatures())
 ifProtocolVersion.uses_prepare_flags(version):
 self.assertEqual(len(io.write.mock_calls), 3)
 else:
@@ -155,7 +156,7 @@ def test_prepare_flag_with_keyspace(self):
 
 forversioninProtocolVersion.SUPPORTED_VERSIONS:
 ifProtocolVersion.uses_keyspace_flag(version):
-message.send_body(io, version)
+message.send_body(io, version, ProtocolFeatures())
 self._check_calls(io, [
  (b'\x00\x00\x00\x01',),
  (b'a',),
@@ -165,15 +166,15 @@ def test_prepare_flag_with_keyspace(self):
  ])
 else:
 withself.assertRaises(UnsupportedOperation):
-message.send_body(io, version)
+message.send_body(io, version, ProtocolFeatures())
 io.reset_mock()
 
 deftest_keyspace_flag_raises_before_v5(self):
 keyspace_message=QueryMessage('a', consistency_level=3, keyspace='ks')
 io=Mock(name='io')
 
 withself.assertRaisesRegex(UnsupportedOperation, 'Keyspaces.*set'):
-keyspace_message.send_body(io, protocol_version=4)
+keyspace_message.send_body(io, protocol_version=4, protocol_features=ProtocolFeatures())
 io.assert_not_called()
 
 deftest_keyspace_written_with_length(self):
@@ -186,7 +187,7 @@ def test_keyspace_written_with_length(self):
  ]
 
 QueryMessage('a', consistency_level=3, keyspace='ks').send_body(
-io, protocol_version=5
+io, protocol_version=5, protocol_features=ProtocolFeatures()
  )
 self._check_calls(io, base_expected+ [
  (b'\x00\x02',), # length of keyspace string
@@ -196,7 +197,7 @@ def test_keyspace_written_with_length(self):
 io.reset_mock()
 
 QueryMessage('a', consistency_level=3, keyspace='keyspace').send_body(
-io, protocol_version=5
+io, protocol_version=5, protocol_features=ProtocolFeatures()
  )
 self._check_calls(io, base_expected+ [
  (b'\x00\x08',), # length of keyspace string
@@ -215,7 +216,7 @@ def test_batch_message_with_keyspace(self):
 consistency_level=3,
 keyspace='ks'
  )
-batch.send_body(io, protocol_version=5)
+batch.send_body(io, protocol_version=5, protocol_features=ProtocolFeatures())
 self._check_calls(io,
  ((b'\x00',), (b'\x00\x03',), (b'\x00',),
  (b'\x00\x00\x00\x06',), (b'stmt a',),