From a59f4e186005c1436e7de04dbf011f985f963b96 Mon Sep 17 00:00:00 2001 From: Yiyuan Chen Date: Thu, 25 Jun 2026 00:30:04 +0000 Subject: [PATCH 1/5] ci: configure Spanner Emulator bypass flag and document limits in README --- tests/datacommons-integration-tests/README.md | 91 +++++++++++++------ .../docker-compose.override.yml.example | 27 ++++++ .../docker-compose.test.yml | 1 + .../local_env.yaml | 5 + .../setup_emulators.sh | 32 +++++++ 5 files changed, 126 insertions(+), 30 deletions(-) create mode 100644 tests/datacommons-integration-tests/docker-compose.override.yml.example create mode 100644 tests/datacommons-integration-tests/local_env.yaml create mode 100755 tests/datacommons-integration-tests/setup_emulators.sh diff --git a/tests/datacommons-integration-tests/README.md b/tests/datacommons-integration-tests/README.md index 03b5cfbe..2a53f42a 100644 --- a/tests/datacommons-integration-tests/README.md +++ b/tests/datacommons-integration-tests/README.md @@ -46,55 +46,86 @@ While running with `--keep-containers`, you can access: - **Ingestion Helper**: `localhost:8081` - **Website Backend**: `localhost:8082` -## Custom Feature Flags - -Local integration tests run with the Spanner multi-entity schema enabled by default in the Mixer backend. This is controlled by mounting [custom_feature_flags.yaml](./custom_feature_flags.yaml) inside the `website` container at `/workspace/deploy/featureflags/custom.yaml`. - -The container entrypoint detects the environment variable `RESOLVE_WITH_SPANNER_EMBEDDINGS=true` (defined in `docker-compose.test.yml`) and starts Mixer using these flags. - -You can modify the flags inside `custom_feature_flags.yaml` if you need to tweak the Mixer server configuration for local debugging. +## Spanner Emulator Local Testing & Limitations + +The local Docker Compose stack uses the Spanner Emulator for offline testing. + +### Emulator Optimizer Bypass +The emulator is started with `--disable_query_null_filtered_index_check`. Required because the emulator rejects queries on null-filtered indexes (e.g., `NodeEmbeddingIndex`) unless all index keys are explicitly filtered for non-null values. + +### Key Drawbacks & Gaps +* **No Real Vector Similarity (NL Search):** The emulator cannot run semantic vector searches (returns empty results). Also, public/production NLP API endpoints do not index custom locally-seeded variables, and the emulator gateway fails to pass `--disable_query_null_filtered_index_check` to the C++ gRPC engine on port `9010`. Local queries must bypass vector search by appending `disable_feature=use_v2_resolve_for_nl_search_vars` to use the text matcher. +* **Local UDF Routing only:** The `--remote_functions_host_port` flag can route `ML.PREDICT` to a local mock container on `localhost`, but connections are restricted to localhost. +* **Performance:** Runs in-memory and does not replicate multi-node Spanner scale or latency. +* **Heuristic Bypass Limitation:** Appending `disable_feature=use_v2_resolve_for_nl_search_vars` disables database-level vector index queries to bypass the emulator limits. **IMPORTANT:** This is a test-only workaround. To ensure that the actual production database index lookup and vector resolving code paths are verified before release, you must run the GCP sandbox integration tests (`run_gcp_integration_test.py`), which cover the active production vector path against a real GCP Spanner instance. + +### Manual Verification of Local NL Search +1. **Boot containers & seed GCS emulator:** + ```bash + docker compose -f docker-compose.test.yml up -d spanner gcs + ./setup_emulators.sh + ``` +2. **Seed Spanner and boot remaining services:** + ```bash + uv run pytest tests/datacommons-integration-tests/run_local_integration_test.py --keep-containers + ``` +3. **Query using the heuristic bypass:** + ```bash + curl -s -X POST "http://localhost:8082/api/explore/detect-and-fulfill?q=Number+of+frogs+in+United+States+of+America&disable_feature=use_v2_resolve_for_nl_search_vars" \ + -H "Content-Type: application/json" \ + -d '{"contextHistory": []}' + ``` + +### Testing Custom NLP Models & Vertex AI + +* **Option A: Custom NLP Service URL:** Point the `NL_SERVICE_ROOT_URL` environment variable of the `website` container to your custom model container port or external model API endpoint (and supply any required API keys in the container env). +* **Option B: Local Real NLP Service Container (CPU/MiniLM):** Run the official model server container locally on CPU using the HuggingFace `all-MiniLM-L6-v2` model: + 1. Enable the override configuration: + ```bash + cp docker-compose.override.yml.example docker-compose.override.yml + ``` + 2. Set your GCP project ID in `docker-compose.override.yml`: + ```yaml + GCP_PROJECT_ID: + ``` + 3. Start the containers using the override: + ```bash + docker compose -f docker-compose.test.yml -f docker-compose.override.yml up -d + ./setup_emulators.sh + uv run pytest tests/datacommons-integration-tests/run_local_integration_test.py --keep-containers + ``` + *(Note: the automated pytest asserts will fail due to real vs. mock output differences, but the local stack will remain running and fully operational for custom testing)*. +* **Option C: Spanner UDF Proxy:** To test database-level `ML.PREDICT` calls, run a local proxy container in the Spanner service's network namespace (`network_mode: "service:spanner"`) to forward requests to Vertex AI. Configure the Spanner emulator container with `--remote_functions_host_port=localhost:`. --- ## GCP Sandbox Integration Tests -In addition to local emulated tests, this package includes a script to run automated, end-to-end integration tests on **real GCP sandbox resources**. This validates Terraform scaffolding, GCS transfers, Spanner database seeding, Cloud Workflow executions, and Cloud Run web server APIs. +This package includes a script to run end-to-end integration tests on real GCP sandbox resources to validate Terraform templates, GCS transfers, Spanner database seeding, and Cloud Run web server APIs. ### Running the GCP Tests - -To run the GCP integration test suite (requires authenticated GCP CLI access): +Requires authenticated GCP CLI access: ```bash python3 tests/datacommons-integration-tests/run_gcp_integration_test.py \ --project-id \ --dc-api-key "YOUR_GOOGLE_DATA_COMMONS_API_KEY" ``` -Alternatively, you can pass the API key via an environment variable: +Alternatively, pass the API key via an environment variable: ```bash DC_API_KEY="YOUR_API_KEY" python3 tests/datacommons-integration-tests/run_gcp_integration_test.py \ --project-id ``` ### Script Arguments +* `--project-id`: GCP Project ID (default: `datcom-ci`). +* `--dc-api-key`: Data Commons API Key. +* `--region`: GCP region (default: `us-central1`). +* `--namespace`: Custom resource naming namespace (default: randomized `itest-XXXX`). +* `--keep-sandbox`: Do not destroy sandbox GCP resources on completion (useful for debugging). +* `--reuse-sandbox`: Reuse existing sandbox resources (requires passing persistent `--namespace` and having run with `--keep-sandbox` previously). +* `--tf-git-ref`: Git reference for the Terraform templates repository (default: `main`). +* `--services-image`, `--preprocessing-image`, `--helper-image`: Override container images deployed during provisioning. -* `--project-id`: The Google Cloud Project ID where the sandbox resources should be provisioned (default: `datcom-ci`). -* `--dc-api-key`: Google Data Commons API Key needed to authenticate and configure sandbox clients. -* `--region`: The GCP region to deploy resources (default: `us-central1`). -* `--namespace`: Custom naming namespace for resources (default: randomized `itest-XXXX`). -* `--keep-sandbox`: Do not destroy sandbox GCP resources on completion/failure. Useful for debugging active instances. -* `--reuse-sandbox`: Reuse existing local workspace and GCP sandbox resources if they exist. Requires passing a persistent, custom `--namespace` (e.g. `--namespace itest-9611`) and having run with `--keep-sandbox` in the previous run. -* `--tf-git-ref`: Git reference branch/tag/commit for the GCP Terraform templates repository (default: `main`). -* `--services-image`, `--preprocessing-image`, `--helper-image`: Override container image tags deployed during provisioning. - -### E2E Verification Stages & Philosophy - -Once the GCP sandbox stack is up and seeded with dataset MCFs, the script triggers the entire API verification suite concurrently. This validates three core capabilities of the deployed Data Commons Platform backend: - -* **Stage A: V2 Observations & Custom Data Retrieval** - * *Philosophy:* Validates E2E query routing and data loading for both standard and custom variables. It queries base DC SDG indicators to verify backward compatibility with global schemas, and asserts observations values on custom seeded frog metrics to verify local database mapping and seeding. -* **Stage B: V2 Bulk Group Info & Hierarchy Resolution** - * *Philosophy:* Verifies category tree resolution and variable grouping APIs. This validates that the serving container's Mixer component correctly processes unconstrained hierarchy crawls and handles location constraints (e.g. comparing available variables for USA vs Vatican City) without backend lookup errors. -* **Stage C: V2 Embeddings & Natural Language Resolution** - * *Philosophy:* Validates custom search index and embeddings generation. This checks that the Ingestion Helper successfully generated vector coordinates from custom MCs, loaded them into Spanner, and that the Resolver can map natural language queries (like "frogs") back to these custom database entities. diff --git a/tests/datacommons-integration-tests/docker-compose.override.yml.example b/tests/datacommons-integration-tests/docker-compose.override.yml.example new file mode 100644 index 00000000..cc02e4c5 --- /dev/null +++ b/tests/datacommons-integration-tests/docker-compose.override.yml.example @@ -0,0 +1,27 @@ +# Example docker-compose override file to run real local NLP (CPU-based MiniLM). +# To use: +# cp docker-compose.override.yml.example docker-compose.override.yml +# +services: + website: + environment: + - NL_SERVICE_ROOT_URL=http://nl-service:6060 + depends_on: + - spanner + - nl-service + + nl-service: + image: gcr.io/datcom-ci/datacommons-nl:latest + container_name: itest-nl-service + volumes: + - ~/.config/gcloud:/root/.config/gcloud + - ./local_env.yaml:/datacommons/nl/env.yaml + - ../../website/deploy/helm_charts/dc_website/nl/catalog.yaml:/datacommons/nl/catalog.yaml + environment: + - GOOGLE_APPLICATION_CREDENTIALS=/root/.config/gcloud/application_default_credentials.json + - GCP_PROJECT_ID= + - NUM_WORKERS=1 + networks: + - test-network + extra_hosts: + - "metadata.google.internal:127.0.0.1" diff --git a/tests/datacommons-integration-tests/docker-compose.test.yml b/tests/datacommons-integration-tests/docker-compose.test.yml index ebe822ee..f29c728f 100644 --- a/tests/datacommons-integration-tests/docker-compose.test.yml +++ b/tests/datacommons-integration-tests/docker-compose.test.yml @@ -5,6 +5,7 @@ services: ports: - "${SPANNER_GRPC_PORT:-9010}:9010" - "${SPANNER_REST_PORT:-9020}:9020" + command: ["./gateway_main", "--hostname", "0.0.0.0", "--disable_query_null_filtered_index_check"] networks: - test-network gcs: diff --git a/tests/datacommons-integration-tests/local_env.yaml b/tests/datacommons-integration-tests/local_env.yaml new file mode 100644 index 00000000..4d3aecfc --- /dev/null +++ b/tests/datacommons-integration-tests/local_env.yaml @@ -0,0 +1,5 @@ +default_indexes: + - medium_ft +enabled_indexes: + - medium_ft +enable_reranking: false diff --git a/tests/datacommons-integration-tests/setup_emulators.sh b/tests/datacommons-integration-tests/setup_emulators.sh new file mode 100755 index 00000000..d293e206 --- /dev/null +++ b/tests/datacommons-integration-tests/setup_emulators.sh @@ -0,0 +1,32 @@ +#!/bin/bash +# Copyright 2026 Google LLC +# Script to wait for local emulators and perform initial GCS seeding. + +set -e + +echo "Waiting for GCS Emulator (port 9099) to start..." +count=0 +until curl -s http://localhost:9099/ > /dev/null || [ $count -eq 60 ]; do + sleep 0.5 + count=$((count + 1)) +done + +if [ $count -eq 60 ]; then + echo "ERROR: GCS Emulator failed to start on port 9099." + exit 1 +fi + + +echo "Seeding GCS emulator bucket and dummy catalog..." +curl -s -X POST "http://localhost:9099/storage/v1/b?project=test-project" \ + -H "Content-Type: application/json" \ + -d '{"name": "test-bucket"}' > /dev/null + +curl -s -X POST "http://localhost:9099/upload/storage/v1/b/test-bucket/o?uploadType=media&name=output/datacommons/nl/embeddings/custom_catalog.yaml" \ + -H "Content-Type: application/x-yaml" \ + -d "version: '1' +models: {} +indexes: {}" > /dev/null + + +echo "Emulators initialized successfully." From 366e395699e567fd8e8e671ae421f9e40ab578d2 Mon Sep 17 00:00:00 2001 From: Yiyuan Chen Date: Tue, 7 Jul 2026 22:35:18 +0000 Subject: [PATCH 2/5] Refactor local test stack to seed GCS natively, remove setup_emulators.sh, and add GCP sandbox cleanup script --- tests/datacommons-integration-tests/README.md | 10 +- .../clean_orphaned_sandboxes.py | 213 ++++++++++++++++++ .../docker-compose.test.yml | 4 +- .../nl/embeddings/custom_catalog.yaml | 3 + .../run_local_integration_test.py | 27 --- .../setup_emulators.sh | 32 --- 6 files changed, 223 insertions(+), 66 deletions(-) create mode 100644 tests/datacommons-integration-tests/clean_orphaned_sandboxes.py create mode 100644 tests/datacommons-integration-tests/gcs_seed/test-bucket/output/datacommons/nl/embeddings/custom_catalog.yaml delete mode 100755 tests/datacommons-integration-tests/setup_emulators.sh diff --git a/tests/datacommons-integration-tests/README.md b/tests/datacommons-integration-tests/README.md index 2a53f42a..0d8e09af 100644 --- a/tests/datacommons-integration-tests/README.md +++ b/tests/datacommons-integration-tests/README.md @@ -60,10 +60,9 @@ The emulator is started with `--disable_query_null_filtered_index_check`. Requir * **Heuristic Bypass Limitation:** Appending `disable_feature=use_v2_resolve_for_nl_search_vars` disables database-level vector index queries to bypass the emulator limits. **IMPORTANT:** This is a test-only workaround. To ensure that the actual production database index lookup and vector resolving code paths are verified before release, you must run the GCP sandbox integration tests (`run_gcp_integration_test.py`), which cover the active production vector path against a real GCP Spanner instance. ### Manual Verification of Local NL Search -1. **Boot containers & seed GCS emulator:** +1. **Boot containers:** ```bash - docker compose -f docker-compose.test.yml up -d spanner gcs - ./setup_emulators.sh + docker compose -f tests/datacommons-integration-tests/docker-compose.test.yml up -d spanner gcs ``` 2. **Seed Spanner and boot remaining services:** ```bash @@ -82,7 +81,7 @@ The emulator is started with `--disable_query_null_filtered_index_check`. Requir * **Option B: Local Real NLP Service Container (CPU/MiniLM):** Run the official model server container locally on CPU using the HuggingFace `all-MiniLM-L6-v2` model: 1. Enable the override configuration: ```bash - cp docker-compose.override.yml.example docker-compose.override.yml + cp tests/datacommons-integration-tests/docker-compose.override.yml.example tests/datacommons-integration-tests/docker-compose.override.yml ``` 2. Set your GCP project ID in `docker-compose.override.yml`: ```yaml @@ -90,8 +89,7 @@ The emulator is started with `--disable_query_null_filtered_index_check`. Requir ``` 3. Start the containers using the override: ```bash - docker compose -f docker-compose.test.yml -f docker-compose.override.yml up -d - ./setup_emulators.sh + docker compose -f tests/datacommons-integration-tests/docker-compose.test.yml -f tests/datacommons-integration-tests/docker-compose.override.yml up -d uv run pytest tests/datacommons-integration-tests/run_local_integration_test.py --keep-containers ``` *(Note: the automated pytest asserts will fail due to real vs. mock output differences, but the local stack will remain running and fully operational for custom testing)*. diff --git a/tests/datacommons-integration-tests/clean_orphaned_sandboxes.py b/tests/datacommons-integration-tests/clean_orphaned_sandboxes.py new file mode 100644 index 00000000..d0fefec7 --- /dev/null +++ b/tests/datacommons-integration-tests/clean_orphaned_sandboxes.py @@ -0,0 +1,213 @@ +#!/usr/bin/env python3 +# Copyright 2026 Google LLC. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Helper script to list and clean up leaked integration test GCP sandboxes. + +It reconstructs local Terraform scaffolding to connect to GCS remote state +and execute 'terraform destroy' for orphaned namespaces. +""" + +import argparse +import os +import re +import shutil +import subprocess +import sys +from pathlib import Path + + +def run_command(args: list[str], cwd: Path | None = None, check: bool = True) -> subprocess.CompletedProcess: + print(f"Running command: {' '.join(args)} (cwd: {cwd or Path.cwd()})", flush=True) + try: + return subprocess.run(args, cwd=cwd, check=check, text=True, capture_output=True) + except subprocess.CalledProcessError as e: + print(f"--- COMMAND FAILED: {' '.join(args)} ---", file=sys.stderr) + if e.stdout: + print(f"STDOUT:\n{e.stdout}", file=sys.stderr) + if e.stderr: + print(f"STDERR:\n{e.stderr}", file=sys.stderr) + print("---------------------------------------", file=sys.stderr) + raise e + + +def get_active_namespaces(tf_state_bucket: str) -> list[str]: + """Lists all prefixes (namespaces) in GCS tf-state-bucket.""" + print(f"Listing active states in gs://{tf_state_bucket}/terraform/state/...") + try: + proc = run_command([ + "gcloud", "storage", "ls", f"gs://{tf_state_bucket}/terraform/state/" + ]) + namespaces = [] + # Expected format: gs://bucket-name/terraform/state/itest-XXXX/ + pattern = rf"gs://{re.escape(tf_state_bucket)}/terraform/state/([^/]+)/" + for line in proc.stdout.splitlines(): + match = re.match(pattern, line.strip()) + if match: + namespaces.append(match.group(1)) + return sorted(list(set(namespaces))) + except Exception as e: + print(f"Error listing state directories: {e}", file=sys.stderr) + return [] + + +def destroy_sandbox(project_id: str, tf_state_bucket: str, namespace: str, workspace_root: Path) -> bool: + """Reconstructs local state and executes terraform destroy.""" + print(f"\n=====================================================================") + print(f" DESTROYING SANDBOX: namespace={namespace} in project={project_id}") + print(f"=====================================================================") + + workspace_dir = Path(f"/tmp/workspace-purge-{namespace}") + sandbox_dir = workspace_dir / namespace + + try: + # Create directories + if workspace_dir.exists(): + shutil.rmtree(workspace_dir) + sandbox_dir.mkdir(parents=True) + + # Copy local infra modules & configs + local_infra_dir = workspace_root / "infra" / "dcp" + shutil.copy(local_infra_dir / "variables.tf", sandbox_dir / "variables.tf") + shutil.copy(local_infra_dir / "outputs.tf", sandbox_dir / "outputs.tf") + shutil.copy(local_infra_dir / "main.tf", sandbox_dir / "main.tf") + shutil.copytree(local_infra_dir / "modules", sandbox_dir / "modules") + + # Construct backend.tf + backend_content = f"""terraform {{ + backend "gcs" {{ + bucket = "{tf_state_bucket}" + prefix = "terraform/state/{namespace}" + }} +}} +""" + (sandbox_dir / "backend.tf").write_text(backend_content, encoding="utf-8") + + # Construct terraform.tfvars + tfvars_content = f"""project_id = "{project_id}" +namespace = "{namespace}" +spanner_create_instance = false +spanner_instance_id = "dcp-integration-test-shared-instance" +spanner_create_database = true +spanner_create_bigquery_reservation = false +datacommons_services_min_instances = 1 +datacommons_services_max_instances = 1 +""" + (sandbox_dir / "terraform.tfvars").write_text(tfvars_content, encoding="utf-8") + + # Init & Destroy + run_command(["terraform", "init"], cwd=sandbox_dir) + run_command(["terraform", "destroy", "-auto-approve"], cwd=sandbox_dir) + + # Delete state file in GCS + print(f"Removing GCS state files for namespace {namespace}...") + run_command([ + "gcloud", "storage", "rm", "--recursive", + f"gs://{tf_state_bucket}/terraform/state/{namespace}/" + ]) + + print(f"Successfully destroyed sandbox resources for namespace: {namespace}") + return True + + except Exception as e: + print(f"Failed to destroy sandbox {namespace}: {e}", file=sys.stderr) + return False + + finally: + if workspace_dir.exists(): + shutil.rmtree(workspace_dir) + + +def main() -> None: + parser = argparse.ArgumentParser(description="Purge orphaned GCP sandbox integration test resources.") + parser.add_argument( + "--project-id", default="datcom-ci", help="GCP project ID (default: datcom-ci)" + ) + parser.add_argument( + "--tf-state-bucket", + default="tf-state-dcp-test-datcom-ci", + help="GCS bucket for remote Terraform state (default: tf-state-dcp-test-datcom-ci)", + ) + parser.add_argument( + "--namespace", + default=None, + help="Specific namespace to destroy (e.g. itest-1234)", + ) + parser.add_argument( + "--list", + action="store_true", + help="List active namespaces in the state bucket and exit", + ) + parser.add_argument( + "--all", + action="store_true", + help="Destroy all active namespaces found in the state bucket", + ) + + args = parser.parse_args() + workspace_root = Path(__file__).resolve().parent.parent.parent + + # 1. Fetch active namespaces + active_namespaces = get_active_namespaces(args.tf_state_bucket) + if not active_namespaces: + print("No active integration test state namespaces found in GCS.") + return + + print(f"Active namespaces found ({len(active_namespaces)}):") + for ns in active_namespaces: + print(f" - {ns}") + + if args.list: + return + + # 2. Determine target namespaces + targets = [] + if args.namespace: + if args.namespace not in active_namespaces: + print(f"Error: Namespace '{args.namespace}' was not found in GCS.", file=sys.stderr) + sys.exit(1) + targets = [args.namespace] + elif args.all: + targets = active_namespaces + else: + print("\nSpecify either --namespace , --all, or --list. See --help for details.") + sys.exit(1) + + # 3. Confirm destroy + print(f"\nWARNING: This will permanently destroy all resources associated with the following namespaces: {', '.join(targets)}") + confirm = input("Are you sure you want to proceed? [y/N]: ").strip().lower() + if confirm not in ["y", "yes"]: + print("Aborted.") + sys.exit(0) + + # 4. Run destroys + successes = 0 + failures = 0 + for ns in targets: + if destroy_sandbox(args.project_id, args.tf_state_bucket, ns, workspace_root): + successes += 1 + else: + failures += 1 + + print(f"\n=====================================================================") + print(f" PURGE SUMMARY") + print(f"=====================================================================") + print(f" Successes: {successes}") + print(f" Failures: {failures}") + if failures > 0: + sys.exit(1) + + +if __name__ == "__main__": + main() diff --git a/tests/datacommons-integration-tests/docker-compose.test.yml b/tests/datacommons-integration-tests/docker-compose.test.yml index f29c728f..08764d74 100644 --- a/tests/datacommons-integration-tests/docker-compose.test.yml +++ b/tests/datacommons-integration-tests/docker-compose.test.yml @@ -13,7 +13,9 @@ services: container_name: itest-gcs ports: - "9099:9099" - command: -scheme http -port 9099 + command: -scheme http -port 9099 -data /data + volumes: + - ./gcs_seed:/data networks: - test-network diff --git a/tests/datacommons-integration-tests/gcs_seed/test-bucket/output/datacommons/nl/embeddings/custom_catalog.yaml b/tests/datacommons-integration-tests/gcs_seed/test-bucket/output/datacommons/nl/embeddings/custom_catalog.yaml new file mode 100644 index 00000000..90d7afe6 --- /dev/null +++ b/tests/datacommons-integration-tests/gcs_seed/test-bucket/output/datacommons/nl/embeddings/custom_catalog.yaml @@ -0,0 +1,3 @@ +version: '1' +models: {} +indexes: {} diff --git a/tests/datacommons-integration-tests/run_local_integration_test.py b/tests/datacommons-integration-tests/run_local_integration_test.py index 7ff6ea54..3fb8e7e8 100755 --- a/tests/datacommons-integration-tests/run_local_integration_test.py +++ b/tests/datacommons-integration-tests/run_local_integration_test.py @@ -164,31 +164,6 @@ def call_helper(action_type: str) -> dict: ) from e # noqa: BLE001 -def seed_gcs_emulator(): - """Create test-bucket and upload dummy catalog yaml to mock GCS emulator.""" - print(">>> Seeding fake GCS emulator bucket and catalogs...", flush=True) - try: - # 1. Create bucket - resp = requests.post( - "http://localhost:9099/storage/v1/b?project=test-project", - json={"name": "test-bucket"}, - timeout=5, - ) - resp.raise_for_status() - - # 2. Upload dummy custom_catalog.yaml - dummy_catalog = "version: '1'\nmodels: {}\nindexes: {}\n" - upload_resp = requests.post( - "http://localhost:9099/upload/storage/v1/b/test-bucket/o?uploadType=media&name=output/datacommons/nl/embeddings/custom_catalog.yaml", - data=dummy_catalog, - headers={"Content-Type": "application/x-yaml"}, - timeout=5, - ) - upload_resp.raise_for_status() - print(" GCS emulator successfully seeded.", flush=True) - except Exception as e: - raise RuntimeError(f"Failed to seed GCS emulator: {e}") from e # noqa: BLE001 - # ============================================================================= # Pytest Orchestration Fixture @@ -255,8 +230,6 @@ def docker_stack(services_image, helper_image, keep_containers): "GCS emulator", ) - seed_gcs_emulator() - # 4. Create database instance print( f">>> Creating Spanner {INSTANCE_ID}/{DATABASE_ID} in emulator...", diff --git a/tests/datacommons-integration-tests/setup_emulators.sh b/tests/datacommons-integration-tests/setup_emulators.sh deleted file mode 100755 index d293e206..00000000 --- a/tests/datacommons-integration-tests/setup_emulators.sh +++ /dev/null @@ -1,32 +0,0 @@ -#!/bin/bash -# Copyright 2026 Google LLC -# Script to wait for local emulators and perform initial GCS seeding. - -set -e - -echo "Waiting for GCS Emulator (port 9099) to start..." -count=0 -until curl -s http://localhost:9099/ > /dev/null || [ $count -eq 60 ]; do - sleep 0.5 - count=$((count + 1)) -done - -if [ $count -eq 60 ]; then - echo "ERROR: GCS Emulator failed to start on port 9099." - exit 1 -fi - - -echo "Seeding GCS emulator bucket and dummy catalog..." -curl -s -X POST "http://localhost:9099/storage/v1/b?project=test-project" \ - -H "Content-Type: application/json" \ - -d '{"name": "test-bucket"}' > /dev/null - -curl -s -X POST "http://localhost:9099/upload/storage/v1/b/test-bucket/o?uploadType=media&name=output/datacommons/nl/embeddings/custom_catalog.yaml" \ - -H "Content-Type: application/x-yaml" \ - -d "version: '1' -models: {} -indexes: {}" > /dev/null - - -echo "Emulators initialized successfully." From b3ab2e01ec0b95fdf9af2173f35b090de9916a7b Mon Sep 17 00:00:00 2001 From: Yiyuan Chen Date: Wed, 8 Jul 2026 23:52:11 +0000 Subject: [PATCH 3/5] Configure local stack to run Spanner Omni with transparent credentials patch --- .../docker-compose.test.yml | 30 +- .../patch_credentials.py | 144 +++ .../run_local_integration_test.py | 82 +- .../spanner_client.py | 917 ------------------ 4 files changed, 210 insertions(+), 963 deletions(-) create mode 100644 tests/datacommons-integration-tests/patch_credentials.py delete mode 100644 tests/datacommons-integration-tests/spanner_client.py diff --git a/tests/datacommons-integration-tests/docker-compose.test.yml b/tests/datacommons-integration-tests/docker-compose.test.yml index 08764d74..64289a9b 100644 --- a/tests/datacommons-integration-tests/docker-compose.test.yml +++ b/tests/datacommons-integration-tests/docker-compose.test.yml @@ -1,11 +1,11 @@ services: spanner: - image: gcr.io/cloud-spanner-emulator/emulator@sha256:caf1bd24c081e005837b5977bae5a250e25cb4da9f25ec1abc91936ad67e4de2 + image: ${SPANNER_IMAGE:-us-docker.pkg.dev/spanner-omni/images/spanner-omni:2026.r1-beta.2} container_name: itest-spanner ports: - - "${SPANNER_GRPC_PORT:-9010}:9010" - - "${SPANNER_REST_PORT:-9020}:9020" - command: ["./gateway_main", "--hostname", "0.0.0.0", "--disable_query_null_filtered_index_check"] + - "${SPANNER_GRPC_PORT:-9010}:${SPANNER_INTERNAL_GRPC_PORT:-15000}" + - "${SPANNER_REST_PORT:-9020}:${SPANNER_INTERNAL_REST_PORT:-15020}" + command: ${SPANNER_COMMAND:-start-single-server} networks: - test-network gcs: @@ -27,15 +27,16 @@ services: ports: - "${HELPER_PORT:-8081}:8080" volumes: - - ./spanner_client.py:/app/clients/spanner.py + - ./patch_credentials.py:/app/patch/sitecustomize.py environment: + - PYTHONPATH=/app/patch - PORT=8080 - - SPANNER_EMULATOR_HOST=spanner:9010 + - SPANNER_EMULATOR_HOST=spanner:${SPANNER_INTERNAL_GRPC_PORT:-15000} - STORAGE_EMULATOR_HOST=http://gcs:9099 - OUTPUT_DIR=gs://test-bucket/output - - PROJECT_ID=test-project - - SPANNER_PROJECT_ID=test-project - - SPANNER_INSTANCE_ID=test-instance + - PROJECT_ID=default + - SPANNER_PROJECT_ID=default + - SPANNER_INSTANCE_ID=default - SPANNER_DATABASE_ID=test-db - SPANNER_GRAPH_DATABASE_ID=test-db depends_on: @@ -54,17 +55,18 @@ services: ports: - "${WEBSITE_PORT:-8082}:8080" volumes: - - ./spanner_client.py:/workspace/spanner_client.py + - ./patch_credentials.py:/workspace/patch/sitecustomize.py - ./custom_feature_flags.yaml:/workspace/deploy/featureflags/custom.yaml environment: + - PYTHONPATH=/workspace/patch - PORT=8080 - - SPANNER_EMULATOR_HOST=spanner:9010 + - SPANNER_EMULATOR_HOST=spanner:${SPANNER_INTERNAL_GRPC_PORT:-15000} - STORAGE_EMULATOR_HOST=http://gcs:9099 - OUTPUT_DIR=gs://test-bucket/output - RESOLVE_WITH_SPANNER_EMBEDDINGS=true - - PROJECT_ID=test-project - - GCP_PROJECT_ID=test-project - - GCP_SPANNER_INSTANCE_ID=test-instance + - PROJECT_ID=default + - GCP_PROJECT_ID=default + - GCP_SPANNER_INSTANCE_ID=default - GCP_SPANNER_DATABASE_NAME=test-db - USE_SPANNER_GRAPH=true - DC_SEARCH_SCOPE=custom_only diff --git a/tests/datacommons-integration-tests/patch_credentials.py b/tests/datacommons-integration-tests/patch_credentials.py new file mode 100644 index 00000000..c71d3189 --- /dev/null +++ b/tests/datacommons-integration-tests/patch_credentials.py @@ -0,0 +1,144 @@ +# ruff: noqa +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Monkeypatch to allow Spanner client libraries to connect to emulators. + +This acts as a temporary workaround. The long-term fix should be implemented in +the ingestion helper service (inside the import repository). +""" + +import os +import sys + +try: + import google.cloud.spanner_admin_database_v1 + import grpc + from google.auth.credentials import AnonymousCredentials + from google.cloud import spanner + from google.cloud.spanner_admin_database_v1 import DatabaseAdminClient + from google.cloud.spanner_admin_database_v1.services.database_admin.transports.grpc import ( + DatabaseAdminGrpcTransport, + ) + + # Avoid attribute errors on AnonymousCredentials + AnonymousCredentials.with_quota_project = lambda self, *args, **kwargs: self + + original_spanner_init = spanner.Client.__init__ + + def patched_spanner_init(self, *args, **kwargs): + emulator_host = os.getenv("SPANNER_EMULATOR_HOST") + if emulator_host: + kwargs["credentials"] = AnonymousCredentials() + original_spanner_init(self, *args, **kwargs) + + spanner.Client.__init__ = patched_spanner_init + + class InsecureDatabaseAdminClient(DatabaseAdminClient): + def __init__(self, *args, **kwargs): + emulator_host = os.getenv("SPANNER_EMULATOR_HOST") + if emulator_host: + channel = grpc.insecure_channel(emulator_host) + kwargs["transport"] = DatabaseAdminGrpcTransport( + channel=channel, credentials=AnonymousCredentials() + ) + if "credentials" in kwargs: + del kwargs["credentials"] + super().__init__(*args, **kwargs) + + def update_database_ddl(self, request=None, *args, **kwargs): + emulator_host = os.getenv("SPANNER_EMULATOR_HOST") + if emulator_host: + req_obj = request if request is not None else kwargs.get("request") + if req_obj is not None: + if hasattr(req_obj, "statements"): + filtered_statements = [] + for stmt in req_obj.statements: + normalized = stmt.strip().upper() + if normalized.startswith( + ("CREATE MODEL", "CREATE VECTOR INDEX") + ): + continue + filtered_statements.append(stmt) + del req_obj.statements[:] + req_obj.statements.extend(filtered_statements) + elif isinstance(req_obj, dict) and "statements" in req_obj: + req_obj["statements"] = [ + s + for s in req_obj["statements"] + if not s.strip() + .upper() + .startswith(("CREATE MODEL", "CREATE VECTOR INDEX")) + ] + elif "statements" in kwargs: + kwargs["statements"] = [ + s + for s in kwargs["statements"] + if not s.strip() + .upper() + .startswith(("CREATE MODEL", "CREATE VECTOR INDEX")) + ] + return super().update_database_ddl(request, *args, **kwargs) + + google.cloud.spanner_admin_database_v1.DatabaseAdminClient = ( + InsecureDatabaseAdminClient + ) + + if "clients.spanner" in sys.modules: + sys.modules["clients.spanner"].DatabaseAdminClient = InsecureDatabaseAdminClient +except ModuleNotFoundError: + pass + +try: + import fs_gcsfs + from google.cloud import storage + + original_bucket_get_blob = storage.Bucket.get_blob + original_makedir = fs_gcsfs._gcsfs.GCSFS.makedir + + # In-memory tracking of directories created during the current process run + created_dirs = set() + + def patched_makedir(self, path, permissions=None, recreate=False): + _path = self.validatepath(path) + _key = self._path_to_dir_key(_path) + res = original_makedir(self, path, permissions, recreate) + created_dirs.add(_key) + return res + + def patched_bucket_get_blob(self, blob_name, client=None, **kwargs): + blob = original_bucket_get_blob(self, blob_name, client, **kwargs) + if blob: + return blob + if isinstance(blob_name, str) and blob_name.endswith("/"): + # Case 1: Check in-memory created directories tracking for empty folders + if blob_name in created_dirs: + return storage.blob.Blob(name=blob_name, bucket=self) + for d in created_dirs: + if d.startswith(blob_name): + return storage.blob.Blob(name=blob_name, bucket=self) + + # Case 2: Check prefix search for pre-existing folders containing files + blobs = list(self.list_blobs(prefix=blob_name, max_results=1)) + if blobs: + # Return a dummy Blob object so getinfo/isdir is satisfied + return storage.blob.Blob(name=blob_name, bucket=self) + return None + + storage.Bucket.get_blob = patched_bucket_get_blob + fs_gcsfs._gcsfs.GCSFS.makedir = patched_makedir +except ModuleNotFoundError: + pass + +__version__ = "0.0.0-dev" diff --git a/tests/datacommons-integration-tests/run_local_integration_test.py b/tests/datacommons-integration-tests/run_local_integration_test.py index 3fb8e7e8..00944180 100755 --- a/tests/datacommons-integration-tests/run_local_integration_test.py +++ b/tests/datacommons-integration-tests/run_local_integration_test.py @@ -25,6 +25,8 @@ import time from pathlib import Path +import grpc +from google.cloud import spanner import pytest import requests @@ -57,8 +59,8 @@ def get_port(env_var: str, default: int) -> int: SPANNER_GRPC_PORT = get_port("SPANNER_GRPC_PORT", 9010) MOCK_NL_PORT = get_port("MOCK_NL_PORT", 6060) -PROJECT_ID = os.getenv("SPANNER_PROJECT_ID", "test-project") -INSTANCE_ID = os.getenv("SPANNER_INSTANCE_ID", "test-instance") +PROJECT_ID = os.getenv("SPANNER_PROJECT_ID", "default") +INSTANCE_ID = os.getenv("SPANNER_INSTANCE_ID", "default") DATABASE_ID = os.getenv("SPANNER_DATABASE_ID", "test-db") # ============================================================================= @@ -144,6 +146,49 @@ def wait_for_service( raise RuntimeError(f"{name} failed to start within {timeout_secs} seconds.") +def wait_for_spanner(timeout_secs: int = 90) -> None: + """Waits for Spanner gRPC endpoint to become ready.""" + + from google.api_core.exceptions import GoogleAPICallError + + print("Waiting for Spanner gRPC endpoint to be ready...", flush=True) + os.environ["SPANNER_EMULATOR_HOST"] = f"127.0.0.1:{SPANNER_GRPC_PORT}" + start_time = time.time() + while time.time() - start_time < timeout_secs: + try: + client = spanner.Client(project=PROJECT_ID) + list(client.list_instances()) + print(" Spanner is ready!", flush=True) + return + except (GoogleAPICallError, grpc.RpcError): + time.sleep(0.5) + raise RuntimeError(f"Spanner failed to start within {timeout_secs} seconds.") + + +def create_spanner_db() -> None: + """Creates Spanner instance and database using client library via gRPC.""" + + os.environ["SPANNER_EMULATOR_HOST"] = f"127.0.0.1:{SPANNER_GRPC_PORT}" + print( + f">>> Creating Spanner instance {INSTANCE_ID} and database {DATABASE_ID} via gRPC client...", + flush=True, + ) + client = spanner.Client(project=PROJECT_ID) + instance = client.instance( + INSTANCE_ID, + configuration_name=f"projects/{PROJECT_ID}/instanceConfigs/emulator-config", + node_count=1, + ) + if not instance.exists(): + op = instance.create() + op.result(timeout=30) + database = instance.database(DATABASE_ID) + if not database.exists(): + op = database.create() + op.result(timeout=30) + print(" Spanner database created.", flush=True) + + def call_helper(action_type: str) -> dict: path = ( "database/initialize" @@ -218,11 +263,8 @@ def docker_stack(services_image, helper_image, keep_containers): env=compose_env, ) - # 3. Wait for Spanner Emulator REST API - wait_for_service( - f"http://localhost:{SPANNER_REST_PORT}/v1/projects/{PROJECT_ID}/instances", - "Spanner emulator", - ) + # 3. Wait for Spanner Emulator ready + wait_for_spanner() # Wait for GCS Emulator ready wait_for_service( @@ -231,31 +273,7 @@ def docker_stack(services_image, helper_image, keep_containers): ) # 4. Create database instance - print( - f">>> Creating Spanner {INSTANCE_ID}/{DATABASE_ID} in emulator...", - flush=True, - ) - inst_resp = requests.post( - f"http://localhost:{SPANNER_REST_PORT}/v1/projects/{PROJECT_ID}/instances", - json={ - "instanceId": INSTANCE_ID, - "instance": { - "name": f"projects/{PROJECT_ID}/instances/{INSTANCE_ID}", - "config": "projects/test-project/instanceConfigs/emulator-config", - "displayName": "Test Instance", - "nodeCount": 1, - }, - }, - timeout=30, - ) - inst_resp.raise_for_status() - - db_resp = requests.post( - f"http://localhost:{SPANNER_REST_PORT}/v1/projects/{PROJECT_ID}/instances/{INSTANCE_ID}/databases", - json={"createStatement": f"CREATE DATABASE `{DATABASE_ID}`"}, - timeout=30, - ) - db_resp.raise_for_status() + create_spanner_db() # 5. Boot Ingestion Helper print(">>> Starting Ingestion Helper container...", flush=True) diff --git a/tests/datacommons-integration-tests/spanner_client.py b/tests/datacommons-integration-tests/spanner_client.py deleted file mode 100644 index 8ecf2a34..00000000 --- a/tests/datacommons-integration-tests/spanner_client.py +++ /dev/null @@ -1,917 +0,0 @@ -# ruff: noqa: LOG015, G004, BLE001, PERF401, E402, ARG002, ANN204, ANN401, PTH118, PTH119, PTH120, PTH123 -# Copyright 2025 Google LLC -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -import logging -import os -from datetime import UTC, datetime -from typing import Any - -from google.cloud import spanner -from google.cloud.spanner_admin_database_v1.types import UpdateDatabaseDdlRequest -from google.cloud.spanner_v1 import Transaction -from google.cloud.spanner_v1.param_types import STRING, Array -from jinja2 import Template - -logging.getLogger().setLevel(logging.INFO) - - -import json - - -class SpannerClient: - """ - Spanner client to handle tasks like acquiring/releasing lock - and getting/updating import statuses. - """ - - _EMBEDDING_MODEL_PATH = "//aiplatform.googleapis.com/projects/{project}/locations/{location}/publishers/google/models/{model}" - _DEFAULT_MODELS = [{"name": "NodeEmbeddingModel", "endpoint": "text-embedding-005"}] - - def __init__( - self, - project_id: str, - instance_id: str, - database_id: str, - graph_database_id: str = None, - location: str = None, - model_id: str = None, - lock_id: str = "global_ingestion_lock", - models: list[dict] = None, - embedding_space: int = 768, - embedding_table: str = "NodeEmbedding", - embedding_index: str = "NodeEmbeddingIndex", - **kwargs: Any, - ) -> None: - """Initializes a Spanner client and connects to a specific database.""" - client_options = {"api_endpoint": "spanner.googleapis.com"} - credentials = None - if os.environ.get("SPANNER_EMULATOR_HOST"): - from google.auth.credentials import AnonymousCredentials - - class MockAnonymousCredentials(AnonymousCredentials): - def with_quota_project(self, quota_project_id): - return self - - credentials = MockAnonymousCredentials() - else: - client_options["quota_project_id"] = project_id - - spanner_client = spanner.Client( - project=project_id, - credentials=credentials, - client_options=client_options, - disable_builtin_metrics=True, - ) - self.spanner_client = spanner_client - instance = spanner_client.instance(instance_id) - database = instance.database(database_id) - logging.info(f"Successfully initialized database: {database.name}") - self.database = database - self.graph_database = database - if graph_database_id: - self.graph_database = instance.database(graph_database_id) - logging.info( - f"Successfully initialized graph database: {self.graph_database.name}" - ) - self.project_id = project_id - self.location = location - self.model_id = model_id - self.lock_id = lock_id - self.embedding_space = embedding_space - self.embedding_table = embedding_table - self.embedding_index = embedding_index - - if not models: - models = self._DEFAULT_MODELS - - self.models = [] - for model in models: - name = model["name"] - endpoint = self._get_embeddings_endpoint(model["endpoint"]) - self.models.append({"name": name, "endpoint": endpoint}) - - def _get_embeddings_endpoint(self, model: str = None) -> str: - """Returns the parameterized embedding model endpoint.""" - return self._EMBEDDING_MODEL_PATH.format( - project=self.project_id, - location=self.location or "us-central1", - model=model or self.model_id or "text-embedding-005", - ) - - def acquire_lock(self, workflow_id: str, timeout: int) -> bool: - """Attempts to acquire the global ingestion lock. - - Args: - workflow_id: The ID of the workflow attempting to acquire the lock. - timeout: The duration in seconds after which a lock is considered stale. - - Returns: - True if the lock was acquired, False otherwise. - """ - logging.info(f"Attempting to acquire lock for {workflow_id}") - - tx_logs = [] - - def _acquire(transaction: Transaction) -> bool: - nonlocal tx_logs - tx_logs.clear() - - sql = "SELECT LockOwner, AcquiredTimestamp FROM IngestionLock WHERE LockID = @lockId" - params = {"lockId": self.lock_id} - param_types = {"lockId": STRING} - - row_found = False - results = transaction.execute_sql(sql, params, param_types) - for row in results: - row_found = True - current_owner, acquired_at = row[0], row[1] - - lock_is_available = False - if not row_found or current_owner is None: - lock_is_available = True - else: - timeout_threshold = datetime.now(UTC) - acquired_at - if timeout_threshold.total_seconds() > timeout: - tx_logs.append( - f"Stale lock found, owned by {current_owner}. Acquiring." - ) - lock_is_available = True - - if lock_is_available: - if not row_found: - sql_statement = """ - INSERT INTO IngestionLock (LockID, LockOwner, AcquiredTimestamp) - VALUES (@lockId, @workflowId, PENDING_COMMIT_TIMESTAMP()) - """ - tx_logs.append( - f"Lock successfully acquired by {workflow_id} (new row created)" - ) - else: - sql_statement = """ - UPDATE IngestionLock - SET LockOwner = @workflowId, AcquiredTimestamp = PENDING_COMMIT_TIMESTAMP() - WHERE LockID = @lockId - """ - tx_logs.append( - f"Lock successfully acquired by {workflow_id} (existing row updated)" - ) - - transaction.execute_update( - sql_statement, - params={"workflowId": workflow_id, "lockId": self.lock_id}, - param_types={"workflowId": STRING, "lockId": STRING}, - ) - return True - tx_logs.append(f"Lock is currently held by {current_owner}") - return False - - try: - result = self.database.run_in_transaction(_acquire) - # Log results outside transaction scope to prevent duplicate logging on retries - for log in tx_logs: - logging.info(log) - return result - except Exception as e: - logging.error(f"Error acquiring lock for {workflow_id}: {e}") - raise - - def release_lock(self, workflow_id: str) -> bool: - """Releases the global lock. - - Args: - workflow_id: The ID of the workflow attempting to release the lock. - - Returns: - True if the lock was released, False otherwise. - """ - logging.info(f"Attempting to release lock for {workflow_id}") - - tx_logs = [] - - def _release(transaction: Transaction) -> bool: - nonlocal tx_logs - tx_logs.clear() - - sql = "SELECT LockOwner, AcquiredTimestamp FROM IngestionLock WHERE LockID = @lockId" - params = {"lockId": self.lock_id} - param_types = {"lockId": STRING} - - current_owner = None - results = transaction.execute_sql(sql, params, param_types) - for row in results: - current_owner = row[0] - - if current_owner == workflow_id: - sql = """ - UPDATE IngestionLock - SET LockOwner = NULL, AcquiredTimestamp = NULL - WHERE LockID = @lockId - """ - transaction.execute_update( - sql, params={"lockId": self.lock_id}, param_types={"lockId": STRING} - ) - tx_logs.append(f"Lock successfully released by {workflow_id}") - return True - tx_logs.append( - f"Lock release skipped. Lock is currently held by {current_owner}" - ) - return False - - try: - result = self.database.run_in_transaction(_release) - for log in tx_logs: - logging.info(log) - return result - except Exception as e: - logging.error(f"Error releasing lock for {workflow_id}: {e}") - raise - - def get_import_info(self, import_list: list) -> list: - """Get the details of imports to ingest. - - If import_list is empty, return info for ready imports (STAGING). - If import_list is not empty, return info for the imports in the list that are in 'STAGING' status. - - Args: - import_list: A list of import names to fetch details for. - - Returns: - A list of dictionaries, where each dictionary contains 'importName', 'latestVersion', and 'graphPath'. - """ - pending_imports = [] - logging.info(f"Fetching imports from import list {import_list}.") - - params = {} - param_types = {} - if import_list: - sql = "SELECT ImportName, LatestVersion, GraphPath FROM ImportStatus WHERE State = 'STAGING' AND ImportName IN UNNEST(@importNames)" - params = {"importNames": import_list} - param_types = {"importNames": Array(STRING)} - else: - sql = "SELECT ImportName, LatestVersion, GraphPath FROM ImportStatus WHERE State = 'STAGING'" - - # Use a read-only snapshot for this query - try: - with self.database.snapshot() as snapshot: - results = snapshot.execute_sql( - sql, params=params, param_types=param_types - ) - for row in results: - import_json = {} - import_json["importName"] = row[0] - import_json["latestVersion"] = os.path.basename(row[1]) - import_json["graphPath"] = ( - f"{row[1].rstrip('/')}/{row[2].lstrip('/')}" - ) - pending_imports.append(import_json) - - logging.info(f"Found {len(pending_imports)} import jobs.") - return pending_imports - except Exception as e: - logging.error(f"Error getting import list: {e}") - raise - - def update_ingestion_status( - self, import_names: list, workflow_id: str, status: str - ): - """Updates the ImportStatus table. - - Args: - import_names: List of import names. - workflow_id: The ID of the workflow. - status: The status of the ingestion. - """ - if not import_names: - return - - logging.info(f"Updated ingestion status for {import_names}") - - def _update(transaction: Transaction): - update_sql = "UPDATE ImportStatus SET State = @importStatus, WorkflowId = @workflowId, StatusUpdateTimestamp = PENDING_COMMIT_TIMESTAMP() WHERE ImportName IN UNNEST(@importNames)" - transaction.execute_update( - update_sql, - params={ - "importNames": import_names, - "workflowId": workflow_id, - "importStatus": status, - }, - param_types={ - "importNames": Array(STRING), - "workflowId": STRING, - "importStatus": STRING, - }, - ) - - try: - self.database.run_in_transaction(_update) - logging.info(f"Marked {len(import_names)} import jobs as {status}.") - except Exception as e: - logging.error(f"Error updating ImportStatus table: {e}") - raise - - def update_ingestion_history( - self, workflow_id: str, job_id: str, ingested_imports: list, metrics: dict - ): - """Updates the IngestionHistory table. - - Args: - workflow_id: The ID of the workflow. - job_id: The Dataflow job ID. - ingested_imports: List of ingested import names. - metrics: A dictionary containing metrics about the ingestion. - """ - - logging.info(f"Updating IngestionHistory table for workflow {workflow_id}") - - def _insert(transaction: Transaction): - columns = [ - "CompletionTimestamp", - "IngestionFailure", - "WorkflowExecutionID", - "DataflowJobId", - "IngestedImports", - "ExecutionTime", - "NodeCount", - "EdgeCount", - "ObservationCount", - ] - values = [ - [ - spanner.COMMIT_TIMESTAMP, - self.check_failed_imports(), - workflow_id, - job_id, - ingested_imports, - metrics["execution_time"], - metrics["node_count"], - metrics["edge_count"], - metrics["obs_count"], - ] - ] - transaction.insert_or_update( - table="IngestionHistory", columns=columns, values=values - ) - - try: - self.database.run_in_transaction(_insert) - logging.info( - f"Updated IngestionHistory table in main database for workflow {workflow_id}" - ) - except Exception as e: - logging.error( - f"Error updating IngestionHistory table in main database: {e}" - ) - raise - - # Decouple dual writes: failures in secondary graph database should not roll back main commit. - if self.graph_database and self.graph_database.name != self.database.name: - try: - self.graph_database.run_in_transaction(_insert) - logging.info( - f"Updated IngestionHistory table in graph database for workflow {workflow_id}" - ) - except Exception as e: - logging.critical( - f"CRITICAL: Failed to update IngestionHistory in graph database (split-brain hazard): {e}. " - f"Main database write succeeded for workflow {workflow_id}." - ) - - def update_import_version_history(self, import_list_json: list, workflow_id: str): - """Updates the ImportVersionHistory table. - - Args: - import_list_json: A list of dictionaries containing import details. - workflow_id: The ID of the workflow. - """ - if not import_list_json: - return - - logging.info(f"Updating ImportVersionHistory table for workflow {workflow_id}") - - def _insert(transaction: Transaction): - version_history_columns = [ - "ImportName", - "Version", - "UpdateTimestamp", - "Comment", - ] - version_history_values = [] - for import_json in import_list_json: - version_history_values.append( - [ - import_json["importName"], - import_json["latestVersion"], - spanner.COMMIT_TIMESTAMP, - "ingestion-workflow:" + workflow_id, - ] - ) - - if version_history_values: - transaction.insert( - table="ImportVersionHistory", - columns=version_history_columns, - values=version_history_values, - ) - - try: - self.database.run_in_transaction(_insert) - logging.info( - f"Updated ImportVersionHistory table for workflow {workflow_id}" - ) - except Exception as e: - logging.error(f"Error updating ImportVersionHistory table: {e}") - raise - - def check_failed_imports(self) -> bool: - """Checks if there are any failed imports.""" - with self.database.snapshot() as snapshot: - results = snapshot.execute_sql( - "SELECT 1 FROM ImportStatus WHERE State = 'PENDING' LIMIT 1" - ) - return any(results) - - def update_import_status(self, params: dict): - """Updates the status for the specified import job. - - Args: - params: A dictionary containing import parameters. - """ - import_name = params["import_name"] - job_id = params["job_id"] - execution_time = params["execution_time"] - data_volume = params["data_volume"] - status = params["status"] - latest_version = params["latest_version"] - next_refresh = datetime.fromisoformat(params["next_refresh"]) - graph_path = params["graph_path"] - logging.info(f"Updating import status in spanner {params}") - - def _record(transaction: Transaction): - columns = [ - "ImportName", - "State", - "JobId", - "ExecutionTime", - "DataVolume", - "NextRefreshTimestamp", - "LatestVersion", - "GraphPath", - "StatusUpdateTimestamp", - ] - - row_values = [ - import_name, - status, - job_id, - execution_time, - data_volume, - next_refresh, - latest_version, - graph_path, - spanner.COMMIT_TIMESTAMP, - ] - - if status == "STAGING": - columns.append("DataImportTimestamp") - row_values.append(spanner.COMMIT_TIMESTAMP) - - transaction.insert_or_update( - table="ImportStatus", columns=columns, values=[row_values] - ) - - try: - self.database.run_in_transaction(_record) - logging.info(f"Marked {import_name} as {status}.") - except Exception as e: - logging.error(f"Error updating import status for {import_name}: {e}") - raise - - def update_version_history(self, import_name: str, version: str, comment: str): - """Updates the version history table. - - Args: - import_name: The name of the import. - version: The version string. - comment: The comment for the update. - """ - import_name = import_name.split(":")[-1] - logging.info(f"Updating version history for {import_name} to {version}") - - def _record(transaction: Transaction): - columns = ["ImportName", "Version", "UpdateTimestamp", "Comment"] - values = [[import_name, version, spanner.COMMIT_TIMESTAMP, comment]] - transaction.insert( - table="ImportVersionHistory", columns=columns, values=values - ) - - try: - self.database.run_in_transaction(_record) - logging.info(f"Added version history entry for {import_name}") - except Exception as e: - logging.error(f"Error updating version history for {import_name}: {e}") - raise - - def initialize_database(self): - """Initializes the database by creating all required tables and proto bundles.""" - logging.info("Initializing database...") - - query = """ - SELECT 'table' as type, table_name as name FROM information_schema.tables WHERE table_schema = '' - UNION ALL - SELECT 'index' as type, index_name as name FROM information_schema.indexes WHERE table_schema = '' AND table_name IN ('NodeEmbedding', 'Edge', 'Observation') - UNION ALL - SELECT 'model' as type, model_name as name FROM information_schema.models WHERE model_schema = '' - """ - - existing_tables = [] - existing_indexes = [] - existing_models = [] - - with self.database.snapshot() as snapshot: - results = snapshot.execute_sql(query) - for row in results: - if len(row) < 2: - logging.warning(f"Invalid row from query: {row}") - continue - obj_type = row[0] - obj_name = row[1] - if obj_type == "table": - existing_tables.append(obj_name) - elif obj_type == "index": - existing_indexes.append(obj_name) - elif obj_type == "model": - existing_models.append(obj_name) - - logging.info(f"Existing tables: {existing_tables}") - logging.info(f"Existing indexes: {existing_indexes}") - logging.info(f"Existing models: {existing_models}") - - required_tables = [ - "Node", - "Edge", - "Observation", - "ImportStatus", - "IngestionHistory", - "ImportVersionHistory", - "IngestionLock", - "Cache", - "NodeEmbedding", - ] - required_indexes = [ - "InEdge", - "VariableMeasuredObservationAbout", - "NodeEmbeddingIndex", - ] - required_models = ["NodeEmbeddingModel"] - - missing_tables = [t for t in required_tables if t not in existing_tables] - missing_indexes = [i for i in required_indexes if i not in existing_indexes] - missing_models = [m for m in required_models if m not in existing_models] - - total_required = ( - len(required_tables) + len(required_indexes) + len(required_models) - ) - total_missing = len(missing_tables) + len(missing_indexes) + len(missing_models) - - if total_missing == 0: - logging.info("Database is properly initialized.") - return - - if total_missing < total_required: - raise RuntimeError( - f"Database inconsistent state. Missing tables: {missing_tables}, missing indexes: {missing_indexes}, missing models: {missing_models}. Please clean up manually." - ) - - logging.info("Creating all tables and proto bundles...") - - schema_path = os.path.join(os.path.dirname(__file__), "schema.sql") - logging.info(f"Reading schema from {schema_path}") - try: - embeddings_endpoint = self._get_embeddings_endpoint() - with open(schema_path) as f: - schema_content = f.read() - - # Remove SQL comment lines starting with '--' - schema_lines = [] - for line in schema_content.splitlines(): - if line.strip().startswith("--"): - continue - schema_lines.append(line) - schema_content = "\n".join(schema_lines) - - schema_content = Template(schema_content).render( - embeddings_endpoint=embeddings_endpoint, - models=self.models, - embedding_space=self.embedding_space, - embedding_table=self.embedding_table, - embedding_index=self.embedding_index, - ) - - if os.environ.get("SPANNER_EMULATOR_HOST"): - import re - - # Note: 'columnar_policy' is a production Enterprise-only feature not supported - # by either standard Cloud Spanner Emulator or local Spanner Omni containers. - # We strip it here along with any resulting empty OPTIONS () clauses to prevent DDL errors. - schema_content = re.sub( - r",\s*columnar_policy\s*=\s*['\"].*?['\"]", "", schema_content - ) - schema_content = re.sub( - r"columnar_policy\s*=\s*['\"].*?['\"]\s*,?", "", schema_content - ) - # Clean up any empty OPTIONS () blocks left behind - schema_content = re.sub(r"OPTIONS\s*\(\s*\)", "", schema_content) - # Clean up any trailing commas before semicolons - schema_content = re.sub(r",\s*;", ";", schema_content) - - # TODO(yiyuanc): Remove this once helper DDL schema is updated to include the 'provenance' column. - schema_content = schema_content.replace( - "provenance_url STRING(1024),", - "provenance_url STRING(1024),\n provenance STRING(1024),", - ) - - ddl_statements = [s.strip() for s in schema_content.split(";") if s.strip()] - except Exception as e: - logging.error(f"Failed to read schema file: {e}") - raise - - database_path = self.database.name - logging.info(f"Updating DDL for {database_path}") - - try: - admin_client = self.spanner_client.database_admin_api - request = UpdateDatabaseDdlRequest( - database=database_path, - statements=ddl_statements, - ) - operation = admin_client.update_database_ddl(request=request) - operation.result() - logging.info("Database initialized successfully.") - except Exception as e: - logging.error(f"Failed to update DDL: {e}") - raise - - def seed_database(self): - """Seeds the database with base empty nodes.""" - logging.info("Seeding database with base nodes...") - - # Collect execution details to log outside transaction scopes - seed_counts = {} - - def _seed(transaction: Transaction): - nonlocal seed_counts - seed_counts.clear() - - # 1. Seed Node Table - nodes = { - # Classes & base nodes - "StatisticalVariable": [ - "StatisticalVariable", - "StatisticalVariable", - "StatisticalVariable", - ["Class"], - spanner.COMMIT_TIMESTAMP, - ], - "StatVarGroup": [ - "StatVarGroup", - "StatVarGroup", - "StatVarGroup", - ["Class"], - spanner.COMMIT_TIMESTAMP, - ], - "StatVarObservation": [ - "StatVarObservation", - "StatVarObservation", - "StatVarObservation", - ["Class"], - spanner.COMMIT_TIMESTAMP, - ], - "Topic": [ - "Topic", - "Topic", - "Topic", - ["Class"], - spanner.COMMIT_TIMESTAMP, - ], - "dc/g/Root": [ - "dc/g/Root", - "Data Commons Variables", - "dc/g/Root", - ["StatVarGroup"], - spanner.COMMIT_TIMESTAMP, - ], - "dc/topic/Root": [ - "dc/topic/Root", - "Statistics", - "dc/topic/Root", - ["Topic"], - spanner.COMMIT_TIMESTAMP, - ], - # Frog custom variables - "Count_Person_Trumal": [ - "Count_Person_Trumal", - "Count of Trumal Person", - "Count of Trumal Person", - ["StatisticalVariable"], - spanner.COMMIT_TIMESTAMP, - ], - "number_of_frogs": [ - "number_of_frogs", - "Number of Frogs", - "Number of Frogs", - ["StatisticalVariable"], - spanner.COMMIT_TIMESTAMP, - ], - "Count_Frog_Green": [ - "Count_Frog_Green", - "Count of Green Frogs", - "Count of Green Frogs", - ["StatisticalVariable"], - spanner.COMMIT_TIMESTAMP, - ], - "Average_Age_Frog": [ - "Average_Age_Frog", - "Average Age of Frogs", - "Average Age of Frogs", - ["StatisticalVariable"], - spanner.COMMIT_TIMESTAMP, - ], - # Places - "country/USA": [ - "country/USA", - "United States of America", - "United States of America", - ["Country", "Place"], - spanner.COMMIT_TIMESTAMP, - ], - "country/CAN": [ - "country/CAN", - "Canada", - "Canada", - ["Country", "Place"], - spanner.COMMIT_TIMESTAMP, - ], - } - - subjects = list(nodes.keys()) - sql = "SELECT subject_id FROM Node WHERE subject_id IN UNNEST(@subjects)" - params = {"subjects": subjects} - param_types = {"subjects": Array(STRING)} - existing = set() - for row in transaction.execute_sql(sql, params, param_types): - existing.add(row[0]) - - node_values = [nodes[subj] for subj in subjects if subj not in existing] - if node_values: - columns = [ - "subject_id", - "name", - "value", - "types", - "last_update_timestamp", - ] - transaction.insert(table="Node", columns=columns, values=node_values) - seed_counts["nodes"] = len(node_values) - - # 2. Seed Edge Table - edges = [ - # Connect custom variables to Root SVG - ("Count_Person_Trumal", "memberOf", "dc/g/Root", "Trumal_Import"), - ("number_of_frogs", "memberOf", "dc/g/Root", "Trumal_Import"), - ("Count_Frog_Green", "memberOf", "dc/g/Root", "Trumal_Import"), - ("Average_Age_Frog", "memberOf", "dc/g/Root", "Trumal_Import"), - # Topic child relation (for browser traversal) - ("dc/topic/Root", "hasMemberSVG", "dc/g/Root", "Trumal_Import"), - ] - edge_columns = ["subject_id", "predicate", "object_id", "provenance"] - transaction.insert_or_update( - table="Edge", columns=edge_columns, values=edges - ) - seed_counts["edges"] = len(edges) - - # 3. Seed TimeSeries and Observation Tables - ts_values = [] - obs_values = [] - - # Facet JSON definition - facet_json = json.dumps( - { - "provenance": "http://example.com/provenance", - "observationPeriod": "", - "unit": "", - "measurementMethod": "", - } - ) - - # Raw observation data to seed - raw_data = [ - # USA - ( - "country/USA", - "Count_Person_Trumal", - {"2025": "4000", "2026": "5000"}, - ), - ("country/USA", "number_of_frogs", {"2025": "40000", "2026": "42000"}), - ("country/USA", "Count_Frog_Green", {"2025": "30000", "2026": "32000"}), - ("country/USA", "Average_Age_Frog", {"2025": "2.5", "2026": "2.7"}), - # CAN - ("country/CAN", "Count_Person_Trumal", {"2025": "800", "2026": "1000"}), - ("country/CAN", "number_of_frogs", {"2025": "10000", "2026": "12000"}), - ("country/CAN", "Count_Frog_Green", {"2025": "8000", "2026": "9000"}), - ("country/CAN", "Average_Age_Frog", {"2025": "3.0", "2026": "3.2"}), - ] - - for entity, var, points in raw_data: - facet_id = "custom_facet" - extra_entities_id = "" - - # Entities JSON field - entities_json = json.dumps({"entity1": entity}) - - # TimeSeries row - ts_values.append( - ( - var, - extra_entities_id, - facet_id, - entities_json, - facet_json, - spanner.COMMIT_TIMESTAMP, - ) - ) - - # Observation child rows - for date, val in points.items(): - obs_values.append( - ( - var, - entity, - extra_entities_id, - facet_id, - date, - val, - spanner.COMMIT_TIMESTAMP, - ) - ) - - # Insert TimeSeries first, then interleaved Observation - ts_columns = [ - "variable_measured", - "extra_entities_id", - "facet_id", - "entities", - "facet", - "last_update_timestamp", - ] - transaction.insert_or_update( - table="TimeSeries", columns=ts_columns, values=ts_values - ) - - obs_columns = [ - "variable_measured", - "entity1", - "extra_entities_id", - "facet_id", - "date", - "value", - "last_update_timestamp", - ] - transaction.insert_or_update( - table="Observation", columns=obs_columns, values=obs_values - ) - seed_counts["observations"] = len(obs_values) - - try: - self.database.run_in_transaction(_seed) - logging.info( - f"Main database seeded successfully: {seed_counts.get('nodes', 0)} nodes, " - f"{seed_counts.get('edges', 0)} edges, and {seed_counts.get('observations', 0)} observations." - ) - except Exception as e: - logging.error(f"Error seeding main database: {e}") - raise - - if self.graph_database and self.graph_database.name != self.database.name: - try: - self.graph_database.run_in_transaction(_seed) - logging.info( - f"Graph database seeded successfully: {seed_counts.get('nodes', 0)} nodes, " - f"{seed_counts.get('edges', 0)} edges, and {seed_counts.get('observations', 0)} observations." - ) - except Exception as e: - logging.critical( - f"CRITICAL: Failed to seed graph database (split-brain hazard): {e}. " - f"Main database seed succeeded." - ) From 67523635c23654e2e0787c2169e038e8aa188129 Mon Sep 17 00:00:00 2001 From: Yiyuan Chen Date: Thu, 9 Jul 2026 16:33:45 +0000 Subject: [PATCH 4/5] chore: fix ruff formatting issues on local emulator compatibility branch --- .../run_local_integration_test.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/tests/datacommons-integration-tests/run_local_integration_test.py b/tests/datacommons-integration-tests/run_local_integration_test.py index 00944180..0d7ec713 100755 --- a/tests/datacommons-integration-tests/run_local_integration_test.py +++ b/tests/datacommons-integration-tests/run_local_integration_test.py @@ -26,9 +26,9 @@ from pathlib import Path import grpc -from google.cloud import spanner import pytest import requests +from google.cloud import spanner # ============================================================================= # Parameterized Configuration Constants @@ -209,7 +209,6 @@ def call_helper(action_type: str) -> dict: ) from e # noqa: BLE001 - # ============================================================================= # Pytest Orchestration Fixture # ============================================================================= From 8803f88cdbf016f8b62d561392563cda809dfaad Mon Sep 17 00:00:00 2001 From: Yiyuan Chen Date: Thu, 9 Jul 2026 16:48:21 +0000 Subject: [PATCH 5/5] chore: fix remaining ruff linter issues in sandbox cleaner script --- .../clean_orphaned_sandboxes.py | 103 +++++++++++------- 1 file changed, 61 insertions(+), 42 deletions(-) diff --git a/tests/datacommons-integration-tests/clean_orphaned_sandboxes.py b/tests/datacommons-integration-tests/clean_orphaned_sandboxes.py index d0fefec7..47399e5a 100644 --- a/tests/datacommons-integration-tests/clean_orphaned_sandboxes.py +++ b/tests/datacommons-integration-tests/clean_orphaned_sandboxes.py @@ -20,7 +20,6 @@ """ import argparse -import os import re import shutil import subprocess @@ -28,10 +27,14 @@ from pathlib import Path -def run_command(args: list[str], cwd: Path | None = None, check: bool = True) -> subprocess.CompletedProcess: +def run_command( + args: list[str], cwd: Path | None = None, *, check: bool = True +) -> subprocess.CompletedProcess: print(f"Running command: {' '.join(args)} (cwd: {cwd or Path.cwd()})", flush=True) try: - return subprocess.run(args, cwd=cwd, check=check, text=True, capture_output=True) + return subprocess.run( # noqa: S603 + args, cwd=cwd, check=check, text=True, capture_output=True + ) except subprocess.CalledProcessError as e: print(f"--- COMMAND FAILED: {' '.join(args)} ---", file=sys.stderr) if e.stdout: @@ -46,9 +49,9 @@ def get_active_namespaces(tf_state_bucket: str) -> list[str]: """Lists all prefixes (namespaces) in GCS tf-state-bucket.""" print(f"Listing active states in gs://{tf_state_bucket}/terraform/state/...") try: - proc = run_command([ - "gcloud", "storage", "ls", f"gs://{tf_state_bucket}/terraform/state/" - ]) + proc = run_command( + ["gcloud", "storage", "ls", f"gs://{tf_state_bucket}/terraform/state/"] + ) namespaces = [] # Expected format: gs://bucket-name/terraform/state/itest-XXXX/ pattern = rf"gs://{re.escape(tf_state_bucket)}/terraform/state/([^/]+)/" @@ -56,34 +59,36 @@ def get_active_namespaces(tf_state_bucket: str) -> list[str]: match = re.match(pattern, line.strip()) if match: namespaces.append(match.group(1)) - return sorted(list(set(namespaces))) - except Exception as e: + return sorted(set(namespaces)) + except Exception as e: # noqa: BLE001 print(f"Error listing state directories: {e}", file=sys.stderr) return [] -def destroy_sandbox(project_id: str, tf_state_bucket: str, namespace: str, workspace_root: Path) -> bool: +def destroy_sandbox( + project_id: str, tf_state_bucket: str, namespace: str, workspace_root: Path +) -> bool: """Reconstructs local state and executes terraform destroy.""" - print(f"\n=====================================================================") + print("\n=====================================================================") print(f" DESTROYING SANDBOX: namespace={namespace} in project={project_id}") - print(f"=====================================================================") - - workspace_dir = Path(f"/tmp/workspace-purge-{namespace}") + print("=====================================================================") + + workspace_dir = Path(f"/tmp/workspace-purge-{namespace}") # noqa: S108 sandbox_dir = workspace_dir / namespace - + try: # Create directories if workspace_dir.exists(): shutil.rmtree(workspace_dir) sandbox_dir.mkdir(parents=True) - + # Copy local infra modules & configs local_infra_dir = workspace_root / "infra" / "dcp" shutil.copy(local_infra_dir / "variables.tf", sandbox_dir / "variables.tf") shutil.copy(local_infra_dir / "outputs.tf", sandbox_dir / "outputs.tf") shutil.copy(local_infra_dir / "main.tf", sandbox_dir / "main.tf") shutil.copytree(local_infra_dir / "modules", sandbox_dir / "modules") - + # Construct backend.tf backend_content = f"""terraform {{ backend "gcs" {{ @@ -93,7 +98,7 @@ def destroy_sandbox(project_id: str, tf_state_bucket: str, namespace: str, works }} """ (sandbox_dir / "backend.tf").write_text(backend_content, encoding="utf-8") - + # Construct terraform.tfvars tfvars_content = f"""project_id = "{project_id}" namespace = "{namespace}" @@ -105,32 +110,39 @@ def destroy_sandbox(project_id: str, tf_state_bucket: str, namespace: str, works datacommons_services_max_instances = 1 """ (sandbox_dir / "terraform.tfvars").write_text(tfvars_content, encoding="utf-8") - + # Init & Destroy run_command(["terraform", "init"], cwd=sandbox_dir) run_command(["terraform", "destroy", "-auto-approve"], cwd=sandbox_dir) - + # Delete state file in GCS print(f"Removing GCS state files for namespace {namespace}...") - run_command([ - "gcloud", "storage", "rm", "--recursive", - f"gs://{tf_state_bucket}/terraform/state/{namespace}/" - ]) - + run_command( + [ + "gcloud", + "storage", + "rm", + "--recursive", + f"gs://{tf_state_bucket}/terraform/state/{namespace}/", + ] + ) + print(f"Successfully destroyed sandbox resources for namespace: {namespace}") return True - - except Exception as e: + + except Exception as e: # noqa: BLE001 print(f"Failed to destroy sandbox {namespace}: {e}", file=sys.stderr) return False - + finally: if workspace_dir.exists(): shutil.rmtree(workspace_dir) def main() -> None: - parser = argparse.ArgumentParser(description="Purge orphaned GCP sandbox integration test resources.") + parser = argparse.ArgumentParser( + description="Purge orphaned GCP sandbox integration test resources." + ) parser.add_argument( "--project-id", default="datcom-ci", help="GCP project ID (default: datcom-ci)" ) @@ -154,43 +166,50 @@ def main() -> None: action="store_true", help="Destroy all active namespaces found in the state bucket", ) - + args = parser.parse_args() workspace_root = Path(__file__).resolve().parent.parent.parent - + # 1. Fetch active namespaces active_namespaces = get_active_namespaces(args.tf_state_bucket) if not active_namespaces: print("No active integration test state namespaces found in GCS.") return - + print(f"Active namespaces found ({len(active_namespaces)}):") for ns in active_namespaces: print(f" - {ns}") - + if args.list: return - + # 2. Determine target namespaces targets = [] if args.namespace: if args.namespace not in active_namespaces: - print(f"Error: Namespace '{args.namespace}' was not found in GCS.", file=sys.stderr) + print( + f"Error: Namespace '{args.namespace}' was not found in GCS.", + file=sys.stderr, + ) sys.exit(1) targets = [args.namespace] elif args.all: targets = active_namespaces else: - print("\nSpecify either --namespace , --all, or --list. See --help for details.") + print( + "\nSpecify either --namespace , --all, or --list. See --help for details." + ) sys.exit(1) - + # 3. Confirm destroy - print(f"\nWARNING: This will permanently destroy all resources associated with the following namespaces: {', '.join(targets)}") + print( + f"\nWARNING: This will permanently destroy all resources associated with the following namespaces: {', '.join(targets)}" + ) confirm = input("Are you sure you want to proceed? [y/N]: ").strip().lower() if confirm not in ["y", "yes"]: print("Aborted.") sys.exit(0) - + # 4. Run destroys successes = 0 failures = 0 @@ -199,10 +218,10 @@ def main() -> None: successes += 1 else: failures += 1 - - print(f"\n=====================================================================") - print(f" PURGE SUMMARY") - print(f"=====================================================================") + + print("\n=====================================================================") + print(" PURGE SUMMARY") + print("=====================================================================") print(f" Successes: {successes}") print(f" Failures: {failures}") if failures > 0: