adding support for setting a query's batch object to None

cqlengine/query.py

@@ -428,8 +428,8 @@ class AbstractQuerySet(object):
         :param batch_obj:
         :return:
         """
-        if not isinstance(batch_obj, BatchQuery):
-            raise CQLEngineException('batch_obj must be a BatchQuery instance')
+        if batch_obj is not None and not isinstance(batch_obj, BatchQuery):
+            raise CQLEngineException('batch_obj must be a BatchQuery instance or None')
         clone = copy.deepcopy(self)
         clone._batch = batch_obj
         return clone
@@ -772,13 +772,13 @@ class DMLQuery(object):
         self.model = model
         self.column_family_name = self.model.column_family_name()
         self.instance = instance
-        self.batch = batch
+        self._batch = batch
         pass
 
     def batch(self, batch_obj):
-        if not isinstance(batch_obj, BatchQuery):
-            raise CQLEngineException('batch_obj must be a BatchQuery instance')
-        self.batch = batch_obj
+        if batch_obj is not None and not isinstance(batch_obj, BatchQuery):
+            raise CQLEngineException('batch_obj must be a BatchQuery instance or None')
+        self._batch = batch_obj
         return self
 
     def save(self):
@@ -852,8 +852,8 @@ class DMLQuery(object):
         # skip query execution if it's empty
         # caused by pointless update queries
         if qs:
-            if self.batch:
-                self.batch.add_query(qs, query_values)
+            if self._batch:
+                self._batch.add_query(qs, query_values)
             else:
                 execute(qs, query_values)
 
@@ -885,8 +885,8 @@ class DMLQuery(object):
 
             qs = ' '.join(qs)
 
-            if self.batch:
-                self.batch.add_query(qs, query_values)
+            if self._batch:
+                self._batch.add_query(qs, query_values)
             else:
                 execute(qs, query_values)
 
@@ -906,8 +906,8 @@ class DMLQuery(object):
         qs += [' AND '.join(where_statements)]
         qs = ' '.join(qs)
 
-        if self.batch:
-            self.batch.add_query(qs, field_values)
+        if self._batch:
+            self._batch.add_query(qs, field_values)
         else:
             execute(qs, field_values)
 

cqlengine/tests/query/test_batch_query.py

@@ -4,7 +4,7 @@ from uuid import uuid4
 import random
 from cqlengine import Model, columns
 from cqlengine.management import delete_table, create_table
-from cqlengine.query import BatchQuery
+from cqlengine.query import BatchQuery, DMLQuery
 from cqlengine.tests.base import BaseCassEngTestCase
 
 class TestMultiKeyModel(Model):
@@ -104,3 +104,22 @@ class BatchQueryTests(BaseCassEngTestCase):
         for m in TestMultiKeyModel.all():
             m.delete()
 
+    def test_none_success_case(self):
+        """ Tests that passing None into the batch call clears any batch object """
+        b = BatchQuery()
+
+        q = TestMultiKeyModel.objects.batch(b)
+        assert q._batch == b
+
+        q = q.batch(None)
+        assert q._batch is None
+
+    def test_dml_none_success_case(self):
+        """ Tests that passing None into the batch call clears any batch object """
+        b = BatchQuery()
+
+        q = DMLQuery(TestMultiKeyModel, batch=b)
+        assert q._batch == b
+
+        q.batch(None)
+        assert q._batch is None