diff --git a/sdks/python/apache_beam/io/gcp/bigtableio_test.py b/sdks/python/apache_beam/io/gcp/bigtableio_test.py index d9ef12a16592..7c371c9b383d 100644 --- a/sdks/python/apache_beam/io/gcp/bigtableio_test.py +++ b/sdks/python/apache_beam/io/gcp/bigtableio_test.py @@ -331,6 +331,61 @@ def test_write_metrics(self): ServiceCallMetric.bigtable_error_code_to_grpc_status_string(OK), 2) + def test_write_batch_error_surfaces_from_async_flush(self): + write_fn = bigtableio._BigTableWriteFn( + self._PROJECT_ID, + self._INSTANCE_ID, + self._TABLE_ID, + flush_count=1, + max_row_bytes=5242880) + write_fn.table = self.table + write_fn.start_bundle() + + direct_rows = [self.generate_row(i) for i in range(5)] + with patch.object(Table, + 'mutate_rows', + side_effect=Exception('batch RPC failed')): + for direct_row in direct_rows: + write_fn.process(direct_row) + with self.assertRaises(Exception): + write_fn.finish_bundle() + + def test_write_batch_error_surfaces_from_buffered_rows(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() + + mock_mutate = MagicMock(side_effect=Exception('batch RPC failed')) + with patch.object(Table, 'mutate_rows', mock_mutate): + write_fn.process(self.generate_row(0)) + with self.assertRaises(Exception): + write_fn.finish_bundle() + self.assertGreater( + mock_mutate.call_count, 0, 'buffered row was never flushed') + + def test_write_close_error_is_surfaced(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() + + with patch.object(MutationsBatcher, + 'close', + side_effect=Exception('error on close')) as mock_close: + write_fn.process(self.generate_row(0)) + with self.assertRaises(Exception): + write_fn.finish_bundle() + mock_close.assert_called_once() + def generate_row(self, index=0): rand = choice(string.ascii_letters + string.digits) value = ''.join(rand for i in range(100)) diff --git a/sdks/python/setup.py b/sdks/python/setup.py index 61dd55cefa21..5b02dc9ab0ef 100644 --- a/sdks/python/setup.py +++ b/sdks/python/setup.py @@ -531,7 +531,9 @@ def get_portability_package_data(): 'google-cloud-bigquery>=2.0.0,<4', 'google-cloud-bigquery-storage>=2.6.3,<3', 'google-cloud-core>=2.0.0,<3', - 'google-cloud-bigtable>=2.19.0,<3', + # 2.42.0 improves MutationsBatcher error handling: it surfaces + # errors raised during async flushes instead of swallowing them. + 'google-cloud-bigtable>=2.42.0,<3', 'google-cloud-build>=3.35.0,<4', 'google-cloud-spanner>=3.0.0,<4', # GCP Packages required by ML functionality