From 1088f4e01d193ecc9e4d20651b00ceb71646a354 Mon Sep 17 00:00:00 2001 From: Josh Stauffer <66793731+joshua-stauffer@users.noreply.github.com> Date: Tue, 4 Aug 2026 16:53:44 +0200 Subject: [PATCH] Sweep stale BigQuery test datasets via datasets.list instead of INFORMATION_SCHEMA INFORMATION_SCHEMA.SCHEMATA is a project-level view, so querying it requires metadata read access across the whole project even though the sweep only needs to see datasets it created itself. It's also region-scoped, so a dataset created in a different location is invisible to the query and never gets cleaned up. The datasets.list API has neither limitation: it returns exactly the datasets the caller can see, works with a credential scoped to just the CI dataset namespace, and isn't tied to a single region. Age filtering now happens client-side against Dataset.created (fetched via get_dataset, since list_datasets doesn't include creation time), and the age threshold is an injectable parameter so it can be exercised in tests without waiting an hour. Deletion uses delete_dataset(..., delete_contents=True) as the equivalent of the previous DROP SCHEMA ... CASCADE, and tolerates a dataset disappearing between listing and deletion. --- scripts/cleanup/cleanup_big_query.py | 100 +++++++++++++++-------- tests/scripts/test_cleanup_big_query.py | 102 ++++++++++++++++++++++++ 2 files changed, 170 insertions(+), 32 deletions(-) create mode 100644 tests/scripts/test_cleanup_big_query.py diff --git a/scripts/cleanup/cleanup_big_query.py b/scripts/cleanup/cleanup_big_query.py index a959a3e5ad3a..18c9d1340750 100644 --- a/scripts/cleanup/cleanup_big_query.py +++ b/scripts/cleanup/cleanup_big_query.py @@ -1,8 +1,10 @@ +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) @@ -10,47 +12,81 @@ 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__": diff --git a/tests/scripts/test_cleanup_big_query.py b/tests/scripts/test_cleanup_big_query.py new file mode 100644 index 000000000000..4d2937cfb2ae --- /dev/null +++ b/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 == []