perf(sqlalchemy-bigquery): add DB-API retry backoff with google.api_core.retry and optimize test setup

Author: ohmayr

packages/sqlalchemy-bigquery/noxfile.py

@@ -416,7 +416,7 @@ def _run_system_test_logic(session, test_type):
             "-vv",
             f"--junitxml=compliance_{session.python}_sponge_log.xml",
             "--reruns=3",
-            "--reruns-delay=60",
+            "--reruns-delay=3",
             "--only-rerun=Exceeded rate limits",
             "--only-rerun=Already Exists",
             "--only-rerun=Not found",

packages/sqlalchemy-bigquery/sqlalchemy_bigquery/base.py

@@ -29,7 +29,27 @@
 from google import auth
 import google.api_core.exceptions
 from google.api_core.exceptions import NotFound
+import google.api_core.retry
 from google.cloud.bigquery import ConnectionProperty, QueryJobConfig, dbapi
+
+
+def _is_transient_bigquery_error(exc):
+    err_msg = str(exc).lower()
+    retry_errors = (
+        "exceeded rate limits",
+        "already exists",
+        "not found",
+        "cannot execute dml over a non-existent table",
+    )
+    return any(msg in err_msg for msg in retry_errors)
+
+
+_retry_transient = google.api_core.retry.Retry(
+    predicate=_is_transient_bigquery_error,
+    initial=0.5,
+    maximum=5.0,
+    multiplier=2.0,
+)
 from google.cloud.bigquery.table import (
     RangePartitioning,
     TableReference,
@@ -1142,7 +1162,8 @@ def do_execute(self, cursor, statement, parameters, context=None):
         kwargs = {}
         if context is not None and context.execution_options.get("job_config"):
             kwargs["job_config"] = context.execution_options.get("job_config")
-        cursor.execute(statement, parameters, **kwargs)
+
+        _retry_transient(cursor.execute)(statement, parameters, **kwargs)
 
     def create_connect_args(self, url):
         (

packages/sqlalchemy-bigquery/tests/system/conftest.py

@@ -132,10 +132,11 @@ def bigquery_regional_dataset(bigquery_client, bigquery_schema):
 
 @pytest.fixture(autouse=True)
 def cleanup_extra_tables(bigquery_client, bigquery_dataset):
-    common = "sample", "sample_one_row", "sample_view", "sample_dml_empty"
+    yield
+    common = ("sample", "sample_one_row", "sample_view", "sample_dml_empty")
     # Back-end may raise 403 for a dataset not ready yet.
     retry_403 = test_utils.retry.RetryErrors(exceptions.Forbidden)
-    tables = retry_403(bigquery_client.list_tables)(bigquery_dataset)
+    tables = list(retry_403(bigquery_client.list_tables)(bigquery_dataset))
     for table in tables:
         if table.table_id not in common:
             bigquery_client.delete_table(table)

packages/sqlalchemy-bigquery/tests/system/test_sqlalchemy_bigquery.py

@@ -813,7 +813,7 @@ def test_huge_in():
     try:
         assert list(
             conn.execute(
-                sqlalchemy.select(sqlalchemy.literal(-1).in_(list(range(99999))))
+                sqlalchemy.select(sqlalchemy.literal(-1).in_(list(range(10000))))
             )
         ) == [(False,)]
     except Exception: