perf(sqlalchemy-bigquery): enable pytest-xdist parallelization with worker dataset isolation

Author: ohmayr

packages/sqlalchemy-bigquery/noxfile.py

@@ -92,6 +92,7 @@
 SYSTEM_TEST_STANDARD_DEPENDENCIES: List[str] = [
     "mock",
     "pytest",
+    "pytest-xdist",
     "google-cloud-testutils",
 ]
 SYSTEM_TEST_EXTERNAL_DEPENDENCIES: List[str] = []
@@ -392,6 +393,7 @@ def _run_system_test_logic(session, test_type):
         session.install(
             "mock",
             "pytest",
+            "pytest-xdist",
             "pytest-rerunfailures",
             "google-cloud-testutils",
             "-c",
@@ -413,6 +415,8 @@ def _run_system_test_logic(session, test_type):
     if test_type == "compliance":
         session.run(
             "py.test",
+            "-n",
+            "auto",
             "-vv",
             f"--junitxml=compliance_{session.python}_sponge_log.xml",
             "--reruns=3",
@@ -426,6 +430,8 @@ def _run_system_test_logic(session, test_type):
     else:
         session.run(
             "py.test",
+            "-n",
+            "auto",
             "--quiet",
             f"--junitxml=system_{session.python}_sponge_log.xml",
             test_path,

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

@@ -70,7 +70,9 @@ def visit_delete(self, delete_stmt, *args, **kw):
 
 
 def pytest_sessionstart(session):
-    dataset_id = prefixer.create_prefix()
+    import os
+    worker_id = os.environ.get("PYTEST_XDIST_WORKER", "master")
+    dataset_id = f"{prefixer.create_prefix()}_{worker_id}"
     session.config.option.dburi = [f"bigquery:///{dataset_id}"]
     with contextlib.closing(google.cloud.bigquery.Client()) as client:
         client.create_dataset(dataset_id)

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

@@ -73,7 +73,8 @@ def bigquery_dataset(
     bigquery_client: bigquery.Client, bigquery_schema: List[bigquery.SchemaField]
 ):
     project_id = bigquery_client.project
-    dataset_id = prefixer.create_prefix()
+    worker_id = os.environ.get("PYTEST_XDIST_WORKER", "master")
+    dataset_id = f"{prefixer.create_prefix()}_{worker_id}"
     dataset = bigquery.Dataset(f"{project_id}.{dataset_id}")
     dataset = bigquery_client.create_dataset(dataset)
     sample_table_id = f"{project_id}.{dataset_id}.sample"
@@ -93,16 +94,6 @@ def bigquery_dataset(
     view.view_query = f"SELECT string FROM `{dataset_id}.sample`"
     bigquery_client.create_table(view)
 
-    # Pre-create session-scoped lake table for geography tests
-    lake_table = bigquery.Table(
-        f"{project_id}.{dataset_id}.lake",
-        schema=[
-            bigquery.SchemaField("name", "STRING"),
-            bigquery.SchemaField("geog", "GEOGRAPHY"),
-        ],
-    )
-    bigquery_client.create_table(lake_table)
-
     yield dataset_id
     bigquery_client.delete_dataset(dataset_id, delete_contents=True)
 
@@ -125,7 +116,8 @@ def bigquery_empty_table(
 @pytest.fixture(scope="session")
 def bigquery_regional_dataset(bigquery_client, bigquery_schema):
     project_id = bigquery_client.project
-    dataset_id = prefixer.create_prefix()
+    worker_id = os.environ.get("PYTEST_XDIST_WORKER", "master")
+    dataset_id = f"{prefixer.create_prefix()}_{worker_id}"
     dataset = bigquery.Dataset(f"{project_id}.{dataset_id}")
     dataset.location = "asia-northeast1"
     dataset = bigquery_client.create_dataset(dataset)

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

@@ -47,6 +47,8 @@ def test_geoalchemy2_core(bigquery_dataset):
         "lake", metadata, Column("name", String), Column("geog", GEOGRAPHY)
     )
 
+    lake_table.create(engine)
+
     # Insertions
 
     conn = engine.connect()