diff --git a/.github/trigger_files/beam_PostCommit_Python_Arm.json b/.github/trigger_files/beam_PostCommit_Python_Arm.json index 1efc8e9e4405..bbdc3a3910ef 100644 --- a/.github/trigger_files/beam_PostCommit_Python_Arm.json +++ b/.github/trigger_files/beam_PostCommit_Python_Arm.json @@ -1,4 +1,4 @@ { "comment": "Modify this file in a trivial way to cause this test suite to run", - "modification": 1 + "modification": 3 } diff --git a/sdks/python/apache_beam/io/gcp/bigtableio.py b/sdks/python/apache_beam/io/gcp/bigtableio.py index cd78deb7466c..f1a19e179f92 100644 --- a/sdks/python/apache_beam/io/gcp/bigtableio.py +++ b/sdks/python/apache_beam/io/gcp/bigtableio.py @@ -62,15 +62,40 @@ from google.cloud.bigtable import Client from google.cloud.bigtable.batcher import MutationsBatcher from google.cloud.bigtable.row import Cell + from google.cloud.bigtable.row import DirectRow from google.cloud.bigtable.row import PartialRowData except ImportError: + DirectRow = None _LOGGER.warning( 'ImportError: from google.cloud.bigtable import Client', exc_info=True) __all__ = ['WriteToBigTable', 'ReadFromBigtable'] +def _restore_direct_row_pb_mutations(row): + # google-cloud-bigtable >= 2.44.0 stores mutations on `_mutations`. + # Older MutationsBatcher reads `_pb_mutations`. Fill that attribute so + # pickle and worker batching both see protobuf mutations. + if hasattr(row, '_pb_mutations'): + return + mutations = getattr(row, '_mutations', None) + if mutations is None: + return + row._pb_mutations = [ + mut._to_pb() if hasattr(mut, '_to_pb') else mut for mut in mutations + ] + + +def _direct_row_getstate(self): + _restore_direct_row_pb_mutations(self) + return self.__dict__ + + +if DirectRow is not None: + DirectRow.__getstate__ = _direct_row_getstate + + class _BigTableWriteFn(beam.DoFn): """ Creates the connector can call and add_row to the batcher using each row in beam pipe line @@ -168,6 +193,7 @@ def process(self, row): # 'field1', # 'value1', # timestamp=datetime.now()) + _restore_direct_row_pb_mutations(row) self.batcher.mutate(row) def finish_bundle(self): diff --git a/sdks/python/apache_beam/io/gcp/bigtableio_test.py b/sdks/python/apache_beam/io/gcp/bigtableio_test.py index 08c33017f9c7..305b22eedc99 100644 --- a/sdks/python/apache_beam/io/gcp/bigtableio_test.py +++ b/sdks/python/apache_beam/io/gcp/bigtableio_test.py @@ -377,6 +377,40 @@ def test_write_batch_error_surfaces_from_buffered_rows(self): self.assertGreater( mock_mutate.call_count, 0, 'buffered row was never flushed') + def test_write_direct_row_without_pb_mutations(self): + write_fn = bigtableio._BigTableWriteFn( + self._PROJECT_ID, + self._INSTANCE_ID, + self._TABLE_ID, + flush_count=1000, + max_row_bytes=5242880) + write_fn.table = self.table + write_fn.start_bundle() + + class Mutation: + def _to_pb(self): + return 'mutation-pb' + + class LegacyDirectRow: + def __init__(self): + self._mutations = [Mutation()] + + row = LegacyDirectRow() + with patch.object(MutationsBatcher, 'mutate') as mock_mutate: + write_fn.process(row) + mock_mutate.assert_called_once_with(row) + self.assertEqual(row._pb_mutations, ['mutation-pb']) + + def test_direct_row_getstate_includes_pb_mutations(self): + row = DirectRow(row_key=b'key-1') + if hasattr(row, '_pb_mutations'): + delattr(row, '_pb_mutations') + mutation = MagicMock() + mutation._to_pb.return_value = 'mutation-pb' + row._mutations = [mutation] + state = bigtableio._direct_row_getstate(row) + self.assertEqual(state['_pb_mutations'], ['mutation-pb']) + def test_write_close_error_is_surfaced(self): write_fn = bigtableio._BigTableWriteFn( self._PROJECT_ID,