from typing import AbstractSet, List, Mapping, Optional, Set, Tuple
from airflow.models.connection import Connection
from airflow.models.dag import DAG
from dagster import (
AssetKey,
AssetsDefinition,
GraphDefinition,
OutputMapping,
TimeWindowPartitionsDefinition,
)
from dagster._core.definitions.graph_definition import create_adjacency_lists
from dagster._utils.schedules import is_valid_cron_schedule
from dagster_airflow.dagster_job_factory import make_dagster_job_from_airflow_dag
from dagster_airflow.utils import (
DagsterAirflowError,
normalized_name,
)
def _build_asset_dependencies(
dag: DAG,
graph: GraphDefinition,
task_ids_by_asset_key: Mapping[AssetKey, AbstractSet[str]],
upstream_dependencies_by_asset_key: Mapping[AssetKey, AbstractSet[AssetKey]],
) -> Tuple[AbstractSet[OutputMapping], Mapping[str, AssetKey], Mapping[str, Set[AssetKey]]]:
"""Builds the asset dependency graph for a given set of airflow task mappings and a dagster graph."""
output_mappings = set()
keys_by_output_name = {}
internal_asset_deps: dict[str, Set[AssetKey]] = {}
visited_nodes: dict[str, bool] = {}
upstream_deps = set()
def find_upstream_dependency(node_name: str) -> None:
"""Uses Depth-Firs-Search to find all upstream asset dependencies
as described in task_ids_by_asset_key.
"""
# node has been visited
if visited_nodes[node_name]:
return
# mark node as visted
visited_nodes[node_name] = True
# traverse upstream nodes
for output_handle in graph.dependency_structure.all_upstream_outputs_from_node(node_name):
forward_node = output_handle.node_name
match = False
# find any assets produced by upstream nodes and add them to the internal asset deps
for asset_key in task_ids_by_asset_key:
if (
forward_node.replace(f"{normalized_name(dag.dag_id)}__", "")
in task_ids_by_asset_key[asset_key]
):
upstream_deps.add(asset_key)
match = True
# don't traverse past nodes that have assets
if not match:
find_upstream_dependency(forward_node)
# iterate through each asset to find all upstream asset dependencies
for asset_key in task_ids_by_asset_key:
asset_upstream_deps = set()
for task_id in task_ids_by_asset_key[asset_key]:
visited_nodes = {s.name: False for s in graph.nodes}
upstream_deps = set()
find_upstream_dependency(normalized_name(dag.dag_id, task_id))
for dep in upstream_deps:
asset_upstream_deps.add(dep)
keys_by_output_name[f"result_{normalized_name(dag.dag_id, task_id)}"] = asset_key
output_mappings.add(
OutputMapping(
graph_output_name=f"result_{normalized_name(dag.dag_id, task_id)}",
mapped_node_name=normalized_name(dag.dag_id, task_id),
mapped_node_output_name="airflow_task_complete", # Default output name
)
)
# the tasks for a given asset should have the same internal deps
for task_id in task_ids_by_asset_key[asset_key]:
if f"result_{normalized_name(dag.dag_id, task_id)}" in internal_asset_deps:
internal_asset_deps[f"result_{normalized_name(dag.dag_id, task_id)}"].update(
asset_upstream_deps
)
else:
internal_asset_deps[f"result_{normalized_name(dag.dag_id, task_id)}"] = (
asset_upstream_deps
)
# add new upstream asset dependencies to the internal deps
for asset_key in upstream_dependencies_by_asset_key:
for key in keys_by_output_name:
if keys_by_output_name[key] == asset_key:
internal_asset_deps[key].update(upstream_dependencies_by_asset_key[asset_key])
return (output_mappings, keys_by_output_name, internal_asset_deps)
[docs]def load_assets_from_airflow_dag(
dag: DAG,
task_ids_by_asset_key: Mapping[AssetKey, AbstractSet[str]] = {},
upstream_dependencies_by_asset_key: Mapping[AssetKey, AbstractSet[AssetKey]] = {},
connections: Optional[List[Connection]] = None,
) -> List[AssetsDefinition]:
"""[Experimental] Construct Dagster Assets for a given Airflow DAG.
Args:
dag (DAG): The Airflow DAG to compile into a Dagster job
task_ids_by_asset_key (Optional[Mapping[AssetKey, AbstractSet[str]]]): A mapping from asset
keys to task ids. Used break up the Airflow Dag into multiple SDAs
upstream_dependencies_by_asset_key (Optional[Mapping[AssetKey, AbstractSet[AssetKey]]]): A
mapping from upstream asset keys to assets provided in task_ids_by_asset_key. Used to
declare new upstream SDA depenencies.
connections (List[Connection]): List of Airflow Connections to be created in the Airflow DB
Returns:
List[AssetsDefinition]
"""
cron_schedule = dag.normalized_schedule_interval
if cron_schedule is not None and not is_valid_cron_schedule(str(cron_schedule)):
raise DagsterAirflowError(f"Invalid cron schedule: {cron_schedule} in DAG {dag.dag_id}")
job = make_dagster_job_from_airflow_dag(dag, connections=connections)
graph = job._graph_def # noqa: SLF001
start_date = dag.start_date if dag.start_date else dag.default_args.get("start_date")
if start_date is None:
raise DagsterAirflowError(f"Invalid start_date: {start_date} in DAG {dag.dag_id}")
# leaf nodes have no downstream nodes
forward_edges, _ = create_adjacency_lists(graph.nodes, graph.dependency_structure)
leaf_nodes = {
node_name.replace(f"{normalized_name(dag.dag_id)}__", "")
for node_name, downstream_nodes in forward_edges.items()
if not downstream_nodes
}
mutated_task_ids_by_asset_key: dict[AssetKey, set[str]] = {}
if task_ids_by_asset_key is None or task_ids_by_asset_key == {}:
# if no mappings are provided the dag becomes a single SDA
task_ids_by_asset_key = {AssetKey(dag.dag_id): leaf_nodes}
else:
# if mappings were provide any unmapped leaf nodes are added to a default asset
used_nodes: set[str] = set()
for key in task_ids_by_asset_key:
used_nodes.update(task_ids_by_asset_key[key])
mutated_task_ids_by_asset_key[AssetKey(dag.dag_id)] = leaf_nodes - used_nodes
for key in task_ids_by_asset_key:
if key not in mutated_task_ids_by_asset_key:
mutated_task_ids_by_asset_key[key] = set(task_ids_by_asset_key[key])
else:
mutated_task_ids_by_asset_key[key].update(task_ids_by_asset_key[key])
output_mappings, keys_by_output_name, internal_asset_deps = _build_asset_dependencies(
dag, graph, mutated_task_ids_by_asset_key, upstream_dependencies_by_asset_key
)
new_graph = graph.copy(
output_mappings=list(output_mappings),
)
asset_def = AssetsDefinition.from_graph(
graph_def=new_graph,
partitions_def=(
TimeWindowPartitionsDefinition(
cron_schedule=str(cron_schedule),
timezone=dag.timezone.name,
start=start_date.strftime("%Y-%m-%dT%H:%M:%S"),
fmt="%Y-%m-%dT%H:%M:%S",
)
if cron_schedule is not None
else None
),
group_name=dag.dag_id,
keys_by_output_name=keys_by_output_name,
internal_asset_deps=internal_asset_deps,
can_subset=True,
)
return [asset_def]