diff --git a/.github/trigger_files/beam_PostCommit_Python_Arm.json b/.github/trigger_files/beam_PostCommit_Python_Arm.json index bbdc3a3910ef..5cf4f475f317 100644 --- a/.github/trigger_files/beam_PostCommit_Python_Arm.json +++ b/.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 } diff --git a/.github/workflows/beam_PostCommit_Python_Arm.yml b/.github/workflows/beam_PostCommit_Python_Arm.yml index eb308c21f7a1..e0471e7ba747 100644 --- a/.github/workflows/beam_PostCommit_Python_Arm.yml +++ b/.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: diff --git a/sdks/python/apache_beam/io/gcp/tests/bigquery_matcher.py b/sdks/python/apache_beam/io/gcp/tests/bigquery_matcher.py index ff48a5644916..6a70850a780e 100644 --- a/sdks/python/apache_beam/io/gcp/tests/bigquery_matcher.py +++ b/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 diff --git a/sdks/python/apache_beam/io/gcp/tests/bigquery_matcher_test.py b/sdks/python/apache_beam/io/gcp/tests/bigquery_matcher_test.py index f0cf14fedced..34657d243456 100644 --- a/sdks/python/apache_beam/io/gcp/tests/bigquery_matcher_test.py +++ b/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,6 +89,16 @@ 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') @@ -94,7 +106,28 @@ def test_bigquery_matcher_query_error_retry(self, mock_bigquery): '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,6 +195,11 @@ 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 @@ -169,7 +207,32 @@ def test__get_query_result_timeout(self, mock_bigquery): 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()