diff --git a/CHANGES.md b/CHANGES.md index f0d5d06b9d9f..92d76fd30f17 100644 --- a/CHANGES.md +++ b/CHANGES.md @@ -98,6 +98,7 @@ * (Prism) Self-checkpointing splittable DoFns now resume after their requested delay instead of immediately, so polling SDFs no longer busy-spin ([#39848](https://github.com/apache/beam/issues/39848)). * (Java) MongoDbIO read splitting now preserves non-ObjectId `_id` types (e.g. string ids) instead of failing to parse the generated range filters ([#39900](https://github.com/apache/beam/issues/39900)). * (Go) Fixed GCS glob matching silently dropping objects when the glob pattern contains multi-byte characters ([#39969](https://github.com/apache/beam/issues/39969)). +* (Python) BigQuery streaming inserts now reset the retry budget and backoff for each batch instead of sharing them across a bundle ([#39995](https://github.com/apache/beam/pull/39995)). ## Security Fixes diff --git a/sdks/python/apache_beam/io/gcp/bigquery.py b/sdks/python/apache_beam/io/gcp/bigquery.py index 40f17bfe9b0e..db7ae71f26ad 100644 --- a/sdks/python/apache_beam/io/gcp/bigquery.py +++ b/sdks/python/apache_beam/io/gcp/bigquery.py @@ -1624,12 +1624,6 @@ def start_bundle(self): minimum_logging_frequency_msec ) = self.streaming_api_logging_frequency_sec * 1000 - self._backoff_calculator = iter( - retry.FuzzedExponentialIntervals( - initial_delay_secs=0.2, - num_retries=self._max_retries, - max_delay_secs=1500)) - def _create_table_if_needed(self, table_reference, schema=None): str_table_reference = '%s:%s.%s' % ( table_reference.projectId, @@ -1808,6 +1802,12 @@ def _flush_batch(self, destination): else: insert_ids = [r[1] for r in rows_and_insert_ids] + backoff_calculator = iter( + retry.FuzzedExponentialIntervals( + initial_delay_secs=0.2, + num_retries=self._max_retries, + max_delay_secs=1500)) + while True: start = time.time() passed, errors = self.bigquery_wrapper.insert_rows( @@ -1824,7 +1824,7 @@ def _flush_batch(self, destination): rows[entry['index']], entry["errors"], window_values[entry['index']]) for entry in errors] failed_insert_ids = [insert_ids[entry['index']] for entry in errors] - retry_backoff = next(self._backoff_calculator, None) + retry_backoff = next(backoff_calculator, None) # If retry_backoff is None, then we will not retry and must log. should_retry = any( diff --git a/sdks/python/apache_beam/io/gcp/bigquery_test.py b/sdks/python/apache_beam/io/gcp/bigquery_test.py index dcadee7f6a1a..69d46ee236f1 100644 --- a/sdks/python/apache_beam/io/gcp/bigquery_test.py +++ b/sdks/python/apache_beam/io/gcp/bigquery_test.py @@ -2283,6 +2283,110 @@ def test_with_batched_input_splits_large_batch(self): @unittest.skipIf(HttpError is None, 'GCP dependencies are not installed') class BigQueryStreamingInsertTransformTests(unittest.TestCase): + def _retry_test_writer(self, responses, **kwargs): + client = mock.Mock() + client.insert_rows_json.side_effect = responses + fn = beam_bq.BigQueryWriteFn( + batch_size=1, + create_disposition=beam_bq.BigQueryDisposition.CREATE_NEVER, + write_disposition=beam_bq.BigQueryDisposition.WRITE_APPEND, + with_batched_input=True, + test_client=client, + **kwargs) + fn.start_bundle() + return fn, client + + def _process_retry_batch( + self, fn, month, destination='project-id:dataset_id.table_id'): + return fn.process( + (destination, [({ + 'month': month + }, 'insertid%d' % month)]), + window_value=beam.transforms.window.GlobalWindows.windowed_value(None)) + + @parameterized.expand([ + ('same_destination', 'project-id:dataset_id.table_id'), + ('different_destination', 'project-id:dataset_id.other_table'), + ]) + @mock.patch('apache_beam.io.gcp.bigquery.time.sleep') + def test_successful_batches_do_not_consume_retries( + self, unused_name, destination, sleep): + errors = [{'index': 0, 'errors': [{'reason': 'backendError'}]}] + fn, client = self._retry_test_writer([[], errors, []], max_retries=1) + + self.assertEqual([], self._process_retry_batch(fn, 1)) + self.assertEqual([], self._process_retry_batch(fn, 2, destination)) + + self.assertEqual([[{ + 'month': 1 + }], [{ + 'month': 2 + }], [{ + 'month': 2 + }]], + [ + call.kwargs['json_rows'] + for call in client.insert_rows_json.call_args_list + ]) + self.assertEqual([['insertid1'], ['insertid2'], ['insertid2']], + [ + call.kwargs['row_ids'] + for call in client.insert_rows_json.call_args_list + ]) + sleep.assert_called_once() + + @mock.patch('apache_beam.io.gcp.bigquery.time.sleep') + def test_exhausted_batch_does_not_consume_next_batch_retries(self, sleep): + errors = [{'index': 0, 'errors': [{'reason': 'backendError'}]}] + fn, client = self._retry_test_writer( + [errors, errors, errors, []], max_retries=1) + + failed = self._process_retry_batch(fn, 1) + self.assertEqual([('project-id:dataset_id.table_id', { + 'month': 1 + })], + [ + output.value.value for output in failed + if output.tag == beam_bq.BigQueryWriteFn.FAILED_ROWS + ]) + self.assertEqual([], self._process_retry_batch(fn, 2)) + self.assertEqual(4, client.insert_rows_json.call_count) + self.assertEqual(2, sleep.call_count) + + @mock.patch('apache_beam.io.gcp.bigquery.time.sleep') + def test_successful_batches_do_not_increase_retry_backoff(self, sleep): + errors = [{'index': 0, 'errors': [{'reason': 'backendError'}]}] + fn, client = self._retry_test_writer([[]] * 20 + [errors, errors, []]) + + for month in range(21): + self.assertEqual([], self._process_retry_batch(fn, month)) + + self.assertEqual(23, client.insert_rows_json.call_count) + self.assertEqual(2, sleep.call_count) + first_delay = sleep.call_args_list[0].args[0] + second_delay = sleep.call_args_list[1].args[0] + self.assertGreaterEqual(first_delay, 0.1) + self.assertLessEqual(first_delay, 0.2) + self.assertGreaterEqual(second_delay, 0.2) + self.assertLessEqual(second_delay, 0.4) + + @mock.patch('apache_beam.io.gcp.bigquery.time.sleep') + def test_zero_retries_attempts_each_batch_once(self, sleep): + errors = [{'index': 0, 'errors': [{'reason': 'backendError'}]}] + fn, client = self._retry_test_writer([errors, []], max_retries=0) + + failed = self._process_retry_batch(fn, 1) + self.assertEqual([('project-id:dataset_id.table_id', { + 'month': 1 + })], + [ + output.value.value for output in failed + if output.tag == beam_bq.BigQueryWriteFn.FAILED_ROWS + ]) + self.assertEqual([], self._process_retry_batch(fn, 2)) + self.assertEqual(2, client.insert_rows_json.call_count) + sleep.assert_not_called() + def test_dofn_client_process_performs_batching(self): client = mock.Mock() client.tables.Get.return_value = bigquery.Table(