diff --git a/airflow-core/src/airflow/api_fastapi/core_api/datamodels/ui/common.py b/airflow-core/src/airflow/api_fastapi/core_api/datamodels/ui/common.py index 702f6997cdf5e..0cccdecc4d3cf 100644 --- a/airflow-core/src/airflow/api_fastapi/core_api/datamodels/ui/common.py +++ b/airflow-core/src/airflow/api_fastapi/core_api/datamodels/ui/common.py @@ -55,6 +55,7 @@ class BaseNodeResponse(BaseModel): "trigger", ] team: str | None = None + asset_condition_type: Literal["or-gate", "and-gate"] | None = None E = TypeVar("E", bound=BaseEdgeResponse) diff --git a/airflow-core/src/airflow/api_fastapi/core_api/datamodels/ui/structure.py b/airflow-core/src/airflow/api_fastapi/core_api/datamodels/ui/structure.py index 5d89b08ee98b6..36f0553d634a4 100644 --- a/airflow-core/src/airflow/api_fastapi/core_api/datamodels/ui/structure.py +++ b/airflow-core/src/airflow/api_fastapi/core_api/datamodels/ui/structure.py @@ -30,7 +30,6 @@ class EdgeResponse(BaseEdgeResponse): is_setup_teardown: bool | None = None label: str | None = None - is_source_asset: bool | None = None class NodeResponse(BaseNodeResponse): @@ -41,7 +40,6 @@ class NodeResponse(BaseNodeResponse): tooltip: str | None = None setup_teardown_type: Literal["setup", "teardown"] | None = None operator: str | None = None - asset_condition_type: Literal["or-gate", "and-gate"] | None = None ui_color: str | None = None ui_fgcolor: str | None = None diff --git a/airflow-core/src/airflow/api_fastapi/core_api/openapi/_private_ui.yaml b/airflow-core/src/airflow/api_fastapi/core_api/openapi/_private_ui.yaml index 00eaa55a5aa87..376844a5479ba 100644 --- a/airflow-core/src/airflow/api_fastapi/core_api/openapi/_private_ui.yaml +++ b/airflow-core/src/airflow/api_fastapi/core_api/openapi/_private_ui.yaml @@ -1030,13 +1030,6 @@ paths: - type: string - type: 'null' title: Root - - name: external_dependencies - in: query - required: false - schema: - type: boolean - default: false - title: External Dependencies - name: version_number in: query required: false @@ -1052,12 +1045,6 @@ paths: application/json: schema: $ref: '#/components/schemas/StructureDataResponse' - '400': - content: - application/json: - schema: - $ref: '#/components/schemas/HTTPExceptionResponse' - description: Bad Request '404': content: application/json: @@ -2146,6 +2133,14 @@ components: - type: string - type: 'null' title: Team + asset_condition_type: + anyOf: + - type: string + enum: + - or-gate + - and-gate + - type: 'null' + title: Asset Condition Type type: object required: - id @@ -2973,11 +2968,6 @@ components: - type: string - type: 'null' title: Label - is_source_asset: - anyOf: - - type: boolean - - type: 'null' - title: Is Source Asset type: object required: - source_id @@ -3626,6 +3616,14 @@ components: - type: string - type: 'null' title: Team + asset_condition_type: + anyOf: + - type: string + enum: + - or-gate + - and-gate + - type: 'null' + title: Asset Condition Type children: anyOf: - items: @@ -3656,14 +3654,6 @@ components: - type: string - type: 'null' title: Operator - asset_condition_type: - anyOf: - - type: string - enum: - - or-gate - - and-gate - - type: 'null' - title: Asset Condition Type ui_color: anyOf: - type: string diff --git a/airflow-core/src/airflow/api_fastapi/core_api/routes/ui/dependencies.py b/airflow-core/src/airflow/api_fastapi/core_api/routes/ui/dependencies.py index d09b8acf0e407..0286b0a126f02 100644 --- a/airflow-core/src/airflow/api_fastapi/core_api/routes/ui/dependencies.py +++ b/airflow-core/src/airflow/api_fastapi/core_api/routes/ui/dependencies.py @@ -69,7 +69,7 @@ def get_dependencies( raise HTTPException(status.HTTP_404_NOT_FOUND, f"Asset with id {asset_id} was not found") return BaseGraphResponse(**data) - data = get_scheduling_dependencies(readable_dags_filter.value) + data = get_scheduling_dependencies(readable_dags_filter.value, session) if node_id is not None: try: diff --git a/airflow-core/src/airflow/api_fastapi/core_api/routes/ui/structure.py b/airflow-core/src/airflow/api_fastapi/core_api/routes/ui/structure.py index 597f44db424bb..07fb23137d58c 100644 --- a/airflow-core/src/airflow/api_fastapi/core_api/routes/ui/structure.py +++ b/airflow-core/src/airflow/api_fastapi/core_api/routes/ui/structure.py @@ -26,13 +26,8 @@ from airflow.api_fastapi.common.router import AirflowRouter from airflow.api_fastapi.core_api.datamodels.ui.structure import StructureDataResponse from airflow.api_fastapi.core_api.openapi.exceptions import create_openapi_http_exception_doc -from airflow.api_fastapi.core_api.security import ReadableDagsFilterDep, requires_access_dag -from airflow.api_fastapi.core_api.services.ui.structure import ( - bind_output_assets_to_tasks, - get_upstream_assets, -) +from airflow.api_fastapi.core_api.security import requires_access_dag from airflow.api_fastapi.core_api.services.ui.task_group import task_group_to_dict -from airflow.models.dag import DagModel from airflow.models.dag_version import DagVersion from airflow.models.serialized_dag import SerializedDagModel from airflow.utils.dag_edges import dag_edges @@ -44,25 +39,21 @@ "/structure_data", responses=create_openapi_http_exception_doc( [ - status.HTTP_400_BAD_REQUEST, status.HTTP_404_NOT_FOUND, ] ), dependencies=[ Depends(requires_access_dag("GET")), - Depends(requires_access_dag("GET", DagAccessEntity.DEPENDENCIES)), Depends(requires_access_dag("GET", DagAccessEntity.TASK_INSTANCE)), ], ) def structure_data( session: SessionDep, dag_id: str, - readable_dags_filter: ReadableDagsFilterDep, include_upstream: QueryIncludeUpstream = False, include_downstream: QueryIncludeDownstream = False, depth: int | None = None, root: str | None = None, - external_dependencies: bool = False, version_number: int | None = None, ) -> StructureDataResponse: """Get Structure Data.""" @@ -79,7 +70,7 @@ def structure_data( select(SerializedDagModel) .join(DagVersion) .where(SerializedDagModel.dag_id == dag_id, DagVersion.version_number == version_number) - .options(joinedload(SerializedDagModel.dag_model).joinedload(DagModel.task_outlet_asset_references)), + .options(joinedload(SerializedDagModel.dag_model)), ) if serialized_dag is None: raise HTTPException( @@ -104,82 +95,4 @@ def structure_data( "edges": edges, } - if external_dependencies: - entry_node_ref = nodes[0] if nodes else None - exit_node_ref = nodes[-1] if nodes else None - - start_edges: list[dict] = [] - end_edges: list[dict] = [] - - readable_dag_ids = readable_dags_filter.value - for dependency_dag_id, dependencies in sorted(SerializedDagModel.get_dag_dependencies().items()): - if readable_dag_ids is not None and dependency_dag_id not in readable_dag_ids: - continue - for dependency in dependencies: - # Dependencies not related to `dag_id` are ignored - if dependency_dag_id != dag_id and dependency.target != dag_id: - continue - # When target is a real Dag ID (not a type label), hide it - # if the caller cannot read that Dag. - if ( - readable_dag_ids is not None - and dependency.target != dependency.dependency_type - and dependency.target not in readable_dag_ids - ): - continue - - # upstream assets are handled by the `get_upstream_assets` function. - if dependency.target != dependency.dependency_type and dependency.dependency_type in [ - "asset-alias", - "asset", - ]: - continue - - # Add edges - # start dependency - if ( - dependency.source == dependency.dependency_type or dependency.target == dag_id - ) and entry_node_ref: - start_edges.append({"source_id": dependency.node_id, "target_id": entry_node_ref["id"]}) - - # end dependency - elif ( - dependency.target == dependency.dependency_type or dependency.source == dag_id - ) and exit_node_ref: - end_edges.append( - { - "source_id": exit_node_ref["id"], - "target_id": dependency.node_id, - "resolved_from_alias": dependency.source.replace("asset-alias:", "", 1) - if dependency.source.startswith("asset-alias:") - else None, - } - ) - - # Add nodes - nodes.append( - { - "id": dependency.node_id, - "label": dependency.label, - "type": dependency.dependency_type, - } - ) - - if (asset_expression := serialized_dag.dag_model.asset_expression) and entry_node_ref: - try: - upstream_asset_nodes, upstream_asset_edges = get_upstream_assets( - asset_expression, entry_node_ref["id"] - ) - except TypeError as e: - raise HTTPException( - status.HTTP_400_BAD_REQUEST, - f"Malformed asset_expression in Dag {dag_id!r} version {version_number}: {e}", - ) from e - data["nodes"] += upstream_asset_nodes - data["edges"] += upstream_asset_edges - - data["edges"] += start_edges + end_edges - - bind_output_assets_to_tasks(data["edges"], serialized_dag, version_number, session) - return StructureDataResponse(**data) diff --git a/airflow-core/src/airflow/api_fastapi/core_api/services/ui/dependencies.py b/airflow-core/src/airflow/api_fastapi/core_api/services/ui/dependencies.py index 20ffc78008a8f..b519a61def3b7 100644 --- a/airflow-core/src/airflow/api_fastapi/core_api/services/ui/dependencies.py +++ b/airflow-core/src/airflow/api_fastapi/core_api/services/ui/dependencies.py @@ -17,15 +17,119 @@ from __future__ import annotations -from collections import defaultdict, deque +from collections import defaultdict from typing import TYPE_CHECKING +import structlog +from sqlalchemy import select + from airflow.models.asset import AssetModel from airflow.models.dag import DagModel if TYPE_CHECKING: from sqlalchemy.orm import Session +log = structlog.get_logger(logger_name=__name__) + +ASSET_UPSTREAM_DEPENDENCY_TYPES = ("asset", "asset-alias", "asset-name-ref", "asset-uri-ref") + + +def _asset_node_id_and_label(asset: dict) -> tuple[str, str]: + """Get the graph node id and label for a single asset/alias/ref entry in an asset_expression.""" + asset_type = asset["type"] + + if asset_type == "asset": + return f"asset:{asset['id']}", asset["name"] + if asset_type in ("asset-alias", "asset-name-ref"): + return f"{asset_type}:{asset['name']}", asset["name"] + if asset_type == "asset-uri-ref": + return f"{asset_type}:{asset['uri']}", asset["uri"] + raise TypeError(f"Unsupported type: {asset_type}") + + +def get_upstream_assets( + asset_expression: dict, entry_node_ref: str, level: int = 0 +) -> tuple[list[dict], list[dict]]: + """Expand a Dag's asset trigger condition into asset-condition (AND/OR gate) nodes and edges.""" + edges: list[dict] = [] + nodes: list[dict] = [] + + asset_expression_type: str | None = None + expr_key = "" + if asset_expression.keys() == {"any"}: + asset_expression_type = "or-gate" + expr_key = "any" + elif asset_expression.keys() == {"all"}: + asset_expression_type = "and-gate" + expr_key = "all" + + if not asset_expression_type: + return nodes, edges + + # include assets, asset-alias, asset-name-refs, asset-uri-refs + assets_info: list[dict] = [] + # nested boolean sub-conditions, each of which becomes its own gate feeding this one + nested_expressions: list[dict] = [] + + for expr in asset_expression[expr_key]: + nested_expr_key = next(iter(expr), None) + if nested_expr_key in ("any", "all"): + nested_expressions.append(expr) + elif nested_expr_key in ("asset", "alias", "asset-name-ref", "asset-uri-ref"): + asset_info = expr[nested_expr_key] + asset_info["type"] = nested_expr_key if nested_expr_key != "alias" else "asset-alias" + assets_info.append(asset_info) + elif nested_expr_key == "asset_ref": + # Asset.ref(...) that hasn't been resolved to a concrete asset yet serializes as + # {"asset_ref": {"name": ...}} or {"asset_ref": {"uri": ...}} -- disambiguate on + # the inner field since, unlike the other branches, the key itself doesn't say which. + ref_info = expr[nested_expr_key] + ref_info["type"] = "asset-name-ref" if "name" in ref_info else "asset-uri-ref" + assets_info.append(ref_info) + else: + raise TypeError(f"Unsupported type: {expr.keys()}") + + branch_count = len(assets_info) + len(nested_expressions) + if branch_count == 0: + return nodes, edges + + # A condition combining exactly one branch isn't a real AND/OR -- connect it directly (or + # recurse straight through the lone nested wrapper) instead of rendering a single-input gate. + if branch_count == 1: + if assets_info: + source_id, label = _asset_node_id_and_label(assets_info[0]) + edges.append({"source_id": source_id, "target_id": entry_node_ref}) + nodes.append({"id": source_id, "label": label, "type": assets_info[0]["type"]}) + return nodes, edges + return get_upstream_assets(nested_expressions[0], entry_node_ref, level=level) + + # Scoped by entry_node_ref (unique per Dag) so gates from different Dags don't collide + # when merged into a single multi-dag graph. + asset_condition_id = f"{entry_node_ref}-{asset_expression_type}-{level}" + edges.append({"source_id": asset_condition_id, "target_id": entry_node_ref}) + nodes.append( + { + "id": asset_condition_id, + "label": asset_condition_id, + "type": "asset-condition", + "asset_condition_type": asset_expression_type, + } + ) + + for asset in assets_info: + source_id, label = _asset_node_id_and_label(asset) + edges.append({"source_id": source_id, "target_id": asset_condition_id}) + nodes.append({"id": source_id, "label": label, "type": asset["type"]}) + + # Each nested sub-condition feeds this gate through its own gate. The branch index keeps + # sibling gates of the same type from colliding on their scoped id. + for index, nested in enumerate(nested_expressions): + nested_nodes, nested_edges = get_upstream_assets(nested, asset_condition_id, level=index) + nodes += nested_nodes + edges += nested_edges + + return nodes, edges + def _dfs_connected_components( temp: list[str], node_id: str, visited: dict[str, bool], adjacency_matrix: dict[str, list[str]] @@ -85,7 +189,7 @@ def extract_single_connected_component( return {"nodes": nodes, "edges": edges} -def get_scheduling_dependencies(readable_dag_ids: set[str] | None = None) -> dict[str, list[dict]]: +def get_scheduling_dependencies(readable_dag_ids: set[str] | None, session: Session) -> dict[str, list[dict]]: """Get scheduling dependencies between Dags.""" from airflow.models.serialized_dag import SerializedDagModel @@ -93,6 +197,40 @@ def get_scheduling_dependencies(readable_dag_ids: set[str] | None = None) -> dic edge_tuples: set[tuple[str, str]] = set() dag_dependencies = SerializedDagModel.get_dag_dependencies() + + relevant_dag_ids = [ + dag for dag in dag_dependencies if readable_dag_ids is None or dag in readable_dag_ids + ] + + # A Dag's asset trigger condition (AND/OR of assets) is expanded into asset-condition + # gate nodes so it renders the same way here as it would per-Dag. A single-asset + # schedule serializes as a bare `{"asset": {...}}`, not `{"any": [...]}`, so + # `get_upstream_assets` returns nothing for it — only Dags with a genuine boolean + # condition are gate-expanded; other Dags keep their flat asset edge untouched. + expression_nodes: dict[str, dict] = {} + expression_edges: set[tuple[str, str]] = set() + gated_dag_ids: set[str] = set() + if relevant_dag_ids: + asset_expressions = session.execute( + select(DagModel.dag_id, DagModel.asset_expression).where(DagModel.dag_id.in_(relevant_dag_ids)) + ) + for dag_id, asset_expression in asset_expressions: + if not asset_expression: + continue + try: + upstream_nodes, upstream_edges = get_upstream_assets(asset_expression, f"dag:{dag_id}") + except TypeError: + # A malformed/unrecognized asset_expression for one Dag must not break this + # endpoint for every Dag -- fall back to that Dag's flat asset edge instead. + log.warning("Could not expand asset_expression into gate nodes", dag_id=dag_id, exc_info=True) + continue + if upstream_nodes: + gated_dag_ids.add(dag_id) + for node in upstream_nodes: + expression_nodes[node["id"]] = node + for edge in upstream_edges: + expression_edges.add((edge["source_id"], edge["target_id"])) + for dag, dependencies in sorted(dag_dependencies.items()): if readable_dag_ids is not None and dag not in readable_dag_ids: continue @@ -113,6 +251,15 @@ def get_scheduling_dependencies(readable_dag_ids: set[str] | None = None) -> dic if not referenced_dag_ids.issubset(readable_dag_ids): continue + # This Dag's upstream asset condition is represented by the gate nodes + # built above instead of a flat edge. + if ( + dag in gated_dag_ids + and dep.target == dag + and dep.dependency_type in ASSET_UPSTREAM_DEPENDENCY_TYPES + ): + continue + # Add nodes nodes_dict[dag_node_id] = {"id": dag_node_id, "label": dag, "type": "dag"} if dep.node_id not in nodes_dict: @@ -135,6 +282,9 @@ def get_scheduling_dependencies(readable_dag_ids: set[str] | None = None) -> dic target = dep.target if ":" in dep.target else f"dag:{dep.target}" edge_tuples.add((source, target)) + nodes_dict.update(expression_nodes) + edge_tuples |= expression_edges + # Create missing ``dag:`` nodes which may have been skipped by the loop above. # A DAG referenced only as a trigger target or a sensor source may have no # scheduling dependencies of its own. Without this loop, these DAGs will not be @@ -163,6 +313,83 @@ def get_scheduling_dependencies(readable_dag_ids: set[str] | None = None) -> dic } +def _entry_task_id_from_serialized_data(data: dict | None) -> str | None: + """ + Pick a Dag's entry-point task id from its serialized JSON, without deserializing the Dag. + + A root task (one nothing else lists as downstream) is where an asset schedule enters the + graph. Task ids and their downstream links are read straight from the serialized payload -- + ``downstream_task_ids`` is stored as a plain list of strings -- so this avoids instantiating + every operator just to read one id. Ties are broken by sorted order for determinism. + """ + tasks = (data or {}).get("dag", {}).get("tasks", []) + task_ids: set[str] = set() + downstream_ids: set[str] = set() + for task in tasks: + var = task.get("__var", {}) + task_id = var.get("task_id") + if task_id is None: + continue + task_ids.add(task_id) + downstream_ids.update(var.get("downstream_task_ids") or []) + + if not task_ids: + return None + roots = task_ids - downstream_ids + return min(roots) if roots else min(task_ids) + + +def _get_dag_entry_points(dag_ids: set[str], session: Session) -> dict[str, tuple[str, dict | None]]: + """ + Batch-resolve each Dag's entry-point task id and its asset_expression. + + Fetches the latest SerializedDagModel row for every dag_id in one query (mirroring the + latest-per-group window-function pattern in + ``SerializedDagModel._prefetch_dag_write_metadata``) instead of one query per Dag, and reads + the entry task from the serialized JSON rather than deserializing the whole Dag -- both matter + when a single asset schedules many Dags across a connected component. + """ + from sqlalchemy import func, select + from sqlalchemy.orm import joinedload + + from airflow.models.dag_version import DagVersion + from airflow.models.serialized_dag import SerializedDagModel + + if not dag_ids: + return {} + + latest_per_dag = ( + select( + SerializedDagModel.id, + func.row_number() + .over(partition_by=SerializedDagModel.dag_id, order_by=DagVersion.version_number.desc()) + .label("rn"), + ) + .join(DagVersion, SerializedDagModel.dag_version_id == DagVersion.id) + .where(SerializedDagModel.dag_id.in_(dag_ids)) + .subquery() + ) + + serialized_dags = session.scalars( + select(SerializedDagModel) + .join(latest_per_dag, SerializedDagModel.id == latest_per_dag.c.id) + .where(latest_per_dag.c.rn == 1) + .options(joinedload(SerializedDagModel.dag_model)) + ).all() + + entry_points: dict[str, tuple[str, dict | None]] = {} + for serialized_dag in serialized_dags: + entry_task_id = _entry_task_id_from_serialized_data(serialized_dag.data) + if entry_task_id is None: + continue + entry_points[serialized_dag.dag_id] = ( + entry_task_id, + serialized_dag.dag_model.asset_expression, + ) + + return entry_points + + def get_data_dependencies( asset_id: int, session: Session, readable_dag_ids: set[str] | None = None ) -> dict[str, list[dict]]: @@ -197,98 +424,225 @@ def get_data_dependencies( nodes_dict: dict[str, dict] = {} edge_set: set[tuple[str, str]] = set() - # BFS to trace full dependencies - assets_to_process: deque[int] = deque([asset_id]) processed_assets: set[int] = set() processed_tasks: set[tuple[str, str]] = set() # (dag_id, task_id) - - while assets_to_process: - current_asset_id = assets_to_process.popleft() - if current_asset_id in processed_assets: - continue - processed_assets.add(current_asset_id) - - # Eagerload producing_tasks and consuming_tasks to avoid lazy queries - asset = session.scalar( + processed_scheduled_dags: set[str] = set() + # Assets linked to an AssetAlias -- these may be produced via the alias with no static + # TaskOutletAssetReference, so their producing task is recovered from AssetEvent afterwards. + alias_linked_asset_ids: set[int] = set() + + # BFS by rounds (a whole frontier of assets at once) rather than one asset at a time, so + # every Dag scheduled by assets in the same round has its entry point resolved through one + # batched query (see _get_dag_entry_points) instead of one query -- and one full Dag + # deserialization -- per Dag. + frontier: set[int] = {asset_id} + + while frontier: + round_asset_ids = frontier - processed_assets + processed_assets |= round_asset_ids + if not round_asset_ids: + break + + # Eagerload producing_tasks, consuming_tasks, and scheduled_dags to avoid lazy queries + assets = session.scalars( select(AssetModel) - .where(AssetModel.id == current_asset_id) + .where(AssetModel.id.in_(round_asset_ids)) .options( selectinload(AssetModel.producing_tasks), selectinload(AssetModel.consuming_tasks), + selectinload(AssetModel.scheduled_dags), + selectinload(AssetModel.aliases), ) - ) - if not asset: - continue - - asset_node_id = f"asset:{current_asset_id}" + ).all() + + next_frontier: set[int] = set() + pending_dag_ids: set[str] = set() + triggering_asset_node_ids_by_dag: dict[str, set[str]] = defaultdict(set) + + for asset in assets: + asset_node_id = f"asset:{asset.id}" + + # Add asset node + if asset_node_id not in nodes_dict: + nodes_dict[asset_node_id] = {"id": asset_node_id, "label": asset.name, "type": "asset"} + + if asset.aliases: + alias_linked_asset_ids.add(asset.id) + + # Process producing tasks (tasks that output this asset) + for ref in asset.producing_tasks: + # Filter out tasks from Dags the user doesn't have access to + if readable_dag_ids is not None and ref.dag_id not in readable_dag_ids: + continue + task_key = (ref.dag_id, ref.task_id) + task_node_id = f"task:{ref.dag_id}{SEPARATOR}{ref.task_id}" + + # Add task node with dag_id.task_id label for disambiguation + if task_node_id not in nodes_dict: + nodes_dict[task_node_id] = { + "id": task_node_id, + "label": f"{ref.dag_id}.{ref.task_id}", + "type": "task", + } - # Add asset node - if asset_node_id not in nodes_dict: - nodes_dict[asset_node_id] = {"id": asset_node_id, "label": asset.name, "type": "asset"} + # Add edge: task → asset + edge_set.add((task_node_id, asset_node_id)) + + # Find other assets this task consumes (inlets) to trace upstream + if task_key not in processed_tasks: + processed_tasks.add(task_key) + inlet_refs = session.scalars( + select(TaskInletAssetReference).where( + TaskInletAssetReference.dag_id == ref.dag_id, + TaskInletAssetReference.task_id == ref.task_id, + ) + ).all() + for inlet_ref in inlet_refs: + if inlet_ref.asset_id not in processed_assets: + next_frontier.add(inlet_ref.asset_id) + + # Process consuming tasks (tasks that input this asset) + for ref in asset.consuming_tasks: + # Filter out tasks from Dags the user doesn't have access to + if readable_dag_ids is not None and ref.dag_id not in readable_dag_ids: + continue + task_key = (ref.dag_id, ref.task_id) + task_node_id = f"task:{ref.dag_id}{SEPARATOR}{ref.task_id}" + + # Add task node with dag_id.task_id label for disambiguation + if task_node_id not in nodes_dict: + nodes_dict[task_node_id] = { + "id": task_node_id, + "label": f"{ref.dag_id}.{ref.task_id}", + "type": "task", + } - # Process producing tasks (tasks that output this asset) - for ref in asset.producing_tasks: - # Filter out tasks from Dags the user doesn't have access to - if readable_dag_ids is not None and ref.dag_id not in readable_dag_ids: + # Add edge: asset → task + edge_set.add((asset_node_id, task_node_id)) + + # Find other assets this task produces (outlets) to trace downstream + if task_key not in processed_tasks: + processed_tasks.add(task_key) + outlet_refs = session.scalars( + select(TaskOutletAssetReference).where( + TaskOutletAssetReference.dag_id == ref.dag_id, + TaskOutletAssetReference.task_id == ref.task_id, + ) + ).all() + for outlet_ref in outlet_refs: + if outlet_ref.asset_id not in processed_assets: + next_frontier.add(outlet_ref.asset_id) + + # Process Dags scheduled by this asset at the Dag level (`schedule=...`). These have + # no task-level inlet reference and would otherwise be a dead end in this graph -- + # collect them here and resolve every entry point for the round together, below. + for ref in asset.scheduled_dags: + if readable_dag_ids is not None and ref.dag_id not in readable_dag_ids: + continue + if ref.dag_id in processed_scheduled_dags: + continue + pending_dag_ids.add(ref.dag_id) + # Record every asset in this round that schedules the Dag -- marking it processed + # is deferred until after the round so a second triggering asset isn't skipped. + triggering_asset_node_ids_by_dag[ref.dag_id].add(asset_node_id) + + processed_scheduled_dags |= pending_dag_ids + + # Route through each Dag's asset_expression (skipping straight to a flat edge for a + # single-asset schedule) into the Dag's topologically-first task, same as the scheduling + # graph does for its `dag:` nodes. + entry_points = _get_dag_entry_points(pending_dag_ids, session) + for dag_id in pending_dag_ids: + entry_point = entry_points.get(dag_id) + if entry_point is None: continue - task_key = (ref.dag_id, ref.task_id) - task_node_id = f"task:{ref.dag_id}{SEPARATOR}{ref.task_id}" + entry_task_id, asset_expression = entry_point + task_key = (dag_id, entry_task_id) + task_node_id = f"task:{dag_id}{SEPARATOR}{entry_task_id}" - # Add task node with dag_id.task_id label for disambiguation if task_node_id not in nodes_dict: nodes_dict[task_node_id] = { "id": task_node_id, - "label": f"{ref.dag_id}.{ref.task_id}", + "label": f"{dag_id}.{entry_task_id}", "type": "task", } - # Add edge: task → asset - edge_set.add((task_node_id, asset_node_id)) + upstream_nodes: list[dict] = [] + upstream_edges: list[dict] = [] + if asset_expression: + try: + upstream_nodes, upstream_edges = get_upstream_assets(asset_expression, task_node_id) + except TypeError: + log.warning( + "Could not expand asset_expression into gate nodes", + dag_id=dag_id, + exc_info=True, + ) - # Find other assets this task consumes (inlets) to trace upstream + if upstream_nodes: + for node in upstream_nodes: + nodes_dict.setdefault(node["id"], node) + # Continue the BFS through any other concrete assets this gate combines + # with, so their own producers are traced too. + if node["id"].startswith("asset:"): + sibling_asset_id = int(node["id"].removeprefix("asset:")) + if sibling_asset_id not in processed_assets: + next_frontier.add(sibling_asset_id) + for edge in upstream_edges: + edge_set.add((edge["source_id"], edge["target_id"])) + else: + # No gate (bare single-asset schedule, or an expression that couldn't be + # expanded) -- link every triggering asset directly to the entry task. + for triggering_asset_node_id in triggering_asset_node_ids_by_dag[dag_id]: + edge_set.add((triggering_asset_node_id, task_node_id)) + + # Find other assets this entry task produces (outlets) to trace downstream. if task_key not in processed_tasks: processed_tasks.add(task_key) - inlet_refs = session.scalars( - select(TaskInletAssetReference).where( - TaskInletAssetReference.dag_id == ref.dag_id, - TaskInletAssetReference.task_id == ref.task_id, + outlet_refs = session.scalars( + select(TaskOutletAssetReference).where( + TaskOutletAssetReference.dag_id == dag_id, + TaskOutletAssetReference.task_id == entry_task_id, ) ).all() - for inlet_ref in inlet_refs: - if inlet_ref.asset_id not in processed_assets: - assets_to_process.append(inlet_ref.asset_id) - - # Process consuming tasks (tasks that input this asset) - for ref in asset.consuming_tasks: - # Filter out tasks from Dags the user doesn't have access to - if readable_dag_ids is not None and ref.dag_id not in readable_dag_ids: + for outlet_ref in outlet_refs: + if outlet_ref.asset_id not in processed_assets: + next_frontier.add(outlet_ref.asset_id) + + frontier = next_frontier + + # Assets produced only through an AssetAlias have no static TaskOutletAssetReference, so the + # producer loop above can't attribute them. Only when the graph actually contains alias-linked + # assets do we consult AssetEvent -- one batched query -- to recover the task that produced each + # via its alias, so alias-produced assets still show their source task. + if alias_linked_asset_ids: + from airflow.models.asset import AssetEvent + + alias_events = session.execute( + select(AssetEvent.asset_id, AssetEvent.source_dag_id, AssetEvent.source_task_id) + .join(AssetEvent.source_aliases) + .where( + AssetEvent.asset_id.in_(alias_linked_asset_ids), + AssetEvent.source_dag_id.is_not(None), + AssetEvent.source_task_id.is_not(None), + ) + .distinct() + ) + for event_asset_id, source_dag_id, source_task_id in alias_events: + if readable_dag_ids is not None and source_dag_id not in readable_dag_ids: continue - task_key = (ref.dag_id, ref.task_id) - task_node_id = f"task:{ref.dag_id}{SEPARATOR}{ref.task_id}" - - # Add task node with dag_id.task_id label for disambiguation + asset_node_id = f"asset:{event_asset_id}" + if asset_node_id not in nodes_dict: + continue + task_node_id = f"task:{source_dag_id}{SEPARATOR}{source_task_id}" if task_node_id not in nodes_dict: nodes_dict[task_node_id] = { "id": task_node_id, - "label": f"{ref.dag_id}.{ref.task_id}", + "label": f"{source_dag_id}.{source_task_id}", "type": "task", } - - # Add edge: asset → task - edge_set.add((asset_node_id, task_node_id)) - - # Find other assets this task produces (outlets) to trace downstream - if task_key not in processed_tasks: - processed_tasks.add(task_key) - outlet_refs = session.scalars( - select(TaskOutletAssetReference).where( - TaskOutletAssetReference.dag_id == ref.dag_id, - TaskOutletAssetReference.task_id == ref.task_id, - ) - ).all() - for outlet_ref in outlet_refs: - if outlet_ref.asset_id not in processed_assets: - assets_to_process.append(outlet_ref.asset_id) + edge_set.add((task_node_id, asset_node_id)) + processed_tasks.add((source_dag_id, source_task_id)) all_dag_ids = list({dag_id for dag_id, _ in processed_tasks}) if all_dag_ids: diff --git a/airflow-core/src/airflow/api_fastapi/core_api/services/ui/structure.py b/airflow-core/src/airflow/api_fastapi/core_api/services/ui/structure.py deleted file mode 100644 index db3d1ba6deac4..0000000000000 --- a/airflow-core/src/airflow/api_fastapi/core_api/services/ui/structure.py +++ /dev/null @@ -1,185 +0,0 @@ -# Licensed to the Apache Software Foundation (ASF) under one -# or more contributor license agreements. See the NOTICE file -# distributed with this work for additional information -# regarding copyright ownership. The ASF licenses this file -# to you under the Apache License, Version 2.0 (the -# "License"); you may not use this file except in compliance -# with the License. You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, -# software distributed under the License is distributed on an -# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY -# KIND, either express or implied. See the License for the -# specific language governing permissions and limitations -# under the License. - -""" -Private service for dag structure. - -:meta private: -""" - -from __future__ import annotations - -from collections import defaultdict - -from sqlalchemy import select -from sqlalchemy.orm import Session - -from airflow.models.asset import AssetAliasModel, AssetEvent -from airflow.models.dag_version import DagVersion -from airflow.models.dagrun import DagRun -from airflow.models.serialized_dag import SerializedDagModel - - -def get_upstream_assets( - asset_expression: dict, entry_node_ref: str, level: int = 0 -) -> tuple[list[dict], list[dict]]: - edges: list[dict] = [] - nodes: list[dict] = [] - asset_expression_type: str | None = None - - # include assets, asset-alias, asset-name-refs, asset-uri-refs - assets_info: list[dict] = [] - - nested_expression: dict = {} - - expr_key = "" - if asset_expression.keys() == {"any"}: - asset_expression_type = "or-gate" - expr_key = "any" - elif asset_expression.keys() == {"all"}: - asset_expression_type = "and-gate" - expr_key = "all" - - if expr_key in asset_expression: - asset_exprs: list[dict] = asset_expression[expr_key] - for expr in asset_exprs: - nested_expr_key = next(iter(expr.keys())) - if nested_expr_key in ("any", "all"): - nested_expression = expr - elif nested_expr_key in ("asset", "alias", "asset-name-ref", "asset-uri-ref"): - asset_info = expr[nested_expr_key] - asset_info["type"] = nested_expr_key if nested_expr_key != "alias" else "asset-alias" - - assets_info.append(asset_info) - else: - raise TypeError(f"Unsupported type: {expr.keys()}") - - if asset_expression_type and assets_info: - asset_condition_id = f"{asset_expression_type}-{level}" - edges.append( - { - "source_id": asset_condition_id, - "target_id": entry_node_ref, - "is_source_asset": level == 0, - } - ) - nodes.append( - { - "id": asset_condition_id, - "label": asset_condition_id, - "type": "asset-condition", - "asset_condition_type": asset_expression_type, - } - ) - - for asset in assets_info: - asset_type = asset["type"] - - if asset_type == "asset": - source_id = str(asset["id"]) - label = asset["name"] - elif asset_type == "asset-alias" or asset_type == "asset-name-ref": - source_id = asset["name"] - label = asset["name"] - elif asset_type == "asset-uri-ref": - source_id = asset["uri"] - label = asset["uri"] - else: - raise TypeError(f"Unsupported type: {asset_type}") - - edges.append( - { - "source_id": source_id, - "target_id": asset_condition_id, - } - ) - nodes.append( - { - "id": source_id, - "label": label, - "type": asset_type, - } - ) - - if nested_expression is not None: - n, e = get_upstream_assets(nested_expression, asset_condition_id, level=level + 1) - - nodes = nodes + n - edges = edges + e - - return nodes, edges - - -def bind_output_assets_to_tasks( - edges: list[dict], serialized_dag: SerializedDagModel, version_number: int, session: Session -) -> None: - """ - Try to bind the downstream assets to the relevant task that produces them. - - This function will mutate the `edges` in place. - """ - # bind normal assets present in the `task_outlet_asset_references` - outlet_asset_references = serialized_dag.dag_model.task_outlet_asset_references - - downstream_asset_edges = [ - edge - for edge in edges - if edge["target_id"].startswith("asset:") and not edge.get("resolved_from_alias") - ] - - for edge in downstream_asset_edges: - # Try to attach the outlet assets to the relevant tasks - asset_id = int(edge["target_id"].replace("asset:", "", 1)) - outlet_asset_reference = next( - outlet_asset_reference - for outlet_asset_reference in outlet_asset_references - if outlet_asset_reference.asset_id == asset_id - ) - edge["source_id"] = outlet_asset_reference.task_id - - # bind assets resolved from aliases, they do not populate the `outlet_asset_references` - downstream_alias_resolved_edges = [ - edge for edge in edges if edge["target_id"].startswith("asset:") and edge.get("resolved_from_alias") - ] - - aliases_names = {edges["resolved_from_alias"] for edges in downstream_alias_resolved_edges} - - result = session.scalars( - select(AssetEvent) - .join(AssetEvent.source_aliases) - .join(AssetEvent.source_dag_run) - # That's a simplification, instead doing `version_number` in `DagRun.dag_versions`. - .join(DagRun.created_dag_version) - .where(AssetEvent.source_aliases.any(AssetAliasModel.name.in_(aliases_names))) - .where(AssetEvent.source_dag_run.has(DagRun.dag_id == serialized_dag.dag_model.dag_id)) - .where(DagVersion.version_number == version_number) - ).unique() - - asset_id_to_task_ids = defaultdict(set) - for asset_event in result: - asset_id_to_task_ids[asset_event.asset_id].add(asset_event.source_task_id) - - for edge in downstream_alias_resolved_edges: - asset_id = int(edge["target_id"].replace("asset:", "", 1)) - task_ids = asset_id_to_task_ids.get(asset_id, set()) - - for index, task_id in enumerate(task_ids): - if index == 0: - edge["source_id"] = task_id - continue - edge_copy = {**edge, "source_id": task_id} - edges.append(edge_copy) diff --git a/airflow-core/src/airflow/ui/openapi-gen/queries/common.ts b/airflow-core/src/airflow/ui/openapi-gen/queries/common.ts index a8f08d3df78c0..a9ac8589e29b6 100644 --- a/airflow-core/src/airflow/ui/openapi-gen/queries/common.ts +++ b/airflow-core/src/airflow/ui/openapi-gen/queries/common.ts @@ -962,15 +962,14 @@ export const UseDeadlinesServiceGetDagDeadlineAlertsKeyFn = ({ dagId, limit, off export type StructureServiceStructureDataDefaultResponse = Awaited>; export type StructureServiceStructureDataQueryResult = UseQueryResult; export const useStructureServiceStructureDataKey = "StructureServiceStructureData"; -export const UseStructureServiceStructureDataKeyFn = ({ dagId, depth, externalDependencies, includeDownstream, includeUpstream, root, versionNumber }: { +export const UseStructureServiceStructureDataKeyFn = ({ dagId, depth, includeDownstream, includeUpstream, root, versionNumber }: { dagId: string; depth?: number; - externalDependencies?: boolean; includeDownstream?: boolean; includeUpstream?: boolean; root?: string; versionNumber?: number; -}, queryKey?: Array) => [useStructureServiceStructureDataKey, ...(queryKey ?? [{ dagId, depth, externalDependencies, includeDownstream, includeUpstream, root, versionNumber }])]; +}, queryKey?: Array) => [useStructureServiceStructureDataKey, ...(queryKey ?? [{ dagId, depth, includeDownstream, includeUpstream, root, versionNumber }])]; export type GridServiceGetDagStructureDefaultResponse = Awaited>; export type GridServiceGetDagStructureQueryResult = UseQueryResult; export const useGridServiceGetDagStructureKey = "GridServiceGetDagStructure"; diff --git a/airflow-core/src/airflow/ui/openapi-gen/queries/ensureQueryData.ts b/airflow-core/src/airflow/ui/openapi-gen/queries/ensureQueryData.ts index 34eff20314070..2c67b7acc4483 100644 --- a/airflow-core/src/airflow/ui/openapi-gen/queries/ensureQueryData.ts +++ b/airflow-core/src/airflow/ui/openapi-gen/queries/ensureQueryData.ts @@ -1936,20 +1936,18 @@ export const ensureUseDeadlinesServiceGetDagDeadlineAlertsData = (queryClient: Q * @param data.includeDownstream * @param data.depth * @param data.root -* @param data.externalDependencies * @param data.versionNumber * @returns StructureDataResponse Successful Response * @throws ApiError */ -export const ensureUseStructureServiceStructureDataData = (queryClient: QueryClient, { dagId, depth, externalDependencies, includeDownstream, includeUpstream, root, versionNumber }: { +export const ensureUseStructureServiceStructureDataData = (queryClient: QueryClient, { dagId, depth, includeDownstream, includeUpstream, root, versionNumber }: { dagId: string; depth?: number; - externalDependencies?: boolean; includeDownstream?: boolean; includeUpstream?: boolean; root?: string; versionNumber?: number; -}) => queryClient.ensureQueryData({ queryKey: Common.UseStructureServiceStructureDataKeyFn({ dagId, depth, externalDependencies, includeDownstream, includeUpstream, root, versionNumber }), queryFn: () => StructureService.structureData({ dagId, depth, externalDependencies, includeDownstream, includeUpstream, root, versionNumber }) }); +}) => queryClient.ensureQueryData({ queryKey: Common.UseStructureServiceStructureDataKeyFn({ dagId, depth, includeDownstream, includeUpstream, root, versionNumber }), queryFn: () => StructureService.structureData({ dagId, depth, includeDownstream, includeUpstream, root, versionNumber }) }); /** * Get Dag Structure * Return dag structure for grid view. diff --git a/airflow-core/src/airflow/ui/openapi-gen/queries/prefetch.ts b/airflow-core/src/airflow/ui/openapi-gen/queries/prefetch.ts index 170a17f9e58ae..52e8944893dc7 100644 --- a/airflow-core/src/airflow/ui/openapi-gen/queries/prefetch.ts +++ b/airflow-core/src/airflow/ui/openapi-gen/queries/prefetch.ts @@ -1936,20 +1936,18 @@ export const prefetchUseDeadlinesServiceGetDagDeadlineAlerts = (queryClient: Que * @param data.includeDownstream * @param data.depth * @param data.root -* @param data.externalDependencies * @param data.versionNumber * @returns StructureDataResponse Successful Response * @throws ApiError */ -export const prefetchUseStructureServiceStructureData = (queryClient: QueryClient, { dagId, depth, externalDependencies, includeDownstream, includeUpstream, root, versionNumber }: { +export const prefetchUseStructureServiceStructureData = (queryClient: QueryClient, { dagId, depth, includeDownstream, includeUpstream, root, versionNumber }: { dagId: string; depth?: number; - externalDependencies?: boolean; includeDownstream?: boolean; includeUpstream?: boolean; root?: string; versionNumber?: number; -}) => queryClient.prefetchQuery({ queryKey: Common.UseStructureServiceStructureDataKeyFn({ dagId, depth, externalDependencies, includeDownstream, includeUpstream, root, versionNumber }), queryFn: () => StructureService.structureData({ dagId, depth, externalDependencies, includeDownstream, includeUpstream, root, versionNumber }) }); +}) => queryClient.prefetchQuery({ queryKey: Common.UseStructureServiceStructureDataKeyFn({ dagId, depth, includeDownstream, includeUpstream, root, versionNumber }), queryFn: () => StructureService.structureData({ dagId, depth, includeDownstream, includeUpstream, root, versionNumber }) }); /** * Get Dag Structure * Return dag structure for grid view. diff --git a/airflow-core/src/airflow/ui/openapi-gen/queries/queries.ts b/airflow-core/src/airflow/ui/openapi-gen/queries/queries.ts index 7efadc2e4f8fa..660bd8372a724 100644 --- a/airflow-core/src/airflow/ui/openapi-gen/queries/queries.ts +++ b/airflow-core/src/airflow/ui/openapi-gen/queries/queries.ts @@ -1936,20 +1936,18 @@ export const useDeadlinesServiceGetDagDeadlineAlerts = = unknown[]>({ dagId, depth, externalDependencies, includeDownstream, includeUpstream, root, versionNumber }: { +export const useStructureServiceStructureData = = unknown[]>({ dagId, depth, includeDownstream, includeUpstream, root, versionNumber }: { dagId: string; depth?: number; - externalDependencies?: boolean; includeDownstream?: boolean; includeUpstream?: boolean; root?: string; versionNumber?: number; -}, queryKey?: TQueryKey, options?: Omit, "queryKey" | "queryFn">) => useQuery({ queryKey: Common.UseStructureServiceStructureDataKeyFn({ dagId, depth, externalDependencies, includeDownstream, includeUpstream, root, versionNumber }, queryKey), queryFn: () => StructureService.structureData({ dagId, depth, externalDependencies, includeDownstream, includeUpstream, root, versionNumber }) as TData, ...options }); +}, queryKey?: TQueryKey, options?: Omit, "queryKey" | "queryFn">) => useQuery({ queryKey: Common.UseStructureServiceStructureDataKeyFn({ dagId, depth, includeDownstream, includeUpstream, root, versionNumber }, queryKey), queryFn: () => StructureService.structureData({ dagId, depth, includeDownstream, includeUpstream, root, versionNumber }) as TData, ...options }); /** * Get Dag Structure * Return dag structure for grid view. diff --git a/airflow-core/src/airflow/ui/openapi-gen/queries/suspense.ts b/airflow-core/src/airflow/ui/openapi-gen/queries/suspense.ts index d4e694554d0ff..cb2b93f6bf1bb 100644 --- a/airflow-core/src/airflow/ui/openapi-gen/queries/suspense.ts +++ b/airflow-core/src/airflow/ui/openapi-gen/queries/suspense.ts @@ -1936,20 +1936,18 @@ export const useDeadlinesServiceGetDagDeadlineAlertsSuspense = = unknown[]>({ dagId, depth, externalDependencies, includeDownstream, includeUpstream, root, versionNumber }: { +export const useStructureServiceStructureDataSuspense = = unknown[]>({ dagId, depth, includeDownstream, includeUpstream, root, versionNumber }: { dagId: string; depth?: number; - externalDependencies?: boolean; includeDownstream?: boolean; includeUpstream?: boolean; root?: string; versionNumber?: number; -}, queryKey?: TQueryKey, options?: Omit, "queryKey" | "queryFn">) => useSuspenseQuery({ queryKey: Common.UseStructureServiceStructureDataKeyFn({ dagId, depth, externalDependencies, includeDownstream, includeUpstream, root, versionNumber }, queryKey), queryFn: () => StructureService.structureData({ dagId, depth, externalDependencies, includeDownstream, includeUpstream, root, versionNumber }) as TData, ...options }); +}, queryKey?: TQueryKey, options?: Omit, "queryKey" | "queryFn">) => useSuspenseQuery({ queryKey: Common.UseStructureServiceStructureDataKeyFn({ dagId, depth, includeDownstream, includeUpstream, root, versionNumber }, queryKey), queryFn: () => StructureService.structureData({ dagId, depth, includeDownstream, includeUpstream, root, versionNumber }) as TData, ...options }); /** * Get Dag Structure * Return dag structure for grid view. diff --git a/airflow-core/src/airflow/ui/openapi-gen/requests/schemas.gen.ts b/airflow-core/src/airflow/ui/openapi-gen/requests/schemas.gen.ts index 6105a80121503..566ce3fccff9f 100644 --- a/airflow-core/src/airflow/ui/openapi-gen/requests/schemas.gen.ts +++ b/airflow-core/src/airflow/ui/openapi-gen/requests/schemas.gen.ts @@ -8608,6 +8608,18 @@ export const $BaseNodeResponse = { } ], title: 'Team' + }, + asset_condition_type: { + anyOf: [ + { + type: 'string', + enum: ['or-gate', 'and-gate'] + }, + { + type: 'null' + } + ], + title: 'Asset Condition Type' } }, type: 'object', @@ -9621,17 +9633,6 @@ export const $EdgeResponse = { } ], title: 'Label' - }, - is_source_asset: { - anyOf: [ - { - type: 'boolean' - }, - { - type: 'null' - } - ], - title: 'Is Source Asset' } }, type: 'object', @@ -10272,6 +10273,18 @@ export const $NodeResponse = { ], title: 'Team' }, + asset_condition_type: { + anyOf: [ + { + type: 'string', + enum: ['or-gate', 'and-gate'] + }, + { + type: 'null' + } + ], + title: 'Asset Condition Type' + }, children: { anyOf: [ { @@ -10331,18 +10344,6 @@ export const $NodeResponse = { ], title: 'Operator' }, - asset_condition_type: { - anyOf: [ - { - type: 'string', - enum: ['or-gate', 'and-gate'] - }, - { - type: 'null' - } - ], - title: 'Asset Condition Type' - }, ui_color: { anyOf: [ { diff --git a/airflow-core/src/airflow/ui/openapi-gen/requests/services.gen.ts b/airflow-core/src/airflow/ui/openapi-gen/requests/services.gen.ts index a53ae4d51cea8..c3d1c3bd17e62 100644 --- a/airflow-core/src/airflow/ui/openapi-gen/requests/services.gen.ts +++ b/airflow-core/src/airflow/ui/openapi-gen/requests/services.gen.ts @@ -4906,7 +4906,6 @@ export class StructureService { * @param data.includeDownstream * @param data.depth * @param data.root - * @param data.externalDependencies * @param data.versionNumber * @returns StructureDataResponse Successful Response * @throws ApiError @@ -4921,11 +4920,9 @@ export class StructureService { include_downstream: data.includeDownstream, depth: data.depth, root: data.root, - external_dependencies: data.externalDependencies, version_number: data.versionNumber }, errors: { - 400: 'Bad Request', 404: 'Not Found', 422: 'Validation Error' } diff --git a/airflow-core/src/airflow/ui/openapi-gen/requests/types.gen.ts b/airflow-core/src/airflow/ui/openapi-gen/requests/types.gen.ts index 901e20ccc13da..84910b563b267 100644 --- a/airflow-core/src/airflow/ui/openapi-gen/requests/types.gen.ts +++ b/airflow-core/src/airflow/ui/openapi-gen/requests/types.gen.ts @@ -2181,6 +2181,7 @@ export type BaseNodeResponse = { label: string; type: 'join' | 'task' | 'asset-condition' | 'asset' | 'asset-alias' | 'asset-name-ref' | 'asset-uri-ref' | 'dag' | 'sensor' | 'trigger'; team?: string | null; + asset_condition_type?: 'or-gate' | 'and-gate' | null; }; export type type = 'join' | 'task' | 'asset-condition' | 'asset' | 'asset-alias' | 'asset-name-ref' | 'asset-uri-ref' | 'dag' | 'sensor' | 'trigger'; @@ -2454,7 +2455,6 @@ export type EdgeResponse = { target_id: string; is_setup_teardown?: boolean | null; label?: string | null; - is_source_asset?: boolean | null; }; /** @@ -2617,12 +2617,12 @@ export type NodeResponse = { label: string; type: 'join' | 'task' | 'asset-condition' | 'asset' | 'asset-alias' | 'asset-name-ref' | 'asset-uri-ref' | 'dag' | 'sensor' | 'trigger'; team?: string | null; + asset_condition_type?: 'or-gate' | 'and-gate' | null; children?: Array | null; is_mapped?: boolean | null; tooltip?: string | null; setup_teardown_type?: 'setup' | 'teardown' | null; operator?: string | null; - asset_condition_type?: 'or-gate' | 'and-gate' | null; ui_color?: string | null; ui_fgcolor?: string | null; }; @@ -4634,7 +4634,6 @@ export type GetDagDeadlineAlertsResponse = DeadlineAlertCollectionResponse; export type StructureDataData = { dagId: string; depth?: number | null; - externalDependencies?: boolean; includeDownstream?: boolean; includeUpstream?: boolean; root?: string | null; @@ -8487,10 +8486,6 @@ export type $OpenApiTs = { * Successful Response */ 200: StructureDataResponse; - /** - * Bad Request - */ - 400: HTTPExceptionResponse; /** * Not Found */ diff --git a/airflow-core/src/airflow/ui/public/i18n/locales/en/dag.json b/airflow-core/src/airflow/ui/public/i18n/locales/en/dag.json index 907fba9e5eac4..3f2916e030c7d 100644 --- a/airflow-core/src/airflow/ui/public/i18n/locales/en/dag.json +++ b/airflow-core/src/airflow/ui/public/i18n/locales/en/dag.json @@ -143,13 +143,7 @@ "label": "Number of Dag Runs" }, "dependencies": { - "label": "Dependencies", - "options": { - "allDagDependencies": "All Dag Dependencies", - "externalConditions": "External conditions", - "onlyTasks": "Only tasks" - }, - "placeholder": "Dependencies" + "allDagDependencies": "All Dag Dependencies" }, "graphDirection": { "label": "Graph Direction" diff --git a/airflow-core/src/airflow/ui/src/components/Graph/Edge.tsx b/airflow-core/src/airflow/ui/src/components/Graph/Edge.tsx index a4afa95c168bf..c2e19448143ad 100644 --- a/airflow-core/src/airflow/ui/src/components/Graph/Edge.tsx +++ b/airflow-core/src/airflow/ui/src/components/Graph/Edge.tsx @@ -39,13 +39,18 @@ const CustomEdge = ({ data, source, target }: Props) => { // don't require the parent to rebuild and pass down a new edges array. // useNodesData subscribes to data changes for these specific node IDs only. const nodesData = useNodesData([source, target]); - const isSelected = nodesData.some((node) => Boolean(node.data.isSelected)); if (data === undefined) { return undefined; } const { rest } = data; + // rest.isSelected covers what node-level selection can't express on its own: a gate + // (asset-condition) node is shared by every asset in its AND/OR condition, so marking the gate + // itself selected would highlight every edge touching it, not just the path relevant to the + // currently selected asset/Dag. See getGatePathEdgeIdsForSelection. + const isSelected = Boolean(rest.isSelected) || nodesData.some((node) => Boolean(node.data.isSelected)); + const edgeStrokeColor = isSelected ? (rest.edgeType === "data" ? dataEdgeColor : blueColor) : strokeColor; return ( diff --git a/airflow-core/src/airflow/ui/src/components/Graph/reactflowUtils.test.ts b/airflow-core/src/airflow/ui/src/components/Graph/reactflowUtils.test.ts new file mode 100644 index 0000000000000..10bc5d81f3c17 --- /dev/null +++ b/airflow-core/src/airflow/ui/src/components/Graph/reactflowUtils.test.ts @@ -0,0 +1,118 @@ +/*! + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ +import { describe, expect, it } from "vitest"; + +import { getGatePathEdgeIdsForSelection } from "./reactflowUtils"; + +const node = (id: string, type: string) => ({ id, type }); +const edge = (id: string, source: string, target: string) => ({ id, source, target }); + +describe("getGatePathEdgeIdsForSelection", () => { + it("returns the edges on the path from a selected asset through a gate to the dag", () => { + const nodes = [node("asset:1", "asset"), node("gate-0", "asset-condition"), node("dag:my_dag", "dag")]; + const edges = [edge("e1", "asset:1", "gate-0"), edge("e2", "gate-0", "dag:my_dag")]; + + const result = getGatePathEdgeIdsForSelection(nodes, edges, (id) => id === "asset:1"); + + expect(result).toEqual(new Set(["e1", "e2"])); + }); + + it("does not highlight a sibling asset's edge in the same AND/OR condition", () => { + // asset:1 and asset:2 both feed and-gate-0, which schedules dag:my_dag. Selecting asset:1 + // must only highlight its own edge into the gate and the gate's single output -- not + // asset:2's edge, even though asset:2 shares the same gate. + const nodes = [ + node("asset:1", "asset"), + node("asset:2", "asset"), + node("and-gate-0", "asset-condition"), + node("dag:my_dag", "dag"), + ]; + const edges = [ + edge("e1", "asset:1", "and-gate-0"), + edge("e2", "asset:2", "and-gate-0"), + edge("e3", "and-gate-0", "dag:my_dag"), + ]; + + const result = getGatePathEdgeIdsForSelection(nodes, edges, (id) => id === "asset:1"); + + expect(result).toEqual(new Set(["e1", "e3"])); + expect(result.has("e2")).toBe(false); + }); + + it("fans out to every input when the Dag (output side) is selected", () => { + // From the Dag's perspective, all of the gate's inputs are genuinely relevant. + const nodes = [ + node("asset:1", "asset"), + node("asset:2", "asset"), + node("and-gate-0", "asset-condition"), + node("dag:my_dag", "dag"), + ]; + const edges = [ + edge("e1", "asset:1", "and-gate-0"), + edge("e2", "asset:2", "and-gate-0"), + edge("e3", "and-gate-0", "dag:my_dag"), + ]; + + const result = getGatePathEdgeIdsForSelection(nodes, edges, (id) => id === "dag:my_dag"); + + expect(result).toEqual(new Set(["e1", "e2", "e3"])); + }); + + it("walks forward through nested gates without fanning into a sibling gate's inputs", () => { + // asset:1 feeds or-gate-1, which is one input (among others) to and-gate-0, which schedules + // the Dag. Selecting asset:1 should light up its single path all the way to the Dag, but not + // and-gate-0's other inputs. + const nodes = [ + node("asset:1", "asset"), + node("asset:2", "asset"), + node("or-gate-1", "asset-condition"), + node("and-gate-0", "asset-condition"), + node("dag:my_dag", "dag"), + ]; + const edges = [ + edge("e1", "asset:1", "or-gate-1"), + edge("e2", "or-gate-1", "and-gate-0"), + edge("e3", "asset:2", "and-gate-0"), + edge("e4", "and-gate-0", "dag:my_dag"), + ]; + + const result = getGatePathEdgeIdsForSelection(nodes, edges, (id) => id === "asset:1"); + + expect(result).toEqual(new Set(["e1", "e2", "e4"])); + expect(result.has("e3")).toBe(false); + }); + + it("returns an empty set when the graph has no gate nodes", () => { + const nodes = [node("asset:1", "asset"), node("dag:my_dag", "dag")]; + const edges = [edge("e1", "asset:1", "dag:my_dag")]; + + const result = getGatePathEdgeIdsForSelection(nodes, edges, (id) => id === "asset:1"); + + expect(result).toEqual(new Set()); + }); + + it("returns an empty set when nothing is selected", () => { + const nodes = [node("asset:1", "asset"), node("gate-0", "asset-condition"), node("dag:my_dag", "dag")]; + const edges = [edge("e1", "asset:1", "gate-0"), edge("e2", "gate-0", "dag:my_dag")]; + + const result = getGatePathEdgeIdsForSelection(nodes, edges, () => false); + + expect(result).toEqual(new Set()); + }); +}); diff --git a/airflow-core/src/airflow/ui/src/components/Graph/reactflowUtils.ts b/airflow-core/src/airflow/ui/src/components/Graph/reactflowUtils.ts index 7a2667a7abc23..93a12e550a0b1 100644 --- a/airflow-core/src/airflow/ui/src/components/Graph/reactflowUtils.ts +++ b/airflow-core/src/airflow/ui/src/components/Graph/reactflowUtils.ts @@ -141,3 +141,108 @@ export const formatFlowEdges = ({ edges }: { edges: Array }): Array, + edges: Array, + isNodeSelected: (nodeId: string) => boolean, +): Set => { + const gateNodeIds = new Set(nodes.filter((node) => node.type === "asset-condition").map((node) => node.id)); + + if (gateNodeIds.size === 0) { + return new Set(); + } + + const outgoingByNode = new Map>(); + const incomingByNode = new Map>(); + const addTo = (map: Map>, key: string, edge: SelectionGraphEdge) => { + const existing = map.get(key); + + if (existing) { + existing.push(edge); + } else { + map.set(key, [edge]); + } + }; + + edges.forEach((edge) => { + addTo(outgoingByNode, edge.source, edge); + addTo(incomingByNode, edge.target, edge); + }); + + const selectedEdgeIds = new Set(); + const visitedForward = new Set(); + const visitedBackward = new Set(); + + const walkForward = (gateId: string) => { + if (visitedForward.has(gateId)) { + return; + } + visitedForward.add(gateId); + + (outgoingByNode.get(gateId) ?? []).forEach((edge) => { + selectedEdgeIds.add(edge.id); + + if (gateNodeIds.has(edge.target)) { + walkForward(edge.target); + } + }); + }; + + const walkBackward = (gateId: string) => { + if (visitedBackward.has(gateId)) { + return; + } + visitedBackward.add(gateId); + + (incomingByNode.get(gateId) ?? []).forEach((edge) => { + selectedEdgeIds.add(edge.id); + + if (gateNodeIds.has(edge.source)) { + walkBackward(edge.source); + } + }); + }; + + nodes.forEach((node) => { + if (!isNodeSelected(node.id)) { + return; + } + + (outgoingByNode.get(node.id) ?? []).forEach((edge) => { + selectedEdgeIds.add(edge.id); + + if (gateNodeIds.has(edge.target)) { + walkForward(edge.target); + } + }); + + (incomingByNode.get(node.id) ?? []).forEach((edge) => { + selectedEdgeIds.add(edge.id); + + if (gateNodeIds.has(edge.source)) { + walkBackward(edge.source); + } + }); + }); + + return selectedEdgeIds; +}; diff --git a/airflow-core/src/airflow/ui/src/constants/localStorage.test.ts b/airflow-core/src/airflow/ui/src/constants/localStorage.test.ts new file mode 100644 index 0000000000000..b485f9420c801 --- /dev/null +++ b/airflow-core/src/airflow/ui/src/constants/localStorage.test.ts @@ -0,0 +1,51 @@ +/*! + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ +import { afterEach, describe, expect, it } from "vitest"; + +import { pruneLegacyDependencyKeys, SHOW_ALL_DEPENDENCIES_KEY } from "./localStorage"; + +describe("pruneLegacyDependencyKeys", () => { + afterEach(() => { + globalThis.localStorage.clear(); + }); + + it("removes only the legacy per-Dag dependency keys, leaving everything else intact", () => { + globalThis.localStorage.setItem("dependencies-my_dag", "all"); + globalThis.localStorage.setItem("dependencies-other_dag", "tasks"); + globalThis.localStorage.setItem(SHOW_ALL_DEPENDENCIES_KEY, "true"); + globalThis.localStorage.setItem("direction-my_dag", "RIGHT"); + + pruneLegacyDependencyKeys(); + + expect(globalThis.localStorage.getItem("dependencies-my_dag")).toBeNull(); + expect(globalThis.localStorage.getItem("dependencies-other_dag")).toBeNull(); + // The global toggle and unrelated Dag-scoped keys are untouched. + expect(globalThis.localStorage.getItem(SHOW_ALL_DEPENDENCIES_KEY)).toBe("true"); + expect(globalThis.localStorage.getItem("direction-my_dag")).toBe("RIGHT"); + }); + + it("is a no-op when there are no legacy keys", () => { + globalThis.localStorage.setItem(SHOW_ALL_DEPENDENCIES_KEY, "false"); + + pruneLegacyDependencyKeys(); + + expect(globalThis.localStorage.getItem(SHOW_ALL_DEPENDENCIES_KEY)).toBe("false"); + expect(globalThis.localStorage.length).toBe(1); + }); +}); diff --git a/airflow-core/src/airflow/ui/src/constants/localStorage.ts b/airflow-core/src/airflow/ui/src/constants/localStorage.ts index 72ff5b8327b4c..9e366796fd6c8 100644 --- a/airflow-core/src/airflow/ui/src/constants/localStorage.ts +++ b/airflow-core/src/airflow/ui/src/constants/localStorage.ts @@ -28,10 +28,10 @@ export const LOG_SHOW_TIMESTAMP_KEY = "log_show_timestamp"; export const LOG_SHOW_SOURCE_KEY = "log_show_source"; export const VERSION_INDICATOR_DISPLAY_MODE_KEY = "version_indicator_display_mode"; export const COLLAPSED_UI_ALERTS_KEY = "collapsed_ui_alerts"; +export const SHOW_ALL_DEPENDENCIES_KEY = "show_all_dependencies"; // Dag-scoped keys export const dagRunsLimitKey = (dagId: string) => `dag_runs_limit-${dagId}`; -export const dependenciesKey = (dagId: string) => `dependencies-${dagId}`; export const directionKey = (dagId: string) => `direction-${dagId}`; export const openGroupsKey = (dagId: string) => `${dagId}/open-groups`; export const allGroupsKey = (dagId: string) => `${dagId}/all-groups`; @@ -45,3 +45,20 @@ export const presetFiltersDefaultKey = (pageName: string) => // SearchBar advanced (substring) toggle, scoped per searchbar via a caller-provided id. export const advancedSearchKey = (id: string) => `advanced_search-${id}`; + +// One-time cleanup of the pre-consolidation per-Dag dependency toggle keys +// (`dependencies-`), now superseded by the global SHOW_ALL_DEPENDENCIES_KEY. Without this +// they linger in every user's localStorage forever. Safe because no current key shares the prefix. +export const pruneLegacyDependencyKeys = (storage: Storage = globalThis.localStorage): void => { + const staleKeys: Array = []; + + for (let index = 0; index < storage.length; index += 1) { + const key = storage.key(index); + + if (key?.startsWith("dependencies-")) { + staleKeys.push(key); + } + } + + staleKeys.forEach((key) => storage.removeItem(key)); +}; diff --git a/airflow-core/src/airflow/ui/src/constants/searchParams.ts b/airflow-core/src/airflow/ui/src/constants/searchParams.ts index b6f4d88e3d0e5..fdde83130a925 100644 --- a/airflow-core/src/airflow/ui/src/constants/searchParams.ts +++ b/airflow-core/src/airflow/ui/src/constants/searchParams.ts @@ -37,7 +37,6 @@ export enum SearchParamsKeys { DEADLINE_TIME_GTE = "deadline_time_gte", DEADLINE_TIME_LTE = "deadline_time_lte", DEADLINE_TIME_RANGE = "deadline_time_range", - DEPENDENCIES = "dependencies", DURATION_GTE = "duration_gte", DURATION_LTE = "duration_lte", END_DATE = "end_date", diff --git a/airflow-core/src/airflow/ui/src/context/groups/GroupsProvider.tsx b/airflow-core/src/airflow/ui/src/context/groups/GroupsProvider.tsx index 2f486f15a2c6a..f467009f05c1c 100644 --- a/airflow-core/src/airflow/ui/src/context/groups/GroupsProvider.tsx +++ b/airflow-core/src/airflow/ui/src/context/groups/GroupsProvider.tsx @@ -21,7 +21,7 @@ import { useDebouncedCallback } from "use-debounce"; import { useLocalStorage } from "usehooks-ts"; import { useStructureServiceStructureData } from "openapi/queries"; -import { allGroupsKey, dependenciesKey, openGroupsKey } from "src/constants/localStorage"; +import { allGroupsKey, openGroupsKey } from "src/constants/localStorage"; import useSelectedVersion from "src/hooks/useSelectedVersion"; import { flattenGraphNodes } from "src/layouts/Details/Grid/utils"; @@ -42,12 +42,10 @@ export const GroupsProvider = ({ children, dagId }: Props) => { }, [allGroupIds]); const selectedVersion = useSelectedVersion(); - const [dependencies] = useLocalStorage<"all" | "immediate" | "tasks">(dependenciesKey(dagId), "tasks"); const { data: structure = { edges: [], nodes: [] } } = useStructureServiceStructureData( { dagId, - externalDependencies: dependencies === "immediate", versionNumber: selectedVersion, }, undefined, diff --git a/airflow-core/src/airflow/ui/src/layouts/Details/Graph/Graph.tsx b/airflow-core/src/airflow/ui/src/layouts/Details/Graph/Graph.tsx index 8c2ff6574530d..73d92d89bfc1a 100644 --- a/airflow-core/src/airflow/ui/src/layouts/Details/Graph/Graph.tsx +++ b/airflow-core/src/airflow/ui/src/layouts/Details/Graph/Graph.tsx @@ -27,9 +27,9 @@ import { useDagRunServiceGetDagRun, useStructureServiceStructureData } from "ope import type { Direction } from "src/components/Graph/DirectionDropdown"; import { DownloadButton } from "src/components/Graph/DownloadButton"; import { edgeTypes, nodeTypes } from "src/components/Graph/graphTypes"; -import type { CustomNodeProps } from "src/components/Graph/reactflowUtils"; +import { getGatePathEdgeIdsForSelection, type CustomNodeProps } from "src/components/Graph/reactflowUtils"; import { useGraphLayout } from "src/components/Graph/useGraphLayout"; -import { dependenciesKey, directionKey } from "src/constants/localStorage"; +import { SHOW_ALL_DEPENDENCIES_KEY, directionKey } from "src/constants/localStorage"; import { useColorMode } from "src/context/colorMode"; import { useGroups } from "src/context/groups"; import useSelectedVersion from "src/hooks/useSelectedVersion"; @@ -70,7 +70,7 @@ export const Graph = () => { const { allGroupIds, openGroupIds, setAllGroupIds } = useGroups(); - const [dependencies] = useLocalStorage<"all" | "immediate" | "tasks">(dependenciesKey(dagId), "tasks"); + const [showAllDependencies] = useLocalStorage(SHOW_ALL_DEPENDENCIES_KEY, false); const [direction] = useLocalStorage(directionKey(dagId), "RIGHT"); const selectedColor = colorMode === "dark" ? selectedDarkColor : selectedLightColor; @@ -78,7 +78,6 @@ export const Graph = () => { { dagId, depth, - externalDependencies: dependencies === "immediate", includeDownstream, includeUpstream, root: hasActiveFilter && filterRoot !== undefined ? filterRoot : undefined, @@ -100,17 +99,17 @@ export const Graph = () => { }, [allGroupIds, graphData.nodes, setAllGroupIds]); const { data: dagDependencies = { edges: [], nodes: [] } } = useDependencyGraph(`dag:${dagId}`, { - enabled: dependencies === "all", + enabled: showAllDependencies, }); - const dagDepEdges = dependencies === "all" ? dagDependencies.edges : []; - const dagDepNodes = dependencies === "all" ? dagDependencies.nodes : []; + const dagDepEdges = showAllDependencies ? dagDependencies.edges : []; + const dagDepNodes = showAllDependencies ? dagDependencies.nodes : []; const layoutEdges = [...graphData.edges, ...dagDepEdges]; const layoutNodes = dagDepNodes.length ? dagDepNodes.map((node) => (node.id === `dag:${dagId}` ? { ...node, children: graphData.nodes } : node)) : graphData.nodes; - const layoutOpenGroupIds = [...openGroupIds, ...(dependencies === "all" ? [`dag:${dagId}`] : [])]; + const layoutOpenGroupIds = [...openGroupIds, ...(showAllDependencies ? [`dag:${dagId}`] : [])]; const { data, isPending } = useGraphLayout({ direction, @@ -131,6 +130,9 @@ export const Graph = () => { }); const gridTISummaries = runId ? summariesByRunId.get(runId) : undefined; + const isNodeSelected = (nodeId: string) => + nodeId === taskId || nodeId === groupId || nodeId === `dag:${dagId}`; + // Add task instances to the node data but without having to recalculate how the graph is laid out const nodesWithTI = data?.nodes.map((node) => { const taskInstance = gridTISummaries?.task_instances.find((ti) => ti.task_id === node.id); @@ -139,7 +141,7 @@ export const Graph = () => { ...node, data: { ...node.data, - isSelected: node.id === taskId || node.id === groupId || node.id === `dag:${dagId}`, + isSelected: isNodeSelected(node.id), taskInstance, }, }; @@ -147,7 +149,7 @@ export const Graph = () => { const baseFilteredNodes = useGraphFilteredNodes(nodesWithTI, graphFilters); - const { edges, nodes } = useFilteredNodesAndEdges({ + const { edges: filteredEdges, nodes } = useFilteredNodesAndEdges({ baseFilteredNodes, dagId, groupId, @@ -155,6 +157,18 @@ export const Graph = () => { taskId, }); + const gatePathEdgeIds = data + ? getGatePathEdgeIdsForSelection(data.nodes, data.edges, isNodeSelected) + : new Set(); + const edges = + gatePathEdgeIds.size > 0 + ? filteredEdges.map((edge) => + gatePathEdgeIds.has(edge.id) + ? { ...edge, data: { ...edge.data, rest: { ...edge.data?.rest, isSelected: true } } } + : edge, + ) + : filteredEdges; + const selectedNodeId = taskId ?? groupId; return ( diff --git a/airflow-core/src/airflow/ui/src/layouts/Details/PanelButtons.tsx b/airflow-core/src/airflow/ui/src/layouts/Details/PanelButtons.tsx index 6109d693190dc..8d8fa30757d89 100644 --- a/airflow-core/src/airflow/ui/src/layouts/Details/PanelButtons.tsx +++ b/airflow-core/src/airflow/ui/src/layouts/Details/PanelButtons.tsx @@ -39,10 +39,10 @@ import { useLocalStorage } from "usehooks-ts"; import { DagVersionSelect } from "src/components/DagVersionSelect"; import { DirectionDropdown } from "src/components/Graph/DirectionDropdown"; import { GraphTaskFilters } from "src/components/GraphTaskFilters"; -import { IconButton } from "src/components/ui"; +import { IconButton, Switch } from "src/components/ui"; import { type ButtonGroupOption, ButtonGroupToggle } from "src/components/ui/ButtonGroupToggle"; import type { DagView } from "src/constants/dagView"; -import { dependenciesKey } from "src/constants/localStorage"; +import { SHOW_ALL_DEPENDENCIES_KEY } from "src/constants/localStorage"; import type { VersionIndicatorOptions } from "src/constants/showVersionIndicatorOptions"; import { SHORTCUTS } from "src/context/keyboardShortcuts"; import { useShortcut } from "src/hooks/useShortcut"; @@ -65,15 +65,6 @@ type Props = { readonly showVersionIndicatorMode: VersionIndicatorOptions; }; -const getOptions = (translate: (key: string) => string) => - createListCollection({ - items: [ - { label: translate("dag:panel.dependencies.options.onlyTasks"), value: "tasks" }, - { label: translate("dag:panel.dependencies.options.externalConditions"), value: "immediate" }, - { label: translate("dag:panel.dependencies.options.allDagDependencies"), value: "all" }, - ], - }); - const getWidthBasedConfig = (width: number, enableResponsiveOptions: boolean) => { const breakpoints = enableResponsiveOptions ? [ @@ -94,10 +85,6 @@ const getWidthBasedConfig = (width: number, enableResponsiveOptions: boolean) => }; }; -const deps = ["all", "immediate", "tasks"]; - -type Dependency = (typeof deps)[number]; - export const PanelButtons = ({ dagView, limit, @@ -111,9 +98,9 @@ export const PanelButtons = ({ const { dagId = "", runId } = useParams(); const { fitView } = useReactFlow(); const shouldShowToggleButtons = Boolean(runId); - const [dependencies, setDependencies, removeDependencies] = useLocalStorage( - dependenciesKey(dagId), - "tasks", + const [showAllDependencies, setShowAllDependencies] = useLocalStorage( + SHOW_ALL_DEPENDENCIES_KEY, + false, ); const containerRef = useRef(null); const containerWidth = useContainerWidth(containerRef); @@ -136,14 +123,6 @@ export const PanelButtons = ({ } }, [defaultLimit, enableResponsiveOptions, limit, setLimit]); - const handleDepsChange = (event: SelectValueChangeDetails<{ label: string; value: Array }>) => { - if (event.value[0] === undefined || event.value[0] === "tasks" || !deps.includes(event.value[0])) { - removeDependencies(); - } else { - setDependencies(event.value[0]); - } - }; - const handleFocus = (view: string) => { if (panelGroupRef.current) { const newLayout = view === "graph" ? [70, 30] : [30, 70]; @@ -231,35 +210,13 @@ export const PanelButtons = ({ - setShowAllDependencies(details.checked)} > - - {translate("dag:panel.dependencies.label")} - - - - - - - - - - - - {getOptions(translate).items.map((option) => ( - - {option.label} - - ))} - - - + {translate("dag:panel.dependencies.allDagDependencies")} + diff --git a/airflow-core/src/airflow/ui/src/main.tsx b/airflow-core/src/airflow/ui/src/main.tsx index e1ef1be5257e9..f3f8f2b0cdbd2 100644 --- a/airflow-core/src/airflow/ui/src/main.tsx +++ b/airflow-core/src/airflow/ui/src/main.tsx @@ -29,6 +29,7 @@ import * as ReactRouterDOM from "react-router-dom"; import * as ReactJSXRuntime from "react/jsx-runtime"; import type { HTTPExceptionResponse } from "openapi/requests/types.gen"; +import { pruneLegacyDependencyKeys } from "src/constants/localStorage"; import { ChakraCustomProvider } from "src/context/ChakraCustomProvider"; import { ColorModeProvider } from "src/context/colorMode"; import { ShortcutRegistryProvider } from "src/context/keyboardShortcuts"; @@ -95,6 +96,8 @@ axios.interceptors.response.use( }, ); +pruneLegacyDependencyKeys(); + createRoot(document.querySelector("#root") as HTMLDivElement).render( diff --git a/airflow-core/src/airflow/ui/src/pages/Asset/AssetGraph.tsx b/airflow-core/src/airflow/ui/src/pages/Asset/AssetGraph.tsx index c0f5218e5c605..e49c4a5cea2fe 100644 --- a/airflow-core/src/airflow/ui/src/pages/Asset/AssetGraph.tsx +++ b/airflow-core/src/airflow/ui/src/pages/Asset/AssetGraph.tsx @@ -26,7 +26,7 @@ import type { AssetResponse } from "openapi/requests/types.gen"; import type { Direction } from "src/components/Graph/DirectionDropdown"; import { DownloadButton } from "src/components/Graph/DownloadButton"; import { edgeTypes, nodeTypes } from "src/components/Graph/graphTypes"; -import type { CustomNodeProps } from "src/components/Graph/reactflowUtils"; +import { getGatePathEdgeIdsForSelection, type CustomNodeProps } from "src/components/Graph/reactflowUtils"; import { useGraphLayout } from "src/components/Graph/useGraphLayout"; import { directionKey } from "src/constants/localStorage"; import { useColorMode } from "src/context/colorMode"; @@ -60,10 +60,16 @@ export const AssetGraph = ({ const selectedColor = colorMode === "dark" ? selectedDarkColor : selectedLightColor; + const isNodeSelected = (nodeId: string) => nodeId === `asset:${assetId}`; + const nodes = layoutData?.nodes.map((node) => - node.id === `asset:${assetId}` ? { ...node, data: { ...node.data, isSelected: true } } : node, + isNodeSelected(node.id) ? { ...node, data: { ...node.data, isSelected: true } } : node, ); + const gatePathEdgeIds = layoutData + ? getGatePathEdgeIdsForSelection(layoutData.nodes, layoutData.edges, isNodeSelected) + : new Set(); + const edges = (layoutData?.edges ?? []).map((edge) => ({ ...edge, data: { @@ -71,7 +77,10 @@ export const AssetGraph = ({ rest: { ...edge.data?.rest, edgeType: dependencyType, - isSelected: `asset:${asset?.id}` === edge.source || `asset:${asset?.id}` === edge.target, + isSelected: + `asset:${asset?.id}` === edge.source || + `asset:${asset?.id}` === edge.target || + gatePathEdgeIds.has(edge.id), }, }, })); diff --git a/airflow-core/tests/unit/api_fastapi/core_api/routes/ui/test_dependencies.py b/airflow-core/tests/unit/api_fastapi/core_api/routes/ui/test_dependencies.py index eca2fa0087541..3c266f856d9a1 100644 --- a/airflow-core/tests/unit/api_fastapi/core_api/routes/ui/test_dependencies.py +++ b/airflow-core/tests/unit/api_fastapi/core_api/routes/ui/test_dependencies.py @@ -23,6 +23,7 @@ from sqlalchemy import select from airflow.models.asset import AssetModel +from airflow.models.dag import DagModel from airflow.models.dagbundle import DagBundleModel from airflow.models.team import Team from airflow.providers.standard.operators.empty import EmptyOperator @@ -116,42 +117,49 @@ def expected_primary_component_response(asset1_id): "label": "downstream", "type": "dag", "team": None, - }, - { - "id": f"asset:{asset1_id}", - "label": "asset1", - "type": "asset", - "team": None, + "asset_condition_type": None, }, { "id": "sensor:other_dag:downstream:external_task_sensor", "label": "external_task_sensor", "type": "sensor", "team": None, + "asset_condition_type": None, }, { "id": "dag:external_trigger_dag_id", "label": "external_trigger_dag_id", "type": "dag", "team": None, + "asset_condition_type": None, }, { "id": "trigger:external_trigger_dag_id:downstream:trigger_dag_run_operator", "label": "trigger_dag_run_operator", "type": "trigger", "team": None, + "asset_condition_type": None, }, { "id": "dag:upstream", "label": "upstream", "type": "dag", "team": None, + "asset_condition_type": None, + }, + { + "id": f"asset:{asset1_id}", + "label": "asset1", + "type": "asset", + "team": None, + "asset_condition_type": None, }, { "id": "dag:other_dag", "label": "other_dag", "type": "dag", "team": None, + "asset_condition_type": None, }, ], } @@ -199,22 +207,25 @@ def expected_secondary_component_response(asset2_id): ], "nodes": [ { - "id": "dag:downstream_secondary", - "label": "downstream_secondary", + "id": "dag:upstream_secondary", + "label": "upstream_secondary", "type": "dag", "team": None, + "asset_condition_type": None, }, { "id": f"asset:{asset2_id}", "label": "asset2", "type": "asset", "team": None, + "asset_condition_type": None, }, { - "id": "dag:upstream_secondary", - "label": "upstream_secondary", + "id": "dag:downstream_secondary", + "label": "downstream_secondary", "type": "dag", "team": None, + "asset_condition_type": None, }, ], } @@ -223,7 +234,7 @@ def expected_secondary_component_response(asset2_id): class TestGetDependencies: @pytest.mark.usefixtures("make_primary_connected_component") def test_should_response_200(self, test_client, expected_primary_component_response): - with assert_queries_count(7): + with assert_queries_count(8): response = test_client.get("/dependencies") assert response.status_code == 200 @@ -259,7 +270,7 @@ def test_delete_dag_should_response_403(self, unauthorized_test_client): @pytest.mark.usefixtures("make_primary_connected_component", "make_secondary_connected_component") def test_with_node_id_filter(self, test_client, node_id, expected_response_fixture, request): expected_response = request.getfixturevalue(expected_response_fixture) - with assert_queries_count(7): + with assert_queries_count(8): response = test_client.get("/dependencies", params={"node_id": node_id}) assert response.status_code == 200 @@ -277,7 +288,7 @@ def test_with_node_id_filter_with_asset( (asset1_id, expected_primary_component_response), (asset2_id, expected_secondary_component_response), ): - with assert_queries_count(7): + with assert_queries_count(8): response = test_client.get("/dependencies", params={"node_id": f"asset:{asset_id}"}) assert response.status_code == 200 @@ -581,3 +592,330 @@ def test_trigger_target_without_scheduling_dependencies_is_materialised(self, da ) actual_edges = {(edge["source_id"], edge["target_id"]) for edge in result["edges"]} assert expected_edge in actual_edges + + def test_scheduling_dependencies_expands_asset_and_gate(self, dag_maker, test_client, session): + """A Dag scheduled by multiple ANDed assets renders an and-gate node instead of flat asset edges.""" + asset_a = Asset(uri="s3://gate-bucket/a", name="gate_asset_a") + asset_b = Asset(uri="s3://gate-bucket/b", name="gate_asset_b") + + with dag_maker(dag_id="gate_upstream", serialized=True, session=session): + EmptyOperator(task_id="produce", outlets=[asset_a, asset_b]) + + with dag_maker( + dag_id="gate_downstream", + schedule=(asset_a & asset_b), + serialized=True, + session=session, + ): + EmptyOperator(task_id="consume") + + dag_maker.sync_dagbag_to_db() + + asset_a_id = session.scalar(select(AssetModel.id).where(AssetModel.name == "gate_asset_a")) + asset_b_id = session.scalar(select(AssetModel.id).where(AssetModel.name == "gate_asset_b")) + + response = test_client.get("/dependencies", params={"node_id": "dag:gate_downstream"}) + assert response.status_code == 200 + + result = response.json() + gate_node = next(node for node in result["nodes"] if node["type"] == "asset-condition") + assert gate_node["asset_condition_type"] == "and-gate" + + edge_tuples = {(edge["source_id"], edge["target_id"]) for edge in result["edges"]} + assert (gate_node["id"], "dag:gate_downstream") in edge_tuples + assert (f"asset:{asset_a_id}", gate_node["id"]) in edge_tuples + assert (f"asset:{asset_b_id}", gate_node["id"]) in edge_tuples + + # The flat asset -> dag edge is replaced by the gate, not duplicated alongside it. + assert (f"asset:{asset_a_id}", "dag:gate_downstream") not in edge_tuples + assert (f"asset:{asset_b_id}", "dag:gate_downstream") not in edge_tuples + + # The producer side (unrelated to the consumer's gate) is untouched. + assert ("dag:gate_upstream", f"asset:{asset_a_id}") in edge_tuples + assert ("dag:gate_upstream", f"asset:{asset_b_id}") in edge_tuples + + def test_scheduling_dependencies_expands_nested_boolean_groups(self, dag_maker, test_client, session): + """`(a & b) | (c & d)` renders both and-gates under the or-gate -- no sibling branch dropped.""" + asset_a = Asset(uri="s3://nested-bucket/a", name="nested_asset_a") + asset_b = Asset(uri="s3://nested-bucket/b", name="nested_asset_b") + asset_c = Asset(uri="s3://nested-bucket/c", name="nested_asset_c") + asset_d = Asset(uri="s3://nested-bucket/d", name="nested_asset_d") + + with dag_maker( + dag_id="nested_downstream", + schedule=((asset_a & asset_b) | (asset_c & asset_d)), + serialized=True, + session=session, + ): + EmptyOperator(task_id="consume") + + dag_maker.sync_dagbag_to_db() + + asset_ids = { + name: session.scalar(select(AssetModel.id).where(AssetModel.name == name)) + for name in ("nested_asset_a", "nested_asset_b", "nested_asset_c", "nested_asset_d") + } + + response = test_client.get("/dependencies", params={"node_id": "dag:nested_downstream"}) + assert response.status_code == 200 + + result = response.json() + gates = [node for node in result["nodes"] if node["type"] == "asset-condition"] + or_gates = [g for g in gates if g["asset_condition_type"] == "or-gate"] + and_gates = [g for g in gates if g["asset_condition_type"] == "and-gate"] + # Both nested and-groups survive: one or-gate, two distinct and-gates. + assert len(or_gates) == 1 + assert len(and_gates) == 2 + assert and_gates[0]["id"] != and_gates[1]["id"] + + edge_tuples = {(edge["source_id"], edge["target_id"]) for edge in result["edges"]} + or_gate_id = or_gates[0]["id"] + assert (or_gate_id, "dag:nested_downstream") in edge_tuples + # Each and-gate feeds the or-gate, and each is fed by exactly its two assets. + for and_gate in and_gates: + assert (and_gate["id"], or_gate_id) in edge_tuples + feeders = {src for src, tgt in edge_tuples if tgt == and_gate["id"]} + assert feeders in ( + {f"asset:{asset_ids['nested_asset_a']}", f"asset:{asset_ids['nested_asset_b']}"}, + {f"asset:{asset_ids['nested_asset_c']}", f"asset:{asset_ids['nested_asset_d']}"}, + ) + # All four assets are represented, none dropped. + all_feeders = {src for src, tgt in edge_tuples if tgt in {g["id"] for g in and_gates}} + assert all_feeders == {f"asset:{aid}" for aid in asset_ids.values()} + + def test_scheduling_dependencies_expands_gate_with_unresolved_asset_ref( + self, dag_maker, test_client, session + ): + """An AND/OR condition mixing a resolved asset with an unresolved Asset.ref(...) must not 500. + + Asset.ref(...) serializes its half of the expression as {"asset_ref": {"name": ...}}, + distinct from the {"asset-name-ref": ...} shape the rest of get_upstream_assets expects. + """ + asset_a = Asset(uri="s3://gate-bucket/ref-a", name="gate_asset_ref_a") + + with dag_maker( + dag_id="gate_downstream_with_ref", + schedule=(asset_a & Asset.ref(name="not_yet_registered_asset")), + serialized=True, + session=session, + ): + EmptyOperator(task_id="consume") + + dag_maker.sync_dagbag_to_db() + + asset_a_id = session.scalar(select(AssetModel.id).where(AssetModel.name == "gate_asset_ref_a")) + + response = test_client.get("/dependencies", params={"node_id": "dag:gate_downstream_with_ref"}) + assert response.status_code == 200 + + result = response.json() + nodes_by_id = {node["id"]: node for node in result["nodes"]} + assert nodes_by_id["asset-name-ref:not_yet_registered_asset"]["type"] == "asset-name-ref" + + edge_tuples = {(edge["source_id"], edge["target_id"]) for edge in result["edges"]} + gate_node_id = next(node["id"] for node in result["nodes"] if node["type"] == "asset-condition") + assert (f"asset:{asset_a_id}", gate_node_id) in edge_tuples + assert ("asset-name-ref:not_yet_registered_asset", gate_node_id) in edge_tuples + + def test_scheduling_dependencies_survives_unrecognized_asset_expression( + self, dag_maker, test_client, session + ): + """A Dag whose asset_expression get_upstream_assets can't parse must not break the whole endpoint. + + The gate expansion falls back to that one Dag's flat asset edge instead of 500ing for + every caller. + """ + asset_a = Asset(uri="s3://gate-bucket/broken", name="gate_asset_broken") + + with dag_maker( + dag_id="gate_downstream_broken", + schedule=[asset_a], + serialized=True, + session=session, + ): + EmptyOperator(task_id="consume") + + dag_maker.sync_dagbag_to_db() + + dag_model = session.scalar(select(DagModel).where(DagModel.dag_id == "gate_downstream_broken")) + dag_model.asset_expression = {"any": [{"weird-op": {}}]} + session.commit() + + asset_a_id = session.scalar(select(AssetModel.id).where(AssetModel.name == "gate_asset_broken")) + + response = test_client.get("/dependencies", params={"node_id": "dag:gate_downstream_broken"}) + assert response.status_code == 200 + + result = response.json() + assert not any(node["type"] == "asset-condition" for node in result["nodes"]) + edge_tuples = {(edge["source_id"], edge["target_id"]) for edge in result["edges"]} + assert (f"asset:{asset_a_id}", "dag:gate_downstream_broken") in edge_tuples + + def test_data_dependencies_routes_through_dag_level_schedule_without_task_inlet( + self, dag_maker, test_client, session + ): + """An asset that schedules a Dag via `schedule=`, with no task declaring it as an inlet, + must still connect to that Dag's entry task in the data-dependencies (Task Dependencies) graph. + + Uses the list form (`schedule=[asset_a]`), since that's the common way a single-asset + schedule is written -- and, like the bare-asset form, it must not render a redundant + single-input AND/OR gate. + """ + asset_a = Asset(uri="s3://data-dep-bucket/a", name="data_dep_asset_a") + + with dag_maker( + dag_id="data_dep_downstream", + schedule=[asset_a], + serialized=True, + session=session, + ): + EmptyOperator(task_id="entry_task") + + dag_maker.sync_dagbag_to_db() + + asset_a_id = session.scalar(select(AssetModel.id).where(AssetModel.name == "data_dep_asset_a")) + + response = test_client.get( + "/dependencies", params={"node_id": f"asset:{asset_a_id}", "dependency_type": "data"} + ) + assert response.status_code == 200 + + result = response.json() + entry_task_node_id = "task:data_dep_downstream__SEPARATOR__entry_task" + nodes_by_id = {node["id"]: node for node in result["nodes"]} + assert nodes_by_id[entry_task_node_id]["type"] == "task" + assert not any(node["type"] == "asset-condition" for node in result["nodes"]) + + edge_tuples = {(edge["source_id"], edge["target_id"]) for edge in result["edges"]} + assert (f"asset:{asset_a_id}", entry_task_node_id) in edge_tuples + + def test_data_dependencies_routes_through_asset_condition_gate(self, dag_maker, test_client, session): + """When the scheduling Dag's condition combines multiple assets, the data-dependencies graph + must show the gate (not a misleading direct edge from just one of the assets). + """ + asset_a = Asset(uri="s3://data-dep-bucket/gate-a", name="data_dep_gate_asset_a") + asset_b = Asset(uri="s3://data-dep-bucket/gate-b", name="data_dep_gate_asset_b") + + with dag_maker(dag_id="data_dep_gate_upstream", serialized=True, session=session): + EmptyOperator(task_id="produce_b", outlets=[asset_b]) + + with dag_maker( + dag_id="data_dep_gate_downstream", + schedule=(asset_a & asset_b), + serialized=True, + session=session, + ): + EmptyOperator(task_id="entry_task") + + dag_maker.sync_dagbag_to_db() + + asset_a_id = session.scalar(select(AssetModel.id).where(AssetModel.name == "data_dep_gate_asset_a")) + asset_b_id = session.scalar(select(AssetModel.id).where(AssetModel.name == "data_dep_gate_asset_b")) + + response = test_client.get( + "/dependencies", params={"node_id": f"asset:{asset_a_id}", "dependency_type": "data"} + ) + assert response.status_code == 200 + + result = response.json() + entry_task_node_id = "task:data_dep_gate_downstream__SEPARATOR__entry_task" + gate_node = next(node for node in result["nodes"] if node["type"] == "asset-condition") + assert gate_node["asset_condition_type"] == "and-gate" + + edge_tuples = {(edge["source_id"], edge["target_id"]) for edge in result["edges"]} + assert (gate_node["id"], entry_task_node_id) in edge_tuples + assert (f"asset:{asset_a_id}", gate_node["id"]) in edge_tuples + assert (f"asset:{asset_b_id}", gate_node["id"]) in edge_tuples + # No misleading direct edge bypassing the gate. + assert (f"asset:{asset_a_id}", entry_task_node_id) not in edge_tuples + + # BFS continues through the sibling asset in the gate to find its producer. + producer_task_node_id = "task:data_dep_gate_upstream__SEPARATOR__produce_b" + assert (producer_task_node_id, f"asset:{asset_b_id}") in edge_tuples + + def test_data_dependencies_batches_entry_point_resolution_across_scheduled_dags( + self, dag_maker, test_client, session + ): + """Resolving each scheduled Dag's entry task must not cost one query (and one full Dag + deserialization) per Dag: the query count must stay flat as the number of Dags an asset + schedules grows, not scale linearly with it. + """ + asset_a = Asset(uri="s3://data-dep-bucket/batch-a", name="data_dep_batch_asset_a") + dag_ids = [f"data_dep_batch_downstream_{i}" for i in range(3)] + + for dag_id in dag_ids: + with dag_maker( + dag_id=dag_id, + schedule=[asset_a], + serialized=True, + session=session, + ): + EmptyOperator(task_id="entry_task") + + dag_maker.sync_dagbag_to_db() + + asset_a_id = session.scalar(select(AssetModel.id).where(AssetModel.name == "data_dep_batch_asset_a")) + + with assert_queries_count(12): + response = test_client.get( + "/dependencies", params={"node_id": f"asset:{asset_a_id}", "dependency_type": "data"} + ) + assert response.status_code == 200 + + result = response.json() + nodes_by_id = {node["id"]: node for node in result["nodes"]} + for dag_id in dag_ids: + entry_task_node_id = f"task:{dag_id}__SEPARATOR__entry_task" + assert nodes_by_id[entry_task_node_id]["type"] == "task" + + def test_data_dependencies_attributes_alias_produced_asset_to_source_task( + self, dag_maker, test_client, session + ): + """An asset produced via an AssetAlias has no static outlet reference, so its producing + task must be recovered from AssetEvent -- otherwise the consumer sees the asset with no + producer. The asset is made visible through a consuming task (the realistic case; an asset + with no reference at all is hidden by the readable-dags gate).""" + from airflow.models.asset import AssetAliasModel, AssetEvent + from airflow.sdk.definitions.asset import AssetAlias + + alias_asset = Asset(uri="s3://alias/resolved", name="alias_resolved_asset") + + with dag_maker(dag_id="alias_producer", serialized=True, session=session): + EmptyOperator(task_id="produce_via_alias", outlets=[AssetAlias("resolving_alias")]) + dr = dag_maker.create_dagrun() + + with dag_maker(dag_id="alias_consumer", serialized=True, session=session): + EmptyOperator(task_id="consume", inlets=[alias_asset]) + + dag_maker.sync_dagbag_to_db() + + asset_id = session.scalar(select(AssetModel.id).where(AssetModel.name == "alias_resolved_asset")) + asset_alias = session.scalar(select(AssetAliasModel).where(AssetAliasModel.name == "resolving_alias")) + asset_alias.assets.append(session.get(AssetModel, asset_id)) + asset_alias.asset_events.append( + AssetEvent( + timestamp=pendulum.datetime(2024, 1, 1, tz="UTC"), + asset_id=asset_id, + source_dag_id="alias_producer", + source_task_id="produce_via_alias", + source_run_id=dr.run_id, + source_map_index=-1, + ) + ) + session.commit() + + response = test_client.get( + "/dependencies", params={"node_id": f"asset:{asset_id}", "dependency_type": "data"} + ) + assert response.status_code == 200 + + result = response.json() + producer_task_node_id = "task:alias_producer__SEPARATOR__produce_via_alias" + consumer_task_node_id = "task:alias_consumer__SEPARATOR__consume" + nodes_by_id = {node["id"]: node for node in result["nodes"]} + assert nodes_by_id[producer_task_node_id]["type"] == "task" + + edge_tuples = {(edge["source_id"], edge["target_id"]) for edge in result["edges"]} + # Producer edge recovered from the alias event, plus the consumer edge that made it visible. + assert (producer_task_node_id, f"asset:{asset_id}") in edge_tuples + assert (f"asset:{asset_id}", consumer_task_node_id) in edge_tuples diff --git a/airflow-core/tests/unit/api_fastapi/core_api/routes/ui/test_structure.py b/airflow-core/tests/unit/api_fastapi/core_api/routes/ui/test_structure.py index 6aaff421ee38e..61b2eed0b3ddc 100644 --- a/airflow-core/tests/unit/api_fastapi/core_api/routes/ui/test_structure.py +++ b/airflow-core/tests/unit/api_fastapi/core_api/routes/ui/test_structure.py @@ -18,21 +18,13 @@ from __future__ import annotations import copy -from unittest import mock import pendulum import pytest -from sqlalchemy import select -from sqlalchemy.orm import Session -from airflow._shared.timezones import timezone -from airflow.models.asset import AssetAliasModel, AssetEvent, AssetModel from airflow.models.dagbag import DBDagBag from airflow.providers.standard.operators.empty import EmptyOperator -from airflow.providers.standard.operators.trigger_dagrun import TriggerDagRunOperator from airflow.providers.standard.sensors.external_task import ExternalTaskSensor -from airflow.sdk import Metadata, task -from airflow.sdk.definitions.asset import Asset, AssetAlias, Dataset from airflow.sdk.definitions.taskgroup import TaskGroup from tests_common.test_utils.asserts import assert_queries_count @@ -41,8 +33,6 @@ pytestmark = pytest.mark.db_test DAG_ID = "dag_with_multiple_versions" -DAG_ID_EXTERNAL_TRIGGER = "external_trigger" -DAG_ID_RESOLVED_ASSET_ALIAS = "dag_with_resolved_asset_alias" DAG_ID_LINEAR_DEPTH = "linear_depth_dag" DAG_ID_NONLINEAR_DEPTH = "nonlinear_depth_dag" LATEST_VERSION_DAG_RESPONSE: dict = { @@ -117,84 +107,20 @@ def clean(): @pytest.fixture -def asset1() -> Asset: - return Asset(uri="s3://bucket/next-run-asset/1", name="asset1") - - -@pytest.fixture -def asset2() -> Asset: - return Asset(uri="s3://bucket/next-run-asset/2", name="asset2") - - -@pytest.fixture -def asset3() -> Dataset: - return Dataset(uri="s3://dataset-bucket/example.csv") - - -@pytest.fixture -def make_dags(dag_maker, session, time_machine, asset1: Asset, asset2: Asset, asset3: Dataset) -> None: - with dag_maker( - dag_id=DAG_ID_EXTERNAL_TRIGGER, - serialized=True, - session=session, - start_date=pendulum.DateTime(2023, 2, 1, 0, 0, 0, tzinfo=pendulum.UTC), - ): - TriggerDagRunOperator(task_id="trigger_dag_run_operator", trigger_dag_id=DAG_ID) - dag_maker.sync_dagbag_to_db() - +def make_dags(dag_maker, session, time_machine) -> None: with dag_maker( dag_id=DAG_ID, serialized=True, session=session, start_date=pendulum.DateTime(2023, 2, 1, 0, 0, 0, tzinfo=pendulum.UTC), - schedule=(asset1 & asset2 & AssetAlias("example-alias")), ): ( - EmptyOperator(task_id="task_1", outlets=[asset3]) + EmptyOperator(task_id="task_1") >> ExternalTaskSensor(task_id="external_task_sensor", external_dag_id=DAG_ID) >> EmptyOperator(task_id="task_2") ) dag_maker.sync_dagbag_to_db() - with dag_maker( - dag_id=DAG_ID_RESOLVED_ASSET_ALIAS, - serialized=True, - session=session, - start_date=pendulum.DateTime(2023, 2, 1, 0, 0, 0, tzinfo=pendulum.UTC), - ): - - @task(outlets=[AssetAlias("example-alias-resolved")]) - def task_1(**context): - yield Metadata( - asset=Asset("resolved_example_asset_alias"), - extra={"k": "v"}, # extra has to be provided, can be {} - alias=AssetAlias("example-alias-resolved"), - ) - - task_1() >> EmptyOperator(task_id="task_2") - - dr = dag_maker.create_dagrun() - asset_alias = session.scalar( - select(AssetAliasModel).where(AssetAliasModel.name == "example-alias-resolved") - ) - asset_model = AssetModel(name="resolved_example_asset_alias") - session.add(asset_model) - session.flush() - asset_alias.assets.append(asset_model) - asset_alias.asset_events.append( - AssetEvent( - id=1, - timestamp=timezone.parse("2021-01-01T00:00:00"), - asset_id=asset_model.id, - source_dag_id=DAG_ID_RESOLVED_ASSET_ALIAS, - source_task_id="task_1", - source_run_id=dr.run_id, - source_map_index=-1, - ) - ) - session.commit() - dag_maker.sync_dagbag_to_db() - # Linear DAG with 5 tasks for depth testing with dag_maker( dag_id=DAG_ID_LINEAR_DEPTH, @@ -232,29 +158,6 @@ def task_1(**context): dag_maker.sync_dagbag_to_db() -def _fetch_asset_id(asset: Asset, session: Session) -> str: - return str( - session.scalar( - select(AssetModel.id).where(AssetModel.name == asset.name, AssetModel.uri == asset.uri) - ) - ) - - -@pytest.fixture -def asset1_id(make_dags, asset1, session: Session) -> str: - return _fetch_asset_id(asset1, session) - - -@pytest.fixture -def asset2_id(make_dags, asset2, session) -> str: - return _fetch_asset_id(asset2, session) - - -@pytest.fixture -def asset3_id(make_dags, asset3, session) -> str: - return _fetch_asset_id(asset3, session) - - class TestStructureDataEndpoint: @pytest.mark.parametrize( ("params", "expected", "expected_queries_count"), @@ -265,14 +168,12 @@ class TestStructureDataEndpoint: "edges": [ { "is_setup_teardown": None, - "is_source_asset": None, "label": None, "source_id": "external_task_sensor", "target_id": "task_2", }, { "is_setup_teardown": None, - "is_source_asset": None, "label": None, "source_id": "task_1", "target_id": "external_task_sensor", @@ -361,51 +262,6 @@ class TestStructureDataEndpoint: }, 7, ), - ( - {"dag_id": DAG_ID_EXTERNAL_TRIGGER, "external_dependencies": True}, - { - "edges": [ - { - "is_source_asset": None, - "is_setup_teardown": None, - "label": None, - "source_id": "trigger_dag_run_operator", - "target_id": "trigger:external_trigger:dag_with_multiple_versions:trigger_dag_run_operator", - } - ], - "nodes": [ - { - "asset_condition_type": None, - "ui_color": "#ffefeb", - "ui_fgcolor": "#000", - "children": None, - "id": "trigger_dag_run_operator", - "is_mapped": None, - "label": "trigger_dag_run_operator", - "tooltip": None, - "setup_teardown_type": None, - "type": "task", - "team": None, - "operator": "TriggerDagRunOperator", - }, - { - "asset_condition_type": None, - "ui_color": None, - "ui_fgcolor": None, - "children": None, - "id": "trigger:external_trigger:dag_with_multiple_versions:trigger_dag_run_operator", - "is_mapped": None, - "label": "trigger_dag_run_operator", - "tooltip": None, - "setup_teardown_type": None, - "type": "trigger", - "team": None, - "operator": None, - }, - ], - }, - 14, - ), ], ) @pytest.mark.usefixtures("make_dags") @@ -415,305 +271,6 @@ def test_should_return_200(self, test_client, params, expected, expected_queries assert response.status_code == 200 assert response.json() == expected - @pytest.mark.usefixtures("make_dags") - def test_should_return_200_with_asset(self, test_client, asset1_id, asset2_id, asset3_id): - params = { - "dag_id": DAG_ID, - "external_dependencies": True, - } - expected = { - "edges": [ - { - "is_setup_teardown": None, - "label": None, - "source_id": "external_task_sensor", - "target_id": "task_2", - "is_source_asset": None, - }, - { - "is_setup_teardown": None, - "label": None, - "source_id": "task_1", - "target_id": "external_task_sensor", - "is_source_asset": None, - }, - { - "is_setup_teardown": None, - "label": None, - "source_id": "and-gate-0", - "target_id": "task_1", - "is_source_asset": True, - }, - { - "is_setup_teardown": None, - "label": None, - "source_id": asset1_id, - "target_id": "and-gate-0", - "is_source_asset": None, - }, - { - "is_setup_teardown": None, - "label": None, - "source_id": asset2_id, - "target_id": "and-gate-0", - "is_source_asset": None, - }, - { - "is_setup_teardown": None, - "label": None, - "source_id": "example-alias", - "target_id": "and-gate-0", - "is_source_asset": None, - }, - { - "is_setup_teardown": None, - "label": None, - "source_id": "sensor:dag_with_multiple_versions:dag_with_multiple_versions:external_task_sensor", - "target_id": "task_1", - "is_source_asset": None, - }, - { - "is_setup_teardown": None, - "label": None, - "source_id": "trigger:external_trigger:dag_with_multiple_versions:trigger_dag_run_operator", - "target_id": "task_1", - "is_source_asset": None, - }, - { - "is_setup_teardown": None, - "label": None, - "source_id": "task_1", - "target_id": f"asset:{asset3_id}", - "is_source_asset": None, - }, - ], - "nodes": [ - { - "children": None, - "id": "task_1", - "is_mapped": None, - "label": "task_1", - "tooltip": None, - "setup_teardown_type": None, - "type": "task", - "team": None, - "operator": "EmptyOperator", - "asset_condition_type": None, - "ui_color": "#e8f7e4", - "ui_fgcolor": "#000", - }, - { - "children": None, - "id": "external_task_sensor", - "is_mapped": None, - "label": "external_task_sensor", - "tooltip": None, - "setup_teardown_type": None, - "type": "task", - "team": None, - "operator": "ExternalTaskSensor", - "asset_condition_type": None, - "ui_color": "#4db7db", - "ui_fgcolor": "#000", - }, - { - "children": None, - "id": "task_2", - "is_mapped": None, - "label": "task_2", - "tooltip": None, - "setup_teardown_type": None, - "type": "task", - "team": None, - "operator": "EmptyOperator", - "asset_condition_type": None, - "ui_color": "#e8f7e4", - "ui_fgcolor": "#000", - }, - { - "children": None, - "id": f"asset:{asset3_id}", - "is_mapped": None, - "label": "s3://dataset-bucket/example.csv", - "tooltip": None, - "setup_teardown_type": None, - "type": "asset", - "team": None, - "operator": None, - "asset_condition_type": None, - "ui_color": None, - "ui_fgcolor": None, - }, - { - "children": None, - "id": "sensor:dag_with_multiple_versions:dag_with_multiple_versions:external_task_sensor", - "is_mapped": None, - "label": "external_task_sensor", - "tooltip": None, - "setup_teardown_type": None, - "type": "sensor", - "team": None, - "operator": None, - "asset_condition_type": None, - "ui_color": None, - "ui_fgcolor": None, - }, - { - "children": None, - "id": "trigger:external_trigger:dag_with_multiple_versions:trigger_dag_run_operator", - "is_mapped": None, - "label": "trigger_dag_run_operator", - "tooltip": None, - "setup_teardown_type": None, - "type": "trigger", - "team": None, - "operator": None, - "asset_condition_type": None, - "ui_color": None, - "ui_fgcolor": None, - }, - { - "children": None, - "id": "and-gate-0", - "is_mapped": None, - "label": "and-gate-0", - "tooltip": None, - "setup_teardown_type": None, - "type": "asset-condition", - "team": None, - "operator": None, - "asset_condition_type": "and-gate", - "ui_color": None, - "ui_fgcolor": None, - }, - { - "children": None, - "id": asset1_id, - "is_mapped": None, - "label": "asset1", - "tooltip": None, - "setup_teardown_type": None, - "type": "asset", - "team": None, - "operator": None, - "asset_condition_type": None, - "ui_color": None, - "ui_fgcolor": None, - }, - { - "children": None, - "id": asset2_id, - "is_mapped": None, - "label": "asset2", - "tooltip": None, - "setup_teardown_type": None, - "type": "asset", - "team": None, - "operator": None, - "asset_condition_type": None, - "ui_color": None, - "ui_fgcolor": None, - }, - { - "children": None, - "id": "example-alias", - "is_mapped": None, - "label": "example-alias", - "tooltip": None, - "setup_teardown_type": None, - "type": "asset-alias", - "team": None, - "operator": None, - "asset_condition_type": None, - "ui_color": None, - "ui_fgcolor": None, - }, - ], - } - - with assert_queries_count(14): - response = test_client.get("/structure/structure_data", params=params) - assert response.status_code == 200 - assert response.json() == expected - - @pytest.mark.usefixtures("make_dags") - def test_should_return_200_with_resolved_asset_alias_attached_to_the_corrrect_producing_task( - self, test_client, session - ): - resolved_asset = session.scalar( - select(AssetModel).where(AssetModel.name == "resolved_example_asset_alias") - ) - params = { - "dag_id": DAG_ID_RESOLVED_ASSET_ALIAS, - "external_dependencies": True, - } - expected = { - "edges": [ - { - "source_id": "task_1", - "target_id": "task_2", - "is_setup_teardown": None, - "label": None, - "is_source_asset": None, - }, - { - "source_id": "task_1", - "target_id": f"asset:{resolved_asset.id}", - "is_setup_teardown": None, - "label": None, - "is_source_asset": None, - }, - ], - "nodes": [ - { - "id": "task_1", - "label": "task_1", - "type": "task", - "team": None, - "children": None, - "is_mapped": None, - "tooltip": None, - "setup_teardown_type": None, - "operator": "@task", - "asset_condition_type": None, - "ui_color": "#ffefeb", - "ui_fgcolor": "#000", - }, - { - "id": "task_2", - "label": "task_2", - "type": "task", - "team": None, - "children": None, - "is_mapped": None, - "tooltip": None, - "setup_teardown_type": None, - "operator": "EmptyOperator", - "asset_condition_type": None, - "ui_color": "#e8f7e4", - "ui_fgcolor": "#000", - }, - { - "id": f"asset:{resolved_asset.id}", - "label": "resolved_example_asset_alias", - "type": "asset", - "team": None, - "children": None, - "is_mapped": None, - "tooltip": None, - "setup_teardown_type": None, - "operator": None, - "asset_condition_type": None, - "ui_color": None, - "ui_fgcolor": None, - }, - ], - } - - response = test_client.get("/structure/structure_data", params=params) - assert response.status_code == 200 - assert response.json() == expected - @pytest.mark.parametrize( ("params", "expected"), [ @@ -753,54 +310,11 @@ def test_delete_dag_should_response_403(self, unauthorized_test_client): response = unauthorized_test_client.get("/structure/structure_data", params={"dag_id": DAG_ID}) assert response.status_code == 403 - @mock.patch( - "airflow.api_fastapi.auth.managers.base_auth_manager.BaseAuthManager.get_authorized_dag_ids", - return_value={DAG_ID_EXTERNAL_TRIGGER}, - ) - @pytest.mark.usefixtures("make_dags") - def test_external_deps_filters_unreadable_dags(self, _, test_client): - response = test_client.get( - "/structure/structure_data", - params={"dag_id": DAG_ID_EXTERNAL_TRIGGER, "external_dependencies": True}, - ) - assert response.status_code == 200 - result = response.json() - node_ids = {node["id"] for node in result["nodes"]} - assert "trigger_dag_run_operator" in node_ids - assert not any(DAG_ID in nid for nid in node_ids if nid != "trigger_dag_run_operator") - edge_targets = {edge["target_id"] for edge in result["edges"]} - assert not any(DAG_ID in tid for tid in edge_targets) - def test_should_return_404(self, test_client): response = test_client.get("/structure/structure_data", params={"dag_id": "not_existing"}) assert response.status_code == 404 assert response.json()["detail"] == "Dag with id not_existing was not found" - @pytest.mark.usefixtures("make_dags") - def test_should_return_400_on_malformed_asset_expression(self, test_client): - """A TypeError from get_upstream_assets surfaces as a 400 with a clear message. - - The asset_expression ultimately comes from user-authored Dag code (via the Task SDK), - so a malformed expression is bad input that ended up persisted -- not a server fault. - Without the try/except wrap, the TypeError propagates uncaught and FastAPI returns a - generic ``{"detail": "Internal Server Error"}`` 500 body with no context about which - Dag triggered it. With the wrap, the response identifies the Dag and version, which - is what a caller needs to fix the upstream Dag definition. - """ - with mock.patch( - "airflow.api_fastapi.core_api.routes.ui.structure.get_upstream_assets", - side_effect=TypeError("Unsupported type: dict_keys(['weird-op'])"), - ): - response = test_client.get( - "/structure/structure_data", - params={"dag_id": DAG_ID, "external_dependencies": True}, - ) - assert response.status_code == 400 - detail = response.json()["detail"] - assert "Malformed asset_expression" in detail - assert DAG_ID in detail - assert "Unsupported type" in detail - def test_should_return_404_when_dag_version_not_found(self, test_client): response = test_client.get( "/structure/structure_data", params={"dag_id": DAG_ID, "version_number": 999}