Python query.BatchQuery类代码示例

    def test_batch_insert_if_not_exists_success(self):
        """ tests that batch insertion with if_not_exists work as expected """

        id = uuid4()

        with BatchQuery() as b:
            TestIfNotExistsModel.batch(b).if_not_exists().create(id=id, count=8, text='123456789')

        b = BatchQuery()
        TestIfNotExistsModel.batch(b).if_not_exists().create(id=id, count=9, text='111111111111')
        with self.assertRaises(LWTException) as assertion:
            b.execute()

        self.assertEqual(assertion.exception.existing, {
            'count': 8,
            'id': id,
            'text': '123456789',
            '[applied]': False,
        })

        q = TestIfNotExistsModel.objects(id=id)
        self.assertEqual(len(q), 1)

        tm = q.first()
        self.assertEqual(tm.count, 8)
        self.assertEqual(tm.text, '123456789')

    def store_cancer_data(self, cancer_data):
        data_group_info = cancer_data.data_group_info

        country = data_group_info.country
        age = data_group_info.age_gender.age
        gender = data_group_info.age_gender.gender

        bq = BatchQuery(consistency=ConsistencyLevel.ONE)
        for cancer_data_record in cancer_data.cancer_data:
            cancer = cancer_data_record.cancer
            deaths = cancer_data_record.deaths
            crude_rate = cancer_data_record.crude_rate
            asr = cancer_data_record.crude_rate
            cumulative_risk = cancer_data_record.cumulative_risk
            CancerDataEntity.batch(bq).create(
                age=age,
                gender=gender,
                country=country,
                cancer=cancer,
                asr=asr,
                crude_rate=crude_rate,
                cumulative_risk=cumulative_risk,
                deaths=deaths,
            )
        bq.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_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 append(self, sequenced_item_or_items):
     if isinstance(sequenced_item_or_items, list):
         if len(sequenced_item_or_items):
             b = BatchQuery()
             for item in sequenced_item_or_items:
                 assert isinstance(item, self.sequenced_item_class), (type(item), self.sequenced_item_class)
                 kwargs = self.get_field_kwargs(item)
                 self.active_record_class.batch(b).if_not_exists().create(**kwargs)
             try:
                 b.execute()
             except LWTException as e:
                 self.raise_sequenced_item_error(sequenced_item_or_items, e)
     else:
         active_record = self.to_active_record(sequenced_item_or_items)
         try:
             active_record.save()
         except LWTException as e:
             self.raise_sequenced_item_error(sequenced_item_or_items, e)

    def test_batch_update_conditional_several_rows(self):
        sync_table(TestUpdateModel)
        self.addCleanup(drop_table, TestUpdateModel)

        first_row = TestUpdateModel.create(partition=1, cluster=1, value=5, text="something")
        second_row = TestUpdateModel.create(partition=1, cluster=2, value=5, text="something")

        b = BatchQuery()
        TestUpdateModel.batch(b).if_not_exists().create(partition=1, cluster=1, value=5, text='something else')
        TestUpdateModel.batch(b).if_not_exists().create(partition=1, cluster=2, value=5, text='something else')
        TestUpdateModel.batch(b).if_not_exists().create(partition=1, cluster=3, value=5, text='something else')

        # The response will be more than two rows because two of the inserts will fail
        with self.assertRaises(LWTException):
            b.execute()

        first_row.delete()
        second_row.delete()
        b.execute()

def parse_csv(filepath, username, trip):
    """
    Parse a csv and import to database
    """
    with open(filepath, 'r') as csvfile:
        reader = csv.DictReader(csvfile.read().split('\n'))
        i = 0
        b = BatchQuery()
        last_dt = None
        for i, line in enumerate(reader):
            if i % 1000 == 0:
                b.execute()
                b = BatchQuery()
            try:
                dt = parser.parse(line['time'])
                if dt == last_dt:
                    continue
                pt = {'lon': line['lon'],
                      'lat': line['lat'],
                      'accurracy': line['accuracy'],
                      'username': username,
                      'created_at': dt,
                      'trip_id': trip}
                last_dt = dt
                Point.batch(b).create(**pt)
            except ValueError:
                continue
        b.execute()
    return username, trip

    def test_batch_update_conditional(self):
        t = TestConditionalModel.create(text='something', count=5)
        id = t.id
        with BatchQuery() as b:
            t.batch(b).iff(count=5).update(text='something else')

        updated = TestConditionalModel.objects(id=id).first()
        self.assertEqual(updated.text, 'something else')

        b = BatchQuery()
        updated.batch(b).iff(count=6).update(text='and another thing')
        with self.assertRaises(LWTException) as assertion:
            b.execute()

        self.assertEqual(assertion.exception.existing, {
            'id': id,
            'count': 5,
            '[applied]': False,
        })

        updated = TestConditionalModel.objects(id=id).first()
        self.assertEqual(updated.text, 'something else')

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_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()

        self.assertEqual(len(call_history), 2)
        self.assertEqual([(), ('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)

        self.assertEqual(batch._callbacks, [
            (my_callback, (), {}),
            (my_callback, (2,), {'named_arg':'value'}),
            (my_callback, (1, 3), {})
        ])

 def __init__(self, **kwargs):
     self.instances = []
     self._batch = BatchQuery()

    def test_empty_batch(self):
        b = BatchQuery()
        b.execute()

        with BatchQuery() as b:
            pass