test: add retry for rate limiting of load jobs

tests/system/__init__.py

@@ -0,0 +1,13 @@
+# Copyright 2021 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.

tests/system/conftest.py

@@ -20,8 +20,10 @@
 
 import pytest
 
-_TABLE_FORMAT = "projects/{}/datasets/{}/tables/{}"
+from . import helpers
+
 
+_TABLE_FORMAT = "projects/{}/datasets/{}/tables/{}"
 _ASSETS_DIR = os.path.join(os.path.abspath(os.path.dirname(__file__)), "assets")
 
 
@@ -98,8 +100,10 @@ def table(project_id, dataset, bq_client):
         bigquery.SchemaField("age", "INTEGER", mode="NULLABLE"),
     ]
 
-    table_id = "{}.{}.{}".format(project_id, dataset.dataset_id, "users")
-    bq_table = bigquery.Table(table_id, schema=schema)
+    unique_suffix = str(uuid.uuid4()).replace("-", "_")
+    table_id = "users_" + unique_suffix
+    table_id_full = f"{project_id}.{dataset.dataset_id}.{table_id}"
+    bq_table = bigquery.Table(table_id_full, schema=schema)
     created_table = bq_client.create_table(bq_table)
 
     yield created_table
@@ -154,7 +158,7 @@ def all_types_table_ref(project_id, dataset, bq_client):
     )
     yield table_ref
 
-    bq_client.delete_table(created_table)
+    helpers.retry_403(bq_client.delete_table)(created_table, not_found_ok=True)
 
 
 @pytest.fixture
@@ -182,7 +186,7 @@ def ingest_partition_table_ref(project_id, dataset, bq_client):
     )
     yield table_ref
 
-    bq_client.delete_table(created_table)
+    helpers.retry_403(bq_client.delete_table)(created_table, not_found_ok=True)
 
 
 @pytest.fixture
@@ -209,29 +213,39 @@ def col_partition_table_ref(project_id, dataset, bq_client):
     )
     yield table_ref
 
-    bq_client.delete_table(created_table)
+    helpers.retry_403(bq_client.delete_table)(created_table, not_found_ok=True)
 
 
 @pytest.fixture
-def table_with_data_ref(dataset, table, bq_client):
+def table_with_data_ref(project_id, dataset, bq_client):
     from google.cloud import bigquery
 
+    unique_suffix = str(uuid.uuid4()).replace("-", "_")
+    table_id = "users_" + unique_suffix
+    table_id_full = f"{project_id}.{dataset.dataset_id}.{table_id}"
+    schema = [
+        bigquery.SchemaField("first_name", "STRING", mode="NULLABLE"),
+        bigquery.SchemaField("last_name", "STRING", mode="NULLABLE"),
+        bigquery.SchemaField("age", "INTEGER", mode="NULLABLE"),
+    ]
+
     job_config = bigquery.LoadJobConfig()
     job_config.source_format = bigquery.SourceFormat.CSV
     job_config.skip_leading_rows = 1
-    job_config.schema = table.schema
+    job_config.schema = schema
 
     filename = os.path.join(_ASSETS_DIR, "people_data.csv")
 
-    with open(filename, "rb") as source_file:
-        job = bq_client.load_table_from_file(source_file, table, job_config=job_config)
+    def create_table():
+        with open(filename, "rb") as source_file:
+            job = bq_client.load_table_from_file(
+                source_file, table_id_full, job_config=job_config
+            )
+        job.result()  # wait for the load to complete
 
-    job.result()  # wait for the load to complete
+    helpers.retry_403(create_table)()
 
-    table_ref = _TABLE_FORMAT.format(table.project, table.dataset_id, table.table_id)
+    table_ref = _TABLE_FORMAT.format(project_id, dataset.dataset_id, table_id)
     yield table_ref
 
-    # truncate table data
-    query = "DELETE FROM {}.{} WHERE 1 = 1".format(dataset.dataset_id, table.table_id)
-    query_job = bq_client.query(query, location="US")
-    query_job.result()
+    helpers.retry_403(bq_client.delete_table)(table_id_full, not_found_ok=True)

tests/system/helpers.py

@@ -0,0 +1,36 @@
+# Copyright 2021 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Test utilities.
+
+Copied from the BigQuery client library
+https://github.com/googleapis/python-bigquery/blob/master/tests/system/helpers.py
+"""
+
+import google.api_core.exceptions
+import test_utils.retry
+
+
+def _rate_limit_exceeded(forbidden):
+    """Predicate: pass only exceptions with 'rateLimitExceeded' as reason."""
+    return any(error["reason"] == "rateLimitExceeded" for error in forbidden._errors)
+
+
+# We need to wait to stay within the rate limits.
+# The alternative outcome is a 403 Forbidden response from upstream, which
+# they return instead of the more appropriate 429.
+# See https://cloud.google.com/bigquery/quota-policy
+retry_403 = test_utils.retry.RetryErrors(
+    google.api_core.exceptions.Forbidden, error_predicate=_rate_limit_exceeded,
+)

tests/system/reader/__init__.py

@@ -0,0 +1,13 @@
+# Copyright 2021 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.

tests/system/reader/test_reader.py

@@ -20,10 +20,11 @@
 import decimal
 import re
 
+from google.cloud import bigquery
 import pytest
 import pytz
 
-from google.cloud import bigquery
+from .. import helpers
 
 
 _TABLE_FORMAT = "projects/{}/datasets/{}/tables/{}"
@@ -184,20 +185,23 @@ def test_column_selection_read(
 def test_snapshot(client_and_types, project_id, table, bq_client, data_format):
     client, types = client_and_types
 
+    def load_json(data):
+        return bq_client.load_table_from_json(data, table).result()
+
     # load original data into the table
     original_data = [
         {"first_name": "OGFoo", "last_name": "Smith", "age": 44},
         {"first_name": "OGBar", "last_name": "Jones", "age": 33},
     ]
-    og_job = bq_client.load_table_from_json(original_data, table).result()
+    og_job = helpers.retry_403(load_json)(original_data)
     og_time = og_job.ended
 
     # load additional data into the table
     new_data = [
         {"first_name": "NewFoo", "last_name": "Smiff", "age": 43},
         {"first_name": "NewBar", "last_name": "Jomes", "age": 34},
     ]
-    new_job = bq_client.load_table_from_json(new_data, table).result()
+    new_job = helpers.retry_403(load_json)(new_data)
     new_time = new_job.ended
 
     # Because we want our snapshot to be between when we loaded the original