Merge pull request apache#497 from datastax/py432

aholmberg · aholmberg · commit a71cd19b40af · 2016-02-29T13:20:10.000-06:00
PYTHON-432: Add IF EXISTS support to UPDATE and DELETE statements of CQLEngine
diff --git a/cassandra/cqlengine/models.py b/cassandra/cqlengine/models.py
@@ -178,7 +178,7 @@ class IfNotExistsDescriptor(object):
     def __get__(self, instance, model):
         if instance:
             # instance method
-            def ifnotexists_setter(ife):
+            def ifnotexists_setter(ife=True):
                 instance._if_not_exists = ife
                 return instance
             return ifnotexists_setter
@@ -189,6 +189,24 @@ def __call__(self, *args, **kwargs):
         raise NotImplementedError
 
 
+class IfExistsDescriptor(object):
+    """
+    return a query set descriptor with a if_exists flag specified
+    """
+    def __get__(self, instance, model):
+        if instance:
+            # instance method
+            def ifexists_setter(ife=True):
+                instance._if_exists = ife
+                return instance
+            return ifexists_setter
+
+        return model.objects.if_exists
+
+    def __call__(self, *args, **kwargs):
+        raise NotImplementedError
+
+
 class ConsistencyDescriptor(object):
     """
     returns a query set descriptor if called on Class, instance if it was an instance call
@@ -303,6 +321,8 @@ class MultipleObjectsReturned(_MultipleObjectsReturned):
 
     if_not_exists = IfNotExistsDescriptor()
 
+    if_exists = IfExistsDescriptor()
+
     # _len is lazily created by __len__
 
     __table_name__ = None
@@ -323,6 +343,8 @@ class MultipleObjectsReturned(_MultipleObjectsReturned):
 
     _if_not_exists = False  # optional if_not_exists flag to check existence before insertion
 
+    _if_exists = False  # optional if_exists flag to check existence before update
+
     _table_name = None  # used internally to cache a derived table name
 
     def __init__(self, **values):
@@ -654,7 +676,8 @@ def save(self):
                           consistency=self.__consistency__,
                           if_not_exists=self._if_not_exists,
                           transaction=self._transaction,
-                          timeout=self._timeout).save()
+                          timeout=self._timeout,
+                          if_exists=self._if_exists).save()
 
         self._set_persisted()
 
@@ -700,7 +723,8 @@ def update(self, **values):
                           timestamp=self._timestamp,
                           consistency=self.__consistency__,
                           transaction=self._transaction,
-                          timeout=self._timeout).update()
+                          timeout=self._timeout,
+                          if_exists=self._if_exists).update()
 
         self._set_persisted()
 
@@ -717,7 +741,8 @@ def delete(self):
                           batch=self._batch,
                           timestamp=self._timestamp,
                           consistency=self.__consistency__,
-                          timeout=self._timeout).delete()
+                          timeout=self._timeout,
+                          if_exists=self._if_exists).delete()
 
     def get_changed_columns(self):
         """
diff --git a/cassandra/cqlengine/query.py b/cassandra/cqlengine/query.py
@@ -38,6 +38,9 @@ class QueryException(CQLEngineException):
 class IfNotExistsWithCounterColumn(CQLEngineException):
     pass
 
+class IfExistsWithCounterColumn(CQLEngineException):
+    pass
+
 
 class LWTException(CQLEngineException):
     """Lightweight transaction exception.
@@ -294,6 +297,7 @@ def __init__(self, model):
         self._timestamp = None
         self._if_not_exists = False
         self._timeout = connection.NOT_SET
+        self._if_exists = False
 
     @property
     def column_family_name(self):
@@ -304,7 +308,7 @@ def _execute(self, q):
             return self._batch.add_query(q)
         else:
             result = connection.execute(q, consistency_level=self._consistency, timeout=self._timeout)
-            if self._transaction:
+            if self._if_not_exists or self._if_exists or self._transaction:
                 check_applied(result)
             return result
 
@@ -780,9 +784,14 @@ def defer(self, fields):
         return self._only_or_defer('defer', fields)
 
     def create(self, **kwargs):
-        return self.model(**kwargs).batch(self._batch).ttl(self._ttl).\
-            consistency(self._consistency).if_not_exists(self._if_not_exists).\
-            timestamp(self._timestamp).save()
+        return self.model(**kwargs) \
+            .batch(self._batch) \
+            .ttl(self._ttl) \
+            .consistency(self._consistency) \
+            .if_not_exists(self._if_not_exists) \
+            .timestamp(self._timestamp) \
+            .if_exists(self._if_exists) \
+            .save()
 
     def delete(self):
         """
@@ -796,7 +805,8 @@ def delete(self):
         dq = DeleteStatement(
             self.column_family_name,
             where=self._where,
-            timestamp=self._timestamp
+            timestamp=self._timestamp,
+            if_exists=self._if_exists
         )
         self._execute(dq)
 
@@ -938,12 +948,29 @@ def timestamp(self, timestamp):
         return clone
 
     def if_not_exists(self):
+        """
+        Check the existence of an object before insertion.
+
+        If the insertion isn't applied, a LWTException is raised.
+        """
         if self.model._has_counter:
             raise IfNotExistsWithCounterColumn('if_not_exists cannot be used with tables containing counter columns')
         clone = copy.deepcopy(self)
         clone._if_not_exists = True
         return clone
 
+    def if_exists(self):
+        """
+        Check the existence of an object before an update or delete.
+
+        If the update or delete isn't applied, a LWTException is raised.
+        """
+        if self.model._has_counter:
+            raise IfExistsWithCounterColumn('if_exists cannot be used with tables containing counter columns')
+        clone = copy.deepcopy(self)
+        clone._if_exists = True
+        return clone
+
     def update(self, **values):
         """
         Performs an update on the row selected by the queryset. Include values to update in the
@@ -1036,7 +1063,7 @@ class Row(Model):
 
         nulled_columns = set()
         us = UpdateStatement(self.column_family_name, where=self._where, ttl=self._ttl,
-                             timestamp=self._timestamp, transactions=self._transaction)
+                             timestamp=self._timestamp, transactions=self._transaction, if_exists=self._if_exists)
         for name, val in values.items():
             col_name, col_op = self._parse_filter_arg(name)
             col = self.model._columns.get(col_name)
@@ -1076,7 +1103,8 @@ class Row(Model):
             self._execute(us)
 
         if nulled_columns:
-            ds = DeleteStatement(self.column_family_name, fields=nulled_columns, where=self._where)
+            ds = DeleteStatement(self.column_family_name, fields=nulled_columns,
+                                 where=self._where, if_exists=self._if_exists)
             self._execute(ds)
 
 
@@ -1092,9 +1120,10 @@ class DMLQuery(object):
     _consistency = None
     _timestamp = None
     _if_not_exists = False
+    _if_exists = False
 
     def __init__(self, model, instance=None, batch=None, ttl=None, consistency=None, timestamp=None,
-                 if_not_exists=False, transaction=None, timeout=connection.NOT_SET):
+                 if_not_exists=False, transaction=None, timeout=connection.NOT_SET, if_exists=False):
         self.model = model
         self.column_family_name = self.model.column_family_name()
         self.instance = instance
@@ -1103,6 +1132,7 @@ def __init__(self, model, instance=None, batch=None, ttl=None, consistency=None,
         self._consistency = consistency
         self._timestamp = timestamp
         self._if_not_exists = if_not_exists
+        self._if_exists = if_exists
         self._transaction = transaction
         self._timeout = timeout
 
@@ -1111,7 +1141,7 @@ def _execute(self, q):
             return self._batch.add_query(q)
         else:
             tmp = connection.execute(q, consistency_level=self._consistency, timeout=self._timeout)
-            if self._if_not_exists or self._transaction:
+            if self._if_not_exists or self._if_exists or self._transaction:
                 check_applied(tmp)
             return tmp
 
@@ -1125,7 +1155,7 @@ def _delete_null_columns(self):
         """
         executes a delete query to remove columns that have changed to null
         """
-        ds = DeleteStatement(self.column_family_name)
+        ds = DeleteStatement(self.column_family_name, if_exists=self._if_exists)
         deleted_fields = False
         for _, v in self.instance._values.items():
             col = v.column
@@ -1159,8 +1189,8 @@ def update(self):
         assert type(self.instance) == self.model
         null_clustering_key = False if len(self.instance._clustering_keys) == 0 else True
         static_changed_only = True
-        statement = UpdateStatement(self.column_family_name, ttl=self._ttl,
-                                    timestamp=self._timestamp, transactions=self._transaction)
+        statement = UpdateStatement(self.column_family_name, ttl=self._ttl, timestamp=self._timestamp,
+                                    transactions=self._transaction, if_exists=self._if_exists)
         for name, col in self.instance._clustering_keys.items():
             null_clustering_key = null_clustering_key and col._val_is_null(getattr(self.instance, name, None))
         # get defined fields and their column names
@@ -1264,7 +1294,7 @@ def delete(self):
         if self.instance is None:
             raise CQLEngineException("DML Query instance attribute is None")
 
-        ds = DeleteStatement(self.column_family_name, timestamp=self._timestamp)
+        ds = DeleteStatement(self.column_family_name, timestamp=self._timestamp, if_exists=self._if_exists)
         for name, col in self.model._primary_keys.items():
             if (not col.partition_key) and (getattr(self.instance, name) is None):
                 continue
diff --git a/cassandra/cqlengine/statements.py b/cassandra/cqlengine/statements.py
@@ -708,7 +708,8 @@ def __init__(self,
                  where=None,
                  ttl=None,
                  timestamp=None,
-                 transactions=None):
+                 transactions=None,
+                 if_exists=False):
         super(UpdateStatement, self). __init__(table,
                                                assignments=assignments,
                                                consistency=consistency,
@@ -721,6 +722,8 @@ def __init__(self,
         for transaction in transactions or []:
             self.add_transaction_clause(transaction)
 
+        self.if_exists = if_exists
+
     def __unicode__(self):
         qs = ['UPDATE', self.table]
 
@@ -744,6 +747,9 @@ def __unicode__(self):
         if len(self.transactions) > 0:
             qs += [self._get_transactions()]
 
+        if self.if_exists:
+            qs += ["IF EXISTS"]
+
         return ' '.join(qs)
 
     def add_transaction_clause(self, clause):
@@ -778,7 +784,7 @@ def update_context_id(self, i):
 class DeleteStatement(BaseCQLStatement):
     """ a cql delete statement """
 
-    def __init__(self, table, fields=None, consistency=None, where=None, timestamp=None):
+    def __init__(self, table, fields=None, consistency=None, where=None, timestamp=None, if_exists=False):
         super(DeleteStatement, self).__init__(
             table,
             consistency=consistency,
@@ -790,6 +796,7 @@ def __init__(self, table, fields=None, consistency=None, where=None, timestamp=N
             fields = [fields]
         for field in fields or []:
             self.add_field(field)
+        self.if_exists = if_exists
 
     def update_context_id(self, i):
         super(DeleteStatement, self).update_context_id(i)
@@ -829,4 +836,7 @@ def __unicode__(self):
         if self.where_clauses:
             qs += [self._where]
 
+        if self.if_exists:
+            qs += ["IF EXISTS"]
+
         return ' '.join(qs)
diff --git a/docs/api/cassandra/cqlengine/models.rst b/docs/api/cassandra/cqlengine/models.rst
@@ -100,6 +100,24 @@ Model
 
         This method is supported on Cassandra 2.0 or later.
 
+    .. method:: if_exists()
+
+        Check the existence of an object before an update or delete. The existence of an
+        object is determined by its primary key(s). And please note using this flag
+        would incur performance cost.
+
+        If the update or delete isn't applied, a :class:`~cassandra.cqlengine.query.LWTException` is raised.
+
+        .. code-block:: python
+
+            try:
+                TestIfExistsModel.objects(id=id).if_exists().update(count=9, text='111111111111')
+            except LWTException as e:
+                # handle failure case
+                pass
+
+        This method is supported on Cassandra 2.0 or later.
+
     .. automethod:: save
 
     .. automethod:: update
diff --git a/docs/api/cassandra/cqlengine/query.rst b/docs/api/cassandra/cqlengine/query.rst
@@ -26,6 +26,10 @@ The methods here are used to filter, order, and constrain results.
 
     .. automethod:: limit
 
+    .. automethod:: if_not_exists
+
+    .. automethod:: if_exists
+
     .. automethod:: order_by
 
     .. automethod:: allow_filtering
diff --git a/tests/integration/cqlengine/test_ifexists.py b/tests/integration/cqlengine/test_ifexists.py
diff --git a/tests/integration/cqlengine/test_ifnotexists.py b/tests/integration/cqlengine/test_ifnotexists.py
