Add built-in Extrenal Link for ExternalTaskMarker operator (apache#23964

)
chethanuk · Jun 4, 2022 · 84c9f4b · 84c9f4b
1 parent 55fd02a
commit 84c9f4b
Show file tree

Hide file tree

Showing 3 changed files with 38 additions and 6 deletions.
diff --git a/airflow/sensors/external_task.py b/airflow/sensors/external_task.py
@@ -18,8 +18,10 @@
 
 import datetime
 import os
+import warnings
 from typing import TYPE_CHECKING, Any, Callable, Collection, FrozenSet, Iterable, Optional, Union
 
+import attr
 from sqlalchemy import func
 
 from airflow.exceptions import AirflowException
@@ -31,10 +33,10 @@
 from airflow.utils.state import State
 
 
-class ExternalTaskSensorLink(BaseOperatorLink):
+class ExternalDagLink(BaseOperatorLink):
     """
-    Operator link for ExternalTaskSensor. It allows users to access
-    DAG waited with ExternalTaskSensor.
+    Operator link for ExternalTaskSensor and ExternalTaskMarker.
+    It allows users to access DAG waited with ExternalTaskSensor or cleared by ExternalTaskMarker.
     """
 
     name = 'External DAG'
@@ -83,7 +85,7 @@ class ExternalTaskSensor(BaseSensorOperator):
     @property
     def operator_extra_links(self):
         """Return operator extra links"""
-        return [ExternalTaskSensorLink()]
+        return [ExternalDagLink()]
 
     def __init__(
         self,
@@ -287,6 +289,11 @@ class ExternalTaskMarker(EmptyOperator):
     # The _serialized_fields are lazily loaded when get_serialized_fields() method is called
     __serialized_fields: Optional[FrozenSet[str]] = None
 
+    @property
+    def operator_extra_links(self):
+        """Return operator extra links"""
+        return [ExternalDagLink()]
+
     def __init__(
         self,
         *,
@@ -318,3 +325,19 @@ def get_serialized_fields(cls):
         if not cls.__serialized_fields:
             cls.__serialized_fields = frozenset(super().get_serialized_fields() | {"recursion_depth"})
         return cls.__serialized_fields
+
+
+@attr.s(auto_attribs=True)
+class ExternalTaskSensorLink(ExternalDagLink):
+    """
+    This external link is deprecated.
+    Please use :class:`airflow.sensors.external_task.ExternalDagLink`.
+    """
+
+    def __attrs_post_init__(self):
+        warnings.warn(
+            "This external link is deprecated. "
+            "Please use :class:`airflow.sensors.external_task.ExternalDagLink`.",
+            DeprecationWarning,
+            stacklevel=2,
+        )
diff --git a/airflow/serialization/serialized_objects.py b/airflow/serialization/serialized_objects.py
@@ -70,8 +70,9 @@
 
 _OPERATOR_EXTRA_LINKS: Set[str] = {
     "airflow.operators.trigger_dagrun.TriggerDagRunLink",
-    "airflow.sensors.external_task.ExternalTaskSensorLink",
+    "airflow.sensors.external_task.ExternalDagLink",
     # Deprecated names, so that existing serialized dags load straight away.
+    "airflow.sensors.external_task.ExternalTaskSensorLink",
     "airflow.operators.dagrun_operator.TriggerDagRunLink",
     "airflow.sensors.external_task_sensor.ExternalTaskSensorLink",
 }

diff --git a/tests/sensors/test_external_task_sensor.py b/tests/sensors/test_external_task_sensor.py
@@ -27,7 +27,7 @@
 from airflow.models.dag import DAG
 from airflow.operators.bash import BashOperator
 from airflow.operators.empty import EmptyOperator
-from airflow.sensors.external_task import ExternalTaskMarker, ExternalTaskSensor
+from airflow.sensors.external_task import ExternalTaskMarker, ExternalTaskSensor, ExternalTaskSensorLink
 from airflow.sensors.time_sensor import TimeSensor
 from airflow.serialization.serialized_objects import SerializedBaseOperator
 from airflow.utils.session import provide_session
@@ -977,3 +977,11 @@ def test_clear_overlapping_external_task_marker(dag_bag_head_tail, session):
         )
         == 30
     )
+
+
+class TestExternalTaskSensorLink:
+    def test_deprecation_warning(self):
+        with pytest.warns(DeprecationWarning) as warnings:
+            ExternalTaskSensorLink()
+            assert len(warnings) == 1
+            assert warnings[0].filename == __file__