diff --git a/simple/stats/runner.py b/simple/stats/runner.py index 38cbe125..9d46c122 100644 --- a/simple/stats/runner.py +++ b/simple/stats/runner.py @@ -57,7 +57,6 @@ from stats.reporter import ImportReporter import stats.schema_constants as sc from stats.svg_cache import generate_svg_cache -from stats.trigger_ingestion_workflow import trigger_ingestion_workflow from stats.validation import MetadataValidator from stats.variable_per_row_importer import VariablePerRowImporter from util.file_match import match @@ -236,15 +235,44 @@ def run(self): store.close() logging.info("File storage closed.") - # Auto-trigger workflow now that all data is guaranteed to be exported and written to GCS - if self.trigger_workflow_info: - trigger_ingestion_workflow(import_list=self.trigger_workflow_info) + # Hand off processed import metadata to workflow or GCS + self._handle_workflow_handoff() except Exception as e: logging.exception("Error updating stats") self.reporter.report_failure(error=str(e)) raise + def _handle_workflow_handoff(self) -> None: + """Writes processed import metadata to a GCS handshake file for the ingestion workflow.""" + if not self.trigger_workflow_info: + return + + workflow_id = os.getenv("WORKFLOW_EXECUTION_ID") + temp_location = os.getenv("TEMP_LOCATION", "") + + if not workflow_id: + logging.warning("WORKFLOW_EXECUTION_ID not set. Skipping GCS handshake.") + return + + if not temp_location.startswith("gs://"): + logging.warning( + "TEMP_LOCATION (%s) is not a GCS path. Skipping GCS handshake.", + temp_location) + return + + handshake_payload = {"importList": json.dumps(self.trigger_workflow_info)} + + output_json_path = f"{temp_location.rstrip('/')}/datacommons/ingestion_records/{workflow_id}.json" + logging.info( + "WORKFLOW_EXECUTION_ID set. Writing preprocessor result to GCS handshake path: %s", + output_json_path) + with create_store(output_json_path, + create_if_missing=True, + treat_as_file=True) as store: + store.as_file().write(json.dumps(handshake_payload, indent=2)) + logging.info("Successfully wrote preprocessor result to GCS.") + def _read_config_from_file(self, config_file_path: str, config_file_dir: Optional[Dir] = None) -> Config: @@ -838,10 +866,10 @@ def _run_imports_and_export_jsonld(self): # Perform strict metadata validation before committing and closing MetadataValidator(self.config, self.db).validate() - # Auto-trigger workflow if output is on GCS + # Populate trigger workflow info if running under ingestion workflow and output is GCS output_path = self.db.jsonld_dir.full_path() import_name = self.db.import_name - if os.getenv("INGESTION_WORKFLOW_NAME") and output_path.startswith("gs://"): + if os.getenv("WORKFLOW_EXECUTION_ID") and output_path.startswith("gs://"): processed_imports = list(self.db._processed_imports) if not processed_imports: processed_imports = [import_name] @@ -852,7 +880,7 @@ def _run_imports_and_export_jsonld(self): self.trigger_workflow_info = import_list else: logging.info( - "Output is local or workflow is missing, skipping auto-trigger of ingestion workflow. Please upload files to GCS and trigger manually." + "Not running under ingestion workflow or output is local. Skipping handshake info." ) diff --git a/simple/stats/trigger_ingestion_workflow.py b/simple/stats/trigger_ingestion_workflow.py deleted file mode 100644 index 394c36a0..00000000 --- a/simple/stats/trigger_ingestion_workflow.py +++ /dev/null @@ -1,102 +0,0 @@ -# Copyright 2026 Google Inc. -# -# 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. -""" -Trigger the Data Commons ingestion workflow. - -Used by the Stats import runner to trigger the ingestion workflow after the Stats import is complete. -""" - -import json -import logging -import os -import re - -import google.auth -import google.auth.transport.requests -import requests - -_MAX_IMPORT_NAME_LENGTH = 64 - - -def _get_env_vars(): - """Helper to get and validate required environment variables.""" - required_env_vars = [ - "GCP_SPANNER_INSTANCE_ID", "GCP_SPANNER_DATABASE_NAME", "PROJECT_ID", - "WORKFLOW_LOCATION", "TEMP_LOCATION", "REGION", "INGESTION_WORKFLOW_NAME" - ] - missing_vars = [var for var in required_env_vars if not os.getenv(var)] - if missing_vars: - logging.error( - f"Cannot trigger ingestion workflow. Missing environment variables: {', '.join(missing_vars)}" - ) - logging.warning("Skipping auto-trigger of ingestion workflow.") - return None - return {var: os.getenv(var) for var in required_env_vars} - - -def trigger_ingestion_workflow(import_list: list[dict]): - """Triggers the Data Commons ingestion workflow via Google Cloud Workflows API.""" - if not import_list: - logging.warning( - "No import list provided. Skipping ingestion workflow trigger.") - return - - logging.info("Attempting to auto-trigger ingestion workflow via API...") - - env_vars = _get_env_vars() - if not env_vars: - return - - raw_import_name = "_".join(item["importName"] for item in import_list) - sanitized_import_name = re.sub(r'[^a-zA-Z0-9_-]', '_', raw_import_name) - sanitized_import_name = sanitized_import_name[:_MAX_IMPORT_NAME_LENGTH] - - data_payload = { - "spannerInstanceId": env_vars["GCP_SPANNER_INSTANCE_ID"], - "spannerDatabaseId": env_vars["GCP_SPANNER_DATABASE_NAME"], - "importName": sanitized_import_name, - "importList": json.dumps(import_list), - "tempLocation": env_vars["TEMP_LOCATION"], - "region": env_vars["REGION"] - } - - try: - credentials, _ = google.auth.default() - auth_request = google.auth.transport.requests.Request() - credentials.refresh(auth_request) - - project_id = env_vars["PROJECT_ID"] - location = env_vars["WORKFLOW_LOCATION"] - workflow_name = env_vars["INGESTION_WORKFLOW_NAME"] - url = f"https://workflowexecutions.googleapis.com/v1/projects/{project_id}/locations/{location}/workflows/{workflow_name}/executions" - - headers = { - "Authorization": f"Bearer {credentials.token}", - "Content-Type": "application/json" - } - - payload = {"argument": json.dumps(data_payload)} - - response = requests.post(url, json=payload, headers=headers, timeout=60) - - if response.status_code == 200: - logging.info("Workflow triggered successfully!") - logging.info(response.json()) - else: - logging.error( - f"Failed to trigger workflow. Status: {response.status_code}") - logging.error(response.text) - - except Exception as e: - logging.error(f"Error triggering workflow via API: {e}") diff --git a/simple/tests/stats/runner_test.py b/simple/tests/stats/runner_test.py index 9935b26d..293dd9eb 100644 --- a/simple/tests/stats/runner_test.py +++ b/simple/tests/stats/runner_test.py @@ -523,8 +523,21 @@ def mock_full_path(self, sub_path=""): return f"gs://my-bucket/jsonld/run_timestamp" return original_full_path(self, sub_path) - with mock.patch.dict(os.environ, - {"INGESTION_WORKFLOW_NAME": "my-workflow"}): + mock_store = mock.MagicMock() + from stats.runner import create_store as real_create_store + + def side_effect(path, *args, **kwargs): + if "ingestion_records" in path: + return mock_store + return real_create_store(path, *args, **kwargs) + + with (mock.patch("stats.runner.create_store", side_effect=side_effect), + mock.patch.dict( + os.environ, { + "INGESTION_WORKFLOW_NAME": "my-workflow", + "WORKFLOW_EXECUTION_ID": "my-execution", + "TEMP_LOCATION": "gs://my-bucket/temp" + })): with mock.patch.object(_StoreWrapper, "full_path", mock_full_path): runner = Runner(config_file_path=None, input_dir_path=input_dir, @@ -558,10 +571,16 @@ def mock_full_path(self, sub_path=""): b.startswith("jsonld/run_timestamp/ilo/node-") for b in called_blobs)) - # Assert that trigger_workflow_info holds paths for both oecd and ilo - self.assertEqual(len(runner.trigger_workflow_info), 2) - trigger_names = sorted( - [t["importName"] for t in runner.trigger_workflow_info]) + # Verify GCS Handshake payload was written correctly + entered_mock = mock_store.__enter__.return_value + entered_mock.as_file.return_value.write.assert_called_once() + written_str = entered_mock.as_file.return_value.write.call_args[0][0] + written_payload = json.loads(written_str) + self.assertIn("importList", written_payload) + + decoded_import_list = json.loads(written_payload["importList"]) + self.assertEqual(len(decoded_import_list), 2) + trigger_names = sorted([t["importName"] for t in decoded_import_list]) self.assertEqual(trigger_names, ["ilo", "oecd"])