Parse LWT flags when creating prepared statement

Author: sylwiaszunejko

cassandra/cluster.py

@@ -3109,7 +3109,9 @@ def prepare(self, query, custom_payload=None, keyspace=None):
         prepared_keyspace = keyspace if keyspace else None
         prepared_statement = PreparedStatement.from_message(
             response.query_id, response.bind_metadata, response.pk_indexes, self.cluster.metadata, query, prepared_keyspace,
-            self._protocol_version, response.column_metadata, response.result_metadata_id, self.cluster.column_encryption_policy)
+            self._protocol_version, response.column_metadata, response.result_metadata_id,
+            future._current_host.lwt_info.is_lwt(response.flags) if future._current_host.lwt_info is not None else False,
+            self.cluster.column_encryption_policy)
         prepared_statement.custom_payload = future.custom_payload
 
         self.cluster.add_prepared(response.query_id, prepared_statement)

cassandra/lwt_info.py

@@ -19,3 +19,6 @@ class _LwtInfo:
 
     def __init__(self, lwt_meta_bit_mask):
         self.lwt_meta_bit_mask = lwt_meta_bit_mask
+
+    def is_lwt(self, flags):
+        return (flags & self.lwt_meta_bit_mask) == self.lwt_meta_bit_mask

cassandra/pool.py

@@ -167,6 +167,8 @@ class Host(object):
 
     sharding_info = None
 
+    lwt_info = None
+
     def __init__(self, endpoint, conviction_policy_factory, datacenter=None, rack=None, host_id=None):
         if endpoint is None:
             raise ValueError("endpoint may not be None")
@@ -438,6 +440,8 @@ def __init__(self, host, host_distance, session):
         if first_connection.features.sharding_info and not self._session.cluster.shard_aware_options.disable:
             self.host.sharding_info = first_connection.features.sharding_info
             self._open_connections_for_all_shards(first_connection.features.shard_id)
+        if first_connection.features.lwt_info is not None:
+            self.host.lwt_info = first_connection.features.lwt_info
         self.tablets_routing_v1 = first_connection.features.tablets_routing_v1
 
         log.debug("Finished initializing connection for host %s", self.host)

cassandra/protocol.py

@@ -686,6 +686,7 @@ class ResultMessage(_MessageType):
     bind_metadata = None
     pk_indexes = None
     schema_change_event = None
+    flags = None
 
     def __init__(self, kind):
         self.kind = kind
@@ -787,6 +788,7 @@ def recv_results_metadata(self, f, user_type_map):
 
     def recv_prepared_metadata(self, f, protocol_version, user_type_map):
         flags = read_int(f)
+        self.flags = flags
         colcount = read_int(f)
         pk_indexes = None
         if protocol_version >= 4:

cassandra/query.py

@@ -454,10 +454,11 @@ class PreparedStatement(object):
     routing_key_indexes = None
     _routing_key_index_set = None
     serial_consistency_level = None  # TODO never used?
+    _is_lwt = False
 
     def __init__(self, column_metadata, query_id, routing_key_indexes, query,
                  keyspace, protocol_version, result_metadata, result_metadata_id,
-                 column_encryption_policy=None):
+                 is_lwt, column_encryption_policy=None):
         self.column_metadata = column_metadata
         self.query_id = query_id
         self.routing_key_indexes = routing_key_indexes
@@ -468,15 +469,16 @@ def __init__(self, column_metadata, query_id, routing_key_indexes, query,
         self.result_metadata_id = result_metadata_id
         self.column_encryption_policy = column_encryption_policy
         self.is_idempotent = False
+        self._is_lwt = is_lwt
 
     @classmethod
     def from_message(cls, query_id, column_metadata, pk_indexes, cluster_metadata,
                      query, prepared_keyspace, protocol_version, result_metadata,
-                     result_metadata_id, column_encryption_policy=None):
+                     result_metadata_id, is_lwt, column_encryption_policy=None):
         if not column_metadata:
             return PreparedStatement(column_metadata, query_id, None,
                                      query, prepared_keyspace, protocol_version, result_metadata,
-                                     result_metadata_id, column_encryption_policy)
+                                     result_metadata_id, is_lwt, column_encryption_policy)
 
         if pk_indexes:
             routing_key_indexes = pk_indexes
@@ -502,7 +504,7 @@ def from_message(cls, query_id, column_metadata, pk_indexes, cluster_metadata,
 
         return PreparedStatement(column_metadata, query_id, routing_key_indexes,
                                  query, prepared_keyspace, protocol_version, result_metadata,
-                                 result_metadata_id, column_encryption_policy)
+                                 result_metadata_id, is_lwt, column_encryption_policy)
 
     def bind(self, values):
         """
@@ -517,6 +519,9 @@ def is_routing_key_index(self, i):
             self._routing_key_index_set = set(self.routing_key_indexes) if self.routing_key_indexes else set()
         return i in self._routing_key_index_set
 
+    def is_lwt(self):
+        return self._is_lwt
+
     def __str__(self):
         consistency = ConsistencyLevel.value_to_name.get(self.consistency_level, 'Not Set')
         return (u'<PreparedStatement query="%s", consistency=%s>' %
@@ -682,6 +687,9 @@ def routing_key(self):
 
         return self._routing_key
 
+    def is_lwt(self):
+        return self.prepared_statement.is_lwt
+
     def __str__(self):
         consistency = ConsistencyLevel.value_to_name.get(self.consistency_level, 'Not Set')
         return (u'<BoundStatement query="%s", values=%s, consistency=%s>' %