From 039ce02f7a3468e0179d382c31d6e91942783a56 Mon Sep 17 00:00:00 2001 From: Shunping Huang Date: Fri, 4 Sep 2026 23:37:46 -0400 Subject: [PATCH] Fix Bigtable DirectRow mutation protobuf retrieval for google-cloud-bigtable >= 2.44.0 --- sdks/python/apache_beam/io/gcp/bigtableio.py | 8 +++++++- .../apache_beam/io/gcp/bigtableio_test.py | 17 +++++++++++++---- 2 files changed, 20 insertions(+), 5 deletions(-) diff --git a/sdks/python/apache_beam/io/gcp/bigtableio.py b/sdks/python/apache_beam/io/gcp/bigtableio.py index 442ba5d68c8b..cd78deb7466c 100644 --- a/sdks/python/apache_beam/io/gcp/bigtableio.py +++ b/sdks/python/apache_beam/io/gcp/bigtableio.py @@ -283,7 +283,13 @@ class _DirectRowMutationsToBeamRow(beam.DoFn): def process(self, direct_row): args = {"key": direct_row.row_key, "mutations": []} # start accumulating mutations in a list - for mutation in direct_row._get_mutations(): + # In google-cloud-bigtable >= 2.44.0, _get_mutations() returns Python + # dataclass objects (RowMutationEntry) instead of protobuf messages. + # Use _get_mutation_pbs() to retrieve Mutation protobuf objects. + mutations = ( + direct_row._get_mutation_pbs() if hasattr( + direct_row, '_get_mutation_pbs') else direct_row._get_mutations()) + for mutation in mutations: if mutation.__contains__("set_cell"): mutation_dict = { "type": b'SetCell', diff --git a/sdks/python/apache_beam/io/gcp/bigtableio_test.py b/sdks/python/apache_beam/io/gcp/bigtableio_test.py index 7c371c9b383d..08c33017f9c7 100644 --- a/sdks/python/apache_beam/io/gcp/bigtableio_test.py +++ b/sdks/python/apache_beam/io/gcp/bigtableio_test.py @@ -119,6 +119,15 @@ def test_beam_row_to_bigtable_row(self): class TestBigtableDirectRowToBeamRow(unittest.TestCase): doFn = bigtableio.WriteToBigTable._DirectRowMutationsToBeamRow() + @staticmethod + def _get_mutation_pbs(direct_row): + # In google-cloud-bigtable >= 2.44.0, _get_mutations() returns Python + # dataclass objects instead of protobuf messages; use _get_mutation_pbs() + # to retrieve Mutation protobuf messages. + if hasattr(direct_row, '_get_mutation_pbs'): + return direct_row._get_mutation_pbs() + return direct_row._get_mutations() + def test_set_cell(self): # create some set cell mutations direct_row: DirectRow = DirectRow('key-1') @@ -144,7 +153,7 @@ def test_set_cell(self): # sort both lists of mutations for convenience beam_row_mutations = sorted(beam_row.mutations, key=lambda m: m['value']) bt_row_mutations = sorted( - direct_row._get_mutations(), key=lambda m: m.set_cell.value) + self._get_mutation_pbs(direct_row), key=lambda m: m.set_cell.value) self.assertEqual(beam_row.key, direct_row.row_key) self.assertEqual(len(beam_row_mutations), len(bt_row_mutations)) @@ -186,7 +195,7 @@ def test_delete_cells(self): beam_row_mutations = sorted( beam_row.mutations, key=lambda m: m['column_qualifier']) bt_row_mutations = sorted( - direct_row._get_mutations(), + self._get_mutation_pbs(direct_row), key=lambda m: m.delete_from_column.column_qualifier) self.assertEqual(beam_row.key, direct_row.row_key) self.assertEqual(len(beam_row_mutations), len(bt_row_mutations)) @@ -232,8 +241,8 @@ def test_delete_column_family(self): beam_row_mutations = sorted( beam_row.mutations, key=lambda m: m['family_name']) bt_row_mutations = sorted( - direct_row._get_mutations(), - key=lambda m: m.delete_from_column.family_name) + self._get_mutation_pbs(direct_row), + key=lambda m: m.delete_from_family.family_name) self.assertEqual(beam_row.key, direct_row.row_key) self.assertEqual(len(beam_row_mutations), len(bt_row_mutations))