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 57f052dbc42a1..cbab0505f3393 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 class StructureDataResponse(BaseGraphResponse[EdgeResponse, NodeResponse]): 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 bcfe8c708ef15..13ce02f6be631 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 @@ -2146,6 +2139,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 +2974,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 +3622,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 +3660,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 type: object required: - id 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..0afb2e4e3cd7a 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 @@ -50,19 +45,16 @@ ), 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 +71,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 +96,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..1320752dd315c 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 @@ -20,12 +20,136 @@ from collections import defaultdict, deque 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 + + # 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) + 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()}") + + if not asset_expression_type: + return nodes, edges + + # A condition combining exactly one branch isn't a real AND/OR -- connect it directly (or + # recurse straight through a nested single-branch wrapper) instead of rendering a gate with + # only one input. + if len(assets_info) + (1 if nested_expression else 0) <= 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"]}) + elif nested_expression: + return get_upstream_assets(nested_expression, entry_node_ref, level=level) + + return nodes, edges + + # 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, + "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: + 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"], + } + ) + + if nested_expression: + n, e = get_upstream_assets(nested_expression, asset_condition_id, level=level + 1) + + nodes = nodes + n + edges = edges + e + + 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 +209,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 +217,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 +271,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 +302,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 +333,38 @@ def get_scheduling_dependencies(readable_dag_ids: set[str] | None = None) -> dic } +def _get_dag_entry_point(dag_id: str, session: Session) -> tuple[str, dict | None] | None: + """Get the id of a Dag's topologically-first task/group and its asset_expression, if any.""" + from sqlalchemy import select + from sqlalchemy.orm import joinedload + + from airflow.api_fastapi.core_api.services.ui.task_group import task_group_to_dict + from airflow.models.dag_version import DagVersion + from airflow.models.serialized_dag import SerializedDagModel + + dag_version = DagVersion.get_latest_version(dag_id, session=session) + if dag_version is None: + return None + + serialized_dag = session.scalar( + select(SerializedDagModel) + .join(DagVersion) + .where( + SerializedDagModel.dag_id == dag_id, + DagVersion.version_number == dag_version.version_number, + ) + .options(joinedload(SerializedDagModel.dag_model)) + ) + if serialized_dag is None: + return None + + entry_points = serialized_dag.dag.task_group.topological_sort() + if not entry_points: + return None + + return task_group_to_dict(entry_points[0])["id"], serialized_dag.dag_model.asset_expression + + def get_data_dependencies( asset_id: int, session: Session, readable_dag_ids: set[str] | None = None ) -> dict[str, list[dict]]: @@ -201,6 +403,7 @@ def get_data_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) + processed_scheduled_dags: set[str] = set() while assets_to_process: current_asset_id = assets_to_process.popleft() @@ -208,13 +411,14 @@ def get_data_dependencies( continue processed_assets.add(current_asset_id) - # Eagerload producing_tasks and consuming_tasks to avoid lazy queries + # Eagerload producing_tasks, consuming_tasks, and scheduled_dags to avoid lazy queries asset = session.scalar( select(AssetModel) .where(AssetModel.id == current_asset_id) .options( selectinload(AssetModel.producing_tasks), selectinload(AssetModel.consuming_tasks), + selectinload(AssetModel.scheduled_dags), ) ) if not asset: @@ -290,6 +494,72 @@ def get_data_dependencies( if outlet_ref.asset_id not in processed_assets: assets_to_process.append(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: route + # through the 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. + 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 + processed_scheduled_dags.add(ref.dag_id) + + entry_point = _get_dag_entry_point(ref.dag_id, session) + if entry_point is None: + continue + entry_task_id, asset_expression = entry_point + task_key = (ref.dag_id, entry_task_id) + task_node_id = f"task:{ref.dag_id}{SEPARATOR}{entry_task_id}" + + if task_node_id not in nodes_dict: + nodes_dict[task_node_id] = { + "id": task_node_id, + "label": f"{ref.dag_id}.{entry_task_id}", + "type": "task", + } + + 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=ref.dag_id, + exc_info=True, + ) + + 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: + assets_to_process.append(sibling_asset_id) + for edge in upstream_edges: + edge_set.add((edge["source_id"], edge["target_id"])) + else: + # No gate (bare single-asset schedule) -- link directly. + edge_set.add((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) + outlet_refs = session.scalars( + select(TaskOutletAssetReference).where( + TaskOutletAssetReference.dag_id == ref.dag_id, + TaskOutletAssetReference.task_id == entry_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) + all_dag_ids = list({dag_id for dag_id, _ in processed_tasks}) if all_dag_ids: dag_id_to_team = DagModel.get_dag_id_to_team_name_mapping(all_dag_ids, session=session) 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 e1880fdd4ea89..1545a0e35b9c3 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 @@ -8597,6 +8597,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', @@ -9610,17 +9622,6 @@ export const $EdgeResponse = { } ], title: 'Label' - }, - is_source_asset: { - anyOf: [ - { - type: 'boolean' - }, - { - type: 'null' - } - ], - title: 'Is Source Asset' } }, type: 'object', @@ -10261,6 +10262,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: [ { @@ -10319,18 +10332,6 @@ export const $NodeResponse = { } ], title: 'Operator' - }, - asset_condition_type: { - anyOf: [ - { - type: 'string', - enum: ['or-gate', 'and-gate'] - }, - { - type: 'null' - } - ], - title: 'Asset Condition Type' } }, type: 'object', 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..4ddd7e9bd49e9 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,7 +4920,6 @@ export class StructureService { include_downstream: data.includeDownstream, depth: data.depth, root: data.root, - external_dependencies: data.externalDependencies, version_number: data.versionNumber }, errors: { 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 24b1fbe5f2e9b..cf8769dfaefa4 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 @@ -2180,6 +2180,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'; @@ -2453,7 +2454,6 @@ export type EdgeResponse = { target_id: string; is_setup_teardown?: boolean | null; label?: string | null; - is_source_asset?: boolean | null; }; /** @@ -2616,12 +2616,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; }; export type OklchColor = string; @@ -4631,7 +4631,6 @@ export type GetDagDeadlineAlertsResponse = DeadlineAlertCollectionResponse; export type StructureDataData = { dagId: string; depth?: number | null; - externalDependencies?: boolean; includeDownstream?: boolean; includeUpstream?: boolean; root?: string | null; 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 67fe2c2d2406d..2911c9c0c92b6 100644 --- a/airflow-core/src/airflow/ui/src/components/Graph/reactflowUtils.ts +++ b/airflow-core/src/airflow/ui/src/components/Graph/reactflowUtils.ts @@ -139,3 +139,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.ts b/airflow-core/src/airflow/ui/src/constants/localStorage.ts index 72ff5b8327b4c..34b1fe2e0f7db 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`; 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/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..1f5be124a4d06 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,194 @@ 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_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 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 0bde80a22dc37..aa07caad77126 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 tests_common.test_utils.asserts import assert_queries_count from tests_common.test_utils.db import clear_db_assets, clear_db_runs @@ -40,8 +32,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 = { @@ -110,84 +100,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, @@ -225,29 +151,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"), @@ -258,14 +161,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", @@ -346,47 +247,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, - "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, - "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") @@ -396,279 +256,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, - }, - { - "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, - }, - { - "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, - }, - { - "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, - }, - { - "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, - }, - { - "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, - }, - { - "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", - }, - { - "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, - }, - { - "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, - }, - { - "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, - }, - ], - } - - 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, - }, - { - "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, - }, - { - "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, - }, - ], - } - - response = test_client.get("/structure/structure_data", params=params) - assert response.status_code == 200 - assert response.json() == expected - @pytest.mark.parametrize( ("params", "expected"), [ @@ -708,54 +295,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}