This is an automated email from the ASF dual-hosted git repository.
bbovenzi pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/airflow.git
The following commit(s) were added to refs/heads/main by this push:
new fd0ec61c503 Consolidate UI Dependency Graphs (#69833)
fd0ec61c503 is described below
commit fd0ec61c503cbebc900dd3e053b7e035cd859f30
Author: Brent Bovenzi <[email protected]>
AuthorDate: Wed Jul 29 11:19:04 2026 -0400
Consolidate UI Dependency Graphs (#69833)
* Consolidate dag dep graphs
* fix prettier
* Batch entry-point resolution for scheduled Dags in the data-dependencies
graph
Resolving each scheduled Dag's entry task previously deserialized the full
Dag once per Dag discovered while tracing an asset's dependencies -- for an
asset feeding many Dags, that's N full Dag deserializations plus N round
trips. Restructures the BFS to process a whole frontier of assets per round
and batch-resolve every scheduled Dag's entry point in a single query.
* Address review feedback on consolidated dependency graphs
Fix nested boolean asset conditions dropping sibling branches so
(a & b) | (c & d) renders both gates; read a scheduled Dag's entry task
from serialized JSON instead of deserializing every Dag while tracing an
asset's lineage; and recover alias-produced assets' source task from
AssetEvent so they still show a producer. Also harden a malformed
expression against a 500 and link a Dag to every triggering asset in the
degraded fallback.
* UI: Prune orphaned per-Dag dependency toggle keys on startup
The dependency-view toggle moved from a per-Dag localStorage key to a
single global one, leaving the old `dependencies-<dag_id>` entries stranded
in every user's browser storage. Clean them up once at load.
---
.../api_fastapi/core_api/datamodels/ui/common.py | 1 +
.../core_api/datamodels/ui/structure.py | 2 -
.../api_fastapi/core_api/openapi/_private_ui.yaml | 42 +-
.../api_fastapi/core_api/routes/ui/dependencies.py | 2 +-
.../api_fastapi/core_api/routes/ui/structure.py | 91 +---
.../core_api/services/ui/dependencies.py | 486 +++++++++++++++++---
.../api_fastapi/core_api/services/ui/structure.py | 185 --------
.../src/airflow/ui/openapi-gen/queries/common.ts | 5 +-
.../ui/openapi-gen/queries/ensureQueryData.ts | 6 +-
.../src/airflow/ui/openapi-gen/queries/prefetch.ts | 6 +-
.../src/airflow/ui/openapi-gen/queries/queries.ts | 6 +-
.../src/airflow/ui/openapi-gen/queries/suspense.ts | 6 +-
.../airflow/ui/openapi-gen/requests/schemas.gen.ts | 47 +-
.../ui/openapi-gen/requests/services.gen.ts | 3 -
.../airflow/ui/openapi-gen/requests/types.gen.ts | 9 +-
.../src/airflow/ui/public/i18n/locales/en/dag.json | 8 +-
.../src/airflow/ui/src/components/Graph/Edge.tsx | 7 +-
.../ui/src/components/Graph/reactflowUtils.test.ts | 118 +++++
.../ui/src/components/Graph/reactflowUtils.ts | 105 +++++
.../airflow/ui/src/constants/localStorage.test.ts | 51 +++
.../src/airflow/ui/src/constants/localStorage.ts | 19 +-
.../src/airflow/ui/src/constants/searchParams.ts | 1 -
.../ui/src/context/groups/GroupsProvider.tsx | 4 +-
.../airflow/ui/src/layouts/Details/Graph/Graph.tsx | 34 +-
.../ui/src/layouts/Details/PanelButtons.tsx | 65 +--
airflow-core/src/airflow/ui/src/main.tsx | 3 +
.../src/airflow/ui/src/pages/Asset/AssetGraph.tsx | 15 +-
.../core_api/routes/ui/test_dependencies.py | 364 ++++++++++++++-
.../core_api/routes/ui/test_structure.py | 490 +--------------------
29 files changed, 1179 insertions(+), 1002 deletions(-)
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 702f6997cdf..0cccdecc4d3 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 5d89b08ee98..36f0553d634 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 317eabd56ee..2237e622d75 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
@@ -1069,13 +1069,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
@@ -1091,12 +1084,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:
@@ -2201,6 +2188,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
@@ -3052,11 +3047,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
@@ -3705,6 +3695,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:
@@ -3735,14 +3733,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 d09b8acf0e4..0286b0a126f 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 a5cb9d1d8d1..876b64dac19 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.parameters import
QueryIncludeDownstream, QueryI
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_router = AirflowRouter(tags=["Structure"],
prefix="/structure")
"/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(
@@ -108,82 +99,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 20ffc78008a..b519a61def3 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 db3d1ba6dea..00000000000
--- 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 ec9da9f0535..64e87bd2a1d 100644
--- a/airflow-core/src/airflow/ui/openapi-gen/queries/common.ts
+++ b/airflow-core/src/airflow/ui/openapi-gen/queries/common.ts
@@ -990,15 +990,14 @@ export const UseDeadlinesServiceGetDagDeadlineAlertsKeyFn
= ({ dagId, limit, off
export type StructureServiceStructureDataDefaultResponse =
Awaited<ReturnType<typeof StructureService.structureData>>;
export type StructureServiceStructureDataQueryResult<TData =
StructureServiceStructureDataDefaultResponse, TError = unknown> =
UseQueryResult<TData, TError>;
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<unknown>) => [useStructureServiceStructureDataKey,
...(queryKey ?? [{ dagId, depth, externalDependencies, includeDownstream,
includeUpstream, root, versionNumber }])];
+}, queryKey?: Array<unknown>) => [useStructureServiceStructureDataKey,
...(queryKey ?? [{ dagId, depth, includeDownstream, includeUpstream, root,
versionNumber }])];
export type GridServiceGetDagStructureDefaultResponse =
Awaited<ReturnType<typeof GridService.getDagStructure>>;
export type GridServiceGetDagStructureQueryResult<TData =
GridServiceGetDagStructureDefaultResponse, TError = unknown> =
UseQueryResult<TData, TError>;
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 9482e9ce6d0..c3c73d7d427 100644
--- a/airflow-core/src/airflow/ui/openapi-gen/queries/ensureQueryData.ts
+++ b/airflow-core/src/airflow/ui/openapi-gen/queries/ensureQueryData.ts
@@ -1908,20 +1908,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 1eaac7e3395..53e93745026 100644
--- a/airflow-core/src/airflow/ui/openapi-gen/queries/prefetch.ts
+++ b/airflow-core/src/airflow/ui/openapi-gen/queries/prefetch.ts
@@ -1908,20 +1908,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 386cd4aca74..49f1fa6d734 100644
--- a/airflow-core/src/airflow/ui/openapi-gen/queries/queries.ts
+++ b/airflow-core/src/airflow/ui/openapi-gen/queries/queries.ts
@@ -1908,20 +1908,18 @@ export const useDeadlinesServiceGetDagDeadlineAlerts =
<TData = Common.Deadlines
* @param data.includeDownstream
* @param data.depth
* @param data.root
-* @param data.externalDependencies
* @param data.versionNumber
* @returns StructureDataResponse Successful Response
* @throws ApiError
*/
-export const useStructureServiceStructureData = <TData =
Common.StructureServiceStructureDataDefaultResponse, TError = unknown,
TQueryKey extends Array<unknown> = unknown[]>({ dagId, depth,
externalDependencies, includeDownstream, includeUpstream, root, versionNumber
}: {
+export const useStructureServiceStructureData = <TData =
Common.StructureServiceStructureDataDefaultResponse, TError = unknown,
TQueryKey extends Array<unknown> = 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<UseQueryOptions<TData, TError>,
"queryKey" | "queryFn">) => useQuery<TData, TError>({ 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<UseQueryOptions<TData, TError>,
"queryKey" | "queryFn">) => useQuery<TData, TError>({ 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 e4834f3ac0a..a67ccaa19ee 100644
--- a/airflow-core/src/airflow/ui/openapi-gen/queries/suspense.ts
+++ b/airflow-core/src/airflow/ui/openapi-gen/queries/suspense.ts
@@ -1908,20 +1908,18 @@ export const
useDeadlinesServiceGetDagDeadlineAlertsSuspense = <TData = Common.D
* @param data.includeDownstream
* @param data.depth
* @param data.root
-* @param data.externalDependencies
* @param data.versionNumber
* @returns StructureDataResponse Successful Response
* @throws ApiError
*/
-export const useStructureServiceStructureDataSuspense = <TData =
Common.StructureServiceStructureDataDefaultResponse, TError = unknown,
TQueryKey extends Array<unknown> = unknown[]>({ dagId, depth,
externalDependencies, includeDownstream, includeUpstream, root, versionNumber
}: {
+export const useStructureServiceStructureDataSuspense = <TData =
Common.StructureServiceStructureDataDefaultResponse, TError = unknown,
TQueryKey extends Array<unknown> = 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<UseQueryOptions<TData, TError>,
"queryKey" | "queryFn">) => useSuspenseQuery<TData, TError>({ 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<UseQueryOptions<TData, TError>,
"queryKey" | "queryFn">) => useSuspenseQuery<TData, TError>({ 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 fedb68983f6..d14b118651c 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
@@ -8743,6 +8743,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',
@@ -9794,17 +9806,6 @@ export const $EdgeResponse = {
}
],
title: 'Label'
- },
- is_source_asset: {
- anyOf: [
- {
- type: 'boolean'
- },
- {
- type: 'null'
- }
- ],
- title: 'Is Source Asset'
}
},
type: 'object',
@@ -10445,6 +10446,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: [
{
@@ -10504,18 +10517,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 1c7365f7ea8..d34910b8276 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
@@ -4901,7 +4901,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
@@ -4916,11 +4915,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 f0b9e09aa46..932326c9d6f 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
@@ -2207,6 +2207,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';
@@ -2489,7 +2490,6 @@ export type EdgeResponse = {
target_id: string;
is_setup_teardown?: boolean | null;
label?: string | null;
- is_source_asset?: boolean | null;
};
/**
@@ -2652,12 +2652,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<NodeResponse> | 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;
@@ -8541,10 +8540,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 f203c21595f..6712444e12e 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
@@ -144,13 +144,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 a4afa95c168..c2e19448143 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 00000000000..10bc5d81f3c
--- /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 7a2667a7abc..93a12e550a0 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<Edge> }): Array<FlowEd
target: edge.targets[0] ?? "",
type: "custom",
}));
+
+type SelectionGraphEdge = { id: string; source: string; target: string };
+type SelectionGraphNode = { id: string; type?: string };
+
+/**
+ * Edges only highlight when one of their two endpoints has node-level
`isSelected` set, or (for
+ * gate-adjacent edges) when this function marks that specific edge id. Gate
edges are directional
+ * -- `source` is always the thing feeding the condition, `target` is what it
schedules -- so:
+ *
+ * - Entering a gate from one of its inputs (an asset, or a nested gate below
it) continues
+ * through the gate's single output edge only. It does NOT fan out to the
gate's other inputs,
+ * which is exactly the over-highlighting bug this replaces: selecting one
asset in an AND/OR
+ * condition must not also light up sibling assets that happen to share the
same gate.
+ * - Entering a gate from its output (e.g. selecting the Dag it schedules)
fans out to every one
+ * of its inputs, since all of them are genuinely part of what makes that
Dag's condition true.
+ *
+ * Returns the set of edge ids to mark selected -- callers OR this into their
own per-edge
+ * `isSelected` computation (see Graph.tsx / AssetGraph.tsx).
+ */
+export const getGatePathEdgeIdsForSelection = (
+ nodes: Array<SelectionGraphNode>,
+ edges: Array<SelectionGraphEdge>,
+ isNodeSelected: (nodeId: string) => boolean,
+): Set<string> => {
+ 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<string, Array<SelectionGraphEdge>>();
+ const incomingByNode = new Map<string, Array<SelectionGraphEdge>>();
+ const addTo = (map: Map<string, Array<SelectionGraphEdge>>, 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<string>();
+ const visitedForward = new Set<string>();
+ const visitedBackward = new Set<string>();
+
+ 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 00000000000..b485f9420c8
--- /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 72ff5b8327b..9e366796fd6 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-<dag_id>`), 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<string> = [];
+
+ 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 cf1174b881e..71f2b763e8a 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 2f486f15a2c..f467009f05c 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 8c2ff657453..73d92d89bfc 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<boolean>(SHOW_ALL_DEPENDENCIES_KEY, false);
const [direction] = useLocalStorage<Direction>(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<string>();
+ 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 6109d693190..8d8fa30757d 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<Dependency>(
- dependenciesKey(dagId),
- "tasks",
+ const [showAllDependencies, setShowAllDependencies] =
useLocalStorage<boolean>(
+ SHOW_ALL_DEPENDENCIES_KEY,
+ false,
);
const containerRef = useRef<HTMLDivElement>(null);
const containerWidth = useContainerWidth(containerRef);
@@ -136,14 +123,6 @@ export const PanelButtons = ({
}
}, [defaultLimit, enableResponsiveOptions, limit, setLimit]);
- const handleDepsChange = (event: SelectValueChangeDetails<{ label: string;
value: Array<string> }>) => {
- 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 = ({
<DagVersionSelect />
<DagRunSelect limit={limit} />
- <Select.Root
- // @ts-expect-error The expected option type is
incorrect
- collection={getOptions(translate)}
- data-testid="dependencies"
- onValueChange={handleDepsChange}
- size="sm"
- value={[dependencies]}
+ <Switch
+ checked={showAllDependencies}
+ data-testid="show-all-dependencies"
+ onCheckedChange={(details) =>
setShowAllDependencies(details.checked)}
>
- <Select.Label fontSize="xs">
- {translate("dag:panel.dependencies.label")}
- </Select.Label>
- <Select.Control>
- <Select.Trigger>
- <Select.ValueText
placeholder={translate("dag:panel.dependencies.label")} />
- </Select.Trigger>
- <Select.IndicatorGroup>
- <Select.Indicator />
- </Select.IndicatorGroup>
- </Select.Control>
- <Select.Positioner>
- <Select.Content>
- {getOptions(translate).items.map((option) => (
- <Select.Item item={option} key={option.value}>
- {option.label}
- </Select.Item>
- ))}
- </Select.Content>
- </Select.Positioner>
- </Select.Root>
+
{translate("dag:panel.dependencies.allDagDependencies")}
+ </Switch>
<DirectionDropdown graphId={dagId} />
</>
diff --git a/airflow-core/src/airflow/ui/src/main.tsx
b/airflow-core/src/airflow/ui/src/main.tsx
index e1ef1be5257..f3f8f2b0cdb 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(
<StrictMode>
<I18nextProvider i18n={i18n}>
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 c0f5218e5c6..e49c4a5cea2 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<string>();
+
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 eca2fa00875..3c266f856d9 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 @@ import pytest
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 @@ class TestGetDependencies:
@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 @@ class TestGetDependencies:
(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 @@ class TestGetDependencies:
)
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 6aaff421ee3..61b2eed0b3d 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 @@ from tests_common.test_utils.db import clear_db_assets,
clear_db_runs
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")
-
-
[email protected]
-def asset2() -> Asset:
- return Asset(uri="s3://bucket/next-run-asset/2", name="asset2")
-
-
[email protected]
-def asset3() -> Dataset:
- return Dataset(uri="s3://dataset-bucket/example.csv")
-
-
[email protected]
-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 make_dags(dag_maker, session, time_machine, asset1:
Asset, asset2: Asset, as
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)
- )
- )
-
-
[email protected]
-def asset1_id(make_dags, asset1, session: Session) -> str:
- return _fetch_asset_id(asset1, session)
-
-
[email protected]
-def asset2_id(make_dags, asset2, session) -> str:
- return _fetch_asset_id(asset2, session)
-
-
[email protected]
-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 @@ class TestStructureDataEndpoint:
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 @@ class TestStructureDataEndpoint:
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}
