[dagster-airlift] make AirflowDefinitionsData aware of jobs

Author: dpeng817

python_modules/libraries/dagster-airlift/dagster_airlift/constants.py

@@ -12,5 +12,6 @@
 DAG_RUN_ID_TAG_KEY = "dagster-airlift/airflow-dag-run-id"
 DAG_ID_TAG_KEY = "dagster-airlift/airflow-dag-id"
 TASK_ID_TAG_KEY = "dagster-airlift/airflow-task-id"
+EXTERNAL_JOB_TAG_KEY = "dagster/external-job"
 
 SOURCE_CODE_METADATA_KEY = "Source Code"

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

@@ -5,6 +5,11 @@
 
 from dagster import AssetKey, AssetsDefinition, AssetSpec
 from dagster._annotations import beta, public
+from dagster._core.definitions.job_definition import JobDefinition
+from dagster._core.definitions.repository_definition.repository_definition import (
+    RepositoryDefinition,
+)
+from dagster._core.definitions.unresolved_asset_job_definition import UnresolvedAssetJobDefinition
 from dagster._record import record
 
 from dagster_airlift.core.airflow_instance import AirflowInstance
@@ -14,7 +19,10 @@
 )
 from dagster_airlift.core.serialization.serialized_data import DagHandle, TaskHandle
 from dagster_airlift.core.utils import (
+    dag_handle_from_job,
     dag_handles_for_spec,
+    get_producing_dag_ids,
+    is_airflow_mapped_job,
     is_dag_mapped_asset_spec,
     is_peered_dag_asset_spec,
     is_task_mapped_asset_spec,
@@ -26,6 +34,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 +53,57 @@ class AirflowDefinitionsData:
     """
 
     airflow_instance: AirflowInstance
-    airflow_mapped_assets: Sequence[MappedAsset]
+    resolved_repository: RepositoryDefinition
+
+    @property
+    def airflow_mapped_asset_specs(self) -> Mapping[AssetKey, AssetSpec]:
+        """The assets that are mapped to Airflow tasks and dags."""
+        return {
+            spec.key: spec
+            for spec in spec_iterator(self.resolved_repository.assets_defs_by_key.values())
+            if _is_mapped_asset(spec)
+        }
+
+    @property
+    def airflow_mapped_jobs(self) -> Sequence[JobDefinition]:
+        """Jobs mapping to Airflow dags."""
+        return [
+            job for job in self.resolved_repository.get_all_jobs() if is_airflow_mapped_job(job)
+        ]
+
+    @property
+    def airflow_mapped_jobs_by_dag_handle(
+        self,
+    ) -> Mapping[DagHandle, Union[JobDefinition, UnresolvedAssetJobDefinition]]:
+        """Jobs mapping to Airflow dags by dag_id."""
+        return {dag_handle_from_job(job): 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.values())
 
     @public
     def task_ids_in_dag(self, dag_id: str) -> set[str]:
@@ -80,7 +128,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 +137,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 +148,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

@@ -16,7 +16,7 @@
     SerializedAirflowDefinitionsData,
     SerializedDagData,
 )
-from dagster_airlift.core.utils import airflow_kind_dict, convert_to_valid_dagster_name
+from dagster_airlift.core.utils import airflow_job_tags, convert_to_valid_dagster_name
 
 
 def construct_dag_jobs(
@@ -49,7 +49,7 @@ 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_job_tags(dag_data.dag_id),
         selection=[asset.key for asset in specs],
     )
 
@@ -61,7 +61,7 @@ def dummy_op():
 
     @job(
         name=convert_to_valid_dagster_name(dag_data.dag_id),
-        tags={**airflow_kind_dict(), **{"dagster/external_job": "airflow"}},
+        tags=airflow_job_tags(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

@@ -109,21 +109,19 @@ def build_airflow_polling_sensor(
     Returns:
         Definitions: A `Definitions` object containing the constructed sensor.
     """
-    airflow_data = AirflowDefinitionsData(
-        airflow_instance=airflow_instance, airflow_mapped_assets=mapped_assets
-    )
 
     @sensor(
-        name=f"{airflow_data.airflow_instance.name}__airflow_dag_status_sensor",
+        name=f"{airflow_instance.name}__airflow_dag_status_sensor",
         minimum_interval_seconds=minimum_interval_seconds,
         default_status=default_sensor_status or DefaultSensorStatus.RUNNING,
         # This sensor will only ever execute asset checks and not asset materializations.
         asset_selection=AssetSelection.all_asset_checks(),
     )
     def airflow_dag_sensor(context: SensorEvaluationContext) -> SensorResult:
         """Sensor to report materialization events for each asset as new runs come in."""
-        context.log.info(
-            f"************Running sensor for {airflow_data.airflow_instance.name}***********"
+        context.log.info(f"************Running sensor for {airflow_instance.name}***********")
+        airflow_data = AirflowDefinitionsData(
+            airflow_instance=airflow_instance, resolved_repository=context.repository_def
         )
         try:
             cursor = (
@@ -416,7 +414,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, Mapping, Sequence
+from typing import TYPE_CHECKING, Any, Optional, Union
 
 from dagster import (
     AssetsDefinition,
@@ -8,13 +8,17 @@
     _check as check,
 )
 from dagster._core.definitions.cacheable_assets import CacheableAssetsDefinition
+from dagster._core.definitions.definitions_class import Definitions
+from dagster._core.definitions.job_definition import JobDefinition
 from dagster._core.definitions.utils import VALID_NAME_REGEX
 from dagster._core.errors import DagsterInvariantViolationError
 from dagster._core.storage.tags import KIND_PREFIX
 
 from dagster_airlift.constants import (
     AIRFLOW_SOURCE_METADATA_KEY_PREFIX,
+    DAG_ID_TAG_KEY,
     DAG_MAPPING_METADATA_KEY,
+    EXTERNAL_JOB_TAG_KEY,
     PEERED_DAG_MAPPING_METADATA_KEY,
     TASK_MAPPING_METADATA_KEY,
 )
@@ -106,3 +110,49 @@ 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 []]
+
+
+def is_airflow_mapped_job(job: JobDefinition) -> bool:
+    return job.tags.get(EXTERNAL_JOB_TAG_KEY) == "airflow"
+
+
+def dag_handle_from_job(job: JobDefinition) -> "DagHandle":
+    from dagster_airlift.core.serialization.serialized_data import DagHandle
+
+    check.invariant(
+        is_airflow_mapped_job(job),
+        "Job is not an Airflow mapped job. Cannot get dag_id.",
+    )
+    return DagHandle(dag_id=job.tags[DAG_ID_TAG_KEY])
+
+
+def airflow_job_tags(dag_id: str) -> Mapping[str, str]:
+    return {
+        **airflow_kind_dict(),
+        EXTERNAL_JOB_TAG_KEY: "airflow",
+        DAG_ID_TAG_KEY: dag_id,
+    }

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

@@ -60,7 +60,7 @@ def asset_spec(asset_str: str, defs: Definitions) -> Optional[AssetSpec]:
     )
 
 
-def get_job(
+def get_job_from_defs(
     name: str, defs: Definitions
 ) -> Optional[Union[JobDefinition, UnresolvedAssetJobDefinition]]:
     """Get the job from the definitions by its name."""