def mutations_batcher(self,
flush_count=FLUSH_COUNT,
max_row_bytes=MAX_ROW_BYTES):
"""Factory to create a mutation batcher associated with this instance.
For example:
.. literalinclude:: snippets_table.py
:start-after: [START bigtable_mutations_batcher]
:end-before: [END bigtable_mutations_batcher]
:type table: class
:param table: class:`~google.cloud.bigtable.table.Table`.
:type flush_count: int
:param flush_count: (Optional) Maximum number of rows per batch. If it
reaches the max number of rows it calls finish_batch() to
mutate the current row batch. Default is FLUSH_COUNT (1000
rows).
:type max_row_bytes: int
:param max_row_bytes: (Optional) Max number of row mutations size to
flush. If it reaches the max number of row mutations size it
calls finish_batch() to mutate the current row batch.
Default is MAX_ROW_BYTES (5 MB).
"""
return MutationsBatcher(self, flush_count, max_row_bytes)
def test_constructor(self):
credentials = _make_credentials()
client = self._make_client(project="project-id",
credentials=credentials,
admin=True)
instance = client.instance(instance_id="instance-id")
table = self._make_table(self.TABLE_ID, instance)
mutation_batcher = MutationsBatcher(table)
self.assertEqual(table, mutation_batcher.table)
def test_mutate_row_with_max_mutations(self):
table = _Table(self.TABLE_NAME)
mutation_batcher = MutationsBatcher(table=table)
row = DirectRow(row_key=b"row_key")
row.set_cell("cf1", b"c1", 1)
row.set_cell("cf1", b"c2", 2)
row.set_cell("cf1", b"c3", 3)
mutation_batcher.mutate(row)
mutation_batcher.flush()
self.assertEqual(table.mutation_calls, 1)
def test_add_row_with_max_flush_count(self):
table = _Table(self.TABLE_NAME)
mutation_batcher = MutationsBatcher(table=table, flush_count=3)
row_1 = DirectRow(row_key=b"row_key_1")
row_2 = DirectRow(row_key=b"row_key_2")
row_3 = DirectRow(row_key=b"row_key_3")
mutation_batcher.mutate(row_1)
mutation_batcher.mutate(row_2)
mutation_batcher.mutate(row_3)
self.assertEqual(table.mutation_calls, 1)
def test_mutate_row_with_max_mutations(self):
table = _Table(self.TABLE_NAME)
mutation_batcher = MutationsBatcher(table=table)
row = DirectRow(row_key=b'row_key')
row.set_cell('cf1', b'c1', 1)
row.set_cell('cf1', b'c2', 2)
row.set_cell('cf1', b'c3', 3)
mutation_batcher.mutate(row)
mutation_batcher.flush()
self.assertEqual(table.mutation_calls, 1)
def test_mutate_row_with_max_mutations_failure(self):
from google.cloud.bigtable.batcher import MaxMutationsError
table = _Table(self.TABLE_NAME)
mutation_batcher = MutationsBatcher(table=table)
row = DirectRow(row_key=b"row_key")
row.set_cell("cf1", b"c1", 1)
row.set_cell("cf1", b"c2", 2)
row.set_cell("cf1", b"c3", 3)
row.set_cell("cf1", b"c4", 4)
with self.assertRaises(MaxMutationsError):
mutation_batcher.mutate(row)
def test_mutate_row_with_max_mutations_failure(self):
from google.cloud.bigtable.batcher import MaxMutationsError
table = _Table(self.TABLE_NAME)
mutation_batcher = MutationsBatcher(table=table)
row = DirectRow(row_key=b'row_key')
row.set_cell('cf1', b'c1', 1)
row.set_cell('cf1', b'c2', 2)
row.set_cell('cf1', b'c3', 3)
row.set_cell('cf1', b'c4', 4)
with self.assertRaises(MaxMutationsError):
mutation_batcher.mutate(row)
def test_mutate_row(self):
table = _Table(self.TABLE_NAME)
mutation_batcher = MutationsBatcher(table=table)
rows = [
DirectRow(row_key=b"row_key"),
DirectRow(row_key=b"row_key_2"),
DirectRow(row_key=b"row_key_3"),
DirectRow(row_key=b"row_key_4"),
]
mutation_batcher.mutate_rows(rows)
mutation_batcher.flush()
self.assertEqual(table.mutation_calls, 1)
def start_bundle(self):
if self.beam_options.credentials is None:
self.client = bigtable.Client(project=self.beam_options.project_id,
admin=True)
else:
self.client = bigtable.Client(
project=self.beam_options.project_id,
credentials=self.beam_options.credentials,
admin=True)
self.instance = self.client.instance(self.beam_options.instance_id)
self.table = self.instance.table(self.beam_options.table_id,
self._app_profile_id)
self.batcher = MutationsBatcher(
self.table, flush_count=self.flush_count,
max_row_bytes=self.max_row_bytes)
def test_mutate_row_with_max_row_bytes(self):
table = _Table(self.TABLE_NAME)
mutation_batcher = MutationsBatcher(table=table,
max_row_bytes=3 * 1024 * 1024)
number_of_bytes = 1 * 1024 * 1024
max_value = b"1" * number_of_bytes
row = DirectRow(row_key=b"row_key")
row.set_cell("cf1", b"c1", max_value)
row.set_cell("cf1", b"c2", max_value)
row.set_cell("cf1", b"c3", max_value)
mutation_batcher.mutate(row)
self.assertEqual(table.mutation_calls, 1)
def test_mutate_row_with_max_row_bytes(self):
table = _Table(self.TABLE_NAME)
mutation_batcher = MutationsBatcher(table=table,
max_row_bytes=3 * 1024 * 1024)
number_of_bytes = 1 * 1024 * 1024
max_value = b'1' * number_of_bytes
row = DirectRow(row_key=b'row_key')
row.set_cell('cf1', b'c1', max_value)
row.set_cell('cf1', b'c2', max_value)
row.set_cell('cf1', b'c3', max_value)
mutation_batcher.mutate(row)
self.assertEqual(table.mutation_calls, 1)
def test_flush_with_no_rows(self):
table = _Table(self.TABLE_NAME)
mutation_batcher = MutationsBatcher(table=table)
mutation_batcher.flush()
self.assertEqual(table.mutation_calls, 0)
def _make_mutation_batcher(table, **kw):
from google.cloud.bigtable.batcher import MutationsBatcher
return MutationsBatcher(table, **kw)