Add support for labelling DAG edges (#15142)

This adds support for putting human-readable labels on edges in the DAG
between Tasks, as well as making the underlying framework for that
generic enough that future metadata could be added if desired.

Author: Andrew Godwin

airflow/models/baseoperator.py

@@ -66,6 +66,7 @@
 from airflow.ti_deps.deps.trigger_rule_dep import TriggerRuleDep
 from airflow.utils import timezone
 from airflow.utils.decorators import apply_defaults
+from airflow.utils.edgemodifier import EdgeModifier
 from airflow.utils.helpers import validate_key
 from airflow.utils.log.logging_mixin import LoggingMixin
 from airflow.utils.operator_resources import Resources
@@ -1205,6 +1206,7 @@ def _set_relatives(
         self,
         task_or_task_list: Union[TaskMixin, Sequence[TaskMixin]],
         upstream: bool = False,
+        edge_modifier: Optional[EdgeModifier] = None,
     ) -> None:
         """Sets relatives for the task or task list."""
         if not isinstance(task_or_task_list, Sequence):
@@ -1259,23 +1261,35 @@ def _set_relatives(
             if upstream:
                 task.add_only_new(task.get_direct_relative_ids(upstream=False), self.task_id)
                 self.add_only_new(self._upstream_task_ids, task.task_id)
+                if edge_modifier:
+                    edge_modifier.add_edge_info(self.dag, task.task_id, self.task_id)
             else:
                 self.add_only_new(self._downstream_task_ids, task.task_id)
                 task.add_only_new(task.get_direct_relative_ids(upstream=True), self.task_id)
+                if edge_modifier:
+                    edge_modifier.add_edge_info(self.dag, self.task_id, task.task_id)
 
-    def set_downstream(self, task_or_task_list: Union[TaskMixin, Sequence[TaskMixin]]) -> None:
+    def set_downstream(
+        self,
+        task_or_task_list: Union[TaskMixin, Sequence[TaskMixin]],
+        edge_modifier: Optional[EdgeModifier] = None,
+    ) -> None:
         """
         Set a task or a task list to be directly downstream from the current
         task. Required by TaskMixin.
         """
-        self._set_relatives(task_or_task_list, upstream=False)
+        self._set_relatives(task_or_task_list, upstream=False, edge_modifier=edge_modifier)
 
-    def set_upstream(self, task_or_task_list: Union[TaskMixin, Sequence[TaskMixin]]) -> None:
+    def set_upstream(
+        self,
+        task_or_task_list: Union[TaskMixin, Sequence[TaskMixin]],
+        edge_modifier: Optional[EdgeModifier] = None,
+    ) -> None:
         """
         Set a task or a task list to be directly upstream from the current
         task. Required by TaskMixin.
         """
-        self._set_relatives(task_or_task_list, upstream=True)
+        self._set_relatives(task_or_task_list, upstream=True, edge_modifier=edge_modifier)
 
     @property
     def output(self):

airflow/models/dag.py

@@ -73,7 +73,7 @@
 from airflow.utils.session import provide_session
 from airflow.utils.sqlalchemy import Interval, UtcDateTime, skip_locked, with_row_locks
 from airflow.utils.state import State
-from airflow.utils.types import DagRunType
+from airflow.utils.types import DagRunType, EdgeInfoType
 
 if TYPE_CHECKING:
     from airflow.utils.task_group import TaskGroup
@@ -349,6 +349,11 @@ def __init__(
         self.on_success_callback = on_success_callback
         self.on_failure_callback = on_failure_callback
 
+        # Keeps track of any extra edge metadata (sparse; will not contain all
+        # edges, so do not iterate over it for that). Outer key is upstream
+        # task ID, inner key is downstream task ID.
+        self.edge_info: Dict[str, Dict[str, EdgeInfoType]] = {}
+
         # To keep it in parity with Serialized DAGs
         # and identify if DAG has on_*_callback without actually storing them in Serialized JSON
         self.has_on_success_callback = self.on_success_callback is not None
@@ -2050,6 +2055,24 @@ def get_serialized_fields(cls):
             }
         return cls.__serialized_fields
 
+    def get_edge_info(self, upstream_task_id: str, downstream_task_id: str) -> EdgeInfoType:
+        """
+        Returns edge information for the given pair of tasks if present, and
+        None if there is no information.
+        """
+        # Note - older serialized DAGs may not have edge_info being a dict at all
+        if self.edge_info:
+            return self.edge_info.get(upstream_task_id, {}).get(downstream_task_id, {})
+        else:
+            return {}
+
+    def set_edge_info(self, upstream_task_id: str, downstream_task_id: str, info: EdgeInfoType):
+        """
+        Sets the given edge information on the DAG. Note that this will overwrite,
+        rather than merge with, existing info.
+        """
+        self.edge_info.setdefault(upstream_task_id, {})[downstream_task_id] = info
+
 
 class DagTag(Base):
     """A tag name per dag, to allow quick filtering in the DAG view."""

airflow/models/xcom_arg.py

@@ -15,12 +15,13 @@
 # specific language governing permissions and limitations
 # under the License.
 
-from typing import Any, Dict, List, Sequence, Union
+from typing import Any, Dict, List, Optional, Sequence, Union
 
 from airflow.exceptions import AirflowException
 from airflow.models.baseoperator import BaseOperator  # pylint: disable=R0401
 from airflow.models.taskmixin import TaskMixin
 from airflow.models.xcom import XCOM_RETURN_KEY
+from airflow.utils.edgemodifier import EdgeModifier
 
 
 class XComArg(TaskMixin):
@@ -111,13 +112,21 @@ def key(self) -> str:
         """Returns keys of this XComArg"""
         return self._key
 
-    def set_upstream(self, task_or_task_list: Union[TaskMixin, Sequence[TaskMixin]]):
+    def set_upstream(
+        self,
+        task_or_task_list: Union[TaskMixin, Sequence[TaskMixin]],
+        edge_modifier: Optional[EdgeModifier] = None,
+    ):
         """Proxy to underlying operator set_upstream method. Required by TaskMixin."""
-        self.operator.set_upstream(task_or_task_list)
+        self.operator.set_upstream(task_or_task_list, edge_modifier)
 
-    def set_downstream(self, task_or_task_list: Union[TaskMixin, Sequence[TaskMixin]]):
+    def set_downstream(
+        self,
+        task_or_task_list: Union[TaskMixin, Sequence[TaskMixin]],
+        edge_modifier: Optional[EdgeModifier] = None,
+    ):
         """Proxy to underlying operator set_downstream method. Required by TaskMixin."""
-        self.operator.set_downstream(task_or_task_list)
+        self.operator.set_downstream(task_or_task_list, edge_modifier)
 
     def resolve(self, context: Dict) -> Any:
         """

airflow/serialization/enums.py

@@ -46,3 +46,4 @@ class DagAttributeTypes(str, Enum):
     TUPLE = 'tuple'
     POD = 'k8s.V1Pod'
     TASK_GROUP = 'taskgroup'
+    EDGE_INFO = 'edgeinfo'

airflow/serialization/schema.json

@@ -102,7 +102,8 @@
         "_task_group": {"anyOf": [
           { "type": "null" },
           { "$ref": "#/definitions/task_group" }
-        ]}
+        ]},
+        "edge_info": { "$ref": "#/definitions/edge_info" }
       },
       "required": [
         "_dag_id",
@@ -217,6 +218,21 @@
         }
       },
       "additionalProperties": false
+    },
+    "edge_info": {
+      "$comment": "Metadata about DAG edges",
+      "type": "object",
+      "additionalProperties": {
+        "type": "object",
+        "additionalProperties": {
+          "type": "object",
+          "properties": {
+            "label": { "type": "string" }
+          },
+          "required": ["label"],
+          "additionalProperties": false
+        }
+      }
     }
   },
 

airflow/serialization/serialized_objects.py

@@ -647,6 +647,9 @@ def serialize_dag(cls, dag: DAG) -> dict:
             serialize_dag["tasks"] = [cls._serialize(task) for _, task in dag.task_dict.items()]
             serialize_dag['_task_group'] = SerializedTaskGroup.serialize_task_group(dag.task_group)
 
+            # Edge info in the JSON exactly matches our internal structure
+            serialize_dag["edge_info"] = dag.edge_info
+
             # has_on_*_callback are only stored if the value is True, as the default is False
             if dag.has_on_success_callback:
                 serialize_dag['has_on_success_callback'] = True
@@ -678,6 +681,9 @@ def deserialize_dag(cls, encoded_dag: Dict[str, Any]) -> 'SerializedDAG':
                 v = cls._deserialize_timedelta(v)
             elif k.endswith("_date"):
                 v = cls._deserialize_datetime(v)
+            elif k == "edge_info":
+                # Value structure matches exactly
+                pass
             elif k in cls._decorated_fields:
                 v = cls._deserialize(v)
             # else use v as it is

airflow/utils/dot_renderer.py

@@ -156,6 +156,10 @@ def render_dag(dag: DAG, tis: Optional[List[TaskInstance]] = None) -> graphviz.D
     _draw_nodes(dag.task_group, dot, states_by_task_id)
 
     for edge in dag_edges(dag):
-        dot.edge(edge["source_id"], edge["target_id"])
+        # Gets an optional label for the edge; this will be None if none is specified.
+        label = dag.get_edge_info(edge["source_id"], edge["target_id"]).get("label")
+        # Add the edge to the graph with optional label
+        # (we can just use the maybe-None label variable directly)
+        dot.edge(edge["source_id"], edge["target_id"], label)
 
     return dot

airflow/utils/edgemodifier.py

@@ -0,0 +1,126 @@
+# 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.
+
+from typing import Sequence, Union
+
+from airflow.models.taskmixin import TaskMixin
+
+
+class EdgeModifier(TaskMixin):
+    """
+    Class that represents edge information to be added between two
+    tasks/operators. Has shorthand factory functions, like Label("hooray").
+
+    Current implementation supports
+        t1 >> Label("Success route") >> t2
+        t2 << Label("Success route") << t2
+
+    Note that due to the potential for use in either direction, this waits
+    to make the actual connection between both sides until both are declared,
+    and will do so progressively if multiple ups/downs are added.
+
+    This and EdgeInfo are related - an EdgeModifier is the Python object you
+    use to add information to (potentially multiple) edges, and EdgeInfo
+    is the representation of the information for one specific edge.
+    """
+
+    def __init__(self, label: str = None):
+        self.label = label
+        self._upstream = []
+        self._downstream = []
+
+    @property
+    def roots(self):
+        """Should return list of root operator List["BaseOperator"]"""
+        return self._downstream
+
+    @property
+    def leaves(self):
+        """Should return list of leaf operator List["BaseOperator"]"""
+        return self._upstream
+
+    def set_upstream(self, task_or_task_list: Union[TaskMixin, Sequence[TaskMixin]], chain: bool = True):
+        """
+        Sets the given task/list onto the upstream attribute, and then checks if
+        we have both sides so we can resolve the relationship.
+
+        Providing this also provides << via TaskMixin.
+        """
+        # Ensure we have a list, even if it's just one item
+        if not isinstance(task_or_task_list, list):
+            task_or_task_list = [task_or_task_list]
+        # Unfurl it into actual operators
+        operators = []
+        for task in task_or_task_list:
+            operators.extend(task.roots)
+        # For each already-declared downstream, pair off with each new upstream
+        # item and store the edge info.
+        for operator in operators:
+            for downstream in self._downstream:
+                self.add_edge_info(operator.dag, operator.task_id, downstream.task_id)
+                if chain:
+                    operator.set_downstream(downstream)
+        # Add the new tasks to our list of ones we've seen
+        self._upstream.extend(operators)
+
+    def set_downstream(self, task_or_task_list: Union[TaskMixin, Sequence[TaskMixin]], chain: bool = True):
+        """
+        Sets the given task/list onto the downstream attribute, and then checks if
+        we have both sides so we can resolve the relationship.
+
+        Providing this also provides >> via TaskMixin.
+        """
+        # Ensure we have a list, even if it's just one item
+        if not isinstance(task_or_task_list, list):
+            task_or_task_list = [task_or_task_list]
+        # Unfurl it into actual operators
+        operators = []
+        for task in task_or_task_list:
+            operators.extend(task.leaves)
+        # Pair them off with existing
+        for operator in operators:
+            for upstream in self._upstream:
+                self.add_edge_info(upstream.dag, upstream.task_id, operator.task_id)
+                if chain:
+                    upstream.set_downstream(operator)
+        # Add the new tasks to our list of ones we've seen
+        self._downstream.extend(operators)
+
+    def update_relative(self, other: "TaskMixin", upstream: bool = True) -> None:
+        """
+        Called if we're not the "main" side of a relationship; we still run the
+        same logic, though.
+        """
+        if upstream:
+            self.set_upstream(other, chain=False)
+        else:
+            self.set_downstream(other, chain=False)
+
+    def add_edge_info(self, dag, upstream_id: str, downstream_id: str):
+        """
+        Adds or updates task info on the DAG for this specific pair of tasks.
+
+        Called either from our relationship trigger methods above, or directly
+        by set_upstream/set_downstream in operators.
+        """
+        dag.set_edge_info(upstream_id, downstream_id, {"label": self.label})
+
+
+# Factory functions
+def Label(label: str):  # pylint: disable=C0103
+    """Creates an EdgeModifier that sets a human-readable label on the edge."""
+    return EdgeModifier(label=label)

airflow/utils/types.py

@@ -15,6 +15,9 @@
 # specific language governing permissions and limitations
 # under the License.
 import enum
+from typing import Optional
+
+from airflow.typing_compat import TypedDict
 
 
 class DagRunType(str, enum.Enum):
@@ -34,3 +37,12 @@ def from_run_id(run_id: str) -> "DagRunType":
             if run_id and run_id.startswith(f"{run_type.value}__"):
                 return run_type
         return DagRunType.MANUAL
+
+
+class EdgeInfoType(TypedDict):
+    """
+    Represents extra metadata that the DAG can store about an edge,
+    usually generated from an EdgeModifier.
+    """
+
+    label: Optional[str]

airflow/www/templates/airflow/graph.html

@@ -686,6 +686,7 @@
             g.setEdge(source_id, target_id, {
               curve: d3.curveBasis,
               arrowheadClass: 'arrowhead',
+              label: edge.label
             });
         })