Support cassandra.query.BatchType with cqlengine BatchQuery

aboudreault · aboudreault · commit c9ef7c39f0f5 · 2017-12-28T10:34:06.000-05:00
diff --git a/CHANGELOG.rst b/CHANGELOG.rst
@@ -4,6 +4,7 @@
 Features
 --------
 * cqlengine: LIKE filter operator (PYTHON-512)
+* Support cassandra.query.BatchType with cqlengine BatchQuery (PYTHON-888)
 
 Bug Fixes
 ---------
diff --git a/cassandra/cqlengine/query.py b/cassandra/cqlengine/query.py
@@ -151,7 +151,7 @@ def __init__(self, batch_type=None, timestamp=None, consistency=None, execute_on
                  timeout=conn.NOT_SET, connection=None):
         """
         :param batch_type: (optional) One of batch type values available through BatchType enum
-        :type batch_type: str or None
+        :type batch_type: BatchType, str or None
         :param timestamp: (optional) A datetime or timedelta object with desired timestamp to be applied
             to the batch conditional.
         :type timestamp: datetime or timedelta or None
@@ -225,7 +225,7 @@ def execute(self):
             self._execute_callbacks()
             return
 
-        opener = 'BEGIN ' + (self.batch_type + ' ' if self.batch_type else '') + ' BATCH'
+        opener = 'BEGIN ' + (str(self.batch_type) + ' ' if self.batch_type else '') + ' BATCH'
         if self.timestamp:
 
             if isinstance(self.timestamp, six.integer_types):
