[dagster-airlift] Ingest metadata from airflow logs

Author: dpeng817

python_modules/dagster/dagster/_core/definitions/metadata/external_metadata.py

@@ -1,5 +1,5 @@
 from collections.abc import Mapping, Sequence
-from typing import TYPE_CHECKING, Any, Literal, TypedDict, Union
+from typing import TYPE_CHECKING, Any, Literal, TypedDict, Union, get_args
 
 import dagster._check as check
 from dagster._core.definitions.asset_key import AssetKey
@@ -48,6 +48,8 @@ class ExternalMetadataValue(TypedDict):
     "table_column_lineage",
     "timestamp",
 ]
+EXTERNAL_METADATA_VALUE_KEYS = frozenset(ExternalMetadataValue.__annotations__.keys())
+EXTERNAL_METADATA_TYPES = frozenset(get_args(ExternalMetadataType))
 
 
 def metadata_map_from_external(

python_modules/libraries/dagster-airlift/dagster_airlift/core/monitoring_job/event_stream.py

@@ -1,8 +1,10 @@
+import asyncio
 from collections.abc import Iterator
 from itertools import chain
 from typing import Optional
 
 from dagster import AssetMaterialization
+from dagster._core.definitions.metadata.metadata_value import MetadataValue
 from dagster._core.events import (
     AssetMaterializationPlannedData,
     DagsterEvent,
@@ -20,6 +22,7 @@
 from dagster_airlift.core.airflow_defs_data import AirflowDefinitionsData
 from dagster_airlift.core.airflow_instance import AirflowInstance
 from dagster_airlift.core.monitoring_job.utils import (
+    extract_metadata_from_logs,
     get_dagster_run_for_airflow_repr,
     structured_log,
 )
@@ -93,6 +96,7 @@ def persist_state(
 @record
 class TaskInstanceCompleted:
     task_instance: TaskInstance
+    metadata: dict[str, MetadataValue]
 
     @property
     def timestamp(self) -> float:
@@ -108,7 +112,9 @@ def persist_state(
         for asset in airflow_data.mapped_asset_keys_by_task_handle[self.task_instance.task_handle]:
             # IMPROVEME: Add metadata to the materialization event.
             _report_materialization(
-                context, corresponding_run, AssetMaterialization(asset_key=asset)
+                context,
+                corresponding_run,
+                AssetMaterialization(asset_key=asset, metadata=self.metadata),
             )
 
 
@@ -225,6 +231,37 @@ def _process_completed_runs(
             break
 
 
+async def _retrieve_logs_for_task_instance(
+    context: OpExecutionContext,
+    airflow_instance: AirflowInstance,
+    task_instance: TaskInstance,
+) -> TaskInstanceCompleted:
+    logs = airflow_instance.get_task_instance_logs(
+        task_instance.dag_id,
+        task_instance.task_id,
+        task_instance.run_id,
+        task_instance.try_number,
+    )
+    metadata = extract_metadata_from_logs(context, logs)
+
+    return TaskInstanceCompleted(task_instance=task_instance, metadata=metadata)
+
+
+async def _async_process_task_instances(
+    context: OpExecutionContext,
+    airflow_instance: AirflowInstance,
+    task_instances: list[TaskInstance],
+) -> Iterator[TaskInstanceCompleted]:
+    results = await asyncio.gather(
+        *(
+            _retrieve_logs_for_task_instance(context, airflow_instance, task_instance)
+            for task_instance in task_instances
+        )
+    )
+
+    return results
+
+
 def _process_task_instances(
     context: OpExecutionContext,
     airflow_data: AirflowDefinitionsData,
@@ -246,9 +283,7 @@ def _process_task_instances(
         context,
         f"Found {len(task_instances)} completed task instances in the time range {datetime_from_timestamp(range_start)} to {datetime_from_timestamp(range_end)}",
     )
-    yield from (
-        TaskInstanceCompleted(task_instance=task_instance) for task_instance in task_instances
-    )
+    yield from asyncio.run(_async_process_task_instances(context, airflow_instance, task_instances))
 
 
 def persist_events(

python_modules/libraries/dagster-airlift/dagster_airlift/core/monitoring_job/utils.py

@@ -1,5 +1,16 @@
-from typing import Optional, Union
+import json
+from collections.abc import Iterable
+from typing import Optional, TypeVar, Union, cast
 
+import dagster._check as check
+from dagster._core.definitions.metadata.external_metadata import (
+    EXTERNAL_METADATA_TYPE_INFER,
+    EXTERNAL_METADATA_TYPES,
+    EXTERNAL_METADATA_VALUE_KEYS,
+    ExternalMetadataValue,
+    metadata_map_from_external,
+)
+from dagster._core.definitions.metadata.metadata_value import MetadataValue
 from dagster._core.execution.context.op_execution_context import OpExecutionContext
 from dagster._core.storage.dagster_run import DagsterRun, DagsterRunStatus, RunsFilter
 from dagster_airlift.constants import DAG_ID_TAG_KEY, DAG_RUN_ID_TAG_KEY
@@ -61,3 +72,51 @@ def get_dagster_run_for_airflow_repr(
         ),
         None,
     )
+
+
+_T = TypeVar("_T")
+
+
+def _assert_param_value(value: _T, expected_values: Iterable[_T]) -> _T:
+    if value not in expected_values:
+        raise Exception(
+            f"Invalid value when translating metadata from logs. Expected one of"
+            f" `{expected_values}`, got `{value}`."
+        )
+    return value
+
+
+def extract_metadata_from_logs(context: OpExecutionContext, logs: str) -> dict[str, MetadataValue]:
+    metadata = {}
+    import re
+
+    matches = re.findall(r"DAGSTER_START(.*?)DAGSTER_END", logs, re.DOTALL)
+    for match in matches:
+        raw_external_metadata_map = json.loads(match)
+        check.mapping_param(raw_external_metadata_map, "raw_external_metadata_map")
+        new_external_metadata_map = {}
+        for key, value in raw_external_metadata_map.items():
+            if not isinstance(key, str):
+                raise Exception(
+                    f"Invalid type when translating metadata from logs. Expected a dict with string"
+                    f" keys, got a key `{key}` of type `{type(key)}`."
+                )
+            elif isinstance(value, dict):
+                if not {*value.keys()} == EXTERNAL_METADATA_VALUE_KEYS:
+                    raise Exception(
+                        f"Invalid type when translating metadata from logs. Expected a dict with"
+                        " string keys and values that are either raw metadata values or dictionaries"
+                        f" with schema `{{raw_value: ..., type: ...}}`. Got a value `{value}`."
+                    )
+                _assert_param_value(value["type"], EXTERNAL_METADATA_TYPES)
+                new_external_metadata_map[key] = cast("ExternalMetadataValue", value)
+            else:
+                new_external_metadata_map[key] = {
+                    "raw_value": value,
+                    "type": EXTERNAL_METADATA_TYPE_INFER,
+                }
+
+        metadata_map = metadata_map_from_external(new_external_metadata_map)
+        metadata.update(metadata_map)
+
+    return metadata

python_modules/libraries/dagster-airlift/dagster_airlift/core/runtime_representations.py

@@ -131,3 +131,7 @@ def dag_handle(self) -> "DagHandle":
         from dagster_airlift.core.serialization.serialized_data import DagHandle
 
         return DagHandle(dag_id=self.dag_id)
+
+    @property
+    def try_number(self) -> int:
+        return self.metadata["try_number"]

python_modules/libraries/dagster-airlift/dagster_airlift/test/airflow_test_instance.py

@@ -41,6 +41,7 @@ def __init__(
         variables: list[dict[str, Any]] = [],
         instance_name: Optional[str] = None,
         max_runs_per_batch: Optional[int] = None,
+        logs: Optional[Mapping[str, Mapping[str, str]]] = None,
     ) -> None:
         self._dag_infos_by_dag_id = {dag_info.dag_id: dag_info for dag_info in dag_infos}
         self._task_infos_by_dag_and_task_id = {
@@ -49,6 +50,11 @@ def __init__(
         self._task_instances_by_dag_and_task_id: dict[tuple[str, str], list[TaskInstance]] = (
             defaultdict(list)
         )
+        self._logs_by_run_id_and_task_id: dict[tuple[str, str], str] = defaultdict(lambda: "")
+        for run_id, task_log_map in (logs or {}).items():
+            for task_id, log in task_log_map.items():
+                self._logs_by_run_id_and_task_id[(run_id, task_id)] = log
+
         for task_instance in task_instances:
             self._task_instances_by_dag_and_task_id[
                 (task_instance.dag_id, task_instance.task_id)
@@ -241,6 +247,11 @@ def get_all_datasets(
             return_datasets.append(dataset)
         return return_datasets
 
+    def get_task_instance_logs(
+        self, dag_id: str, task_id: str, run_id: str, try_number: int
+    ) -> str:
+        return self._logs_by_run_id_and_task_id[(run_id, task_id)]
+
 
 def make_dag_info(
     instance_name: str, dag_id: str, file_token: Optional[str], dag_props: Mapping[str, Any]
@@ -281,6 +292,7 @@ def make_task_instance(
             "start_date": start_date.isoformat(),
             "end_date": end_date.isoformat(),
             "logical_date": logical_date.isoformat() if logical_date else start_date.isoformat(),
+            "try_number": 1,
         },
     )
 
@@ -352,6 +364,7 @@ def make_instance(
     max_runs_per_batch: Optional[int] = None,
     dag_props: dict[str, Any] = {},
     task_instances: Optional[list[TaskInstance]] = None,
+    logs: Optional[Mapping[str, Mapping[str, str]]] = None,
 ) -> AirflowInstanceFake:
     """Constructs DagInfo, TaskInfo, and TaskInstance objects from provided data.
 
@@ -415,4 +428,5 @@ def make_instance(
         instance_name=instance_name,
         max_runs_per_batch=max_runs_per_batch,
         datasets=datasets,
+        logs=logs,
     )

python_modules/libraries/dagster-airlift/dagster_airlift_tests/unit_tests/conftest.py

@@ -1,5 +1,5 @@
 from collections import defaultdict
-from collections.abc import Generator, Sequence
+from collections.abc import Generator, Mapping, Sequence
 from datetime import datetime, timedelta
 from typing import Any, Optional, Union
 
@@ -73,6 +73,7 @@ def create_defs_and_instance(
     dag_level_asset_overrides: Optional[dict[str, list[str]]] = None,
     seeded_runs: Optional[list[DagRun]] = None,
     seeded_task_instances: Optional[list[TaskInstance]] = None,
+    seeded_logs: Optional[Mapping[str, Mapping[str, str]]] = None,
 ) -> tuple[Definitions, AirflowInstance]:
     assets = []
     dag_and_task_structure = defaultdict(list)
@@ -132,6 +133,7 @@ def _asset():
         dag_runs=runs,
         dataset_construction_info=dataset_construction_info or [],
         task_instances=seeded_task_instances,
+        logs=seeded_logs,
     )
     defs = Definitions.merge(
         additional_defs,

python_modules/libraries/dagster-airlift/dagster_airlift_tests/unit_tests/core_tests/test_monitoring_job.py

@@ -1,3 +1,4 @@
+import json
 from datetime import datetime, timedelta, timezone
 
 from dagster import AssetKey, DagsterInstance
@@ -32,6 +33,10 @@ def test_monitoring_job_execution(init_load_context: None, instance: DagsterInst
     """Test that monitoring job correctly represents state in Dagster."""
     freeze_datetime = datetime(2021, 1, 1, tzinfo=timezone.utc)
 
+    raw_metadata = {
+        "foo": "bar",
+        "my_timestamp": {"raw_value": 111, "type": "timestamp"},
+    }
     with freeze_time(freeze_datetime):
         defs, af_instance = create_defs_and_instance(
             assets_per_task={
@@ -77,6 +82,9 @@ def test_monitoring_job_execution(init_load_context: None, instance: DagsterInst
                     end_date=freeze_datetime,
                 ),
             ],
+            seeded_logs={
+                "run-dag": {"task": f"DAGSTER_START{json.dumps(raw_metadata)}DAGSTER_END"}
+            },
         )
         defs = build_job_based_airflow_defs(
             airflow_instance=af_instance,

python_modules/libraries/dagster-airlift/kitchen-sink/kitchen_sink/airflow_dags/dataset_dags.py

@@ -19,8 +19,20 @@ def print_fn() -> None:
     import json
 
     print("Hello")  # noqa: T201
-    data = json.dumps({"foo": "bar"})
+    data = json.dumps(
+        {
+            "foo": "bar",
+            "my_timestamp": {"raw_value": 111, "type": "timestamp"},
+        }
+    )
     print(f"DAGSTER_START{data}DAGSTER_END")  # noqa: T201
+    another_data = json.dumps(
+        {
+            "foo": "baz",
+            "my_other_timestamp": {"raw_value": 113, "type": "timestamp"},
+        }
+    )
+    print(f"DAGSTER_START{another_data}DAGSTER_END")  # noqa: T201
 
 
 # Inter-dag structure as follows:

python_modules/libraries/dagster-airlift/kitchen-sink/kitchen_sink_tests/integration_tests/test_e2e_jobs.py

@@ -1,4 +1,5 @@
 from dagster._core.definitions.asset_key import AssetKey
+from dagster._core.definitions.metadata.metadata_value import TimestampMetadataValue
 from dagster._core.event_api import EventRecordsFilter
 from dagster._core.events import DagsterEventType
 from dagster._core.instance_for_test import instance_for_test
@@ -78,3 +79,18 @@ def test_job_based_defs(
         assert {materialized_records.run_id for materialized_records in materialized_records} == {
             producer_run.run_id
         }
+        for key in ["example1", "example2"]:
+            key_record = next(
+                materialized_records
+                for materialized_records in materialized_records
+                if materialized_records.asset_key == AssetKey(key)
+            )
+            assert key_record.asset_materialization.metadata
+            assert key_record.asset_materialization.metadata[
+                "my_timestamp"
+            ] == TimestampMetadataValue(value=111.0)
+            assert key_record.asset_materialization.metadata[
+                "my_other_timestamp"
+            ] == TimestampMetadataValue(value=113.0)
+            # It gets overridden by the second print
+            assert key_record.asset_materialization.metadata["foo"].value == "baz"