[Python] Reset BigQuery retries for each batch

CHANGES.md

@@ -89,6 +89,7 @@
 
 ## Bugfixes
 
+* (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)).
 * (Go) Fixed a data race on the Prism runner's artifact cache map in JobServices ([#32656](https://github.com/apache/beam/issues/32656)).
 * (Go) Fixed the harness leaking Data/State gRPC streams after the worker stops, and a deadlock when Send returns EOF ([#40260](https://github.com/apache/beam/issues/40260)).
 * (Go) Fixed a deadlock recreating a data channel while holding the channel lock ([#40414](https://github.com/apache/beam/issues/40414)).

sdks/python/apache_beam/io/gcp/bigquery.py

@@ -1640,12 +1640,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,
@@ -1824,6 +1818,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(
@@ -1840,7 +1840,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(

sdks/python/apache_beam/io/gcp/bigquery_test.py

@@ -2335,6 +2335,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(