[dagster-airlift] make AirflowDefinitionsData aware of jobs

Author: dpeng817

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

@@ -3,8 +3,16 @@
 from functools import cached_property
 from typing import AbstractSet, Union  # noqa: UP035
 
-from dagster import AssetKey, AssetsDefinition, AssetSpec
+from dagster import (
+    AssetKey,
+    AssetsDefinition,
+    AssetSpec,
+    _check as check,
+)
 from dagster._annotations import beta, public
+from dagster._core.definitions.definitions_class import Definitions
+from dagster._core.definitions.job_definition import JobDefinition
+from dagster._core.definitions.unresolved_asset_job_definition import UnresolvedAssetJobDefinition
 from dagster._record import record
 
 from dagster_airlift.core.airflow_instance import AirflowInstance
@@ -15,6 +23,7 @@
 from dagster_airlift.core.serialization.serialized_data import DagHandle, TaskHandle
 from dagster_airlift.core.utils import (
     dag_handles_for_spec,
+    get_producing_dag_ids,
     is_dag_mapped_asset_spec,
     is_peered_dag_asset_spec,
     is_task_mapped_asset_spec,
@@ -26,6 +35,14 @@
 MappedAsset = Union[AssetSpec, AssetsDefinition]
 
 
+def _is_mapped_asset(asset: MappedAsset) -> bool:
+    return (
+        is_task_mapped_asset_spec(asset)
+        or is_dag_mapped_asset_spec(asset)
+        or is_peered_dag_asset_spec(asset)
+    )
+
+
 @beta
 @record
 class AirflowDefinitionsData:
@@ -37,25 +54,61 @@ class AirflowDefinitionsData:
     """
 
     airflow_instance: AirflowInstance
-    airflow_mapped_assets: Sequence[MappedAsset]
+    defs: Definitions
+
+    @property
+    def airflow_mapped_asset_specs(self) -> Mapping[AssetKey, AssetSpec]:
+        """The assets that are mapped to Airflow tasks and dags."""
+        result = {}
+        for asset in self.defs.assets:
+            if not isinstance(asset, (AssetsDefinition, AssetSpec)):
+                continue
+            for spec in spec_iterator([asset]):
+                if _is_mapped_asset(spec):
+                    result[spec.key] = spec
+        return result
+
+    @property
+    def airflow_mapped_jobs(self) -> Sequence[Union[JobDefinition, UnresolvedAssetJobDefinition]]:
+        """Jobs mapping to Airflow dags."""
+        return [job for job in self.defs.jobs if job.tags.get("dagster/external_job") == "airflow"]
+
+    @property
+    def airflow_mapped_jobs_by_dag_handle(
+        self,
+    ) -> Mapping[DagHandle, Union[JobDefinition, UnresolvedAssetJobDefinition]]:
+        """Jobs mapping to Airflow dags by dag_id."""
+        return {
+            DagHandle(dag_id=check.not_none(job.tags)["dagster-airlift/dag_id"]): job
+            for job in self.airflow_mapped_jobs
+        }
+
+    @property
+    def assets_per_job(self) -> Mapping[str, AbstractSet[AssetKey]]:
+        """Assets per job mapping to Airflow dags."""
+        return {
+            job.name: self.assets_produced_by_dags[dag_handle.dag_id]
+            for dag_handle, job in self.airflow_mapped_jobs_by_dag_handle.items()
+        }
+
+    @property
+    def assets_produced_by_dags(self) -> Mapping[str, AbstractSet[AssetKey]]:
+        """Assets produced by Airflow dags."""
+        result = defaultdict(set)
+        for spec in self.airflow_mapped_asset_specs.values():
+            for dag_id in get_producing_dag_ids(spec):
+                result[dag_id].add(spec.key)
+        return result
 
     @public
     @property
     def instance_name(self) -> str:
         """The name of the Airflow instance."""
         return self.airflow_instance.name
 
-    @cached_property
-    def all_asset_specs(self) -> Sequence[AssetSpec]:
-        return list(spec_iterator(self.airflow_mapped_assets))
-
     @cached_property
     def mapping_info(self) -> AirliftMetadataMappingInfo:
-        return build_airlift_metadata_mapping_info(self.airflow_mapped_assets)
-
-    @cached_property
-    def all_asset_specs_by_key(self) -> Mapping[AssetKey, AssetSpec]:
-        return {spec.key: spec for spec in self.all_asset_specs}
+        return build_airlift_metadata_mapping_info(self.airflow_mapped_asset_specs)
 
     @public
     def task_ids_in_dag(self, dag_id: str) -> set[str]:
@@ -80,7 +133,7 @@ def dag_ids_with_mapped_asset_keys(self) -> AbstractSet[str]:
     @cached_property
     def mapped_asset_keys_by_task_handle(self) -> Mapping[TaskHandle, AbstractSet[AssetKey]]:
         asset_keys_per_handle = defaultdict(set)
-        for spec in self.all_asset_specs:
+        for spec in self.airflow_mapped_asset_specs.values():
             if is_task_mapped_asset_spec(spec):
                 task_handles = task_handles_for_spec(spec)
                 for task_handle in task_handles:
@@ -89,8 +142,9 @@ def mapped_asset_keys_by_task_handle(self) -> Mapping[TaskHandle, AbstractSet[As
 
     @cached_property
     def mapped_asset_keys_by_dag_handle(self) -> Mapping[DagHandle, AbstractSet[AssetKey]]:
+        """Assets specifically mapped to each dag."""
         asset_keys_per_handle = defaultdict(set)
-        for spec in self.all_asset_specs:
+        for spec in self.airflow_mapped_asset_specs.values():
             if is_dag_mapped_asset_spec(spec):
                 dag_handles = dag_handles_for_spec(spec)
                 for dag_handle in dag_handles:
@@ -99,8 +153,9 @@ def mapped_asset_keys_by_dag_handle(self) -> Mapping[DagHandle, AbstractSet[Asse
 
     @cached_property
     def peered_dag_asset_keys_by_dag_handle(self) -> Mapping[DagHandle, AbstractSet[AssetKey]]:
+        """Autogenerated "peered" dag assets."""
         asset_keys_per_handle = defaultdict(set)
-        for spec in self.all_asset_specs:
+        for spec in self.airflow_mapped_asset_specs.values():
             if is_peered_dag_asset_spec(spec):
                 dag_handles = peered_dag_handles_for_spec(spec)
                 for dag_handle in dag_handles:

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

@@ -49,7 +49,10 @@ def dag_asset_job(
     return define_asset_job(
         name=convert_to_valid_dagster_name(dag_data.dag_id),
         metadata=dag_asset_metadata(dag_data.dag_info),
-        tags={**airflow_kind_dict(), **{"dagster/external_job": "airflow"}},
+        tags={
+            **airflow_kind_dict(),
+            **{"dagster/external_job": "airflow", "dagster-airlift/dag_id": dag_data.dag_id},
+        },
         selection=[asset.key for asset in specs],
     )
 
@@ -61,7 +64,10 @@ def dummy_op():
 
     @job(
         name=convert_to_valid_dagster_name(dag_data.dag_id),
-        tags={**airflow_kind_dict(), **{"dagster/external_job": "airflow"}},
+        tags={
+            **airflow_kind_dict(),
+            **{"dagster/external_job": "airflow", "dagster-airlift/dag_id": dag_data.dag_id},
+        },
     )
     def dummy_job():
         dummy_op()

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

@@ -1,14 +1,9 @@
 from collections import defaultdict
 from collections.abc import Iterable, Mapping, Sequence
 from dataclasses import dataclass
-from typing import Any, Callable, Optional, Union, cast
+from typing import Callable, Optional, Union, cast
 
-from dagster import (
-    AssetsDefinition,
-    AssetSpec,
-    Definitions,
-    _check as check,
-)
+from dagster import AssetsDefinition, AssetSpec, Definitions
 from dagster._annotations import beta
 from dagster._core.definitions.asset_key import AssetKey
 from dagster._core.definitions.asset_spec import map_asset_specs
@@ -19,7 +14,6 @@
 from dagster._core.definitions.sensor_definition import DefaultSensorStatus
 from dagster._core.definitions.unresolved_asset_job_definition import UnresolvedAssetJobDefinition
 
-from dagster_airlift.core.airflow_defs_data import MappedAsset
 from dagster_airlift.core.airflow_instance import AirflowInstance
 from dagster_airlift.core.filter import AirflowFilter
 from dagster_airlift.core.job_builder import construct_dag_jobs
@@ -41,12 +35,14 @@
     SerializedAirflowDefinitionsData,
 )
 from dagster_airlift.core.utils import (
+    MappedAsset,
     dag_handles_for_spec,
     get_metadata_key,
     is_dag_mapped_asset_spec,
     is_task_mapped_asset_spec,
     spec_iterator,
     task_handles_for_spec,
+    type_narrow_defs_assets,
 )
 
 
@@ -230,7 +226,7 @@ def only_include_dag(dag_info: DagInfo) -> bool:
 
     """
     defs = defs or Definitions()
-    mapped_assets = _type_narrow_defs_assets(defs)
+    mapped_assets = type_narrow_defs_assets(defs)
     serialized_airflow_data = AirflowInstanceDefsLoader(
         airflow_instance=airflow_instance,
         mapped_assets=mapped_assets,
@@ -272,18 +268,6 @@ def only_include_dag(dag_info: DagInfo) -> bool:
     )
 
 
-def _type_check_asset(asset: Any) -> MappedAsset:
-    return check.inst(
-        asset,
-        (AssetSpec, AssetsDefinition),
-        "Expected passed assets to all be AssetsDefinitions or AssetSpecs.",
-    )
-
-
-def _type_narrow_defs_assets(defs: Definitions) -> Sequence[MappedAsset]:
-    return [_type_check_asset(asset) for asset in defs.assets or []]
-
-
 def _apply_airflow_data_to_specs(
     assets: Sequence[MappedAsset],
     serialized_data: SerializedAirflowDefinitionsData,
@@ -456,7 +440,7 @@ def build_job_based_airflow_defs(
     mapped_defs: Optional[Definitions] = None,
 ) -> Definitions:
     mapped_defs = mapped_defs or Definitions()
-    mapped_assets = _type_narrow_defs_assets(mapped_defs)
+    mapped_assets = type_narrow_defs_assets(mapped_defs)
     serialized_airflow_data = AirflowInstanceDefsLoader(
         airflow_instance=airflow_instance,
         mapped_assets=mapped_assets,

python_modules/libraries/dagster-airlift/dagster_airlift/core/sensor/event_translation.py

@@ -40,7 +40,7 @@ def default_event_transformer(
     """The default event transformer function, which attaches a partition key to materializations which are from time-window partitioned assets."""
     cached_partition_calculations = defaultdict(dict)
     for mat in materializations:
-        asset_spec = airflow_data.all_asset_specs_by_key[mat.asset_key]
+        asset_spec = airflow_data.airflow_mapped_asset_specs[mat.asset_key]
         if not asset_spec.partitions_def or not isinstance(
             asset_spec.partitions_def, TimeWindowPartitionsDefinition
         ):

python_modules/libraries/dagster-airlift/dagster_airlift/core/sensor/sensor_builder.py

@@ -17,6 +17,7 @@
 from dagster._annotations import beta
 from dagster._core.definitions.asset_check_evaluation import AssetCheckEvaluation
 from dagster._core.definitions.asset_selection import AssetSelection
+from dagster._core.definitions.definitions_class import Definitions
 from dagster._core.definitions.events import AssetObservation
 from dagster._core.definitions.repository_definition.repository_definition import (
     RepositoryDefinition,
@@ -110,7 +111,7 @@ def build_airflow_polling_sensor(
         Definitions: A `Definitions` object containing the constructed sensor.
     """
     airflow_data = AirflowDefinitionsData(
-        airflow_instance=airflow_instance, airflow_mapped_assets=mapped_assets
+        airflow_instance=airflow_instance, defs=Definitions(assets=mapped_assets)
     )
 
     @sensor(
@@ -416,7 +417,7 @@ def automapped_tasks_asset_keys(
     asset_keys_to_emit = set()
     asset_keys = airflow_data.asset_keys_in_task(dag_run.dag_id, task_instance.task_id)
     for asset_key in asset_keys:
-        spec = airflow_data.all_asset_specs_by_key[asset_key]
+        spec = airflow_data.airflow_mapped_asset_specs[asset_key]
         if spec.metadata.get(AUTOMAPPED_TASK_METADATA_KEY):
             asset_keys_to_emit.add(asset_key)
     return asset_keys_to_emit

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

@@ -1,5 +1,5 @@
-from collections.abc import Iterable, Iterator
-from typing import TYPE_CHECKING, Optional, Union
+from collections.abc import Iterable, Iterator, Sequence
+from typing import TYPE_CHECKING, Any, Optional, Union
 
 from dagster import (
     AssetsDefinition,
@@ -8,6 +8,7 @@
     _check as check,
 )
 from dagster._core.definitions.cacheable_assets import CacheableAssetsDefinition
+from dagster._core.definitions.definitions_class import Definitions
 from dagster._core.definitions.utils import VALID_NAME_REGEX
 from dagster._core.errors import DagsterInvariantViolationError
 from dagster._core.storage.tags import KIND_PREFIX
@@ -106,3 +107,27 @@ def peered_dag_handles_for_spec(spec: AssetSpec) -> set["DagHandle"]:
         DagHandle(dag_id=dag_handle_dict["dag_id"])
         for dag_handle_dict in spec.metadata[PEERED_DAG_MAPPING_METADATA_KEY]
     }
+
+
+def get_producing_dag_ids(spec: AssetSpec) -> set[str]:
+    if is_dag_mapped_asset_spec(spec):
+        return {dag_handle.dag_id for dag_handle in dag_handles_for_spec(spec)}
+    if is_peered_dag_asset_spec(spec):
+        return {dag_handle.dag_id for dag_handle in peered_dag_handles_for_spec(spec)}
+    else:
+        return {task_handle.dag_id for task_handle in task_handles_for_spec(spec)}
+
+
+MappedAsset = Union[AssetSpec, AssetsDefinition]
+
+
+def _type_check_asset(asset: Any) -> MappedAsset:
+    return check.inst(
+        asset,
+        (AssetSpec, AssetsDefinition),
+        "Expected passed assets to all be AssetsDefinitions or AssetSpecs.",
+    )
+
+
+def type_narrow_defs_assets(defs: Definitions) -> Sequence[MappedAsset]:
+    return [_type_check_asset(asset) for asset in defs.assets or []]

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

@@ -42,6 +42,7 @@
 )
 from dagster_airlift.core.serialization.defs_construction import make_default_dag_asset_key
 from dagster_airlift.core.serialization.serialized_data import (
+    DagHandle,
     SerializedAirflowDefinitionsData,
     TaskHandle,
 )
@@ -844,3 +845,21 @@ def test_load_job_defs() -> None:
     assert isinstance(get_job("producer2", defs), UnresolvedAssetJobDefinition)
     assert isinstance(get_job("consumer1", defs), UnresolvedAssetJobDefinition)
     assert isinstance(get_job("consumer2", defs), JobDefinition)
+
+    airflow_defs_data = AirflowDefinitionsData(
+        airflow_instance=af_instance,
+        defs=defs,
+    )
+
+    assert airflow_defs_data.airflow_mapped_jobs_by_dag_handle == {
+        DagHandle(dag_id="producer1"): get_job("producer1", defs),
+        DagHandle(dag_id="producer2"): get_job("producer2", defs),
+        DagHandle(dag_id="consumer1"): get_job("consumer1", defs),
+        DagHandle(dag_id="consumer2"): get_job("consumer2", defs),
+    }
+    assert airflow_defs_data.assets_per_job == {
+        "producer1": {AssetKey("example1"), AssetKey("a")},
+        "producer2": {AssetKey("example1")},
+        "consumer1": {AssetKey("example2")},
+        "consumer2": set(),
+    }