From f90cfb046a7e727a160e4e331d799064d1588c06 Mon Sep 17 00:00:00 2001
From: Blake Eggleston <bdeggleston@gmail.com>
Date: Thu, 27 Jun 2013 13:10:16 -0700
Subject: [PATCH] adding support for setting a query's batch object to None

---
 cqlengine/query.py                        | 24 +++++++++++------------
 cqlengine/tests/query/test_batch_query.py | 21 +++++++++++++++++++-
 2 files changed, 32 insertions(+), 13 deletions(-)

diff --git a/cqlengine/query.py b/cqlengine/query.py
index b8b4926d..d86b504b 100644
--- a/cqlengine/query.py
+++ b/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)
 
diff --git a/cqlengine/tests/query/test_batch_query.py b/cqlengine/tests/query/test_batch_query.py
index 5b31826a..37e8d28d 100644
--- a/cqlengine/tests/query/test_batch_query.py
+++ b/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