Fix PostCommit Python tests related comments

Author: akashorabek

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

@@ -132,8 +132,8 @@ def _setup_new_types_env(self):
         projectId=self.project, datasetId=self.dataset_id, table=table)
     self.bigquery_client.client.tables.Insert(request)
 
-    # Wait for table creation to propagate.
-    time.sleep(10)
+    # Call get_table so that we wait until the table is visible.
+    _ = self.bigquery_client.get_table(self.project, self.dataset_id, NEW_TYPES_INPUT_TABLE)
 
     table_data = [{
         'bytes': b'xyw', 'date': '2011-01-01', 'time': '23:59:59.999999'

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

@@ -24,7 +24,7 @@
 import datetime
 import logging
 import secrets
-import time as mod_time
+import time
 import unittest
 import uuid
 from decimal import Decimal
@@ -99,7 +99,7 @@ def setUpClass(cls):
 
     cls.bigquery_client = BigQueryWrapper()
     cls.dataset_id = '%s%d%s' % (
-        cls.BIG_QUERY_DATASET_ID, int(mod_time.time()), secrets.token_hex(3))
+        cls.BIG_QUERY_DATASET_ID, int(time.time()), secrets.token_hex(3))
     cls.bigquery_client.get_or_create_dataset(cls.project, cls.dataset_id)
     _LOGGER.info(
         "Created dataset %s in project %s", cls.dataset_id, cls.project)
@@ -158,9 +158,11 @@ def create_table(cls, table_name):
     request = bigquery.BigqueryTablesInsertRequest(
         projectId=cls.project, datasetId=cls.dataset_id, table=table)
     cls.bigquery_client.client.tables.Insert(request)
+    # Call get_table so that we wait until the table is visible.
+    _ = cls.bigquery_client.get_table(
+        cls.project, cls.dataset_id, table_name)
     cls.bigquery_client.insert_rows(
         cls.project, cls.dataset_id, table_name, cls.TABLE_DATA)
-    mod_time.sleep(10)
 
   @skip(['PortableRunner', 'FlinkRunner'])
   @pytest.mark.it_postcommit
@@ -333,7 +335,7 @@ class ReadUsingStorageApiTests(BigQueryReadIntegrationTests):
   def setUpClass(cls):
     super(ReadUsingStorageApiTests, cls).setUpClass()
     cls.table_name = '%s%d%s' % (
-        cls.BIG_QUERY_DATASET_ID, int(mod_time.time()), secrets.token_hex(3))
+        cls.BIG_QUERY_DATASET_ID, int(time.time()), secrets.token_hex(3))
     cls._create_table(cls.table_name)
 
     table_id = '{}.{}'.format(cls.dataset_id, cls.table_name)
@@ -396,9 +398,11 @@ def _create_table(cls, table_name):
     request = bigquery.BigqueryTablesInsertRequest(
         projectId=cls.project, datasetId=cls.dataset_id, table=table)
     cls.bigquery_client.client.tables.Insert(request)
+    # Call get_table so that we wait until the table is visible.
+    _ = cls.bigquery_client.get_table(
+        cls.project, cls.dataset_id, table_name)
     cls.bigquery_client.insert_rows(
         cls.project, cls.dataset_id, table_name, cls.TABLE_DATA)
-    mod_time.sleep(10)
 
   @classmethod
   def _setup_temporary_dataset(cls, project, query):
@@ -411,7 +415,7 @@ def _execute_query(cls, project, query):
         'materializing_table_before_reading',
         str(uuid.uuid4())[0:10],
         bigquery_tools.BigQueryJobTypes.QUERY,
-        '%d_%s' % (int(mod_time.time()), secrets.token_hex(3)))
+        '%d_%s' % (int(time.time()), secrets.token_hex(3)))
     cls._setup_temporary_dataset(cls.project, cls.query)
     job = cls.bigquery_client._start_query_job(
         project,
@@ -658,6 +662,9 @@ def create_table(cls, table_name):
     request = bigquery.BigqueryTablesInsertRequest(
         projectId=cls.project, datasetId=cls.dataset_id, table=table)
     cls.bigquery_client.client.tables.Insert(request)
+    # Call get_table so that we wait until the table is visible.
+    _ = cls.bigquery_client.get_table(
+        cls.project, cls.dataset_id, table_name)
     row_data = {
         'float': 0.33,
         'numeric': Decimal('10'),
@@ -676,7 +683,6 @@ def create_table(cls, table_name):
 
     cls.bigquery_client.insert_rows(
         cls.project, cls.dataset_id, table_name, table_data)
-    mod_time.sleep(10)
 
   def get_expected_data(self, native=True):
     byts = b'\xab\xac'
@@ -782,9 +788,11 @@ def create_table(cls, table_name, data, table_schema):
     request = bigquery.BigqueryTablesInsertRequest(
         projectId=cls.project, datasetId=cls.dataset_id, table=table)
     cls.bigquery_client.client.tables.Insert(request)
+    # Call get_table so that we wait until the table is visible.
+    _ = cls.bigquery_client.get_table(
+        cls.project, cls.dataset_id, table_name)
     cls.bigquery_client.insert_rows(
         cls.project, cls.dataset_id, table_name, data)
-    mod_time.sleep(10)
     return table_schema
 
   @classmethod

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

@@ -60,6 +60,9 @@
   NotFound = None
 
 
+# Number of seconds to wait for bucket deletion to propagate.
+WAIT_DELETE_BUCKET_PROPAGATION_SECONDS = 10
+
 @unittest.skipIf(gcsio is None, 'GCP dependencies are not installed')
 @parameterized_class(
     ('no_gcsio_throttling_counter', 'enable_gcsio_blob_generation'),
@@ -222,7 +225,7 @@ def test_create_default_bucket(self, mock_default_gcs_bucket_name):
     if existing_bucket:
       try:
         existing_bucket.delete()
-        time.sleep(10)
+        time.sleep(WAIT_DELETE_BUCKET_PROPAGATION_SECONDS)
       except NotFound:
         # Bucket existence check from get_bucket may be inaccurate due to gcs
         # cache or delay

sdks/python/apache_beam/ml/gcp/recommendations_ai.py

@@ -24,7 +24,6 @@
 from typing import Sequence
 from typing import Tuple
 
-import tenacity
 from google.api_core.retry import Retry
 
 from apache_beam import pvalue
@@ -34,6 +33,7 @@
 from apache_beam.transforms import ParDo
 from apache_beam.transforms import PTransform
 from apache_beam.transforms.util import GroupIntoBatches
+from apache_beam.utils import retry
 from cachetools.func import ttl_cache
 
 # pylint: disable=wrong-import-order, wrong-import-position, ungrouped-imports
@@ -54,6 +54,7 @@
 ]
 
 FAILED_CATALOG_ITEMS = "failed_catalog_items"
+MAX_RETRIES=5
 
 
 @ttl_cache(maxsize=128, ttl=3600)
@@ -155,10 +156,9 @@ def process(self, element):
     request = recommendationengine.CreateCatalogItemRequest(
         parent=self.parent, catalog_item=catalog_item)
 
-    @tenacity.retry(
-        wait=tenacity.wait_exponential(multiplier=1, min=1, max=10),
-        stop=tenacity.stop_after_attempt(5),
-        reraise=True)
+    @retry.with_exponential_backoff(
+        num_retries=MAX_RETRIES,
+        retry_filter=retry.retry_on_server_errors_timeout_or_quota_issues_filter)
     def create_item():
       return self._client.create_catalog_item(
           request=request,