pierrejeambrun commented on code in PR #69833:
URL: https://github.com/apache/airflow/pull/69833#discussion_r3614405001
##########
airflow-core/src/airflow/api_fastapi/core_api/services/ui/dependencies.py:
##########
@@ -20,12 +20,136 @@
from collections import defaultdict, deque
from typing import TYPE_CHECKING
+import structlog
+from sqlalchemy import select
+
from airflow.models.asset import AssetModel
from airflow.models.dag import DagModel
if TYPE_CHECKING:
from sqlalchemy.orm import Session
+log = structlog.get_logger(logger_name=__name__)
+
+ASSET_UPSTREAM_DEPENDENCY_TYPES = ("asset", "asset-alias", "asset-name-ref",
"asset-uri-ref")
+
+
+def _asset_node_id_and_label(asset: dict) -> tuple[str, str]:
+ """Get the graph node id and label for a single asset/alias/ref entry in
an asset_expression."""
+ asset_type = asset["type"]
+
+ if asset_type == "asset":
+ return f"asset:{asset['id']}", asset["name"]
+ if asset_type in ("asset-alias", "asset-name-ref"):
+ return f"{asset_type}:{asset['name']}", asset["name"]
+ if asset_type == "asset-uri-ref":
+ return f"{asset_type}:{asset['uri']}", asset["uri"]
+ raise TypeError(f"Unsupported type: {asset_type}")
+
+
+def get_upstream_assets(
+ asset_expression: dict, entry_node_ref: str, level: int = 0
+) -> tuple[list[dict], list[dict]]:
+ """Expand a Dag's asset trigger condition into asset-condition (AND/OR
gate) nodes and edges."""
+ edges: list[dict] = []
+ nodes: list[dict] = []
+ asset_expression_type: str | None = None
+
+ # include assets, asset-alias, asset-name-refs, asset-uri-refs
+ assets_info: list[dict] = []
+
+ nested_expression: dict = {}
+
+ expr_key = ""
+ if asset_expression.keys() == {"any"}:
+ asset_expression_type = "or-gate"
+ expr_key = "any"
+ elif asset_expression.keys() == {"all"}:
+ asset_expression_type = "and-gate"
+ expr_key = "all"
+
+ if expr_key in asset_expression:
+ asset_exprs: list[dict] = asset_expression[expr_key]
+ for expr in asset_exprs:
+ nested_expr_key = next(iter(expr.keys()))
+ if nested_expr_key in ("any", "all"):
+ nested_expression = expr
Review Comment:
`nested_expression` is overwritten each iteration, so a condition with 2+
nested gates at the same level (e.g. `(A|B) & (C|D)` →
`{"all":[{"any":…},{"any":…}]}`) keeps only the last — the other branch is
silently dropped. Collect nested expressions into a list and expand each.
##########
airflow-core/src/airflow/api_fastapi/core_api/services/ui/dependencies.py:
##########
@@ -163,6 +333,38 @@ def get_scheduling_dependencies(readable_dag_ids: set[str]
| None = None) -> dic
}
+def _get_dag_entry_point(dag_id: str, session: Session) -> tuple[str, dict |
None] | None:
+ """Get the id of a Dag's topologically-first task/group and its
asset_expression, if any."""
+ from sqlalchemy import select
+ from sqlalchemy.orm import joinedload
+
+ from airflow.api_fastapi.core_api.services.ui.task_group import
task_group_to_dict
+ from airflow.models.dag_version import DagVersion
+ from airflow.models.serialized_dag import SerializedDagModel
+
+ dag_version = DagVersion.get_latest_version(dag_id, session=session)
+ if dag_version is None:
+ return None
+
+ serialized_dag = session.scalar(
+ select(SerializedDagModel)
+ .join(DagVersion)
+ .where(
+ SerializedDagModel.dag_id == dag_id,
+ DagVersion.version_number == dag_version.version_number,
+ )
+ .options(joinedload(SerializedDagModel.dag_model))
+ )
+ if serialized_dag is None:
+ return None
+
+ entry_points = serialized_dag.dag.task_group.topological_sort()
Review Comment:
This deserializes the full DAG for every scheduled Dag in the BFS just to
read its first task — N full-DAG loads for an asset feeding many Dags. Consider
batching the serialized-dag fetch or deriving the entry task without full
deserialization.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
