perf(sqlalchemy-bigquery): optimize system test setup, DB-API retries, and compliance suite scope

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,28 @@
 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",
+        "job exceeded rate limits",
+        "quota exceeded",
+        "backend error",
+        "service unavailable",
+    )
+    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 +1163,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/sqlalchemy_dialect_compliance/test_dialect_compliance.py

@@ -643,3 +643,30 @@ def test_no_results_for_non_returning_insert(cls):
 del (
     WindowFunctionTest.test_window_rows_between
 )  # test expects BQ to return sorted results
+
+# Deleting compliance suites for features BigQuery (OLAP) explicitly does not support:
+UNSUPPORTED_SUITES = [
+    "IsolationLevelTest",
+    "TransactionTest",
+    "AutocommitIsolationTest",
+    "ComputedColumnTest",
+    "ComputedReflectionTest",
+    "TempTableElementsTest",
+    "TableNoColumnsTest",
+    "NativeUUIDTest",
+    "JSONLegacyStringCastIndexTest",
+    "DateTimeTZTest",
+    "TimeTZTest",
+    "IntervalTest",
+    "HasSequenceTest",
+    "HasSequenceTestEmpty",
+    "SequenceCompilerTest",
+    "SequenceTest",
+    "IdentityColumnTest",
+    "IdentityReflectionTest",
+    "ServerSideCursorsTest",
+]
+
+for _suite_name in UNSUPPORTED_SUITES:
+    if _suite_name in globals():
+        del globals()[_suite_name]

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 = retry_403(lambda: list(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: