From d3703d1b52cbfa6d33cceb244159df8d7cd46e80 Mon Sep 17 00:00:00 2001 From: Vishal Gupta Date: Fri, 11 Sep 2026 08:14:06 +0000 Subject: [PATCH] Airflow implementation of import automation workflow --- .gitignore | 3 +- .../cloudbuild/cloudbuild.workflow.yaml | 113 +++ ...estion-golden-verification.cloudbuild.yaml | 127 +++ import-automation/workflow/README.md | 32 + .../workflow/build_manifest_catalog.py | 153 ++++ import-automation/workflow/e2e_dag_test.py | 435 +++++++++++ .../workflow/golden_verification.py | 517 +++++++++++++ .../workflow/import_automation_workflow.py | 732 ++++++++++++++++++ .../workflow/import_dags_factory.py | 101 +++ import-automation/workflow/manifest.json | 12 + 10 files changed, 2224 insertions(+), 1 deletion(-) create mode 100644 import-automation/cloudbuild/cloudbuild.workflow.yaml create mode 100644 import-automation/cloudbuild/ingestion-golden-verification.cloudbuild.yaml create mode 100644 import-automation/workflow/README.md create mode 100644 import-automation/workflow/build_manifest_catalog.py create mode 100644 import-automation/workflow/e2e_dag_test.py create mode 100644 import-automation/workflow/golden_verification.py create mode 100644 import-automation/workflow/import_automation_workflow.py create mode 100644 import-automation/workflow/import_dags_factory.py create mode 100644 import-automation/workflow/manifest.json diff --git a/.gitignore b/.gitignore index 87f3d76b24..af4044b563 100644 --- a/.gitignore +++ b/.gitignore @@ -33,4 +33,5 @@ import-automation/executor/config_override.json # Ignore generated troubleshooting post-mortem documents agents/troubleshooting/ - +# Ignore compiled Airflow import catalog (generated at build time) +import-automation/workflow/imports_catalog.json diff --git a/import-automation/cloudbuild/cloudbuild.workflow.yaml b/import-automation/cloudbuild/cloudbuild.workflow.yaml new file mode 100644 index 0000000000..1030b5e1d1 --- /dev/null +++ b/import-automation/cloudbuild/cloudbuild.workflow.yaml @@ -0,0 +1,113 @@ +# 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. +# +# Cloud Build configuration to compile manifest.json files, sync Airflow DAGs to Cloud Composer, +# and run an End-to-End (E2E) DAG verification test. +# +# Usage: +# gcloud builds submit . \ +# --config=import-automation/cloudbuild/cloudbuild.workflow.yaml \ +# --project=datcom-ci \ +# --substitutions=_PROJECT_ID=datcom-import-automation-prod,_COMPOSER_ENV_NAME=import-automation-airflow,_LOCATION=us-central1 + +substitutions: + _PROJECT_ID: 'datcom-import-automation-prod' + _LOCATION: 'us-central1' + _COMPOSER_ENV_NAME: 'import-automation-airflow' + _DAG_SUBDIR: 'datacommons_airflow' + _DAG_BUCKET: '' + _ALLOWLIST: '' + _RUN_E2E_TEST: 'true' + _E2E_DAG_ID: 'USFed_ConstantMaturityRates_Test' + _E2E_SKIP_IMPORT_JOB: 'false' + +steps: + # 1. Compile manifest.json files into imports_catalog.json + - id: 'build-catalog' + name: 'python:3.11-slim' + entrypoint: 'bash' + args: + - '-c' + - | + ARGS="--data-dir=. --output=import-automation/workflow/imports_catalog.json" + if [ -n "${_ALLOWLIST}" ]; then + ARGS="$${ARGS} --allowlist ${_ALLOWLIST}" + fi + python import-automation/workflow/build_manifest_catalog.py $${ARGS} + + # 2. Validate DAG, factory, and E2E test syntax + - id: 'validate-dags' + name: 'python:3.11-slim' + entrypoint: 'bash' + args: + - '-c' + - | + python -m py_compile import-automation/workflow/import_automation_workflow.py + python -m py_compile import-automation/workflow/golden_verification.py + python -m py_compile import-automation/workflow/import_dags_factory.py + python -m py_compile import-automation/workflow/e2e_dag_test.py + echo "DAG syntax validation passed." + + # 3. Deploy/Sync DAGs, factory, and catalog to Cloud Composer GCS bucket + - id: 'deploy-dags' + name: 'gcr.io/cloud-builders/gcloud' + entrypoint: 'bash' + args: + - '-c' + - | + if [ -n "${_DAG_BUCKET}" ]; then + TARGET_GCS="gs://${_DAG_BUCKET}/dags/${_DAG_SUBDIR}" + else + DAG_PREFIX=$(gcloud composer environments describe ${_COMPOSER_ENV_NAME} \ + --location=${_LOCATION} \ + --project=${_PROJECT_ID} \ + --format="value(config.dagGcsPrefix)") + TARGET_GCS="$${DAG_PREFIX}/${_DAG_SUBDIR}" + fi + + gcloud composer environments describe ${_COMPOSER_ENV_NAME} \ + --location=${_LOCATION} \ + --project=${_PROJECT_ID} \ + --format="value(config.airflowUri)" > /workspace/composer_webserver_url.txt || true + + echo "Deploying Airflow DAGs to: $${TARGET_GCS}" + gcloud storage cp import-automation/workflow/import_automation_workflow.py "$${TARGET_GCS}/import_automation_workflow.py" + gcloud storage cp import-automation/workflow/golden_verification.py "$${TARGET_GCS}/golden_verification.py" + gcloud storage cp import-automation/workflow/import_dags_factory.py "$${TARGET_GCS}/import_dags_factory.py" + gcloud storage cp import-automation/workflow/imports_catalog.json "$${TARGET_GCS}/imports_catalog.json" + echo "Successfully deployed Airflow DAGs." + + # 4. Run End-to-End DAG Verification Test in Cloud Composer + - id: 'e2e-test-dag' + name: 'python:3.11-slim' + entrypoint: 'bash' + args: + - '-c' + - | + if [ "${_RUN_E2E_TEST}" != "true" ]; then + echo "Skipping E2E DAG test (_RUN_E2E_TEST=${_RUN_E2E_TEST})." + exit 0 + fi + pip install --quiet google-auth requests + E2E_ARGS="--project-id=${_PROJECT_ID} --location=${_LOCATION} --composer-env=${_COMPOSER_ENV_NAME} --dag-id=${_E2E_DAG_ID}" + if [ -s /workspace/composer_webserver_url.txt ]; then + E2E_ARGS="$${E2E_ARGS} --webserver-url=$(cat /workspace/composer_webserver_url.txt)" + fi + if [ "${_E2E_SKIP_IMPORT_JOB}" = "true" ]; then + E2E_ARGS="$${E2E_ARGS} --skip-import-job" + fi + python3 import-automation/workflow/e2e_dag_test.py $${E2E_ARGS} + +options: + logging: CLOUD_LOGGING_ONLY diff --git a/import-automation/cloudbuild/ingestion-golden-verification.cloudbuild.yaml b/import-automation/cloudbuild/ingestion-golden-verification.cloudbuild.yaml new file mode 100644 index 0000000000..5e313f4a51 --- /dev/null +++ b/import-automation/cloudbuild/ingestion-golden-verification.cloudbuild.yaml @@ -0,0 +1,127 @@ +# 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. +# +# Cloud Build configuration for staging schema verification gate. +# Clears Staging Redis cache, runs golden tests in update mode, generates diffs, +# creates a GitHub PR if changes are detected, and outputs summary for human approval. + +substitutions: + _GITHUB_ORG: 'datacommonsorg' + _GITHUB_REPO: 'website' + _GITHUB_BRANCH: 'master' + _GITHUB_AUTHOR: 'datacommons-robot-author' + _PR_REVIEWER: 'datacommonsorg/website-reviewers' + _DIFF_BUCKET: 'datcom-ci-test' + +availableSecrets: + secretManager: + - versionName: projects/879489846695/secrets/GH_PAT/versions/latest + env: 'ghsecret' + +steps: + # 1. Initialize mixer submodule so helm configs for mixer are available + - id: 'init-mixer-submodule' + name: 'gcr.io/cloud-builders/git' + args: ['submodule', 'update', '--init', '--depth', '1', 'mixer'] + + # 2. Clear Redis cache on Staging Mixer + - id: 'clear-staging-mixer-cache' + name: 'gcr.io/datcom-ci/datacommons-script-runner:latest' + entrypoint: 'bash' + args: + - 'tools/clearcache/run.sh' + - 'mixer' + - 'staging' + + # 3. Run Explore & NL golden tests against staging backend in update mode (-g) + - id: 'run-golden-tests' + name: 'python:3.11.3' + entrypoint: '/bin/sh' + args: + - '-c' + - | + pip install uv + ./run_test.sh --explore -g || true + ./run_test.sh --nl -g || true + + # 4. Detect diffs, create PR if needed, and write diff_summary.json to GCS + - id: 'create-pr-and-summary' + name: 'gcr.io/google.com/cloudsdktool/cloud-sdk:slim' + entrypoint: 'bash' + secretEnv: ['ghsecret'] + args: + - '-c' + - | + apt-get update && apt-get install -y git gh jq + + DIFF_TARGETS="server/integration_tests/test_data/ server/config/nl_page/" + BUILD_ID_SAFE=$(echo "${BUILD_ID}" | tr -cd '[:alnum:]-') + + # Check for diffs in golden directories + if git diff --exit-code --quiet -- $${DIFF_TARGETS}; then + echo "NO_DIFF_DETECTED" + cat < /workspace/diff_summary.json + { + "has_diff": false, + "pr_url": "", + "build_id": "${BUILD_ID}", + "message": "No golden diffs detected" + } + EOF + else + echo "DIFF_DETECTED: Generating branch and opening PR..." + BRANCH_NAME="schema-golden-diff-$${BUILD_ID_SAFE}" + PR_URL="" + + if [ -n "$$ghsecret" ]; then + git config user.name "${_GITHUB_AUTHOR}" + git config user.email "${_GITHUB_AUTHOR}@users.noreply.github.com" + git checkout -b "$${BRANCH_NAME}" + git add $${DIFF_TARGETS} + git commit -m "test(goldens): automated schema staging diff update (build $${BUILD_ID_SAFE})" + + export GH_TOKEN="$$ghsecret" + git push "https://${_GITHUB_AUTHOR}:$${ghsecret}@github.com/${_GITHUB_ORG}/${_GITHUB_REPO}.git" "HEAD:$${BRANCH_NAME}" || true + + PR_URL=$(gh pr create \ + --repo "${_GITHUB_ORG}/${_GITHUB_REPO}" \ + --title "Staging Schema Golden Diff: Build $${BUILD_ID_SAFE}" \ + --body "Automated golden diff from staging schema import. Please review and merge. [Build Log](https://console.cloud.google.com/cloud-build/builds/${BUILD_ID}?project=${PROJECT_ID})" \ + --base "${_GITHUB_BRANCH}" \ + --reviewer "${_PR_REVIEWER}" \ + --head "$${BRANCH_NAME}" \ + --label "golden-diff-gate" 2>/dev/null || echo "") + echo "AUTOMATED_PR_URL:$${PR_URL}" + else + echo "GH_PAT secret not provided; skipped git push/PR creation." + fi + + cat < /workspace/diff_summary.json + { + "has_diff": true, + "pr_url": "$${PR_URL}", + "build_id": "${BUILD_ID}", + "branch_name": "$${BRANCH_NAME}", + "message": "Golden diff detected" + } + EOF + fi + + echo "Uploading diff_summary.json to gs://${_DIFF_BUCKET}/golden_diffs/${BUILD_ID}/diff_summary.json..." + gcloud storage cp /workspace/diff_summary.json "gs://${_DIFF_BUCKET}/golden_diffs/${BUILD_ID}/diff_summary.json" || true + cat /workspace/diff_summary.json + +options: + machineType: 'E2_HIGHCPU_32' +timeout: '2400s' diff --git a/import-automation/workflow/README.md b/import-automation/workflow/README.md new file mode 100644 index 0000000000..ef0b099c3e --- /dev/null +++ b/import-automation/workflow/README.md @@ -0,0 +1,32 @@ +# Import Automation Workflow (Airflow / Cloud Composer) + +This directory contains the Apache Airflow DAG definitions and dynamic factory for automating Data Commons data imports. + +## Architecture + +1. **`build_manifest_catalog.py`**: + Scans all `manifest.json` files in the repository (excluding `scripts/entities`) and compiles them into `imports_catalog.json`. + To run manually: + ```bash + python3 import-automation/workflow/build_manifest_catalog.py + ``` + +2. **`imports_catalog.json`**: + The compiled catalog of all import configurations, cron schedules, curator emails, and resource allocations. + +3. **`import_dags_factory.py`**: + Airflow dynamic DAG factory that reads `imports_catalog.json` and registers an independent DAG for each import specification. + All DAGs are created paused by default (`is_paused_upon_creation=True`, `catchup=False`). + +4. **`import_automation_workflow.py`**: + Core Airflow DAG definition that defines `build_dag` and executes the 4-stage pipeline: + - **Cloud Batch Job**: Runs `dc-import-executor` container. + - **Staging Ingestion**: Updates staging version via `import-helper-service-staging`, triggers ingestion via Spanner Cloud Workflow, and polls until completion. + - **Production Ingestion**: Triggers fire-and-forget production Spanner ingestion upon staging success. + - **Workflow Summary**: Aggregates execution status across stages and reports errors. + +## Continuous Deployment + +Cloud Build automatically updates the catalog and syncs DAGs to Cloud Composer: +- Config: `import-automation/cloudbuild/cloudbuild.workflow.yaml` +- Target: `gs:///dags/datacommons_airflow/` diff --git a/import-automation/workflow/build_manifest_catalog.py b/import-automation/workflow/build_manifest_catalog.py new file mode 100644 index 0000000000..ce0a775fcf --- /dev/null +++ b/import-automation/workflow/build_manifest_catalog.py @@ -0,0 +1,153 @@ +# 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. +"""Builds a consolidated JSON catalog of import specifications from manifest.json files. + +Usage: + # Scan all manifests in data repository: + python3 build_manifest_catalog.py + + # Scan a specific subfolder (e.g. scripts/us_fed): + python3 build_manifest_catalog.py --subfolder=scripts/us_fed +""" + +from __future__ import annotations + +import argparse +import glob +import json +import os +import re +from typing import Any + + +def strip_leading_scripts(path_str: str) -> str: + """Safely removes only the leading 'scripts' directory component from a path.""" + parts = os.path.normpath(path_str).split(os.sep) + if parts and parts[0] == "scripts": + parts = parts[1:] + return "/".join(parts) + + + + +def build_catalog( + data_dir: str, + subfolder: str | None = None, + output_path: str | None = None, + allowlist: list[str] | set[str] | None = None, +) -> list[dict[str, Any]]: + """Scans manifest.json files and compiles all import specifications into a list.""" + base_search = os.path.join(data_dir, subfolder) if subfolder else data_dir + search_pattern = os.path.join(base_search, "**/manifest.json") + manifest_files = sorted(glob.glob(search_pattern, recursive=True)) + + catalog: list[dict[str, Any]] = [] + seen_dag_ids: dict[str, str] = {} + allowlist_set = set(allowlist) if allowlist else None + + for manifest_path in manifest_files: + rel_dir = os.path.relpath(os.path.dirname(manifest_path), data_dir) + # Skip dated execution / output directories + if re.search(r"\b\d{4}-\d{2}-\d{2}\b", rel_dir): + continue + + try: + with open(manifest_path, "r", encoding="utf-8") as f: + data = json.load(f) + except Exception as e: + print(f"Warning: Failed to load {manifest_path}: {e}") + continue + + import_specs = data.get("import_specifications", []) + for spec in import_specs: + import_name = spec.get("import_name") + if not import_name: + continue + + dag_id = import_name + if allowlist_set and (dag_id not in allowlist_set and import_name not in allowlist_set): + continue + + if dag_id in seen_dag_ids: + prev_manifest = seen_dag_ids[dag_id] + raise ValueError( + f"Duplicate DAG ID '{dag_id}' detected in '{manifest_path}'. " + f"This DAG ID is already defined in '{prev_manifest}'. " + "Each import specification must have a globally unique import_name / DAG ID." + ) + seen_dag_ids[dag_id] = manifest_path + + subpath = strip_leading_scripts(rel_dir) + category = subpath.split("/")[0] if subpath else "general" + + entry = { + "dag_id": dag_id, + "import_name": import_name, + "full_import_name": f"{rel_dir}:{import_name}", + "script_dir": rel_dir, + "cron_schedule": spec.get("cron_schedule"), + "curator_emails": spec.get("curator_emails", ["support@datacommons.org"]), + "provenance_description": spec.get("provenance_description", ""), + "provenance_url": spec.get("provenance_url", ""), + "config_override": spec.get("config_override", {}), + "resource_limits": spec.get("resource_limits", {}), + "category": category, + } + catalog.append(entry) + + if output_path: + os.makedirs(os.path.dirname(os.path.abspath(output_path)), exist_ok=True) + with open(output_path, "w", encoding="utf-8") as out: + json.dump(catalog, out, indent=2) + print(f"Wrote {len(catalog)} import specifications to {output_path}") + + return catalog + + +def main() -> None: + parser = argparse.ArgumentParser(description="Compile manifest.json files into imports_catalog.json") + parser.add_argument( + "--data-dir", + default=os.path.abspath(os.path.join(os.path.dirname(__file__), "../..")), + help="Path to the data repository root", + ) + parser.add_argument( + "--subfolder", + default=None, + help="Optional subfolder to restrict scan (e.g. scripts/us_fed)", + ) + parser.add_argument( + "--allowlist", + nargs="*", + default=None, + help="Optional list of specific DAG IDs / import names to include", + ) + parser.add_argument( + "--output", + default=os.path.join(os.path.dirname(__file__), "imports_catalog.json"), + help="Output path for the compiled catalog JSON", + ) + args = parser.parse_args() + + catalog = build_catalog( + data_dir=args.data_dir, + subfolder=args.subfolder, + output_path=args.output, + allowlist=args.allowlist, + ) + print(f"Catalog build complete: {len(catalog)} DAG entries compiled.") + + +if __name__ == "__main__": + main() diff --git a/import-automation/workflow/e2e_dag_test.py b/import-automation/workflow/e2e_dag_test.py new file mode 100644 index 0000000000..85376efba0 --- /dev/null +++ b/import-automation/workflow/e2e_dag_test.py @@ -0,0 +1,435 @@ +#!/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. +"""End-to-End (E2E) Cloud Composer DAG Test Runner. + +Triggers an Airflow DAG run in Google Cloud Composer via the Airflow Stable REST API, +monitors task execution states, simulates human-in-the-loop approval if applicable, +verifies the final workflow_summary XCom output, and dumps task logs upon failure. + +Usage: + python3 e2e_dag_test.py \ + --project-id=datcom-import-automation-prod \ + --location=us-central1 \ + --composer-env=import-automation-airflow \ + --dag-id=USFed_ConstantMaturityRates_Test +""" + +from __future__ import annotations + +import argparse +import json +import logging +import os +import sys +import time +import uuid +from datetime import datetime, timezone +from typing import Any + +import google.auth +from google.auth.transport.requests import AuthorizedSession + +logging.basicConfig( + level=logging.INFO, + format="%(asctime)s [%(levelname)s] %(message)s", + datefmt="%Y-%m-%d %H:%M:%S", +) + +DEFAULT_PROJECT_ID = "datcom-import-automation-prod" +DEFAULT_LOCATION = "us-central1" +DEFAULT_COMPOSER_ENV = "import-automation-airflow" +DEFAULT_DAG_ID = "USFed_ConstantMaturityRates_Test" + + +class GcloudCliCredentials(google.auth.credentials.Credentials): + """Credentials backed by `gcloud auth print-access-token`.""" + + def __init__(self) -> None: + super().__init__() + self.refresh(None) + + def refresh(self, request: Any) -> None: + import subprocess + + self.token = subprocess.check_output( + ["gcloud", "auth", "print-access-token"], + text=True, + timeout=30, + ).strip() + + +def get_authorized_session(project_id: str = DEFAULT_PROJECT_ID) -> AuthorizedSession: + """Creates an AuthorizedSession with Google Cloud Platform scopes.""" + session: AuthorizedSession | None = None + try: + credentials, _ = google.auth.default( + scopes=["https://www.googleapis.com/auth/cloud-platform"] + ) + from google.auth.transport.requests import Request + credentials.refresh(Request()) + session = AuthorizedSession(credentials) + except Exception as adc_err: + logging.info("ADC refresh failed (%s); using gcloud auth print-access-token...", adc_err) + session = AuthorizedSession(GcloudCliCredentials()) + + if project_id: + session.headers.update({"x-goog-user-project": project_id}) + return session + + +def resolve_composer_webserver_url( + session: AuthorizedSession, + project_id: str, + location: str, + composer_env: str, +) -> str: + """Queries the Cloud Composer API to discover the Airflow Webserver URI.""" + env_url = ( + f"https://composer.googleapis.com/v1/projects/{project_id}" + f"/locations/{location}/environments/{composer_env}" + ) + logging.info("Resolving Composer webserver URI from %s...", env_url) + try: + resp = session.get(env_url, timeout=30) + resp.raise_for_status() + data = resp.json() + airflow_uri = data.get("config", {}).get("airflowUri", "") + if airflow_uri: + return airflow_uri.rstrip("/") + except Exception as rest_err: + logging.warning("Composer REST API lookup failed (%s); falling back to gcloud CLI...", rest_err) + import subprocess + cmd = [ + "gcloud", "composer", "environments", "describe", composer_env, + f"--location={location}", + f"--project={project_id}", + "--format=value(config.airflowUri)", + ] + airflow_uri = subprocess.check_output(cmd, text=True, timeout=30).strip() + if airflow_uri: + return airflow_uri.rstrip("/") + raise RuntimeError(f"Could not resolve Composer webserver URI via REST or CLI: {rest_err}") from rest_err + + raise RuntimeError(f"Could not find config.airflowUri in Composer environment response: {data}") + + +def check_dag_import_errors(session: AuthorizedSession, webserver_url: str) -> None: + """Verifies that there are no DAG import/syntax errors in Airflow.""" + err_url = f"{webserver_url}/api/v1/importErrors" + resp = session.get(err_url, timeout=30) + resp.raise_for_status() + data = resp.json() + import_errors = data.get("import_errors", []) + if import_errors: + logging.error("Airflow reported %d DAG import error(s):", len(import_errors)) + for err in import_errors: + logging.error( + "File: %s\nTimestamp: %s\nStack Trace:\n%s\n%s", + err.get("filename"), + err.get("timestamp"), + err.get("stack_trace"), + "-" * 60, + ) + raise RuntimeError(f"Aborting E2E test due to {len(import_errors)} DAG import error(s) in Airflow.") + logging.info("Pre-flight check passed: 0 DAG import errors in Airflow.") + + +def ensure_dag_available_and_unpaused( + session: AuthorizedSession, + webserver_url: str, + dag_id: str, + max_wait_sec: int = 120, +) -> dict[str, Any]: + """Waits for DAG to appear in Airflow and ensures it is unpaused.""" + dag_url = f"{webserver_url}/api/v1/dags/{dag_id}" + start_time = time.time() + + while time.time() - start_time < max_wait_sec: + resp = session.get(dag_url, timeout=30) + if resp.status_code == 200: + dag_info = resp.json() + if dag_info.get("is_paused"): + logging.info("DAG '%s' is currently paused. Unpausing for E2E test...", dag_id) + patch_resp = session.patch(dag_url, json={"is_paused": False}, timeout=30) + patch_resp.raise_for_status() + dag_info = patch_resp.json() + logging.info("DAG '%s' is active and ready (is_paused=%s).", dag_id, dag_info.get("is_paused")) + return dag_info + logging.info("Waiting for DAG '%s' to be registered in Airflow (status %d)...", dag_id, resp.status_code) + time.sleep(10) + + raise TimeoutError(f"DAG '{dag_id}' was not found in Airflow after {max_wait_sec}s.") + + +def trigger_dag_run( + session: AuthorizedSession, + webserver_url: str, + dag_id: str, + run_id: str, + conf: dict[str, Any], +) -> dict[str, Any]: + """Triggers a new DAG run via the Airflow REST API.""" + runs_url = f"{webserver_url}/api/v1/dags/{dag_id}/dagRuns" + payload = { + "dag_run_id": run_id, + "conf": conf, + } + logging.info("Triggering DAG '%s' with run_id='%s' and conf=%s", dag_id, run_id, json.dumps(conf)) + resp = session.post(runs_url, json=payload, timeout=30) + resp.raise_for_status() + return resp.json() + + +def dump_failed_task_logs( + session: AuthorizedSession, + webserver_url: str, + dag_id: str, + run_id: str, + task_instances: list[dict[str, Any]], +) -> None: + """Fetches and prints Airflow logs for any failed tasks.""" + for ti in task_instances: + t_id = ti.get("task_id") + state = ti.get("state") + try_num = ti.get("try_number", 1) + if state in ("failed", "upstream_failed") and t_id: + log_url = f"{webserver_url}/api/v1/dags/{dag_id}/dagRuns/{run_id}/taskInstances/{t_id}/logs/{max(1, try_num)}" + try: + resp = session.get(log_url, headers={"Accept": "text/plain"}, timeout=30) + logging.error( + "\n%s\nTask Failure Log: %s (state=%s, try=%s)\n%s\n%s\n%s", + "=" * 80, + t_id, + state, + try_num, + "-" * 80, + resp.text[-4000:] if resp.text else "", + "=" * 80, + ) + except Exception as ex: + logging.warning("Could not retrieve log for failed task '%s': %s", t_id, ex) + + +def run_e2e_test( + project_id: str, + location: str, + composer_env: str, + dag_id: str, + webserver_url: str = "", + skip_import_job: bool = False, + skip_staging_ingestion: bool = False, + run_golden_tests: bool = False, + simulate_hitl_approval: bool = False, + sync_wait_sec: int = 20, + poll_interval_sec: int = 15, + timeout_sec: int = 1800, +) -> dict[str, Any]: + """Executes the E2E DAG test and returns the final workflow summary.""" + session = get_authorized_session(project_id) + + if not webserver_url: + webserver_url = resolve_composer_webserver_url( + session=session, + project_id=project_id, + location=location, + composer_env=composer_env, + ) + logging.info("Using Airflow Webserver URL: %s", webserver_url) + + if sync_wait_sec > 0: + logging.info("Waiting %ds for Composer DAG processor to sync latest GCS files...", sync_wait_sec) + time.sleep(sync_wait_sec) + + # 1. Pre-flight check for DAG import errors + check_dag_import_errors(session, webserver_url) + + # 2. Ensure target DAG exists and is unpaused + ensure_dag_available_and_unpaused(session, webserver_url, dag_id) + + # 3. Trigger DAG run + build_tag = os.environ.get("BUILD_ID", uuid.uuid4().hex[:8]) + timestamp_str = datetime.now(timezone.utc).strftime("%Y%m%d-%H%M%S") + run_id = f"e2e-test-{build_tag}-{timestamp_str}" + + conf = { + "skipImportJob": skip_import_job, + "skipStagingIngestion": skip_staging_ingestion, + "skipProdIngestion": True, + "runGoldenTests": run_golden_tests, + "autoApproveGoldenDiff": not simulate_hitl_approval, + } + + trigger_dag_run(session, webserver_url, dag_id, run_id, conf) + + # 4. Monitor DAG run progress + run_url = f"{webserver_url}/api/v1/dags/{dag_id}/dagRuns/{run_id}" + ti_url = f"{run_url}/taskInstances" + var_url = f"{webserver_url}/api/v1/variables" + + start_time = time.time() + approval_variable_set = False + dag_state = "queued" + task_instances: list[dict[str, Any]] = [] + + try: + while time.time() - start_time < timeout_sec: + resp = session.get(run_url, timeout=30) + resp.raise_for_status() + dag_state = resp.json().get("state", "unknown") + + ti_resp = session.get(ti_url, timeout=30) + ti_resp.raise_for_status() + task_instances = ti_resp.json().get("task_instances", []) + + task_summary = ", ".join( + f"{ti.get('task_id')}:{ti.get('state') or 'none'}" + for ti in sorted(task_instances, key=lambda x: x.get("task_id", "")) + ) + elapsed = int(time.time() - start_time) + logging.info("[%04ds] DAG State: %-10s | Tasks: %s", elapsed, dag_state.upper(), task_summary) + + # Check if HumanApprovalSensor is paused waiting for approval + for ti in task_instances: + if ti.get("task_id") == "await_human_approval" and ti.get("state") == "up_for_reschedule": + if simulate_hitl_approval and not approval_variable_set: + logging.info( + "Task 'await_human_approval' is paused in reschedule mode. " + "Injecting Airflow Variable PROD_APPROVE_ALL='true' to unblock..." + ) + session.post( + var_url, + json={"key": "PROD_APPROVE_ALL", "value": "true"}, + timeout=30, + ) + approval_variable_set = True + + if dag_state in ("success", "failed"): + break + + time.sleep(poll_interval_sec) + else: + raise TimeoutError(f"DAG run '{run_id}' timed out after {timeout_sec}s (last state: {dag_state}).") + + finally: + if approval_variable_set: + logging.info("Cleaning up temporary Airflow Variable 'PROD_APPROVE_ALL'...") + try: + session.delete(f"{var_url}/PROD_APPROVE_ALL", timeout=30) + except Exception as ex: + logging.warning("Failed to delete temporary variable PROD_APPROVE_ALL: %s", ex) + + # 5. Evaluate final state & fetch XCom summary + if dag_state != "success": + logging.error("E2E DAG run '%s' FAILED with state '%s'!", run_id, dag_state) + dump_failed_task_logs(session, webserver_url, dag_id, run_id, task_instances) + raise RuntimeError(f"E2E test failed: DAG '{dag_id}' run '{run_id}' finished with state '{dag_state}'.") + + # Fetch workflow_summary XCom return_value + xcom_url = f"{run_url}/taskInstances/workflow_summary/xcomEntries/return_value" + summary_data: dict[str, Any] = {} + try: + xcom_resp = session.get(xcom_url, timeout=30) + if xcom_resp.ok: + raw_val = xcom_resp.json().get("value") + if isinstance(raw_val, str): + try: + summary_data = json.loads(raw_val) + except Exception: + summary_data = {"raw": raw_val} + elif isinstance(raw_val, dict): + summary_data = raw_val + except Exception as ex: + logging.warning("Could not fetch workflow_summary XCom: %s", ex) + + logging.info( + "\n%s\nE2E DAG TEST SUCCEEDED: %s (%s)\nWorkflow Summary XCom:\n%s\n%s", + "=" * 80, + dag_id, + run_id, + json.dumps(summary_data, indent=2), + "=" * 80, + ) + return summary_data + + +def main() -> None: + parser = argparse.ArgumentParser(description="Run E2E test for an Airflow DAG in Cloud Composer.") + parser.add_argument("--project-id", default=os.environ.get("PROJECT_ID", DEFAULT_PROJECT_ID)) + parser.add_argument("--location", default=os.environ.get("LOCATION", DEFAULT_LOCATION)) + parser.add_argument("--composer-env", default=os.environ.get("COMPOSER_ENV_NAME", DEFAULT_COMPOSER_ENV)) + parser.add_argument("--dag-id", default=os.environ.get("E2E_DAG_ID", DEFAULT_DAG_ID)) + parser.add_argument("--webserver-url", default=os.environ.get("COMPOSER_WEBSERVER_URL", "")) + parser.add_argument( + "--skip-import-job", + action="store_true", + default=os.environ.get("E2E_SKIP_IMPORT_JOB", "").lower() in ("true", "1", "yes"), + help="Skip the Cloud Batch import job step to test orchestration quickly.", + ) + parser.add_argument( + "--skip-staging-ingestion", + action="store_true", + default=os.environ.get("E2E_SKIP_STAGING_INGESTION", "").lower() in ("true", "1", "yes"), + help="Skip the Staging Spanner ingestion step.", + ) + parser.add_argument( + "--run-golden-tests", + action="store_true", + default=os.environ.get("E2E_RUN_GOLDEN_TESTS", "").lower() in ("true", "1", "yes"), + help="Force running staging golden verification Cloud Build.", + ) + parser.add_argument( + "--simulate-hitl-approval", + action="store_true", + default=False, + help="Test human-in-the-loop pause and variable-based approval.", + ) + parser.add_argument( + "--sync-wait", + type=int, + default=int(os.environ.get("E2E_SYNC_WAIT_SEC", "20")), + help="Seconds to wait for Composer DAG sync before triggering.", + ) + parser.add_argument( + "--poll-interval", + type=int, + default=15, + help="Polling interval in seconds.", + ) + parser.add_argument( + "--timeout", + type=int, + default=1800, + help="Maximum DAG run timeout in seconds.", + ) + + args = parser.parse_args() + run_e2e_test( + project_id=args.project_id, + location=args.location, + composer_env=args.composer_env, + dag_id=args.dag_id, + webserver_url=args.webserver_url, + skip_import_job=args.skip_import_job, + skip_staging_ingestion=args.skip_staging_ingestion, + run_golden_tests=args.run_golden_tests, + simulate_hitl_approval=args.simulate_hitl_approval, + sync_wait_sec=args.sync_wait, + poll_interval_sec=args.poll_interval, + timeout_sec=args.timeout, + ) + + +if __name__ == "__main__": + main() diff --git a/import-automation/workflow/golden_verification.py b/import-automation/workflow/golden_verification.py new file mode 100644 index 0000000000..77d222c748 --- /dev/null +++ b/import-automation/workflow/golden_verification.py @@ -0,0 +1,517 @@ +# 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. +"""Generic staging golden test verification gate for Data Commons imports. + +Evaluates an allowlist of sensitive imports (e.g. Schema, Place). When an import +is in the allowlist, it clears the Staging Redis cache and runs Explore/NL golden +tests via Cloud Build in diff mode to gate production promotion. +All other imports skip this gate and auto-approve immediately. +""" + +from __future__ import annotations + +from datetime import datetime, timezone +import json +import logging +import os +import re +import sys +import time +from typing import Any + +from airflow.decorators import task +from airflow.exceptions import AirflowFailException +from airflow.models import Variable +from airflow.sensors.base import BaseSensorOperator +from airflow.utils.trigger_rule import TriggerRule + +current_dir = os.path.dirname(os.path.abspath(__file__)) +if current_dir not in sys.path: + sys.path.insert(0, current_dir) + +DEFAULT_GOLDEN_TEST_IMPORTS = frozenset({ + "Schema", + "Place", +}) + +DEFAULT_IMPORT_TRIGGERS: dict[str, str] = { + "schema": "ingestion-golden-verification", + "place": "ingestion-golden-verification", +} +FALLBACK_TRIGGER_ID = "ingestion-golden-verification" +DEFAULT_BUILD_PROJECT = "datcom-ci" +DEFAULT_BRANCH = "master" +DEFAULT_DIFF_BUCKET = "datcom-ci-test" + + +def get_golden_test_imports() -> set[str]: + """Retrieves allowlist from Airflow Variable, OS env, or defaults.""" + fallback = os.environ.get("GOLDEN_TEST_IMPORTS", "") + try: + raw = Variable.get("GOLDEN_TEST_IMPORTS", default_var=fallback) + except Exception: + raw = fallback + + if raw and isinstance(raw, str): + try: + parsed = json.loads(raw) + if isinstance(parsed, list): + return {item.strip() for item in parsed if isinstance(item, str) and item.strip()} + except Exception: + pass + return {item.strip() for item in re.split(r"[,;\s]+", raw) if item.strip()} + + return set(DEFAULT_GOLDEN_TEST_IMPORTS) + + +def is_golden_test_import(import_name: str, allowlist: set[str] | None = None) -> bool: + """Checks whether the import matches the golden test allowlist.""" + if not import_name: + return False + targets = allowlist if allowlist is not None else get_golden_test_imports() + short_name = import_name.split(":")[-1] + normalized = {t.lower() for t in targets} + return import_name.lower() in normalized or short_name.lower() in normalized + + +def get_golden_trigger_id(import_name: str, custom_trigger: str = "") -> str: + """Resolves the Cloud Build trigger ID for an import.""" + if custom_trigger: + return custom_trigger + short_name = import_name.split(":")[-1].lower() + return DEFAULT_IMPORT_TRIGGERS.get(short_name, FALLBACK_TRIGGER_ID) + + +def get_diff_bucket(custom_bucket: str = "") -> str: + """Resolves the GCS bucket storing golden diff summaries.""" + if custom_bucket: + return custom_bucket + fallback = os.environ.get("GOLDEN_DIFF_BUCKET", DEFAULT_DIFF_BUCKET) + try: + return Variable.get("GOLDEN_DIFF_BUCKET", default_var=fallback) + except Exception: + return fallback + + +def _fetch_diff_summary( + build_id: str, + bucket_name: str = DEFAULT_DIFF_BUCKET, + max_retries: int = 3, +) -> dict[str, Any]: + """Fetches diff_summary.json generated by Cloud Build from GCS.""" + if not build_id: + return {"has_diff": False, "pr_url": "", "message": "No build ID provided"} + + object_path = f"golden_diffs/{build_id}/diff_summary.json" + logging.info("Fetching diff summary from gs://%s/%s...", bucket_name, object_path) + + for attempt in range(1, max_retries + 1): + # 1. Try GCSHook + try: + from airflow.providers.google.cloud.hooks.gcs import GCSHook + hook = GCSHook(gcp_conn_id="google_cloud_default") + content = hook.download(bucket_name=bucket_name, object_name=object_path) + if isinstance(content, bytes): + content = content.decode("utf-8") + parsed = json.loads(content) + logging.info("Fetched diff summary via GCSHook on attempt %d: %s", attempt, parsed) + return parsed + except Exception as ex: + logging.debug("GCSHook download attempt %d failed: %s", attempt, ex) + + # 2. Try google.cloud.storage client + try: + from google.cloud import storage + client = storage.Client() + bucket = client.bucket(bucket_name) + blob = bucket.blob(object_path) + if blob.exists(): + content = blob.download_as_text() + parsed = json.loads(content) + logging.info("Fetched diff summary via google.cloud.storage on attempt %d: %s", attempt, parsed) + return parsed + except Exception as ex: + logging.debug("google.cloud.storage attempt %d failed: %s", attempt, ex) + + # 3. Direct REST API via AuthorizedSession + try: + import google.auth + from google.auth.transport.requests import AuthorizedSession + + credentials, _ = google.auth.default( + scopes=[ + "https://www.googleapis.com/auth/devstorage.read_only", + "https://www.googleapis.com/auth/cloud-platform", + ] + ) + session = AuthorizedSession(credentials) + url = f"https://storage.googleapis.com/storage/v1/b/{bucket_name}/o/{object_path.replace('/', '%2F')}?alt=media" + resp = session.get(url, timeout=30) + if resp.ok: + parsed = resp.json() + logging.info("Fetched diff summary via REST API on attempt %d: %s", attempt, parsed) + return parsed + except Exception as ex: + logging.debug("REST API attempt %d failed: %s", attempt, ex) + + if attempt < max_retries: + time.sleep(2) + + logging.warning( + "Could not retrieve diff_summary.json from gs://%s/%s after %d attempts", + bucket_name, object_path, max_retries + ) + return { + "has_diff": False, + "pr_url": "", + "message": f"Could not retrieve diff_summary.json from gs://{bucket_name}/{object_path}", + } + + +def _run_cloud_build_verification( + project_id: str, + trigger_id: str, + branch_name: str = "master", + timeout: int = 2400, + poll_interval: int = 30, +) -> dict[str, Any]: + """Triggers and polls Cloud Build trigger for golden test verification.""" + try: + from airflow.providers.google.cloud.hooks.cloud_build import CloudBuildHook + hook = CloudBuildHook(gcp_conn_id="google_cloud_default") + try: + build = hook.run_build_trigger( + trigger_id=trigger_id, + source={"branch_name": branch_name}, + project_id=project_id, + wait=True, + ) + except TypeError: + build = hook.run_build_trigger( + request={ + "project_id": project_id, + "trigger_id": trigger_id, + "source": {"branch_name": branch_name}, + } + ) + b_status = getattr(build, "status", None) or (build.get("status") if isinstance(build, dict) else "") + b_id = getattr(build, "id", None) or (build.get("id") if isinstance(build, dict) else "") + b_log = getattr(build, "log_url", None) or getattr(build, "logUrl", None) or (build.get("logUrl") if isinstance(build, dict) else "") + status_str = getattr(b_status, "name", str(b_status)) + return {"id": str(b_id), "status": status_str, "logUrl": str(b_log)} + except Exception as ex: + logging.warning("CloudBuildHook execution error, falling back to direct API: %s", ex) + + # Fallback to direct REST API + import google.auth + from google.auth.transport.requests import AuthorizedSession + + credentials, _ = google.auth.default(scopes=["https://www.googleapis.com/auth/cloud-platform"]) + authed_session = AuthorizedSession(credentials) + + run_url = f"https://cloudbuild.googleapis.com/v1/projects/{project_id}/triggers/{trigger_id}:run" + resp = authed_session.post(run_url, json={"branchName": branch_name}, timeout=60) + resp.raise_for_status() + op = resp.json() + + metadata = op.get("metadata", {}) + build_info = metadata.get("build", {}) + build_id = build_info.get("id") or op.get("name", "").split("/")[-1] + log_url = build_info.get("logUrl", "") + + logging.info("Started Cloud Build '%s' via trigger '%s'. Waiting for completion...", build_id, trigger_id) + start_time = time.time() + build_url = f"https://cloudbuild.googleapis.com/v1/projects/{project_id}/builds/{build_id}" + while time.time() - start_time < timeout: + time.sleep(poll_interval) + poll_resp = authed_session.get(build_url, timeout=60) + if not poll_resp.ok: + continue + b_data = poll_resp.json() + b_status = b_data.get("status", "") + log_url = b_data.get("logUrl", log_url) + if b_status in ("SUCCESS", "FAILURE", "INTERNAL_ERROR", "TIMEOUT", "CANCELLED", "EXPIRED"): + return {"id": build_id, "status": b_status, "logUrl": log_url} + + raise TimeoutError(f"Cloud Build verification '{build_id}' timed out after {timeout} seconds.") + + +@task(task_id="verify_golden_tests", trigger_rule=TriggerRule.NONE_FAILED) +def verify_golden_tests(**context) -> dict[str, Any]: + """Runs staging verification gate for allowlisted imports (e.g. Schema, Place).""" + # Lazy import to avoid circular dependency + from import_automation_workflow import resolve_workflow_context, _report_import_failure + + cfg = resolve_workflow_context(context) + ti = context.get("ti") + staging_res = ti.xcom_pull(task_ids="trigger_staging_ingestion") if ti else None + + dag_run = context.get("dag_run") + conf = dag_run.conf if dag_run and dag_run.conf else {} + params = context.get("params") or {} + + def _p(key: str, default: Any = None) -> Any: + if key in conf and conf[key] is not None: + return conf[key] + return params.get(key, default) + + explicit_skip = bool(_p("skipGoldenTests", False)) + explicit_run = bool(_p("runGoldenTests", False)) + import_name = cfg["importName"] + + # 1. Eligibility Check: Allowlist or explicit run flag + if not (explicit_run or is_golden_test_import(import_name)) or explicit_skip: + logging.info("Import '%s' is not configured for golden verification; skipping.", import_name) + return { + "status": "SKIPPED", + "hasDiff": False, + "prUrl": "", + "reason": f"Import '{import_name}' not in golden test allowlist", + } + + # 2. Check staging prerequisite + if ( + cfg["skipStagingIngestion"] + or not staging_res + or staging_res.get("status") not in ("SUBMITTED", "SUCCESS") + ): + logging.info("Staging ingestion was not successful; skipping golden gate.") + return { + "status": "SKIPPED", + "hasDiff": False, + "prUrl": "", + "reason": "Staging ingestion was not SUCCESS", + } + + # 3. Resolve Trigger ID, Project, Branch, and Diff Bucket + trigger_id = get_golden_trigger_id( + import_name, + custom_trigger=_p("goldenTestTriggerId") or os.environ.get("GOLDEN_TEST_TRIGGER_ID", ""), + ) + project_id = _p("goldenTestProjectId") or os.environ.get("GOLDEN_TEST_PROJECT_ID", DEFAULT_BUILD_PROJECT) + branch_name = _p("goldenTestBranch") or os.environ.get("GOLDEN_TEST_BRANCH", DEFAULT_BRANCH) + diff_bucket = get_diff_bucket(_p("goldenDiffBucket") or os.environ.get("GOLDEN_DIFF_BUCKET", "")) + + logging.info("Triggering staging golden verification Cloud Build '%s' on %s (%s)...", trigger_id, project_id, branch_name) + start_time = time.time() + try: + build_res = _run_cloud_build_verification( + project_id=project_id, + trigger_id=trigger_id, + branch_name=branch_name, + ) + build_id = build_res.get("id", "") + log_url = build_res.get("logUrl", "") + b_status = build_res.get("status", "") + duration = int(time.time() - start_time) + + if b_status == "SUCCESS": + diff_summary = _fetch_diff_summary(build_id=build_id, bucket_name=diff_bucket) + has_diff = bool(diff_summary.get("has_diff", False)) + pr_url = str(diff_summary.get("pr_url", "")) + branch = str(diff_summary.get("branch_name", "")) + logging.info( + "Golden verification succeeded in %ss (build: %s, hasDiff: %s, prUrl: %s, log: %s).", + duration, build_id, has_diff, pr_url, log_url + ) + return { + "status": "SUCCESS", + "buildId": build_id, + "logUrl": log_url, + "executionTime": duration, + "hasDiff": has_diff, + "prUrl": pr_url, + "branchName": branch, + "diffSummary": diff_summary, + } + + # Failure: Report failure to import-helper and raise exception to block prod + exec_dt = context.get("logical_date") or context.get("execution_date") + exec_time = int(time.time() - (exec_dt.timestamp() if exec_dt else start_time)) + helper_url = cfg["helperUrlFn"](cfg["importHelperService"], "-staging") + _report_import_failure(helper_url, cfg["jobId"], cfg["importName"], cfg["gcsImportBucket"], exec_time) + + err_msg = f"Golden gate failed with status '{b_status}' (build: {build_id}, logs: {log_url}). Blocking prod promotion." + logging.error(err_msg) + raise AirflowFailException(err_msg) + + except AirflowFailException: + raise + except Exception as e: + exec_dt = context.get("logical_date") or context.get("execution_date") + exec_time = int(time.time() - (exec_dt.timestamp() if exec_dt else start_time)) + helper_url = cfg["helperUrlFn"](cfg["importHelperService"], "-staging") + _report_import_failure(helper_url, cfg["jobId"], cfg["importName"], cfg["gcsImportBucket"], exec_time) + raise AirflowFailException(f"Golden verification error: {e}") from e + + +class HumanApprovalSensor(BaseSensorOperator): + """Airflow sensor that pauses DAG execution awaiting human approval if golden diffs were detected. + + Auto-approves immediately if: + - Golden test verification was skipped, or + - Golden test verification succeeded and detected no diffs (hasDiff is False), or + - autoApproveGoldenDiff parameter is set to True. + + Waits for approval if: + - Golden test verification detected diffs (hasDiff is True). + + Approval options: + - Set Airflow Variable `PROD_APPROVE_` = 'true' + - Set Airflow Variable `PROD_APPROVE_` = 'true' + - Set Airflow Variable `PROD_APPROVE_ALL` = 'true' + - Click 'Mark Success' on this task in the Airflow UI. + + Rejection options: + - Set Airflow Variable `PROD_REJECT_` = 'true' + - Set Airflow Variable `PROD_REJECT_` = 'true' + - Click 'Mark Failed' on this task in the Airflow UI. + """ + + template_fields = ("job_id", "import_name") + + def __init__( + self, + job_id: str = "", + import_name: str = "", + mode: str = "reschedule", + poke_interval: int = 60, + timeout: int = 86400, # 24 hours + soft_fail: bool = False, + trigger_rule: str = TriggerRule.NONE_FAILED, + **kwargs, + ): + super().__init__( + mode=mode, + poke_interval=poke_interval, + timeout=timeout, + soft_fail=soft_fail, + trigger_rule=trigger_rule, + **kwargs, + ) + self.job_id = job_id + self.import_name = import_name + + def poke(self, context: dict[str, Any]) -> bool: + dag_run = context.get("dag_run") + conf = dag_run.conf if dag_run and dag_run.conf else {} + params = context.get("params") or {} + + def _p(key: str, default: Any = None) -> Any: + if key in conf and conf[key] is not None: + return conf[key] + return params.get(key, default) + + if bool(_p("autoApproveGoldenDiff", False)): + logging.info("autoApproveGoldenDiff parameter is True; auto-approving prod promotion.") + return True + + ti = context.get("ti") + golden_res = ti.xcom_pull(task_ids="verify_golden_tests") if ti else None + + # Auto-approve if golden check was skipped, not run, or had no diffs + if not golden_res or golden_res.get("status") != "SUCCESS" or not golden_res.get("hasDiff"): + logging.info( + "HumanApprovalSensor: No golden diffs detected (or gate skipped). Auto-approving prod promotion." + ) + return True + + from import_automation_workflow import resolve_workflow_context + try: + cfg = resolve_workflow_context(context) + except Exception: + cfg = {} + job_id = self.job_id or cfg.get("jobId", "") + import_name = self.import_name or cfg.get("importName", "") + short_import = import_name.split(":")[-1] if import_name else "" + + pr_url = golden_res.get("prUrl", "") + log_url = golden_res.get("logUrl", "") + build_id = golden_res.get("buildId", "") + + # Check for rejection signals first + reject_keys = [ + f"PROD_REJECT_{job_id}", + f"PROD_REJECT_{import_name}", + f"PROD_REJECT_{short_import}", + ] + for rk in reject_keys: + if not rk.strip(): + continue + try: + r_val = Variable.get(rk, default_var="") + if isinstance(r_val, str) and r_val.strip().lower() in ("true", "1", "yes"): + msg = f"Prod promotion for '{import_name}' (job: {job_id}) was REJECTED via Airflow Variable '{rk}'." + logging.error(msg) + raise AirflowFailException(msg) + except AirflowFailException: + raise + except Exception: + pass + + # Check for approval signals + approve_keys = [ + f"PROD_APPROVE_{job_id}", + f"PROD_APPROVE_{import_name}", + f"PROD_APPROVE_{short_import}", + "PROD_APPROVE_ALL", + ] + for ak in approve_keys: + if not ak.strip(): + continue + try: + a_val = Variable.get(ak, default_var="") + if isinstance(a_val, str) and a_val.strip().lower() in ("true", "1", "yes"): + logging.info( + "Prod promotion for '%s' (job: %s) APPROVED via Airflow Variable '%s'.", + import_name, job_id, ak + ) + return True + except Exception: + pass + + # Diff detected and awaiting human approval + logging.info( + "\n" + "================================================================================\n" + " PAUSED FOR HUMAN APPROVAL: STAGING GOLDEN DIFF DETECTED\n" + "================================================================================\n" + " Import: %s\n" + " Job ID: %s\n" + " Build ID: %s\n" + " PR URL: %s\n" + " Build Log: %s\n" + "--------------------------------------------------------------------------------\n" + " To APPROVE prod promotion:\n" + " 1. In Airflow UI, select task 'await_human_approval' and click 'Mark Success', OR\n" + " 2. Set Airflow Variable 'PROD_APPROVE_%s' = 'true'\n" + "\n" + " To REJECT prod promotion:\n" + " 1. In Airflow UI, select task 'await_human_approval' and click 'Mark Failed', OR\n" + " 2. Set Airflow Variable 'PROD_REJECT_%s' = 'true'\n" + "================================================================================\n", + import_name, job_id, build_id, pr_url, log_url, job_id, job_id + ) + return False + + def execute(self, context: dict[str, Any]) -> dict[str, Any]: + super().execute(context) + ti = context.get("ti") + golden_res = ti.xcom_pull(task_ids="verify_golden_tests") if ti else {} + return { + "status": "APPROVED", + "hasDiff": bool(golden_res.get("hasDiff", False)), + "prUrl": golden_res.get("prUrl", ""), + "approvedAt": datetime.now(timezone.utc).isoformat(), + } diff --git a/import-automation/workflow/import_automation_workflow.py b/import-automation/workflow/import_automation_workflow.py new file mode 100644 index 0000000000..edb2fe3b7e --- /dev/null +++ b/import-automation/workflow/import_automation_workflow.py @@ -0,0 +1,732 @@ +# 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. +"""Airflow DAG for Data Commons Import Automation. + +Orchestrates Cloud Batch data imports, updates metadata via import-helper Cloud Run, +and triggers Spanner ingestion Cloud Workflows for staging and production. +""" + +from __future__ import annotations + +import json +import logging +import os +import re +import sys +import time +from datetime import datetime, timezone +from typing import Any, Callable + +from airflow import DAG +from airflow.decorators import task +from airflow.exceptions import AirflowException, AirflowFailException, AirflowSkipException +from airflow.models import Variable +from airflow.models.param import Param +from airflow.providers.google.cloud.hooks.cloud_batch import CloudBatchHook +from airflow.providers.google.cloud.sensors.workflows import WorkflowExecutionSensor +from airflow.utils.trigger_rule import TriggerRule + +# Enable Jinja templating for project_id on WorkflowExecutionSensor +WorkflowExecutionSensor.template_fields = (*WorkflowExecutionSensor.template_fields, "project_id") + +current_dir = os.path.dirname(os.path.abspath(__file__)) +if current_dir not in sys.path: + sys.path.insert(0, current_dir) + +from golden_verification import ( + HumanApprovalSensor, + get_golden_test_imports, + is_golden_test_import, + verify_golden_tests, +) + + +# ----------------------------------------------------------------------------- +# Configuration & Helpers +# ----------------------------------------------------------------------------- + +DAG_ID = os.environ.get("IMPORT_AUTOMATION_DAG_ID", "import_automation_workflow") +DEFAULT_IMAGE_URI = "us-docker.pkg.dev/datcom-ci/gcr.io/dc-import-executor:stable" +DEFAULT_RESOURCES = {"machine": "n2-standard-8", "cpu": 8000, "memory": 32768, "disk": 100} + +PROD_DENYLIST = frozenset({ + "CDC500", + "CDC_OzoneCounty", + "CDC_PM25County", + "CensusCountyBusinessPatterns", + "CensusSAHIE", + "EIA_Electricity", + "EPA_EJSCREEN", + "EPA_GHGRP", + "FARS_CrashData", + "FBIGovCrime", + "IndiaNSS_HealthAilments", + "India_RBIStateDomesticProduct", + "NASA_VIIRSActiveFiresEvents", + "NCES_PrivateSchool", + "NCES_PublicSchool", + "NCES_SchoolDistrict", + "NOAA_GPCC_StandardardizedPrecipitationIndex", + "NOAA_GlobalForecastSystem", + "OECDRegionalDemography_Population", + "UNEnergy", + "USCensusPEP_AgeSexRaceHispanicOrigin", + "USDA_AgricultureCensus", + "USFed_ConstantMaturityRates_Test", + "USNationalPrisonerStatistics", + "WorldBankDatasets", +}) + + +def is_prod_denylisted(import_name: str) -> bool: + short_name = import_name.split(":")[-1] + return import_name in PROD_DENYLIST or short_name in PROD_DENYLIST + + +def get_config_var(key: str, default: str = "") -> str: + """Reads configuration from Airflow Variable or OS environment.""" + fallback = os.environ.get(key, default) + try: + return Variable.get(key, default_var=fallback) + except Exception: + return fallback + + +def generate_job_id(import_name: str, timestamp: int | None = None) -> str: + """Generates a valid RFC 1035 Cloud Batch job ID.""" + cleaned = re.sub(r"[^a-z0-9-]", "-", import_name.split(":")[-1][:50].lower()).strip("-") + job_id = f"{cleaned}-{timestamp or int(time.time())}" + return job_id if job_id[0].isalpha() else f"job-{job_id}"[:63] + + +def _get_oidc_token(audience: str) -> str | None: + """Fetches an OIDC identity token for invoking Cloud Run services.""" + from google.auth.transport.requests import Request + req = Request() + try: + from google.oauth2 import id_token + return id_token.fetch_id_token(req, audience) + except Exception: + try: + import google.auth + creds, _ = google.auth.default() + creds.refresh(req) + return getattr(creds, "id_token", None) or getattr(creds, "token", None) + except Exception as ex: + logging.error("OIDC token fetch failed: %s", ex) + return None + + +def _make_http_post(url: str, body: dict[str, Any], timeout: int = 60) -> dict[str, Any]: + """Posts JSON payload to an authenticated Cloud Run endpoint with retries.""" + import requests + from requests.adapters import HTTPAdapter + from urllib3.util import Retry + + headers = {"Content-Type": "application/json"} + token = _get_oidc_token(url) + if token: + headers["Authorization"] = f"Bearer {token}" + + session = requests.Session() + retries = Retry(total=3, backoff_factor=2, status_forcelist=[500, 502, 503, 504]) + session.mount("https://", HTTPAdapter(max_retries=retries)) + + resp = session.post(url, json=body, headers=headers, timeout=timeout) + resp.raise_for_status() + return resp.json() if resp.content else {} + + +def _report_import_failure(helper_url: str, job_id: str, import_name: str, bucket: str, exec_time: int) -> None: + """Notifies import-helper of a batch job failure to update Spanner metadata.""" + try: + _make_http_post(f"{helper_url}/imports/status", { + "jobId": job_id, "executionTime": exec_time, + "imports": [{"importName": import_name, "status": "FAILURE", "latestVersion": f"gs://{bucket}/{import_name.replace(':', '/')}"}], + }) + except Exception as e: + logging.error("Failed reporting import failure: %s", e) + + +# ----------------------------------------------------------------------------- +# Cloud Batch Job Spec & TaskFlow Operator +# ----------------------------------------------------------------------------- + +def _build_batch_job_spec( + image_uri: str, import_name: str, import_config: str, job_id: str, + resources: dict[str, Any], gcs_mount_bucket: str, gcs_mount_path: str = "/tmp/gcs", +) -> dict[str, Any]: + """Builds a dictionary-based Cloud Batch job descriptor.""" + sa_email = get_config_var("CLOUD_BATCH_SERVICE_ACCOUNT") + task_spec = { + "runnables": [{ + "container": {"imageUri": image_uri, "commands": [f"--import_name={import_name}", f"--import_config={import_config}"]}, + "environment": {"variables": {"IMPORT_NAME": import_name, "BATCH_JOB_NAME": job_id}}, + }], + "computeResource": {"cpuMilli": int(resources.get("cpu", 8000)), "memoryMib": int(resources.get("memory", 32768))}, + **({"volumes": [{"gcs": {"remotePath": gcs_mount_bucket}, "mountPath": gcs_mount_path}]} if gcs_mount_bucket else {}), + } + policy = { + "instances": [{ + "policy": { + "machineType": str(resources.get("machine", "n2-standard-8")), + "provisioningModel": "STANDARD", + "bootDisk": {"image": "projects/debian-cloud/global/images/family/debian-12", "sizeGb": int(resources.get("disk", 100))}, + }, + "installOpsAgent": True, + }], + **({"serviceAccount": {"email": sa_email}} if sa_email else {}), + } + return { + "taskGroups": [{"taskSpec": task_spec, "taskCount": 1, "parallelism": 1}], + "allocationPolicy": policy, + "logsPolicy": {"destination": "CLOUD_LOGGING"}, + } + + +@task(task_id="run_import_job") +def run_import_job(import_name: str = "", **context) -> dict[str, Any]: + """Submits and monitors Google Cloud Batch import jobs using CloudBatchHook.""" + cfg = resolve_workflow_context(context, default_import_name=import_name) + if cfg["skipImportJob"]: + logging.info("skipImportJob is True; skipping Cloud Batch import job.") + return {"status": "SKIPPED", "message": "Import job skipped by configuration"} + + hook = CloudBatchHook(gcp_conn_id="google_cloud_default") + job_spec = _build_batch_job_spec( + cfg["imageUri"], cfg["importName"], cfg["batchImportConfig"], + cfg["jobId"], cfg["resources"], cfg["gcsMountBucket"], cfg["gcsMountPath"], + ) + + start_time = time.time() + logging.info("Submitting Cloud Batch job '%s' (%s, %s)...", cfg["jobId"], cfg["projectId"], cfg["region"]) + try: + submitted_job = hook.submit_batch_job( + job_name=cfg["jobId"], + job=job_spec, + region=cfg["region"], + project_id=cfg["projectId"], + ) + job_resource_name = ( + getattr(submitted_job, "name", None) + or f"projects/{cfg['projectId']}/locations/{cfg['region']}/jobs/{cfg['jobId']}" + ) + job = hook.wait_for_job( + job_name=job_resource_name, + timeout=604800, + ) + exec_time = int(time.time() - start_time) + logging.info("Cloud Batch job '%s' succeeded in %ss.", cfg["jobId"], exec_time) + return { + "status": "SUCCESS", + "jobId": cfg["jobId"], + "name": job_resource_name, + "executionTime": exec_time, + "result": str(job), + } + except Exception as e: + exec_time = int(time.time() - start_time) + helper_url = cfg["helperUrlFn"](cfg["importHelperService"], "-staging") + _report_import_failure(helper_url, cfg["jobId"], cfg["importName"], cfg["gcsImportBucket"], exec_time) + raise AirflowException(f"Cloud Batch import job failed: {e}") from e + + +def _run_cloud_run_job( + project_id: str, + region: str, + job_name: str, + args: list[str], + timeout_seconds: int = 7200, + poll_interval: int = 15, +) -> dict[str, Any]: + """Triggers a Cloud Run v2 Job execution and polls the LRO until completion.""" + import google.auth + from google.auth.transport.requests import AuthorizedSession + + creds, _ = google.auth.default(scopes=["https://www.googleapis.com/auth/cloud-platform"]) + session = AuthorizedSession(creds) + + run_url = f"https://run.googleapis.com/v2/projects/{project_id}/locations/{region}/jobs/{job_name}:run" + payload = { + "overrides": { + "containerOverrides": [ + { + "args": args, + } + ] + } + } + logging.info("Triggering Cloud Run job '%s' via %s with args: %s", job_name, run_url, args) + resp = session.post(run_url, json=payload, timeout=60) + resp.raise_for_status() + operation = resp.json() + op_name = operation.get("name") + if not op_name: + raise RuntimeError(f"Cloud Run jobs.run did not return an operation name: {operation}") + + logging.info("Waiting for Cloud Run job operation '%s' to complete (timeout: %ss)...", op_name, timeout_seconds) + op_url = f"https://run.googleapis.com/v2/{op_name}" + start_time = time.time() + + while True: + if operation.get("done"): + if "error" in operation: + err = operation["error"] + raise RuntimeError(f"Cloud Run job '{job_name}' failed with error: {err}") + response_data = operation.get("response", {}) + failed_count = response_data.get("failedCount", 0) + cancelled_count = response_data.get("cancelledCount", 0) + if failed_count > 0 or cancelled_count > 0: + raise RuntimeError( + f"Cloud Run job '{job_name}' execution failed " + f"(failedCount={failed_count}, cancelledCount={cancelled_count}): {response_data}" + ) + logging.info("Cloud Run job '%s' completed successfully.", job_name) + return response_data or operation + + if time.time() - start_time > timeout_seconds: + raise TimeoutError(f"Cloud Run job '{job_name}' operation '{op_name}' timed out after {timeout_seconds}s.") + + time.sleep(poll_interval) + poll_resp = session.get(op_url, timeout=60) + poll_resp.raise_for_status() + operation = poll_resp.json() + + +@task(task_id="run_validation_job") +def run_validation_job(import_name: str = "", **context) -> dict[str, Any]: + """Executes the Cloud Run validation job after the Batch import job.""" + cfg = resolve_workflow_context(context, default_import_name=import_name) + if cfg["skipImportJob"] or cfg.get("skipValidationJob"): + logging.info( + "Skipping Cloud Run validation job (skipImportJob=%s, skipValidationJob=%s).", + cfg["skipImportJob"], + cfg.get("skipValidationJob"), + ) + return {"status": "SKIPPED", "message": "Validation job skipped by configuration"} + + start_time = time.time() + try: + res = _run_cloud_run_job( + project_id=cfg["projectId"], + region=cfg["region"], + job_name=cfg["validationJobName"], + args=[ + f"--import_name={cfg['importName']}", + f"--import_config={cfg['importConfig']}", + ], + timeout_seconds=7200, + ) + exec_time = int(time.time() - start_time) + return { + "status": "SUCCESS", + "jobName": cfg["validationJobName"], + "executionTime": exec_time, + "result": str(res), + } + except Exception as e: + exec_time = int(time.time() - start_time) + helper_url = cfg["helperUrlFn"](cfg["importHelperService"], "-staging") + _report_import_failure(helper_url, cfg["jobId"], cfg["importName"], cfg["gcsImportBucket"], exec_time) + raise AirflowException(f"Cloud Run validation job failed: {e}") from e + + +# ----------------------------------------------------------------------------- +# Cloud Workflows & Ingestion Helpers +# ----------------------------------------------------------------------------- + +def trigger_environment_ingestion(cfg: dict[str, Any], env_suffix: str) -> dict[str, Any]: + """Updates version via import-helper and triggers Spanner ingestion via ingestion-helper.""" + import_helper_url = cfg["helperUrlFn"](cfg["importHelperService"], env_suffix) + ingestion_helper_url = cfg["helperUrlFn"](cfg["ingestionHelperService"], env_suffix) + + version_res = _make_http_post(f"{import_helper_url}/imports/version", { + "imports": [cfg["importName"]], "version": "STAGING", "override": False, "comment": f"import-workflow:{cfg['runId']}", + }) + + imports_res = version_res.get("imports", []) + if not imports_res or imports_res[0].get("status") not in ("STAGING", "SKIP"): + msg = version_res.get("message", "Status not STAGING or SKIP") + logging.info("Import status is not STAGING or SKIP (%s); skipping Spanner ingestion.", msg) + return {"status": "SKIPPED", "message": f"Skipped: {msg}"} + + target = imports_res[0] + import_entry = { + "importName": target.get("importName", cfg["importName"]).split(":")[-1], + "latestVersion": target.get("latestVersion", ""), + } + + ingest_res = _make_http_post(f"{ingestion_helper_url}/imports/ingest", { + "importList": [import_entry], + "dryRun": cfg.get("dryRunIngestion", False), + "forceIngestion": cfg.get("forceIngestion", False), + }) + + if ingest_res.get("status") == "SKIPPED": + logging.info("Ingestion helper skipped ingestion: %s", ingest_res.get("message")) + return {"status": "SKIPPED", "message": ingest_res.get("message", "Skipped")} + + if ingest_res.get("status") != "SUBMITTED": + raise RuntimeError(f"Ingestion helper returned unexpected status: {ingest_res.get('status')}") + + exec_name = ingest_res.get("executionName", "") + if not exec_name: + raise RuntimeError("Ingestion helper returned status 'SUBMITTED' but executionName was empty.") + + logging.info("Triggered ingestion workflow execution: %s", exec_name) + + parts = exec_name.split("/") + # Cloud Workflows execution format: projects/{project}/locations/{location}/workflows/{workflow}/executions/{execution_id} + project_id = parts[1] if len(parts) >= 2 else cfg["projectId"] + location = parts[3] if len(parts) >= 4 else cfg["region"] + workflow_id = parts[5] if len(parts) >= 6 else cfg["spannerWorkflowName"] + execution_id = parts[7] if len(parts) >= 8 else "" + + return { + "status": "SUBMITTED", + "executionName": exec_name, + "projectId": project_id, + "location": location, + "workflowId": workflow_id, + "executionId": execution_id, + } + + +# ----------------------------------------------------------------------------- +# Workflow Context Resolver +# ----------------------------------------------------------------------------- + +def resolve_workflow_context(context: dict[str, Any], default_import_name: str = "") -> dict[str, Any]: + """Normalizes runtime parameters, environment variables, and fallback defaults.""" + dag_run, dag = context.get("dag_run"), context.get("dag") + conf = dag_run.conf if dag_run and dag_run.conf else {} + params = context.get("params") or {} + dag_params = getattr(dag, "params", {}) if dag else {} + + def get_val(key: str, default: Any = None) -> Any: + for src in (conf, params, dag_params): + if hasattr(src, "get"): + v = src.get(key) + if v is not None and v != "": + return getattr(v, "default", v) + return default + + def cfg_val(key: str, default: str, *env_keys: str) -> str: + v = get_val(key) + if v: + return str(v) + for ek in env_keys: + ev = get_config_var(ek) + if ev: + return ev + return default + + def to_bool(key: str, default: bool = False) -> bool: + v = get_val(key, default) + return v is True or str(v).lower() == "true" + + import_name = get_val("importName", default_import_name) or getattr(dag, "dag_id", "") + if not import_name: + raise ValueError("Parameter 'importName' is required to execute the import workflow.") + + project_id = cfg_val("projectId", "datcom-ci", "GCP_PROJECT_ID", "PROJECT_ID") + region = cfg_val("region", "us-central1", "CLOUD_BATCH_REGION", "LOCATION") + default_proj_num = "965988403328" if project_id == "datcom-import-automation-prod" else "879489846695" + project_number = cfg_val("projectNumber", default_proj_num, "PROJECT_NUMBER") + gcs_mount_bucket = cfg_val("gcsMountBucket", "datcom-ci-test", "GCS_MOUNT_BUCKET") + gcs_import_bucket = cfg_val("gcsImportBucket", "datcom-ci-test", "GCS_BUCKET_ID") + import_helper = cfg_val("importHelperService", "import-helper-service", "IMPORT_HELPER_SERVICE") + ingestion_helper = cfg_val("ingestionHelperService", "ingestion-helper-service", "INGESTION_HELPER_SERVICE") + spanner_workflow = cfg_val("spannerWorkflowName", "spanner-ingestion-workflow", "SPANNER_INGESTION_WORKFLOW_NAME") + + env_suffix = cfg_val("envSuffix", "", "ENV_SUFFIX") + validation_job_name = cfg_val("validationJobName", f"import-validator-job{env_suffix}", "VALIDATION_JOB_NAME") + + import_config = get_val("importConfig") + if not import_config or import_config == "{}": + import_config_dict: dict[str, Any] = { + "gcp_project_id": project_id, + "gcs_project_id": project_id, + "storage_prod_bucket_name": gcs_import_bucket, + "gcs_bucket_volume_mount": gcs_mount_bucket, + } + elif isinstance(import_config, str): + try: + import_config_dict = json.loads(import_config) + except Exception: + import_config_dict = {} + elif isinstance(import_config, dict): + import_config_dict = dict(import_config) + else: + import_config_dict = {} + + batch_config_dict = dict(import_config_dict) + batch_config_dict["invoke_import_validation"] = False + batch_config_dict["invoke_differ_tool"] = False + + import_config_str = json.dumps(import_config_dict) + batch_import_config_str = json.dumps(batch_config_dict) + + exec_dt = context.get("logical_date") or context.get("execution_date") + run_ts = int(exec_dt.timestamp()) if exec_dt else int(time.time()) + run_id = dag_run.run_id if dag_run else f"manual__{datetime.now(timezone.utc).isoformat()}" + + def helper_url(service_name: str, suffix: str = "") -> str: + svc = f"{service_name}{suffix}" + return f"https://{svc}-{project_number}.{region}.run.app" if project_number else f"https://{svc}.{region}.run.app" + + return { + "projectId": project_id, "region": region, "projectNumber": project_number, + "importName": import_name, "jobId": get_val("jobId") or generate_job_id(import_name, run_ts), + "imageUri": get_val("imageUri", DEFAULT_IMAGE_URI), + "importConfig": import_config_str, + "batchImportConfig": batch_import_config_str, + "envSuffix": env_suffix, + "validationJobName": validation_job_name, + "skipValidationJob": to_bool("skipValidationJob", False), + "gcsMountBucket": gcs_mount_bucket, "gcsImportBucket": gcs_import_bucket, + "gcsMountPath": get_val("gcsMountPath", "/tmp/gcs"), + "helperUrlFn": helper_url, + "importHelperService": import_helper, "ingestionHelperService": ingestion_helper, + "spannerWorkflowName": spanner_workflow, + "skipImportJob": to_bool("skipImportJob"), + "skipStagingIngestion": to_bool("skipStagingIngestion"), + "skipProdIngestion": is_prod_denylisted(import_name) or to_bool("skipProdIngestion", True), + "dryRunIngestion": to_bool("dryRunIngestion"), + "forceIngestion": to_bool("forceIngestion"), + "resources": {**DEFAULT_RESOURCES, **(get_val("resources") if isinstance(get_val("resources"), dict) else {})}, + "runId": run_id, + } + + +# ----------------------------------------------------------------------------- +# Airflow Tasks +# ----------------------------------------------------------------------------- + +@task(task_id="trigger_staging_ingestion") +def trigger_staging_ingestion(**context) -> dict[str, Any]: + """Updates staging version and triggers Spanner ingestion via ingestion-helper.""" + cfg = resolve_workflow_context(context) + if cfg["skipStagingIngestion"]: + raise AirflowSkipException("Staging ingestion skipped by configuration (skipStagingIngestion=True).") + + res = trigger_environment_ingestion(cfg, env_suffix="-staging") + if res.get("status") == "SKIPPED": + raise AirflowSkipException(f"Staging ingestion skipped: {res.get('message', 'Skipped')}") + + return res + + +@task(task_id="ingest_prod", trigger_rule=TriggerRule.NONE_FAILED) +def ingest_prod(**context) -> dict[str, Any]: + """Updates prod version and triggers fire-and-forget prod Spanner ingestion.""" + cfg = resolve_workflow_context(context) + ti = context.get("ti") + staging_res = ti.xcom_pull(task_ids="trigger_staging_ingestion") if ti else None + + # Check if an upstream golden test verification or pre-prod gate reported failure + for gate_task_id in ("verify_golden_tests", "verify_schema_golden_gate", "pre_prod_gate"): + gate_res = ti.xcom_pull(task_ids=gate_task_id) if ti else None + if isinstance(gate_res, dict) and gate_res.get("status") in ("FAILURE", "FAILED"): + logging.error("Pre-prod gate '%s' failed: %s. Blocking production promotion.", gate_task_id, gate_res) + return {"status": "BLOCKED", "message": f"Blocked by pre-prod gate {gate_task_id}: {gate_res.get('message', 'FAILURE')}"} + + if cfg["skipProdIngestion"]: + logging.info("skipProdIngestion is True; skipping production ingestion.") + return {"status": "SKIPPED", "message": "Production ingestion skipped"} + + if not cfg["skipStagingIngestion"] and not staging_res: + logging.info("Staging ingestion was not triggered or was skipped; skipping production ingestion.") + return {"status": "SKIPPED", "message": "Staging was not executed; skipping production ingestion."} + + return trigger_environment_ingestion(cfg, env_suffix="") + + +@task(task_id="workflow_summary", trigger_rule=TriggerRule.ALL_DONE) +def workflow_summary(**context) -> dict[str, Any]: + """Aggregates workflow outputs and fails the DAG run if any upstream task failed.""" + cfg, ti = resolve_workflow_context(context), context.get("ti") + default_res = {"status": "SKIPPED"} + dag = context.get("dag") + dag_task_ids = set(dag.task_ids) if dag and hasattr(dag, "task_ids") else None + has_golden = ( + ("verify_golden_tests" in dag_task_ids) + if dag_task_ids is not None + else is_golden_test_import(cfg["importName"]) + ) + stages = [ + ("import", "run_import_job"), + ("validation", "run_validation_job"), + ("staging_trigger", "trigger_staging_ingestion"), + ("staging_wait", "wait_staging_ingestion"), + *( + [ + ("golden_check", "verify_golden_tests"), + ("human_approval", "await_human_approval"), + ] + if has_golden + else [] + ), + ("prod", "ingest_prod"), + ] + results = { + name: default_res if val is None else (val if isinstance(val, dict) else {"status": "SUCCESS", "result": val}) + for name, tid in stages + for val in [(ti.xcom_pull(task_ids=tid) if ti else default_res)] + } + for gate_task_id in ("verify_schema_golden_gate", "pre_prod_gate"): + gate_val = ti.xcom_pull(task_ids=gate_task_id) if ti else None + if gate_val is not None: + results[gate_task_id] = gate_val if isinstance(gate_val, dict) else {"status": "SUCCESS", "result": gate_val} + + summary = {"jobId": cfg["jobId"], "importName": cfg["importName"], **results} + logging.info("Workflow summary: %s", json.dumps(summary, indent=2)) + + dag_run = context.get("dag_run") + failed = [ + t.task_id + for t in (dag_run.get_task_instances() if dag_run else []) + if t.task_id != "workflow_summary" and t.state in ("failed", "upstream_failed") + ] + failed += [ + f"{k} ({v.get('status')})" + for k, v in results.items() + if isinstance(v, dict) and v.get("status") in ("FAILURE", "FAILED") + ] + if failed: + error_msg = f"Workflow failed in upstream stage(s): {', '.join(dict.fromkeys(failed))}" + logging.error(error_msg) + raise AirflowFailException(error_msg) + + return summary + + +# ----------------------------------------------------------------------------- +# DAG Factory +# ----------------------------------------------------------------------------- + +default_args = { + "owner": "data-commons", + "depends_on_past": False, + "retries": 0, + "email_on_failure": False, + "email_on_retry": False, +} + +base_dag_params = { + "importName": Param(default="", type="string", description="Full import name"), + "imageUri": Param(default=DEFAULT_IMAGE_URI, type="string", description="Executor container image"), + "importConfig": Param(default="{}", type=["string", "object"], description="Import configuration JSON"), + "skipImportJob": Param(default=False, type="boolean", description="Skip Batch import job"), + "skipValidationJob": Param(default=False, type="boolean", description="Skip Cloud Run validation job"), + "validationJobName": Param(default="", type="string", description="Cloud Run validation job name override"), + "skipStagingIngestion": Param(default=False, type="boolean", description="Skip staging ingestion"), + "skipProdIngestion": Param(default=True, type="boolean", description="Skip prod ingestion"), + "dryRunIngestion": Param(default=False, type="boolean", description="Dry run ingestion"), + "forceIngestion": Param(default=False, type="boolean", description="Force ingestion even if already SUCCESS"), + "resources": Param(default=DEFAULT_RESOURCES, type="object", description="Batch job resources"), +} + +golden_dag_params = { + "runGoldenTests": Param(default=False, type="boolean", description="Force staging golden verification gate"), + "skipGoldenTests": Param(default=False, type="boolean", description="Skip staging golden verification gate"), + "goldenTestTriggerId": Param(default="", type="string", description="Cloud Build trigger ID override"), + "goldenTestBranch": Param(default="master", type="string", description="Cloud Build branch override"), + "goldenTestProjectId": Param(default="datcom-ci", type="string", description="Cloud Build project ID override"), + "goldenDiffBucket": Param(default="datcom-ci-test", type="string", description="GCS bucket containing golden diff summaries"), + "autoApproveGoldenDiff": Param(default=False, type="boolean", description="Auto-approve staging golden diffs without waiting for human input"), +} + +dag_params = { + **base_dag_params, + **golden_dag_params, +} + + +def build_dag( + dag_id: str = DAG_ID, + schedule: str | None = None, + import_name: str = "", + curator_emails: list[str] | None = None, + config_override: dict[str, Any] | None = None, + resource_limits: dict[str, Any] | None = None, + extra_tags: list[str] | None = None, + is_paused_upon_creation: bool = True, + pre_prod_gate_factory: Callable[[], Any] | None = None, + golden_test_allowlist: set[str] | None = None, +) -> DAG: + """Builds and returns an Airflow DAG instance for import automation.""" + allowlist = golden_test_allowlist if golden_test_allowlist is not None else get_golden_test_imports() + has_golden_check = ( + is_golden_test_import(import_name, allowlist=allowlist) + or is_golden_test_import(dag_id, allowlist=allowlist) + ) + + params = { + **base_dag_params, + **(golden_dag_params if has_golden_check else {}), + **({"importName": Param(import_name, type="string", description="Full import name")} if import_name else {}), + **({"importConfig": Param(config_override, type=["string", "object"], description="Import configuration JSON")} if config_override else {}), + **({"resources": Param({**DEFAULT_RESOURCES, **resource_limits}, type="object", description="Batch job resources")} if resource_limits else {}), + } + + dag_instance = DAG( + dag_id=dag_id, + default_args={**default_args, **({"email": curator_emails} if curator_emails else {})}, + description=f"Orchestrates import automation for {import_name or dag_id}", + schedule=schedule, + start_date=datetime(2025, 1, 1, tzinfo=timezone.utc), + catchup=False, + max_active_runs=10, + params=params, + tags=["data-commons", "import-automation"] + (extra_tags or []), + is_paused_upon_creation=is_paused_upon_creation, + ) + + with dag_instance: + staging_wait = WorkflowExecutionSensor( + task_id="wait_staging_ingestion", + project_id="{{ (ti.xcom_pull(task_ids='trigger_staging_ingestion') or {}).get('projectId', '') }}", + location="{{ (ti.xcom_pull(task_ids='trigger_staging_ingestion') or {}).get('location', '') }}", + workflow_id="{{ (ti.xcom_pull(task_ids='trigger_staging_ingestion') or {}).get('workflowId', '') }}", + execution_id="{{ (ti.xcom_pull(task_ids='trigger_staging_ingestion') or {}).get('executionId', '') }}", + mode="reschedule", + poke_interval=60, + timeout=21600, + ) + batch_task = run_import_job(import_name=import_name) + validation_task = run_validation_job(import_name=import_name) + staging_task = trigger_staging_ingestion() + prod_task = ingest_prod() + summary_task = workflow_summary() + + pipeline: list[Any] = [batch_task, validation_task, staging_task, staging_wait] + if has_golden_check: + golden_task = verify_golden_tests() + approval_task = HumanApprovalSensor( + task_id="await_human_approval", + import_name=import_name, + ) + pipeline.extend([golden_task, approval_task]) + + if pre_prod_gate_factory: + custom_gate = pre_prod_gate_factory() + pipeline.append(custom_gate) + + pipeline.extend([prod_task, summary_task]) + + for upstream, downstream in zip(pipeline, pipeline[1:]): + upstream >> downstream + + return dag_instance diff --git a/import-automation/workflow/import_dags_factory.py b/import-automation/workflow/import_dags_factory.py new file mode 100644 index 0000000000..462226cda6 --- /dev/null +++ b/import-automation/workflow/import_dags_factory.py @@ -0,0 +1,101 @@ +# 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. +"""Dynamic DAG Factory for Data Commons Imports. + +Reads imports_catalog.json (compiled from manifest.json files in the data repository) +and dynamically registers an independent Airflow DAG for each import specification. +""" + +from __future__ import annotations + +import json +import logging +import os +import sys +from typing import Any + +current_dir = os.path.dirname(os.path.abspath(__file__)) +if current_dir not in sys.path: + sys.path.insert(0, current_dir) + +import importlib +import golden_verification as _golden_verification +import import_automation_workflow as _import_automation_workflow + +importlib.reload(_golden_verification) +importlib.reload(_import_automation_workflow) + +from golden_verification import get_golden_test_imports +from import_automation_workflow import build_dag + + +def find_catalog_file() -> str | None: + """Locates the imports_catalog.json file across known search locations.""" + candidates = [ + os.environ.get("IMPORTS_CATALOG_PATH"), + os.path.join(current_dir, "imports_catalog.json"), + os.path.join(current_dir, "datacommons_airflow", "imports_catalog.json"), + "/home/airflow/gcs/dags/datacommons_airflow/imports_catalog.json", + "/home/airflow/gcs/dags/imports_catalog.json", + ] + return next((p for p in candidates if p and os.path.isfile(p)), None) + + +def load_catalog_and_register_dags(target_globals: dict[str, Any]) -> int: + """Loads imports_catalog.json and registers DAG instances into target_globals.""" + catalog_path = find_catalog_file() + if not catalog_path: + logging.info("imports_catalog.json not found; dynamic DAG factory skipped.") + return 0 + + try: + with open(catalog_path, "r", encoding="utf-8") as f: + catalog = json.load(f) + except Exception as e: + logging.error("Failed to load %s: %s", catalog_path, e) + return 0 + + golden_allowlist = get_golden_test_imports() + seen_dag_ids: set[str] = set() + count = 0 + for entry in catalog: + dag_id, full_name = entry.get("dag_id"), entry.get("full_import_name") + if not dag_id or not full_name: + continue + if dag_id in seen_dag_ids: + raise ValueError( + f"Duplicate DAG ID '{dag_id}' encountered in {catalog_path}; " + "cannot overwrite an already registered DAG." + ) + seen_dag_ids.add(dag_id) + cron = entry.get("cron_schedule") + target_globals[dag_id] = build_dag( + dag_id=dag_id, + schedule=cron, + import_name=full_name, + curator_emails=entry.get("curator_emails"), + config_override=entry.get("config_override"), + resource_limits=entry.get("resource_limits"), + extra_tags=["scheduled" if cron else "manual", entry.get("category", "data-commons")], + is_paused_upon_creation=True, + golden_test_allowlist=golden_allowlist, + ) + count += 1 + + logging.info("Registered %d import DAGs from %s", count, catalog_path) + return count + + +# Automatically register DAGs in global namespace when imported by Airflow +load_catalog_and_register_dags(globals()) diff --git a/import-automation/workflow/manifest.json b/import-automation/workflow/manifest.json new file mode 100644 index 0000000000..f47fc2b83f --- /dev/null +++ b/import-automation/workflow/manifest.json @@ -0,0 +1,12 @@ +{ + "import_specifications": [ + { + "import_name": "import_automation_workflow", + "curator_emails": [ + "support@datacommons.org" + ], + "provenance_url": "https://datacommons.org", + "provenance_description": "Generic Import Automation Workflow for ad-hoc or unregistered imports" + } + ] +}