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 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(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(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(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(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 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 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)
class WriteToBigtable(beam.DoFn):
""" Creates the connector can call and add_row to the batcher using each
row in beam pipe line
:type beam_options: class:`~bigtable_configuration.BigtableConfiguration`
:param beam_options: Class `~bigtable_configuration.BigtableConfiguration`.
:type flush_count: int
:param flush_count: (Optional) Max number of rows to flush. 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_mutations: int
:param max_mutations: (Optional) Max number of row mutations to flush.
If it reaches the max number of row mutations it calls finish_batch() to
mutate the current row batch. Default is MAX_MUTATIONS (100000 mutations).
: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).
:type app_profile_id: str
:param app_profile_id: (Optional) The unique name of the AppProfile.
"""
def __init__(self, beam_options, flush_count=None, max_row_bytes=None,
app_profile_id=None):
super(WriteToBigtable, self).__init__(beam_options)
self.beam_options = beam_options
self.client = None
self.instance = None
self.table = None
self.batcher = None
self._app_profile_id = app_profile_id
self.flush_count = flush_count
self.max_row_bytes = max_row_bytes
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 process(self, row):
# row.table = self.table
self.batcher.mutate(row)
def finish_bundle(self):
return self.batcher.flush()