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,
+            response.lwt_info.is_lwt(response.flags) if response.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/protocol.py

@@ -686,19 +686,21 @@ class ResultMessage(_MessageType):
     bind_metadata = None
     pk_indexes = None
     schema_change_event = None
+    flags = None
+    lwt_info = None
 
     def __init__(self, kind):
         self.kind = kind
 
-    def recv(self, f, protocol_version, user_type_map, result_metadata, column_encryption_policy):
+    def recv(self, f, protocol_version, protocol_features, user_type_map, result_metadata, column_encryption_policy):
         if self.kind == RESULT_KIND_VOID:
             return
         elif self.kind == RESULT_KIND_ROWS:
             self.recv_results_rows(f, protocol_version, user_type_map, result_metadata, column_encryption_policy)
         elif self.kind == RESULT_KIND_SET_KEYSPACE:
             self.new_keyspace = read_string(f)
         elif self.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)
         elif self.kind == RESULT_KIND_SCHEMA_CHANGE:
             self.recv_results_schema_change(f, protocol_version)
         else:
@@ -708,7 +710,7 @@ def recv(self, f, protocol_version, user_type_map, result_metadata, column_encry
     def recv_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)
         return msg
 
     def recv_results_rows(self, f, protocol_version, user_type_map, result_metadata, column_encryption_policy):
@@ -741,8 +743,9 @@ def decode_row(row):
                                                                                                      col_md[3].cql_parameterized_type(),
                                                                                                      str(e)))
 
-    def recv_results_prepared(self, f, protocol_version, user_type_map):
+    def recv_results_prepared(self, f, protocol_version, protocol_features, user_type_map):
         self.query_id = read_binary_string(f)
+        self.lwt_info = protocol_features.lwt_info
         if ProtocolVersion.uses_prepared_metadata(protocol_version):
             self.result_metadata_id = read_binary_string(f)
         else:
@@ -787,6 +790,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

@@ -345,6 +345,9 @@ def _set_serial_consistency_level(self, serial_consistency_level):
     def _del_serial_consistency_level(self):
         self._serial_consistency_level = None
 
+    def is_lwt(self):
+        return False
+
     serial_consistency_level = property(
         _get_serial_consistency_level,
         _set_serial_consistency_level,
@@ -454,10 +457,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=None, column_encryption_policy=None):
         self.column_metadata = column_metadata
         self.query_id = query_id
         self.routing_key_indexes = routing_key_indexes
@@ -468,15 +472,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 +507,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 +522,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 +690,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>' %