updating batch query to work with 1.2

cqlengine/query.py

@@ -126,22 +126,20 @@ class LessThanOrEqualOperator(QueryOperator):
     symbol = "LTE"
     cql_symbol = '<='
 
-class Consistency(object):
-    ANY = 'ANY'
-    ONE = 'ONE'
-    QUORUM = 'QUORUM'
-    LOCAL_QUORUM = 'LOCAL_QUORUM'
-    EACH_QUORUM = 'EACH_QUORUM'
-    ALL = 'ALL'
+class BatchType(object):
+    Unlogged    = 'UNLOGGED'
+    Counter     = 'COUNTER'
 
 class BatchQuery(object):
     """
     Handles the batching of queries
+
+    http://www.datastax.com/docs/1.2/cql_cli/cql/BATCH
     """
 
-    def __init__(self, consistency=None, timestamp=None):
+    def __init__(self, batch_type=None, timestamp=None):
         self.queries = []
-        self.consistency = consistency
+        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
@@ -150,13 +148,11 @@ class BatchQuery(object):
         self.queries.append((query, params))
 
     def execute(self):
-        opener = 'BEGIN BATCH'
-        if self.consistency:
-            opener += ' USING CONSISTENCY {}'.format(self.consistency)
+        opener = 'BEGIN ' + (self.batch_type + ' ' if self.batch_type else '') + ' BATCH'
         if self.timestamp:
             epoch = datetime(1970, 1, 1)
             ts = long((self.timestamp - epoch).total_seconds() * 1000)
-            opener += ' TIMESTAMP {}'.format(ts)
+            opener += ' USING TIMESTAMP {}'.format(ts)
 
         query_list = [opener]
         parameters = {}

cqlengine/tests/test_batch_query.py

@@ -3,7 +3,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, Consistency
+from cqlengine.query import BatchQuery
 from cqlengine.tests.base import BaseCassEngTestCase
 
 class TestMultiKeyModel(Model):