feat(ingestion): move to dcp-tools 1.1.0 for Data Commons Platform 1.1.6

.env.example

@@ -15,7 +15,6 @@ DUCKDB_PATH=path/to/repo/data/dfo.duckdb
 # DCP_PROJECT_ID=
 # DCP_REGION=
 # DCP_BUCKET_NAME=
-# DCP_PREP_JOB_NAME=
 # DCP_WORKFLOW_NAME=
 
 # The GitHub repository the DFO agent skills are fetched from. Only the skills assets read it.

README.md

@@ -43,6 +43,12 @@ Automation runs through one sensor, `automation_condition_sensor`. On a new loca
 
 On a fresh environment, materialize the dbt seeds, and any other asset with no parents, once by hand. `eager()` never builds an asset with no parents, so nothing downstream of them automates until they exist.
 
+### Ingestion
+
+Each `*_dc_ingestion` asset syncs its export's bundle to the ingestion bucket, starts the Data Commons Platform ingestion workflow, and watches the execution for two minutes. An execution that fails or is cancelled in that time fails the asset. One still running is left to finish, and the asset records its resource name as `execution_name` metadata, the name to look it up by in Cloud Workflows.
+
+The identity that starts the workflow also needs `roles/workflows.viewer`. The asset reads the workflow's definition before it uploads the bundle and stops if the workflow does not come from a Data Commons Platform 1.1.6 stack, and the asset reads the execution's state while it watches.
+
 ## Running dbt
 
 The dbt models can also be built without Dagster. Run this from the repository root:

pyproject.toml

@@ -22,7 +22,8 @@ dependencies = [
     "universal-pathlib>=0.3.10",
     "pandas>=3.0.5",
     "imf-reader>=2.2.0",
-    "dcp-tools>=1.0.1",
+    "dcp-tools>=1.1.0",
+    "pydantic-settings>=2.9.1,<3.0.0",
     "dfo-ontology",
     "oda-data==2.8.1",
     "pyyaml>=6.0",

src/dfo_pipelines/data_commons/export.py

@@ -53,10 +53,9 @@
 # The provenance property `ExportSpec.conversion_recipe` is written onto.
 CONVERSION_RECIPE_PROPERTY = "conversionRecipe"
 
-# The provenance model column every pipeline emits, and the property it is written
-# onto.
+# The provenance model column every pipeline emits, written onto the provenance as
+# its `importedSourceReleaseDate`.
 SOURCE_RELEASE_DATE_COLUMN = "source_release_date"
-SOURCE_RELEASE_DATE_PROPERTY = "importedSourceReleaseDate"
 
 
 class ConversionRecipe(StrEnum):
@@ -157,7 +156,7 @@ class ExportSpec:
     def __post_init__(self) -> None:
         reserved = DOCUMENTATION_FILES.keys() | {
             CONVERSION_RECIPE_PROPERTY,
-            SOURCE_RELEASE_DATE_PROPERTY,
+            "importedSourceReleaseDate",
             "lastDataRefreshDate",
         }
         clashes = sorted(set(self.provenance_column_to_property.values()) & reserved)
@@ -280,17 +279,16 @@ def export_bundle(
 
     log.info("Registering Provenance nodes")
     (provenance,) = get_model_rows(connection, spec.provenance_model)
-    column_to_property = spec.provenance_column_to_property | {
-        SOURCE_RELEASE_DATE_COLUMN: SOURCE_RELEASE_DATE_PROPERTY
-    }
-    if missing := column_to_property.keys() - provenance.keys():
+    required = spec.provenance_column_to_property.keys() | {SOURCE_RELEASE_DATE_COLUMN}
+    if missing := required - provenance.keys():
         raise ValueError(
             f"{spec.provenance_model} lacks the mapped columns {sorted(missing)}"
         )
     mapped_properties = {
         property_name: _provenance_property_value(column, provenance.pop(column))
-        for column, property_name in column_to_property.items()
+        for column, property_name in spec.provenance_column_to_property.items()
     }
+    source_release_date = provenance.pop(SOURCE_RELEASE_DATE_COLUMN)
     recipe = (
         {}
         if spec.conversion_recipe is None
@@ -306,6 +304,7 @@ def export_bundle(
     manager.add_provenance(
         **provenance,
         dcid=spec.provenance_dcid,
+        imported_source_release_date=source_release_date,
         additional_properties=recipe | mapped_properties | documentation,
     )
 

src/dfo_pipelines/data_commons/ingestion.py

@@ -5,39 +5,129 @@
 workflow that reads it. An ingestion rebuilds every node and edge carrying the
 import's provenance and leaves the rest of the graph alone, which is why the
 folder holds that import's current bundle and nothing else.
+
+The workflow runs for as long as the ingestion takes, which is longer than an
+asset should wait. Its execution is watched for its first minutes only, which is
+when a bad request or a failed launch shows.
 """
 
 import logging
+import time
+from collections.abc import Callable
+from dataclasses import dataclass
 from pathlib import Path
 
 from dcp_tools.gcp_utilities.clients import get_gcs_client
-from dcp_tools.gcp_utilities.ingestion import run_ingestion_workflow
+from dcp_tools.gcp_utilities.ingestion import (
+    IngestionExecutionState,
+    check_ingestion_workflow_contract,
+    get_ingestion_execution_state,
+    run_ingestion_workflow,
+)
 from dcp_tools.gcp_utilities.settings import KGSettings
 from dcp_tools.gcp_utilities.storage import sync_directory_to_gcs
 from upath import UPath
 
+# Workflow execution states that end an execution without ingesting it.
+FAILED_STATES = frozenset({"FAILED", "CANCELLED"})
+
+
+class IngestionExecutionError(RuntimeError):
+    """The ingestion workflow's execution failed or was cancelled."""
+
+
+@dataclass(frozen=True, kw_only=True)
+class Ingestion:
+    """An ingestion that has started.
+
+    Attributes:
+        destination: URI of the folder the bundle was synced to.
+        execution_name: Resource name of the workflow execution ingesting it.
+    """
+
+    destination: str
+    execution_name: str
+
+
+def watch_execution_start(
+    settings: KGSettings,
+    execution_name: str,
+    *,
+    log: logging.Logger,
+    watch_seconds: float,
+    poll_seconds: float,
+    sleep: Callable[[float], None] = time.sleep,
+    clock: Callable[[], float] = time.monotonic,
+) -> IngestionExecutionState:
+    """Poll an ingestion's execution until it succeeds or the watch ends.
+
+    Args:
+        settings: Knowledge graph the execution ingests into.
+        execution_name: Resource name of the execution.
+        log: Logger to use, e.g. an asset's `context.log`.
+        watch_seconds: How long to poll for before leaving the execution running.
+        poll_seconds: How long to wait between polls.
+        sleep: Waits the given number of seconds.
+        clock: Monotonic time in seconds.
+
+    Returns:
+        IngestionExecutionState: The last state read, `SUCCEEDED` or still running.
+
+    Raises:
+        IngestionExecutionError: If the execution fails or is cancelled while
+            watched.
+    """
+    deadline = clock() + watch_seconds
+    while True:
+        execution = get_ingestion_execution_state(settings, execution_name)
+        if execution.state in FAILED_STATES:
+            raise IngestionExecutionError(
+                f"Ingestion execution {execution_name} ended {execution.state}: "
+                f"{execution.error}"
+            )
+        if execution.state == "SUCCEEDED" or clock() >= deadline:
+            log.info(
+                f"Ingestion execution {execution_name} is {execution.state}"
+                f" at steps {list(execution.current_steps)}"
+            )
+            return execution
+        sleep(poll_seconds)
+
 
 def ingest_bundle(
-    bundle: UPath, *, import_name: str, settings: KGSettings, log: logging.Logger
-) -> str:
+    bundle: UPath,
+    *,
+    import_name: str,
+    settings: KGSettings,
+    log: logging.Logger,
+    watch_seconds: float = 120,
+    poll_seconds: float = 5,
+) -> Ingestion:
     """Sync a bundle into the import's folder in the ingestion bucket and ingest it.
 
     The sync removes anything in the folder that the bundle does not contain, so
     the graph is built from what the export wrote rather than from that plus
-    whatever an earlier bundle left behind. The workflow is only started: it runs
-    on for as long as the ingestion takes, after this returns.
+    whatever an earlier bundle left behind. The workflow is checked before the
+    sync, so a workflow dcp-tools cannot start leaves the folder untouched. Its
+    execution is watched for `watch_seconds` and runs on after this returns.
 
     Args:
         bundle: Local directory the export wrote the bundle to.
         import_name: Name of the import, which names its folder in the bucket.
         settings: Knowledge graph the bundle is ingested into.
         log: Logger to use, e.g. an asset's `context.log`.
+        watch_seconds: How long to watch the execution for after starting it.
+        poll_seconds: How long to wait between reads of its state.
 
     Returns:
-        str: URI of the folder the bundle was synced to.
+        Ingestion: Where the bundle was synced to and the execution ingesting it.
 
     Raises:
         ValueError: If the bundle is not on local disk, or holds no files to upload.
+        WorkflowContractError: If the workflow is not one dcp-tools can start,
+            raised before anything is synced.
+        IngestionExecutionError: If the execution fails or is cancelled while
+            watched.
     """
     # dcp-tools uploads from local disk.
     if bundle.protocol not in ("", "file"):
@@ -52,12 +142,25 @@ def ingest_bundle(
     ):
         raise ValueError(f"Found no bundle to ingest at {bundle}")
 
+    check_ingestion_workflow_contract(settings)
+
     folder = f"{settings.gcs_input_folder_path}/{import_name}"
-    bucket = get_gcs_client(settings.gcp_credentials).bucket(settings.gcs_bucket_name)
+    client = get_gcs_client(settings.gcp_credentials, project=settings.gcp_project_id)
+    bucket = client.bucket(settings.gcs_bucket_name)
     sync_directory_to_gcs(bucket, directory, gcs_folder_name=folder)
     destination = f"gs://{settings.gcs_bucket_name}/{folder}"
     log.info(f"Synced {bundle} to {destination}")
 
     log.info(f"Starting the ingestion workflow for {import_name}")
-    run_ingestion_workflow(settings, imports=import_name)
-    return destination
+    execution_name = run_ingestion_workflow(
+        settings, imports=import_name, skip_workflow_contract_check=True
+    )
+    log.info(f"Started ingestion execution {execution_name}")
+    watch_execution_start(
+        settings,
+        execution_name,
+        log=log,
+        watch_seconds=watch_seconds,
+        poll_seconds=poll_seconds,
+    )
+    return Ingestion(destination=destination, execution_name=execution_name)

src/dfo_pipelines/definitions.py

@@ -26,7 +26,6 @@ def defs():
                     project_id=dg.EnvVar("DCP_PROJECT_ID"),
                     region=dg.EnvVar("DCP_REGION"),
                     bucket_name=dg.EnvVar("DCP_BUCKET_NAME"),
-                    prep_job_name=dg.EnvVar("DCP_PREP_JOB_NAME"),
                     workflow_name=dg.EnvVar("DCP_WORKFLOW_NAME"),
                 ),
                 "skills_repository": SkillsRepositoryResource(

src/dfo_pipelines/defs/dfo/ontology/assets.py

@@ -205,15 +205,16 @@ def dfo_ontology_dc_ingestion(
     Materialized by hand only: an ingestion rewrites the shared graph and runs
     for as long as it takes after the workflow is started.
     """
-    destination = ingest_bundle(
+    ingestion = ingest_bundle(
         UPath(dfo_ontology_dc_export),
         import_name=IMPORT_NAME,
         settings=knowledge_graph.settings,
         log=context.log,
     )
     return dg.MaterializeResult(
         metadata={
-            "dagster/uri": dg.MetadataValue.text(destination),
+            "dagster/uri": dg.MetadataValue.text(ingestion.destination),
             "import_name": dg.MetadataValue.text(IMPORT_NAME),
+            "execution_name": dg.MetadataValue.text(ingestion.execution_name),
         },
     )

src/dfo_pipelines/defs/dfo/shared/assets.py

@@ -87,15 +87,16 @@ def dfo_shared_dc_ingestion(
     for as long as it takes after the workflow is started. It follows the
     ontology's ingestion, which defines the source the provenance names.
     """
-    destination = ingest_bundle(
+    ingestion = ingest_bundle(
         UPath(dfo_shared_dc_export),
         import_name=IMPORT_NAME,
         settings=knowledge_graph.settings,
         log=context.log,
     )
     return dg.MaterializeResult(
         metadata={
-            "dagster/uri": dg.MetadataValue.text(destination),
+            "dagster/uri": dg.MetadataValue.text(ingestion.destination),
             "import_name": dg.MetadataValue.text(IMPORT_NAME),
+            "execution_name": dg.MetadataValue.text(ingestion.execution_name),
         },
     )

src/dfo_pipelines/defs/dfo/skills/assets.py

@@ -693,15 +693,16 @@ def dfo_skills_dc_ingestion(
     ontology's ingestion, which defines the source, the skill classes and
     properties, and the nodes skills link to.
     """
-    destination = ingest_bundle(
+    ingestion = ingest_bundle(
         UPath(dfo_skills_dc_export),
         import_name=IMPORT_NAME,
         settings=knowledge_graph.settings,
         log=context.log,
     )
     return dg.MaterializeResult(
         metadata={
-            "dagster/uri": dg.MetadataValue.text(destination),
+            "dagster/uri": dg.MetadataValue.text(ingestion.destination),
             "import_name": dg.MetadataValue.text(IMPORT_NAME),
+            "execution_name": dg.MetadataValue.text(ingestion.execution_name),
         },
     )

src/dfo_pipelines/defs/imf/weo/assets.py

@@ -273,15 +273,16 @@ def imf_weo_dc_ingestion(
     Materialized by hand only: an ingestion rewrites the shared graph and runs
     for as long as it takes after the workflow is started.
     """
-    destination = ingest_bundle(
+    ingestion = ingest_bundle(
         UPath(imf_weo_dc_export),
         import_name=IMF_WEO_EXPORT_SPEC.import_name,
         settings=knowledge_graph.settings,
         log=context.log,
     )
     return dg.MaterializeResult(
         metadata={
-            "dagster/uri": dg.MetadataValue.text(destination),
+            "dagster/uri": dg.MetadataValue.text(ingestion.destination),
             "import_name": dg.MetadataValue.text(IMF_WEO_EXPORT_SPEC.import_name),
+            "execution_name": dg.MetadataValue.text(ingestion.execution_name),
         },
     )

src/dfo_pipelines/defs/oecd/dac1/assets.py

@@ -298,15 +298,16 @@ def oecd_dac1_dc_ingestion(
     Materialized by hand only: an ingestion rewrites the shared graph and runs
     for as long as it takes after the workflow is started.
     """
-    destination = ingest_bundle(
+    ingestion = ingest_bundle(
         UPath(oecd_dac1_dc_export),
         import_name=OECD_DAC1_EXPORT_SPEC.import_name,
         settings=knowledge_graph.settings,
         log=context.log,
     )
     return dg.MaterializeResult(
         metadata={
-            "dagster/uri": dg.MetadataValue.text(destination),
+            "dagster/uri": dg.MetadataValue.text(ingestion.destination),
             "import_name": dg.MetadataValue.text(OECD_DAC1_EXPORT_SPEC.import_name),
+            "execution_name": dg.MetadataValue.text(ingestion.execution_name),
         },
     )

src/dfo_pipelines/defs/oecd/dac2a/assets.py

@@ -287,15 +287,16 @@ def oecd_dac2a_dc_ingestion(
     Materialized by hand only: an ingestion rewrites the shared graph and runs
     for as long as it takes after the workflow is started.
     """
-    destination = ingest_bundle(
+    ingestion = ingest_bundle(
         UPath(oecd_dac2a_dc_export),
         import_name=OECD_DAC2A_EXPORT_SPEC.import_name,
         settings=knowledge_graph.settings,
         log=context.log,
     )
     return dg.MaterializeResult(
         metadata={
-            "dagster/uri": dg.MetadataValue.text(destination),
+            "dagster/uri": dg.MetadataValue.text(ingestion.destination),
             "import_name": dg.MetadataValue.text(OECD_DAC2A_EXPORT_SPEC.import_name),
+            "execution_name": dg.MetadataValue.text(ingestion.execution_name),
         },
     )

src/dfo_pipelines/defs/oecd/dac2b/assets.py

@@ -281,15 +281,16 @@ def oecd_dac2b_dc_ingestion(
     Materialized by hand only: an ingestion rewrites the shared graph and runs
     for as long as it takes after the workflow is started.
     """
-    destination = ingest_bundle(
+    ingestion = ingest_bundle(
         UPath(oecd_dac2b_dc_export),
         import_name=OECD_DAC2B_EXPORT_SPEC.import_name,
         settings=knowledge_graph.settings,
         log=context.log,
     )
     return dg.MaterializeResult(
         metadata={
-            "dagster/uri": dg.MetadataValue.text(destination),
+            "dagster/uri": dg.MetadataValue.text(ingestion.destination),
             "import_name": dg.MetadataValue.text(OECD_DAC2B_EXPORT_SPEC.import_name),
+            "execution_name": dg.MetadataValue.text(ingestion.execution_name),
         },
     )