def test_spanner_update(self): _prefex = 'test_update' # adding data to perform test instance = self._SPANNER_INSTANCE database = instance.database(self.TEST_DATABASE) data = [ (_prefex + '1', _prefex + 'inset-1'), (_prefex + '2', _prefex + 'inset-2'), (_prefex + '3', _prefex + 'inset-3'), ] with database.batch() as batch: batch.insert(table='Users', columns=('UserId', 'Key'), values=data) mutations_update = [ WriteMutation.update('Users', ('UserId', 'Key'), [(_prefex + '1', _prefex + 'update-1')]), WriteMutation.update('Users', ('UserId', 'Key'), [(_prefex + '2', _prefex + 'update-2')]), WriteMutation.delete('Users', spanner.KeySet(keys=[[_prefex + '3']])) ] p = beam.Pipeline(argv=self.args) _ = (p | beam.Create(mutations_update) | WriteToSpanner(project_id=self.project, instance_id=self.instance, database_id=self.TEST_DATABASE)) res = p.run() res.wait_until_finish() self.assertEqual(self._count_data(_prefex), 2)
def test_spanner_write_mutation_groups(self, mock_batch_snapshot_class, mock_batch_checkout): ks = spanner.KeySet(keys=[[1233], [1234]]) mutation_groups = [ MutationGroup([ WriteMutation.insert("roles", ("key", "rolename"), [('9001233', "mutations-inset-1233")]), WriteMutation.insert("roles", ("key", "rolename"), [('9001234', "mutations-inset-1234")]) ]), MutationGroup([ WriteMutation.update( "roles", ("key", "rolename"), [('9001234', "mutations-inset-9001233-updated")]) ]), MutationGroup([WriteMutation.delete("roles", ks)]) ] p = TestPipeline() _ = (p | beam.Create(mutation_groups) | WriteToSpanner(project_id=TEST_PROJECT_ID, instance_id=TEST_INSTANCE_ID, database_id=_generate_database_name(), max_batch_size_bytes=100)) res = p.run() res.wait_until_finish() metric_results = res.metrics().query( MetricsFilter().with_name('SpannerBatches')) batches_counter = metric_results['counters'][0] self.assertEqual(batches_counter.committed, 3) self.assertEqual(batches_counter.attempted, 3)
def test_spanner_error(self): mutations_update = [ WriteMutation.update('Users', ('UserId', 'Key'), [('INVALD_ID', 'Error-error')]), ] with self.assertRaises(Exception): p = beam.Pipeline(argv=self.args) _ = (p | beam.Create(mutations_update) | WriteToSpanner(project_id=self.project, instance_id=self.instance, database_id=self.TEST_DATABASE)) p.run()