Negotiate the TABLETS_ROUTING_V2 protocol extension

Add per-connection negotiation of the TABLETS_ROUTING_V2 extension, the
successor to TABLETS_ROUTING_V1. When the server advertises it in the
SUPPORTED response, the driver echoes it back during STARTUP to opt in;
a driver that negotiates v2 does not negotiate v1.

While the feature is experimental the wire name carries the
`_EXPERIMENTAL` suffix (TABLETS_ROUTING_V2_EXPERIMENTAL), and the server
only advertises it when started with the `strongly-consistent-tables`
experimental feature enabled.

Also add the optional trailing tablet_version_block byte to the EXECUTE
message body, written only when set, so later commits can carry the
cached tablet version to the server.

cassandra/protocol.py

@@ -632,9 +632,11 @@ class ExecuteMessage(_QueryMessage):
     def __init__(self, query_id, query_params, consistency_level,
                  serial_consistency_level=None, fetch_size=None,
                  paging_state=None, timestamp=None, skip_meta=False,
-                 continuous_paging_options=None, result_metadata_id=None):
+                 continuous_paging_options=None, result_metadata_id=None,
+                 tablet_version_block=None):
         self.query_id = query_id
         self.result_metadata_id = result_metadata_id
+        self.tablet_version_block = tablet_version_block
         super(ExecuteMessage, self).__init__(query_params, consistency_level, serial_consistency_level, fetch_size,
                                              paging_state, timestamp, skip_meta, continuous_paging_options)
 
@@ -646,6 +648,8 @@ def send_body(self, f, protocol_version):
         if ProtocolVersion.uses_prepared_metadata(protocol_version):
             write_string(f, self.result_metadata_id)
         self._write_query_params(f, protocol_version)
+        if self.tablet_version_block is not None:
+            write_byte(f, self.tablet_version_block)
 
 
 CUSTOM_TYPE = object()

cassandra/protocol_features.py

@@ -10,28 +10,36 @@
 LWT_OPTIMIZATION_META_BIT_MASK = "LWT_OPTIMIZATION_META_BIT_MASK"
 RATE_LIMIT_ERROR_EXTENSION = "SCYLLA_RATE_LIMIT_ERROR"
 TABLETS_ROUTING_V1 = "TABLETS_ROUTING_V1"
+# The server advertises and expects this exact extension name in SUPPORTED/STARTUP
+# (see scylladb transport/cql_protocol_extension.cc). While the feature is gated
+# behind the server's `strongly-consistent-tables` experimental flag, the wire
+# name carries the `_EXPERIMENTAL` suffix.
+TABLETS_ROUTING_V2 = "TABLETS_ROUTING_V2_EXPERIMENTAL"
 
 class ProtocolFeatures(object):
     rate_limit_error = None
     shard_id = 0
     sharding_info = None
     tablets_routing_v1 = False
+    tablets_routing_v2 = False
     lwt_info = None
 
-    def __init__(self, rate_limit_error=None, shard_id=0, sharding_info=None, tablets_routing_v1=False, lwt_info=None):
+    def __init__(self, rate_limit_error=None, shard_id=0, sharding_info=None, tablets_routing_v1=False, tablets_routing_v2=False, lwt_info=None):
         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.tablets_routing_v2 = tablets_routing_v2
         self.lwt_info = lwt_info
 
     @staticmethod
     def parse_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)
+        tablets_routing_v2 = ProtocolFeatures.parse_tablets_v2_info(supported)
         lwt_info = ProtocolFeatures.parse_lwt_info(supported)
-        return ProtocolFeatures(rate_limit_error, shard_id, sharding_info, tablets_routing_v1, lwt_info)
+        return ProtocolFeatures(rate_limit_error, shard_id, sharding_info, tablets_routing_v1, tablets_routing_v2, lwt_info)
 
     @staticmethod
     def maybe_parse_rate_limit_error(supported):
@@ -53,7 +61,11 @@ def get_cql_extension_field(vals, key):
     def add_startup_options(self, options):
         if self.rate_limit_error is not None:
             options[RATE_LIMIT_ERROR_EXTENSION] = ""
-        if self.tablets_routing_v1:
+        # Only one of TABLETS_ROUTING_V{1,2} should be negotiated
+        # per connection. Hence the if-else branch.
+        if self.tablets_routing_v2:
+            options[TABLETS_ROUTING_V2] = ""
+        elif self.tablets_routing_v1:
             options[TABLETS_ROUTING_V1] = ""
         if self.lwt_info is not None:
             options[LWT_ADD_METADATA_MARK] = str(self.lwt_info.lwt_meta_bit_mask)
@@ -81,6 +93,10 @@ def parse_sharding_info(options):
     def parse_tablets_info(options):
         return TABLETS_ROUTING_V1 in options
 
+    @staticmethod
+    def parse_tablets_v2_info(options):
+        return TABLETS_ROUTING_V2 in options
+
     @staticmethod
     def parse_lwt_info(options):
         value_list = options.get(LWT_ADD_METADATA_MARK, [None])

tests/unit/test_protocol_features.py

@@ -2,7 +2,7 @@
 
 import logging
 
-from cassandra.protocol_features import ProtocolFeatures
+from cassandra.protocol_features import ProtocolFeatures, TABLETS_ROUTING_V1, TABLETS_ROUTING_V2
 
 LOGGER = logging.getLogger(__name__)
 
@@ -22,3 +22,45 @@ class OptionsHolder(object):
         assert protocol_features.rate_limit_error == 123
         assert protocol_features.shard_id == 0
         assert protocol_features.sharding_info is None
+
+    def test_tablets_routing_v2_negotiation(self):
+        """V2 is detected from SUPPORTED and subsumes V1 in STARTUP options."""
+        options = {
+            TABLETS_ROUTING_V1: [''],
+            TABLETS_ROUTING_V2: [''],
+        }
+        features = ProtocolFeatures.parse_from_supported(options)
+        assert features.tablets_routing_v1 is True
+        assert features.tablets_routing_v2 is True
+
+        # V2 subsumes V1: only TABLETS_ROUTING_V2 should appear in startup.
+        startup = {}
+        features.add_startup_options(startup)
+        assert TABLETS_ROUTING_V2 in startup
+        assert TABLETS_ROUTING_V1 not in startup
+
+    def test_tablets_routing_v1_only(self):
+        """When server only advertises V1, only V1 is negotiated."""
+        options = {
+            TABLETS_ROUTING_V1: [''],
+        }
+        features = ProtocolFeatures.parse_from_supported(options)
+        assert features.tablets_routing_v1 is True
+        assert features.tablets_routing_v2 is False
+
+        startup = {}
+        features.add_startup_options(startup)
+        assert TABLETS_ROUTING_V1 in startup
+        assert TABLETS_ROUTING_V2 not in startup
+
+    def test_no_tablets_routing(self):
+        """When server advertises neither V1 nor V2."""
+        options = {}
+        features = ProtocolFeatures.parse_from_supported(options)
+        assert features.tablets_routing_v1 is False
+        assert features.tablets_routing_v2 is False
+
+        startup = {}
+        features.add_startup_options(startup)
+        assert TABLETS_ROUTING_V1 not in startup
+        assert TABLETS_ROUTING_V2 not in startup