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_insert_success_case(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() TestMultiKeyModel.get(partition=self.pkey, cluster=2)
def test_delete_success_case(self): inst = TestMultiKeyModel.create(partition=self.pkey, cluster=2, count=3, text='4') b = BatchQuery() inst.batch(b).delete() TestMultiKeyModel.get(partition=self.pkey, cluster=2) b.execute() with self.assertRaises(TestMultiKeyModel.DoesNotExist): TestMultiKeyModel.get(partition=self.pkey, cluster=2)
def test_update_success_case(self): inst = TestMultiKeyModel.create(partition=self.pkey, cluster=2, count=3, text='4') b = BatchQuery() inst.count = 4 inst.batch(b).save() inst2 = TestMultiKeyModel.get(partition=self.pkey, cluster=2) assert inst2.count == 3 b.execute() inst3 = TestMultiKeyModel.get(partition=self.pkey, cluster=2) assert inst3.count == 4
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
def test_batch_execute_no_timeout(self): with mock.patch.object(Session, 'execute', autospec=True) as mock_execute: with BatchQuery() as b: BatchQueryLogModel.batch(b).create(k=2, v=2) mock_execute.assert_called_once_with(mock.ANY, mock.ANY, mock.ANY, timeout=NOT_SET)
def test_context_manager(self): with BatchQuery() as b: for i in range(5): TestMultiKeyModel.batch(b).create(partition=self.pkey, cluster=i, count=3, text='4') for i in range(5): with self.assertRaises(TestMultiKeyModel.DoesNotExist): TestMultiKeyModel.get(partition=self.pkey, cluster=i) for i in range(5): TestMultiKeyModel.get(partition=self.pkey, cluster=i)
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
class Batch(BatchQuery): ''' Performs a batch of insert queries using async connections ''' def __init__(self, **kwargs): self.instances = [] self._batch = BatchQuery() def batch_insert(self, model_instance): self.instances.append(model_instance) def __enter__(self): return self def add_query(self, query): self._batch.add_query(query) def add_callback(self, fn, *args, **kwargs): raise TypeError('not supported') def execute(self): promises = [] session = get_session() for instance in self.instances: query = instance.__dmlquery__(instance.__class__, instance) query.batch(self._batch) query.save() for query in self._batch.queries: statement = SimpleStatement(str(query)) params = query.get_context() promises.append(session.execute_async(statement, params)) return [r.result() for r in promises] def __exit__(self, exc_type, exc_val, exc_tb): self.execute()
def test_bulk_delete_success_case(self): for i in range(1): for j in range(5): TestMultiKeyModel.create(partition=i, cluster=j, count=i*j, text='{}:{}'.format(i,j)) with BatchQuery() as b: TestMultiKeyModel.objects.batch(b).filter(partition=0).delete() assert TestMultiKeyModel.filter(partition=0).count() == 5 assert TestMultiKeyModel.filter(partition=0).count() == 0 #cleanup for m in TestMultiKeyModel.all(): m.delete()
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_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_batch_execute_on_exception_succeeds(self): # makes sure if execute_on_exception == True we still apply the batch drop_table(BatchQueryLogModel) sync_table(BatchQueryLogModel) obj = BatchQueryLogModel.objects(k=1) self.assertEqual(0, len(obj)) try: with BatchQuery(execute_on_exception=True) as b: BatchQueryLogModel.batch(b).create(k=1, v=1) raise Exception("Blah") except: pass obj = BatchQueryLogModel.objects(k=1) # should be 1 because the batch should execute self.assertEqual(1, len(obj))
def test_batch_execute_on_exception_skips_if_not_specified(self): # makes sure if execute_on_exception == True we still apply the batch drop_table(BatchQueryLogModel) sync_table(BatchQueryLogModel) obj = BatchQueryLogModel.objects(k=2) self.assertEqual(0, len(obj)) try: with BatchQuery() as b: BatchQueryLogModel.batch(b).create(k=2, v=2) raise Exception("Blah") except: pass obj = BatchQueryLogModel.objects(k=2) # should be 0 because the batch should not execute self.assertEqual(0, len(obj))
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_empty_batch(self): b = BatchQuery() b.execute() with BatchQuery() as b: pass
def __init__(self, **kwargs): self.instances = [] self._batch = BatchQuery()