[MAINTENANCE] Discover stale BigQuery test datasets via datasets.list instead of INFORMATION_SCHEMA

scripts/cleanup/cleanup_big_query.py

@@ -1,56 +1,92 @@
+import datetime
 import logging
+import re
 import sys
 
+from great_expectations.compatibility.google import NotFound, python_bigquery
 from great_expectations.compatibility.pydantic import BaseSettings
-from great_expectations.compatibility.sqlalchemy import TextClause, create_engine
 
 logger = logging.getLogger(__name__)
 logger.setLevel(logging.INFO)
 logger.addHandler(logging.StreamHandler(sys.stdout))
 
 
 class BigQueryConnectionConfig(BaseSettings):
-    """Environment variables for BigQuery connection.
+    """Environment variables for BigQuery access.
     These are injected in via CI, but when running locally, you may use your own credentials.
-    GOOGLE_APPLICATION_CREDENTIALS must be kept secret
+    GOOGLE_APPLICATION_CREDENTIALS must be kept secret. It is not read directly by this script;
+    Application Default Credentials picks it up automatically.
     """
 
     GE_TEST_GCP_PROJECT: str
-    GE_TEST_BIGQUERY_DATASET: str
     GOOGLE_APPLICATION_CREDENTIALS: str
 
-    @property
-    def connection_string(self) -> str:
-        return f"bigquery://{self.GE_TEST_GCP_PROJECT}/{self.GE_TEST_BIGQUERY_DATASET}?credentials_path={self.GOOGLE_APPLICATION_CREDENTIALS}"
-
 
 # Schema patterns for different test types
 SCHEMA_PATTERN_TEST = "^gx_ci_test_[a-f0-9]{10}$"  # General SQL testing framework
 SCHEMA_PATTERN_PY_VERSION = "^py3[0-9]{1,2}_i[a-f0-9]{32}$"  # Python version-specific test schemas
-SCHEMA_FORMAT = f"{SCHEMA_PATTERN_TEST}|{SCHEMA_PATTERN_PY_VERSION}"
-
-
-def cleanup_big_query(config: BigQueryConnectionConfig) -> None:
-    engine = create_engine(url=config.connection_string)
-    with engine.connect() as conn, conn.begin():
-        results = conn.execute(
-            TextClause(
-                """
-                SELECT 'DROP SCHEMA ' || schema_name || ' CASCADE;'
-                FROM INFORMATION_SCHEMA.SCHEMATA
-                WHERE REGEXP_CONTAINS(schema_name, :schema_format)
-                AND creation_time < TIMESTAMP_SUB(CURRENT_TIMESTAMP(), INTERVAL 1 HOUR);
-                """
-            ),
-            {"schema_format": SCHEMA_FORMAT},
-        ).fetchall()
-        if results:
-            to_run = TextClause("\n".join([row[0] for row in results]))
-            conn.execute(to_run)
-            logger.info(f"Cleaned up {len(results)} BigQuery schema(s)")
-        else:
-            logger.info("No BigQuery schemas to clean up!")
-    engine.dispose()
+SCHEMA_FORMAT = re.compile(f"{SCHEMA_PATTERN_TEST}|{SCHEMA_PATTERN_PY_VERSION}")
+
+# Only sweep datasets older than this. Kept small enough that a dataset from a run that is
+# still in progress is never deleted out from under it.
+DEFAULT_MAX_AGE = datetime.timedelta(hours=1)
+
+
+def find_stale_dataset_ids(
+    client: python_bigquery.Client, max_age: datetime.timedelta = DEFAULT_MAX_AGE
+) -> list[str]:
+    """Find test dataset ids old enough to be cleaned up.
+
+    Uses the `datasets.list` API rather than querying `INFORMATION_SCHEMA.SCHEMATA`:
+    - `datasets.list` only returns datasets the caller can already see, so a credential scoped
+      to just the CI dataset namespace can run this sweep. A project-level
+      `INFORMATION_SCHEMA.SCHEMATA` query requires permission to read dataset metadata across
+      the whole project, which is more access than a CI credential should need.
+    - `INFORMATION_SCHEMA` is region-scoped: it only sees datasets in the region the query runs
+      in, so a dataset created in a different location would be silently invisible to it.
+      `datasets.list` is not region-scoped.
+    """
+    now = datetime.datetime.now(datetime.timezone.utc)
+    stale_ids = []
+    for dataset_item in client.list_datasets():
+        dataset_id = dataset_item.dataset_id
+        if not SCHEMA_FORMAT.match(dataset_id):
+            continue
+
+        try:
+            # `list_datasets` results don't include creation time; `get_dataset` does.
+            dataset = client.get_dataset(dataset_item.reference)
+        except NotFound:
+            # Dataset was deleted between listing and inspecting it.
+            continue
+
+        created = dataset.created
+        if created is not None and now - created > max_age:
+            stale_ids.append(dataset_id)
+
+    return stale_ids
+
+
+def cleanup_big_query(
+    config: BigQueryConnectionConfig, max_age: datetime.timedelta = DEFAULT_MAX_AGE
+) -> None:
+    client = python_bigquery.Client(project=config.GE_TEST_GCP_PROJECT)
+
+    stale_ids = find_stale_dataset_ids(client, max_age=max_age)
+    if not stale_ids:
+        logger.info("No BigQuery datasets to clean up!")
+        return
+
+    cleaned_up = 0
+    for dataset_id in stale_ids:
+        try:
+            client.delete_dataset(dataset_id, delete_contents=True)
+            cleaned_up += 1
+        except NotFound:
+            # Dataset was deleted (e.g. by a concurrent sweep) between listing and deleting it.
+            logger.info(f"Dataset {dataset_id} was already deleted")
+
+    logger.info(f"Cleaned up {cleaned_up} BigQuery dataset(s)")
 
 
 if __name__ == "__main__":

tests/scripts/test_cleanup_big_query.py

@@ -0,0 +1,102 @@
+import datetime
+
+import pytest
+from pytest_mock import MockerFixture
+from scripts.cleanup.cleanup_big_query import find_stale_dataset_ids
+
+from great_expectations.compatibility.google import NotFound
+
+pytestmark = pytest.mark.unit
+
+
+def _dataset_item(mocker: MockerFixture, dataset_id: str):
+    item = mocker.Mock()
+    item.dataset_id = dataset_id
+    item.reference = dataset_id
+    return item
+
+
+def _dataset(mocker: MockerFixture, created: datetime.datetime | None):
+    dataset = mocker.Mock()
+    dataset.created = created
+    return dataset
+
+
+def _ago(**kwargs) -> datetime.datetime:
+    return datetime.datetime.now(datetime.timezone.utc) - datetime.timedelta(**kwargs)
+
+
+def test_matches_gx_ci_test_pattern_and_is_old_enough(mocker: MockerFixture):
+    dataset_id = "gx_ci_test_" + "a" * 10
+    client = mocker.Mock()
+    client.list_datasets.return_value = [_dataset_item(mocker, dataset_id)]
+    client.get_dataset.return_value = _dataset(mocker, _ago(hours=2))
+
+    result = find_stale_dataset_ids(client, max_age=datetime.timedelta(hours=1))
+
+    assert result == [dataset_id]
+
+
+def test_matches_py_version_pattern(mocker: MockerFixture):
+    dataset_id = "py312_i" + "a" * 32
+    client = mocker.Mock()
+    client.list_datasets.return_value = [_dataset_item(mocker, dataset_id)]
+    client.get_dataset.return_value = _dataset(mocker, _ago(hours=2))
+
+    result = find_stale_dataset_ids(client, max_age=datetime.timedelta(hours=1))
+
+    assert result == [dataset_id]
+
+
+def test_ignores_dataset_that_does_not_match_any_pattern(mocker: MockerFixture):
+    client = mocker.Mock()
+    client.list_datasets.return_value = [_dataset_item(mocker, "great_expectations_ci")]
+
+    result = find_stale_dataset_ids(client, max_age=datetime.timedelta(hours=1))
+
+    assert result == []
+    # A dataset that doesn't match the naming pattern should never even be inspected,
+    # since inspecting it is an extra API call this credential may not even have access to.
+    client.get_dataset.assert_not_called()
+
+
+def test_excludes_dataset_younger_than_max_age(mocker: MockerFixture):
+    dataset_id = "gx_ci_test_" + "b" * 10
+    client = mocker.Mock()
+    client.list_datasets.return_value = [_dataset_item(mocker, dataset_id)]
+    client.get_dataset.return_value = _dataset(mocker, _ago(minutes=1))
+
+    result = find_stale_dataset_ids(client, max_age=datetime.timedelta(hours=1))
+
+    assert result == []
+
+
+def test_zero_max_age_includes_freshly_created_dataset(mocker: MockerFixture):
+    dataset_id = "gx_ci_test_" + "c" * 10
+    client = mocker.Mock()
+    client.list_datasets.return_value = [_dataset_item(mocker, dataset_id)]
+    client.get_dataset.return_value = _dataset(mocker, _ago(seconds=1))
+
+    result = find_stale_dataset_ids(client, max_age=datetime.timedelta(seconds=0))
+
+    assert result == [dataset_id]
+
+
+def test_dataset_deleted_between_list_and_get_is_skipped(mocker: MockerFixture):
+    client = mocker.Mock()
+    client.list_datasets.return_value = [_dataset_item(mocker, "gx_ci_test_" + "d" * 10)]
+    client.get_dataset.side_effect = NotFound("gone")
+
+    result = find_stale_dataset_ids(client, max_age=datetime.timedelta(hours=1))
+
+    assert result == []
+
+
+def test_dataset_with_no_creation_time_is_skipped(mocker: MockerFixture):
+    client = mocker.Mock()
+    client.list_datasets.return_value = [_dataset_item(mocker, "gx_ci_test_" + "e" * 10)]
+    client.get_dataset.return_value = _dataset(mocker, None)
+
+    result = find_stale_dataset_ids(client, max_age=datetime.timedelta(hours=1))
+
+    assert result == []