Perform LWT metadata protocol handshake

Author: sylwiaszunejko

cassandra/lwt_info.py

@@ -0,0 +1,21 @@
+# Copyright 2020 ScyllaDB, Inc.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+class _LwtInfo:
+    """
+    Holds LWT-related information parsed from the server's supported features.
+    """
+
+    def __init__(self, lwt_meta_bit_mask):
+        self.lwt_meta_bit_mask = lwt_meta_bit_mask

cassandra/protocol_features.py

@@ -1,10 +1,13 @@
 import logging
 
 from cassandra.shard_info import _ShardingInfo
+from cassandra.lwt_info import _LwtInfo
 
 log = logging.getLogger(__name__)
 
 
+LWT_ADD_METADATA_MARK = "SCYLLA_LWT_ADD_METADATA_MARK"
+LWT_OPTIMIZATION_META_BIT_MASK = "LWT_OPTIMIZATION_META_BIT_MASK"
 RATE_LIMIT_ERROR_EXTENSION = "SCYLLA_RATE_LIMIT_ERROR"
 TABLETS_ROUTING_V1 = "TABLETS_ROUTING_V1"
 
@@ -13,19 +16,22 @@ class ProtocolFeatures(object):
     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):
+    def __init__(self, rate_limit_error=None, shard_id=0, sharding_info=None, tablets_routing_v1=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.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)
-        return ProtocolFeatures(rate_limit_error, shard_id, sharding_info, tablets_routing_v1)
+        lwt_info = ProtocolFeatures.parse_lwt_info(supported)
+        return ProtocolFeatures(rate_limit_error, shard_id, sharding_info, tablets_routing_v1, lwt_info)
 
     @staticmethod
     def maybe_parse_rate_limit_error(supported):
@@ -49,6 +55,8 @@ def add_startup_options(self, options):
             options[RATE_LIMIT_ERROR_EXTENSION] = ""
         if 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)
 
     @staticmethod
     def parse_sharding_info(options):
@@ -72,3 +80,18 @@ def parse_sharding_info(options):
     @staticmethod
     def parse_tablets_info(options):
         return TABLETS_ROUTING_V1 in options
+
+    @staticmethod
+    def parse_lwt_info(options):
+        value_list = options.get(LWT_ADD_METADATA_MARK, [None])
+        for value in value_list:
+            if value is None or not value.startswith(LWT_OPTIMIZATION_META_BIT_MASK + "="):
+                continue
+            try:
+                lwt_meta_bit_mask = int(value[len(LWT_OPTIMIZATION_META_BIT_MASK + "="):])
+                return _LwtInfo(lwt_meta_bit_mask)
+            except Exception as e:
+                log.exception(f"Error while parsing {LWT_ADD_METADATA_MARK}: {e}")
+                return None
+
+        return None