[DCP Preprocessor] Move preprocessor to workflow - part 1

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."
       )
 
 

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"])