Merge branch 'pull-178'

This commit is contained in:
Blake Eggleston
2014-04-17 15:28:46 -07:00
3 changed files with 164 additions and 12 deletions

View File

@@ -78,8 +78,25 @@ class BatchQuery(object):
"""
_consistency = None
def __init__(self, batch_type=None, timestamp=None, consistency=None, execute_on_exception=False):
"""
:param batch_type: (optional) One of batch type values available through BatchType enum
:type batch_type: str or None
:param timestamp: (optional) A datetime or timedelta object with desired timestamp to be applied
to the batch transaction.
:type timestamp: datetime or timedelta or None
:param consistency: (optional) One of consistency values ("ANY", "ONE", "QUORUM" etc)
:type consistency: str or None
:param execute_on_exception: (Defaults to False) Indicates that when the BatchQuery instance is used
as a context manager the queries accumulated within the context must be executed despite
encountering an error within the context. By default, any exception raised from within
the context scope will cause the batched queries not to be executed.
:type execute_on_exception: bool
:param callbacks: A list of functions to be executed after the batch executes. Note, that if the batch
does not execute, the callbacks are not executed. This, thus, effectively is a list of "on success"
callback handlers. If defined, must be a collection of callables.
:type callbacks: list or set or tuple
"""
self.queries = []
self.batch_type = batch_type
if timestamp is not None and not isinstance(timestamp, (datetime, timedelta)):
@@ -87,6 +104,7 @@ class BatchQuery(object):
self.timestamp = timestamp
self._consistency = consistency
self._execute_on_exception = execute_on_exception
self._callbacks = []
def add_query(self, query):
if not isinstance(query, BaseCQLStatement):
@@ -96,9 +114,35 @@ class BatchQuery(object):
def consistency(self, consistency):
self._consistency = consistency
def _execute_callbacks(self):
for callback, args, kwargs in self._callbacks:
callback(*args, **kwargs)
# trying to clear up the ref counts for objects mentioned in the set
del self._callbacks
def add_callback(self, fn, *args, **kwargs):
"""Add a function and arguments to be passed to it to be executed after the batch executes.
A batch can support multiple callbacks.
Note, that if the batch does not execute, the callbacks are not executed.
A callback, thus, is an "on batch success" handler.
:param fn: Callable object
:type fn: callable
:param *args: Positional arguments to be passed to the callback at the time of execution
:param **kwargs: Named arguments to be passed to the callback at the time of execution
"""
if not callable(fn):
raise ValueError("Value for argument 'fn' is {} and is not a callable object.".format(type(fn)))
self._callbacks.append((fn, args, kwargs))
def execute(self):
if len(self.queries) == 0:
# Empty batch is a no-op
# except for callbacks
self._execute_callbacks()
return
opener = 'BEGIN ' + (self.batch_type + ' ' if self.batch_type else '') + ' BATCH'
@@ -130,6 +174,7 @@ class BatchQuery(object):
execute('\n'.join(query_list), parameters, self._consistency)
self.queries = []
self._execute_callbacks()
def __enter__(self):
return self

View File

@@ -17,6 +17,7 @@ class TestMultiKeyModel(Model):
count = columns.Integer(required=False)
text = columns.Text(required=False)
class BatchQueryTests(BaseCassEngTestCase):
@classmethod
@@ -41,24 +42,13 @@ class BatchQueryTests(BaseCassEngTestCase):
b = BatchQuery()
inst = TestMultiKeyModel.batch(b).create(partition=self.pkey, cluster=2, count=3, text='4')
with self.assertRaises(TestMultiKeyModel.DoesNotExist):
TestMultiKeyModel.get(partition=self.pkey, cluster=2)
b.execute()
TestMultiKeyModel.get(partition=self.pkey, cluster=2)
def test_batch_is_executed(self):
b = BatchQuery()
inst = TestMultiKeyModel.batch(b).create(partition=self.pkey, cluster=2, count=3, text='4')
with self.assertRaises(TestMultiKeyModel.DoesNotExist):
TestMultiKeyModel.get(partition=self.pkey, cluster=2)
b.execute()
def test_update_success_case(self):
inst = TestMultiKeyModel.create(partition=self.pkey, cluster=2, count=3, text='4')
@@ -125,3 +115,80 @@ class BatchQueryTests(BaseCassEngTestCase):
with BatchQuery() as b:
pass
class BatchQueryCallbacksTests(BaseCassEngTestCase):
def test_API_managing_callbacks(self):
# Callbacks can be added at init and after
def my_callback(*args, **kwargs):
pass
# adding on init:
batch = BatchQuery()
batch.add_callback(my_callback)
batch.add_callback(my_callback, 2, named_arg='value')
batch.add_callback(my_callback, 1, 3)
assert batch._callbacks == [
(my_callback, (), {}),
(my_callback, (2,), {'named_arg':'value'}),
(my_callback, (1, 3), {})
]
def test_callbacks_properly_execute_callables_and_tuples(self):
call_history = []
def my_callback(*args, **kwargs):
call_history.append(args)
# adding on init:
batch = BatchQuery()
batch.add_callback(my_callback)
batch.add_callback(my_callback, 'more', 'args')
batch.execute()
assert len(call_history) == 2
assert [(), ('more', 'args')] == call_history
def test_callbacks_tied_to_execute(self):
"""Batch callbacks should NOT fire if batch is not executed in context manager mode"""
call_history = []
def my_callback(*args, **kwargs):
call_history.append(args)
with BatchQuery() as batch:
batch.add_callback(my_callback)
pass
assert len(call_history) == 1
class SomeError(Exception):
pass
with self.assertRaises(SomeError):
with BatchQuery() as batch:
batch.add_callback(my_callback)
# this error bubbling up through context manager
# should prevent callback runs (along with b.execute())
raise SomeError
# still same call history. Nothing added
assert len(call_history) == 1
# but if execute ran, even with an error bubbling through
# the callbacks also would have fired
with self.assertRaises(SomeError):
with BatchQuery(execute_on_exception=True) as batch:
batch.add_callback(my_callback)
# this error bubbling up through context manager
# should prevent callback runs (along with b.execute())
raise SomeError
# still same call history
assert len(call_history) == 2

View File

@@ -274,6 +274,9 @@ Batch Queries
cqlengine now supports batch queries using the BatchQuery class. Batch queries can be started and stopped manually, or within a context manager. To add queries to the batch object, you just need to precede the create/save/delete call with a call to batch, and pass in the batch object.
Batch Query General Use Pattern
-------------------------------
You can only create, update, and delete rows with a batch query, attempting to read rows out of the database with a batch query will fail.
.. code-block:: python
@@ -324,6 +327,43 @@ Batch Queries
If an exception is thrown somewhere in the block, any statements that have been added to the batch will still be executed. This is useful for some logging situations.
Batch Query Execution Callbacks
-------------------------------
In order to allow secondary tasks to be chained to the end of batch, BatchQuery instances allow callbacks to be
registered with the batch, to be executed immediately after the batch executes.
Multiple callbacks can be attached to same BatchQuery instance, they are executed in the same order that they
are added to the batch.
The callbacks attached to a given batch instance are executed only if the batch executes. If the batch is used as a
context manager and an exception is raised, the queued up callbacks will not be run.
.. code-block:: python
def my_callback(*args, **kwargs):
pass
batch = BatchQuery()
batch.add_callback(my_callback)
batch.add_callback(my_callback, 'positional arg', named_arg='named arg value')
# if you need reference to the batch within the callback,
# just trap it in the arguments to be passed to the callback:
batch.add_callback(my_callback, cqlengine_batch=batch)
# once the batch executes...
batch.execute()
# the effect of the above scheduled callbacks will be similar to
my_callback()
my_callback('positional arg', named_arg='named arg value')
my_callback(cqlengine_batch=batch)
Failure in any of the callbacks does not affect the batch's execution, as the callbacks are started after the execution
of the batch is complete.
QuerySet method reference