From 6b1391749edaccbbbd83767a16b8cacb8903a92c Mon Sep 17 00:00:00 2001
From: Jon Haddad <jon@jonhaddad.com>
Date: Thu, 24 Oct 2013 18:19:04 -0700
Subject: [PATCH] consistency with batch verified

---
 cqlengine/query.py                  |  9 +++++++--
 cqlengine/tests/test_consistency.py | 17 ++++++++++++++++-
 2 files changed, 23 insertions(+), 3 deletions(-)

diff --git a/cqlengine/query.py b/cqlengine/query.py
index d67b45d4..0e9395de 100644
--- a/cqlengine/query.py
+++ b/cqlengine/query.py
@@ -208,17 +208,22 @@ class BatchQuery(object):
 
     http://www.datastax.com/docs/1.2/cql_cli/cql/BATCH
     """
+    _consistency = None
 
-    def __init__(self, batch_type=None, timestamp=None):
+    def __init__(self, batch_type=None, timestamp=None, consistency=None):
         self.queries = []
         self.batch_type = batch_type
         if timestamp is not None and not isinstance(timestamp, datetime):
             raise CQLEngineException('timestamp object must be an instance of datetime')
         self.timestamp = timestamp
+        self._consistency = consistency
 
     def add_query(self, query, params):
         self.queries.append((query, params))
 
+    def consistency(self, consistency):
+        self._consistency = consistency
+
     def execute(self):
         if len(self.queries) == 0:
             # Empty batch is a no-op
@@ -238,7 +243,7 @@ class BatchQuery(object):
 
         query_list.append('APPLY BATCH;')
 
-        execute('\n'.join(query_list), parameters)
+        execute('\n'.join(query_list), parameters, self._consistency)
 
         self.queries = []
 
diff --git a/cqlengine/tests/test_consistency.py b/cqlengine/tests/test_consistency.py
index 27c455ef..f4c5986b 100644
--- a/cqlengine/tests/test_consistency.py
+++ b/cqlengine/tests/test_consistency.py
@@ -5,7 +5,7 @@ from uuid import uuid4
 from cqlengine import columns
 import mock
 from cqlengine.connection import ConnectionPool
-from cqlengine import ALL
+from cqlengine import ALL, BatchQuery
 
 class TestConsistencyModel(Model):
     id      = columns.UUID(primary_key=True, default=lambda:uuid4())
@@ -50,3 +50,18 @@ class TestConsistency(BaseConsistencyTest):
         self.assertEqual(ALL, args[0][2])
 
 
+    def test_batch_consistency(self):
+
+        with mock.patch.object(ConnectionPool, 'execute') as m:
+            with BatchQuery(consistency=ALL) as b:
+                TestConsistencyModel.batch(b).create(text="monkey")
+
+        args = m.call_args
+        self.assertEqual(ALL, args[0][2])
+
+        with mock.patch.object(ConnectionPool, 'execute') as m:
+            with BatchQuery() as b:
+                TestConsistencyModel.batch(b).create(text="monkey")
+
+        args = m.call_args
+        self.assertNotEqual(ALL, args[0][2])