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)