poll BigQuery matchers on missing tables

.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": 3
+    "modification": 7
 }

.github/workflows/beam_PostCommit_Python_Arm.yml

@@ -58,7 +58,7 @@ jobs:
     # apache-beam-testing Snapshots image tagged with sdk_version, plus an
     # aarch64 wheel from this checkout. Keep headroom for the IT suite and
     # qemu wheel build.
-    timeout-minutes: 180
+    timeout-minutes: 210
     strategy:
       fail-fast: false
       matrix:

sdks/python/apache_beam/io/gcp/tests/bigquery_matcher.py

@@ -48,7 +48,13 @@
 
 
 def retry_on_http_timeout_and_value_error(exception):
-  """Filter allowing retries on Bigquery errors and value error."""
+  """Filter allowing retries on Bigquery errors and value error.
+
+  A missing table (HTTP 404) is not retried here. Streaming inserts create
+  the table on the first write, and matchers with a timeout poll for it.
+  """
+  if getattr(exception, 'code', None) == 404:
+    return False
   return isinstance(
       exception,
       (GoogleCloudError, ValueError, concurrent.futures.TimeoutError))
@@ -87,12 +93,18 @@ def __init__(self, project, query, checksum, timeout_secs=0):
   def _matches(self, _):
     @retry.with_exponential_backoff(
         num_retries=1000,
-        initial_delay_secs=0.5,
+        initial_delay_secs=5,
         max_delay_secs=30,
         stop_after_secs=self.timeout_secs,
     )
     def get_checksum():
-      response = self._query_with_retry()
+      try:
+        response = self._query_with_retry()
+      except GoogleCloudError as exc:
+        # Keep polling until timeout_secs. The query retry does not retry 404.
+        if self.timeout_secs and getattr(exc, 'code', None) == 404:
+          raise ValueError('BigQuery table is not visible yet: %s' % exc)
+        raise
       _LOGGER.info(
           'Read from given query (%s), total rows %d',
           self.query,
@@ -196,11 +208,19 @@ def __init__(self, project, query, data, timeout=DEFAULT_TIMEOUT):
   def _get_query_result(self):
     start_time = time.time()
     while time.time() - start_time <= self.timeout:
-      response = self._query_with_retry()
+      try:
+        response = self._query_with_retry()
+      except GoogleCloudError as exc:
+        # The destination table does not exist until the first streaming write.
+        if getattr(exc, 'code', None) != 404:
+          raise
+        _LOGGER.info('BigQuery table is not visible yet: %s', exc)
+        response = []
       if len(response) >= len(self.expected_data):
         return response
       _LOGGER.debug('Query result contains %d rows' % len(response))
-      time.sleep(1)
+      # Space out polls so repeated 404s do not hit BigQuery too often.
+      time.sleep(5)
     raise TimeoutError('Timeout exceeded for matcher.')  # noqa: F821
 
 

sdks/python/apache_beam/io/gcp/tests/bigquery_matcher_test.py

@@ -33,9 +33,11 @@
 # pylint: disable=wrong-import-order, wrong-import-position
 try:
   from google.cloud import bigquery
+  from google.cloud.exceptions import InternalServerError
   from google.cloud.exceptions import NotFound
 except ImportError:
   bigquery = None
+  InternalServerError = None
   NotFound = None
 # pylint: enable=wrong-import-order, wrong-import-position
 
@@ -87,14 +89,45 @@ def test_bigquery_matcher_success_streaming_retry(self, mock_bigquery):
     self.assertEqual(2, mock_query.call_count)
 
   def test_bigquery_matcher_query_error_retry(self, mock_bigquery):
+    mock_query = mock_bigquery.return_value.query
+    mock_query.side_effect = InternalServerError('internal')
+
+    matcher = bq_verifier.BigqueryMatcher(
+        'mock_project', 'mock_query', 'mock_checksum')
+    with self.assertRaises(InternalServerError):
+      hc_assert_that(self._mock_result, matcher)
+    self.assertEqual(bq_verifier.MAX_RETRIES + 1, mock_query.call_count)
+
+  def test_bigquery_matcher_not_found_without_timeout(self, mock_bigquery):
     mock_query = mock_bigquery.return_value.query
     mock_query.side_effect = NotFound('table not found')
 
     matcher = bq_verifier.BigqueryMatcher(
         'mock_project', 'mock_query', 'mock_checksum')
     with self.assertRaises(NotFound):
       hc_assert_that(self._mock_result, matcher)
-    self.assertEqual(bq_verifier.MAX_RETRIES + 1, mock_query.call_count)
+    self.assertEqual(1, mock_query.call_count)
+
+  def test_bigquery_matcher_not_found_retries_until_visible(
+      self, mock_bigquery):
+    mock_query_result = [mock.Mock(), mock.Mock(), mock.Mock()]
+    mock_query_result[0].values.return_value = []
+    mock_query_result[1].values.return_value = None
+    mock_query_result[2].values.return_value = None
+    ready = mock.Mock()
+    ready.result.return_value = mock_query_result
+
+    mock_query = mock_bigquery.return_value.query
+    mock_query.side_effect = [NotFound('table not found'), ready]
+
+    matcher = bq_verifier.BigqueryMatcher(
+        'mock_project',
+        'mock_query',
+        '59f9d6bdee30d67ea73b8aded121c3a0280f9cd8',
+        timeout_secs=5,
+    )
+    hc_assert_that(self._mock_result, matcher)
+    self.assertEqual(2, mock_query.call_count)
 
   def test_bigquery_matcher_query_error_checksum(self, mock_bigquery):
     empty_query_result = []
@@ -162,14 +195,44 @@ def test_bigquery_table_matcher_query_error_retry(self, mock_bigquery):
 class BigqueryFullResultStreamingMatcherTest(unittest.TestCase):
   def setUp(self):
     self.timeout = 0.01
+    # The matcher waits between polls. Skip that wait here.
+    sleep_patcher = mock.patch(
+        'apache_beam.io.gcp.tests.bigquery_matcher.time.sleep')
+    sleep_patcher.start()
+    self.addCleanup(sleep_patcher.stop)
 
   def test__get_query_result_timeout(self, mock_bigquery):
     mock_query = mock_bigquery.return_value.query
     mock_query.return_value.result.return_value = []
 
     matcher = bq_verifier.BigqueryFullResultStreamingMatcher(
         'some-project', 'some-query', [1, 2, 3], timeout=self.timeout)
-    with self.assertRaises(TimeoutError):  # noqa: F821
+    with self.assertRaises(TimeoutError):
+      matcher._get_query_result()
+
+  def test_get_query_result_retries_missing_table(self, mock_bigquery):
+    rows = []
+    for value in (1, 2, 3):
+      row = mock.Mock()
+      row.values.return_value = (value, )
+      rows.append(row)
+    ready = mock.Mock()
+    ready.result.return_value = rows
+    mock_query = mock_bigquery.return_value.query
+    mock_query.side_effect = [NotFound('table not found'), ready]
+
+    matcher = bq_verifier.BigqueryFullResultStreamingMatcher(
+        'some-project', 'some-query', [(1, ), (2, ), (3, )], timeout=5)
+    self.assertEqual([(1, ), (2, ), (3, )], matcher._get_query_result())
+    self.assertEqual(2, mock_query.call_count)
+
+  def test_get_query_result_missing_table_timeout(self, mock_bigquery):
+    mock_query = mock_bigquery.return_value.query
+    mock_query.side_effect = NotFound('table not found')
+
+    matcher = bq_verifier.BigqueryFullResultStreamingMatcher(
+        'some-project', 'some-query', [1, 2, 3], timeout=self.timeout)
+    with self.assertRaises(TimeoutError):
       matcher._get_query_result()