From 946ce2bc67ce55a00b6aecca9a13e87b558d09c0 Mon Sep 17 00:00:00 2001
From: Ephraim Anierobi <splendidzigy24@gmail.com>
Date: Sat, 6 Aug 2022 00:29:01 +0100
Subject: [PATCH 1/6] Fix mapped sensor with reschedule mode

There are two issues with mapped sensor with `reschedule` mode. First, the reschedule table is being populated with a default map_index of -1 even when the map_index is not -1. Secondly, MappedOperator does not have the `ReadyToReschedule` dependency.
This PR is an attempt to fix this
---
 airflow/decorators/base.py                  |  1 +
 airflow/models/baseoperator.py              |  2 ++
 airflow/models/mappedoperator.py            | 10 ++++++++++
 airflow/models/taskinstance.py              |  1 +
 airflow/models/taskreschedule.py            |  1 +
 airflow/sensors/base.py                     |  5 -----
 airflow/serialization/serialized_objects.py |  1 +
 7 files changed, 16 insertions(+), 5 deletions(-)

diff --git a/airflow/decorators/base.py b/airflow/decorators/base.py
index cfc94b525ed25..0fa28fd48a993 100644
--- a/airflow/decorators/base.py
+++ b/airflow/decorators/base.py
@@ -405,6 +405,7 @@ def _expand(self, expand_input: ExpandInput, *, strict: bool) -> XComArg:
             ui_color=self.operator_class.ui_color,
             ui_fgcolor=self.operator_class.ui_fgcolor,
             is_empty=False,
+            is_sensor=False,
             task_module=self.operator_class.__module__,
             task_type=self.operator_class.__name__,
             dag=dag,
diff --git a/airflow/models/baseoperator.py b/airflow/models/baseoperator.py
index 395e1d3b676e3..d8358a42e9336 100644
--- a/airflow/models/baseoperator.py
+++ b/airflow/models/baseoperator.py
@@ -83,6 +83,7 @@
 from airflow.ti_deps.deps.not_in_retry_period_dep import NotInRetryPeriodDep
 from airflow.ti_deps.deps.not_previously_skipped_dep import NotPreviouslySkippedDep
 from airflow.ti_deps.deps.prev_dagrun_dep import PrevDagrunDep
+from airflow.ti_deps.deps.ready_to_reschedule import ReadyToRescheduleDep
 from airflow.ti_deps.deps.trigger_rule_dep import TriggerRuleDep
 from airflow.triggers.base import BaseTrigger
 from airflow.utils import timezone
@@ -1075,6 +1076,7 @@ def has_dag(self):
             PrevDagrunDep(),
             TriggerRuleDep(),
             NotPreviouslySkippedDep(),
+            ReadyToRescheduleDep(),
         }
     )
     """
diff --git a/airflow/models/mappedoperator.py b/airflow/models/mappedoperator.py
index f62d1c687b58e..52856908a60ed 100644
--- a/airflow/models/mappedoperator.py
+++ b/airflow/models/mappedoperator.py
@@ -202,6 +202,7 @@ def expand_kwargs(self, kwargs: "XComArg", *, strict: bool = True) -> "MappedOpe
 
     def _expand(self, expand_input: ExpandInput, *, strict: bool) -> "MappedOperator":
         from airflow.operators.empty import EmptyOperator
+        from airflow.sensors.base import BaseSensorOperator
 
         self._expand_called = True
         ensure_xcomarg_return_value(expand_input.value)
@@ -228,6 +229,7 @@ def _expand(self, expand_input: ExpandInput, *, strict: bool) -> "MappedOperator
             ui_color=self.operator_class.ui_color,
             ui_fgcolor=self.operator_class.ui_fgcolor,
             is_empty=issubclass(self.operator_class, EmptyOperator),
+            is_sensor=issubclass(self.operator_class, BaseSensorOperator),
             task_module=self.operator_class.__module__,
             task_type=self.operator_class.__name__,
             dag=dag,
@@ -277,6 +279,7 @@ class MappedOperator(AbstractOperator):
     ui_color: str
     ui_fgcolor: str
     _is_empty: bool
+    _is_sensor: bool
     _task_module: str
     _task_type: str
 
@@ -383,6 +386,13 @@ def inherits_from_empty_operator(self) -> bool:
         """Implementing Operator."""
         return self._is_empty
 
+    @property
+    def reschedule(self) -> bool:
+        """Check if the operator is a sensor and has mode reschedule"""
+        if self._is_sensor:
+            return self.partial_kwargs['mode'] == 'reschedule'
+        raise AttributeError("reschedule is not defined for MappedOperator")
+
     @property
     def roots(self) -> Sequence[AbstractOperator]:
         """Implementing DAGNode."""
diff --git a/airflow/models/taskinstance.py b/airflow/models/taskinstance.py
index 7930d91fb8a64..ebdf07a3820e1 100644
--- a/airflow/models/taskinstance.py
+++ b/airflow/models/taskinstance.py
@@ -1806,6 +1806,7 @@ def _handle_reschedule(
                 actual_start_date,
                 self.end_date,
                 reschedule_exception.reschedule_date,
+                self.map_index,
             )
         )
 
diff --git a/airflow/models/taskreschedule.py b/airflow/models/taskreschedule.py
index 48e91783454fc..3dc0e6d609c00 100644
--- a/airflow/models/taskreschedule.py
+++ b/airflow/models/taskreschedule.py
@@ -112,6 +112,7 @@ def query_for_task_instance(task_instance, descending=False, session=None, try_n
             TR.dag_id == task_instance.dag_id,
             TR.task_id == task_instance.task_id,
             TR.run_id == task_instance.run_id,
+            TR.map_index == task_instance.map_index,
             TR.try_number == try_number,
         )
         if descending:
diff --git a/airflow/sensors/base.py b/airflow/sensors/base.py
index a3621bf1a49c1..f5f4e65c1b360 100644
--- a/airflow/sensors/base.py
+++ b/airflow/sensors/base.py
@@ -34,7 +34,6 @@
 from airflow.models.baseoperator import BaseOperator
 from airflow.models.skipmixin import SkipMixin
 from airflow.models.taskreschedule import TaskReschedule
-from airflow.ti_deps.deps.ready_to_reschedule import ReadyToRescheduleDep
 from airflow.utils import timezone
 from airflow.utils.context import Context
 
@@ -101,10 +100,6 @@ class BaseSensorOperator(BaseOperator, SkipMixin):
     ui_color = '#e6f1f2'  # type: str
     valid_modes = ['poke', 'reschedule']  # type: Iterable[str]
 
-    # Adds one additional dependency for all sensor operators that checks if a
-    # sensor task instance can be rescheduled.
-    deps = BaseOperator.deps | {ReadyToRescheduleDep()}
-
     def __init__(
         self,
         *,
diff --git a/airflow/serialization/serialized_objects.py b/airflow/serialization/serialized_objects.py
index 9ddec956c5425..2e00aeeeb11ba 100644
--- a/airflow/serialization/serialized_objects.py
+++ b/airflow/serialization/serialized_objects.py
@@ -845,6 +845,7 @@ def deserialize_operator(cls, encoded_op: Dict[str, Any]) -> Operator:
                 ui_color=BaseOperator.ui_color,
                 ui_fgcolor=BaseOperator.ui_fgcolor,
                 is_empty=False,
+                is_sensor=False,
                 task_module=encoded_op["_task_module"],
                 task_type=encoded_op["_task_type"],
                 dag=None,

From 6131dbabf6d9cd6fb13d70d09d3520288c63e44f Mon Sep 17 00:00:00 2001
From: Ephraim Anierobi <splendidzigy24@gmail.com>
Date: Mon, 8 Aug 2022 21:03:58 +0100
Subject: [PATCH 2/6] Move logic to ready_to_reschedule dep

---
 airflow/decorators/base.py                         |  1 -
 airflow/models/mappedoperator.py                   | 10 ----------
 airflow/serialization/serialized_objects.py        |  1 -
 airflow/ti_deps/deps/ready_to_reschedule.py        |  6 +++++-
 tests/serialization/test_dag_serialization.py      |  5 ++++-
 tests/ti_deps/deps/test_ready_to_reschedule_dep.py |  2 ++
 6 files changed, 11 insertions(+), 14 deletions(-)

diff --git a/airflow/decorators/base.py b/airflow/decorators/base.py
index 0fa28fd48a993..cfc94b525ed25 100644
--- a/airflow/decorators/base.py
+++ b/airflow/decorators/base.py
@@ -405,7 +405,6 @@ def _expand(self, expand_input: ExpandInput, *, strict: bool) -> XComArg:
             ui_color=self.operator_class.ui_color,
             ui_fgcolor=self.operator_class.ui_fgcolor,
             is_empty=False,
-            is_sensor=False,
             task_module=self.operator_class.__module__,
             task_type=self.operator_class.__name__,
             dag=dag,
diff --git a/airflow/models/mappedoperator.py b/airflow/models/mappedoperator.py
index 52856908a60ed..f62d1c687b58e 100644
--- a/airflow/models/mappedoperator.py
+++ b/airflow/models/mappedoperator.py
@@ -202,7 +202,6 @@ def expand_kwargs(self, kwargs: "XComArg", *, strict: bool = True) -> "MappedOpe
 
     def _expand(self, expand_input: ExpandInput, *, strict: bool) -> "MappedOperator":
         from airflow.operators.empty import EmptyOperator
-        from airflow.sensors.base import BaseSensorOperator
 
         self._expand_called = True
         ensure_xcomarg_return_value(expand_input.value)
@@ -229,7 +228,6 @@ def _expand(self, expand_input: ExpandInput, *, strict: bool) -> "MappedOperator
             ui_color=self.operator_class.ui_color,
             ui_fgcolor=self.operator_class.ui_fgcolor,
             is_empty=issubclass(self.operator_class, EmptyOperator),
-            is_sensor=issubclass(self.operator_class, BaseSensorOperator),
             task_module=self.operator_class.__module__,
             task_type=self.operator_class.__name__,
             dag=dag,
@@ -279,7 +277,6 @@ class MappedOperator(AbstractOperator):
     ui_color: str
     ui_fgcolor: str
     _is_empty: bool
-    _is_sensor: bool
     _task_module: str
     _task_type: str
 
@@ -386,13 +383,6 @@ def inherits_from_empty_operator(self) -> bool:
         """Implementing Operator."""
         return self._is_empty
 
-    @property
-    def reschedule(self) -> bool:
-        """Check if the operator is a sensor and has mode reschedule"""
-        if self._is_sensor:
-            return self.partial_kwargs['mode'] == 'reschedule'
-        raise AttributeError("reschedule is not defined for MappedOperator")
-
     @property
     def roots(self) -> Sequence[AbstractOperator]:
         """Implementing DAGNode."""
diff --git a/airflow/serialization/serialized_objects.py b/airflow/serialization/serialized_objects.py
index 2e00aeeeb11ba..9ddec956c5425 100644
--- a/airflow/serialization/serialized_objects.py
+++ b/airflow/serialization/serialized_objects.py
@@ -845,7 +845,6 @@ def deserialize_operator(cls, encoded_op: Dict[str, Any]) -> Operator:
                 ui_color=BaseOperator.ui_color,
                 ui_fgcolor=BaseOperator.ui_fgcolor,
                 is_empty=False,
-                is_sensor=False,
                 task_module=encoded_op["_task_module"],
                 task_type=encoded_op["_task_type"],
                 dag=None,
diff --git a/airflow/ti_deps/deps/ready_to_reschedule.py b/airflow/ti_deps/deps/ready_to_reschedule.py
index 9086822ceac8d..224f50ee7eeb6 100644
--- a/airflow/ti_deps/deps/ready_to_reschedule.py
+++ b/airflow/ti_deps/deps/ready_to_reschedule.py
@@ -40,7 +40,8 @@ def _get_dep_statuses(self, ti, session, dep_context):
         considered as passed. This dependency fails if the latest reschedule
         request's reschedule date is still in future.
         """
-        if not getattr(ti.task, "reschedule", False):
+        is_mapped = ti.task.is_mapped
+        if not is_mapped and not getattr(ti.task, "reschedule", False):
             yield self._passing_status(reason="Task is not in reschedule mode.")
             return
 
@@ -62,6 +63,9 @@ def _get_dep_statuses(self, ti, session, dep_context):
             .first()
         )
         if not task_reschedule:
+            if is_mapped:
+                yield self._passing_status(reason="The task is mapped and not in reschedule mode")
+                return
             yield self._passing_status(reason="There is no reschedule request for this task instance.")
             return
 
diff --git a/tests/serialization/test_dag_serialization.py b/tests/serialization/test_dag_serialization.py
index 37da287cf1de1..6c5e306b07402 100644
--- a/tests/serialization/test_dag_serialization.py
+++ b/tests/serialization/test_dag_serialization.py
@@ -1270,6 +1270,7 @@ def test_deps_sorted(self):
             task1 >> task2
 
         serialize_op = SerializedBaseOperator.serialize_operator(dag.task_dict["task1"])
+
         deps = serialize_op["deps"]
         assert deps == [
             'airflow.ti_deps.deps.not_in_retry_period_dep.NotInRetryPeriodDep',
@@ -1324,6 +1325,7 @@ class DummyTask(BaseOperator):
             'airflow.ti_deps.deps.not_in_retry_period_dep.NotInRetryPeriodDep',
             'airflow.ti_deps.deps.not_previously_skipped_dep.NotPreviouslySkippedDep',
             'airflow.ti_deps.deps.prev_dagrun_dep.PrevDagrunDep',
+            'airflow.ti_deps.deps.ready_to_reschedule.ReadyToRescheduleDep',
             'airflow.ti_deps.deps.trigger_rule_dep.TriggerRuleDep',
             'test_plugin.CustomTestTriggerRule',
         ]
@@ -1334,6 +1336,7 @@ class DummyTask(BaseOperator):
             '<TIDep(Not In Retry Period)>',
             '<TIDep(Not Previously Skipped)>',
             '<TIDep(Previous Dagrun State)>',
+            '<TIDep(Ready To Reschedule)>',
             '<TIDep(Trigger Rule)>',
         ]
 
@@ -1605,7 +1608,7 @@ def poke(self, context: Context):
         op = DummySensor(task_id='dummy', mode=mode, poke_interval=23)
 
         blob = SerializedBaseOperator.serialize_operator(op)
-        assert "deps" in blob
+        assert op.deps == BaseOperator.deps
 
         serialized_op = SerializedBaseOperator.deserialize_operator(blob)
         assert serialized_op.reschedule == (mode == "reschedule")
diff --git a/tests/ti_deps/deps/test_ready_to_reschedule_dep.py b/tests/ti_deps/deps/test_ready_to_reschedule_dep.py
index 99416bbbc8927..6448a0ec98423 100644
--- a/tests/ti_deps/deps/test_ready_to_reschedule_dep.py
+++ b/tests/ti_deps/deps/test_ready_to_reschedule_dep.py
@@ -32,11 +32,13 @@ class TestNotInReschedulePeriodDep(unittest.TestCase):
     def _get_task_instance(self, state):
         dag = DAG('test_dag')
         task = Mock(dag=dag, reschedule=True)
+        task.is_mapped = False
         ti = TaskInstance(task=task, state=state, run_id=None)
         return ti
 
     def _get_task_reschedule(self, reschedule_date):
         task = Mock(dag_id='test_dag', task_id='test_task')
+        task.is_mapped = False
         reschedule = TaskReschedule(
             task=task,
             run_id=None,

From 47f1fcd8de61eb1cf1c1442b0c4e6ae8778dbac3 Mon Sep 17 00:00:00 2001
From: Ephraim Anierobi <splendidzigy24@gmail.com>
Date: Tue, 9 Aug 2022 11:51:57 +0100
Subject: [PATCH 3/6] Add test to expose reschedule not in mapped sensor
 serialization

---
 airflow/models/baseoperator.py                |  2 --
 airflow/sensors/base.py                       |  5 +++
 airflow/serialization/serialized_objects.py   |  3 +-
 tests/serialization/test_dag_serialization.py | 33 +++++++++++++++++--
 4 files changed, 36 insertions(+), 7 deletions(-)

diff --git a/airflow/models/baseoperator.py b/airflow/models/baseoperator.py
index d8358a42e9336..395e1d3b676e3 100644
--- a/airflow/models/baseoperator.py
+++ b/airflow/models/baseoperator.py
@@ -83,7 +83,6 @@
 from airflow.ti_deps.deps.not_in_retry_period_dep import NotInRetryPeriodDep
 from airflow.ti_deps.deps.not_previously_skipped_dep import NotPreviouslySkippedDep
 from airflow.ti_deps.deps.prev_dagrun_dep import PrevDagrunDep
-from airflow.ti_deps.deps.ready_to_reschedule import ReadyToRescheduleDep
 from airflow.ti_deps.deps.trigger_rule_dep import TriggerRuleDep
 from airflow.triggers.base import BaseTrigger
 from airflow.utils import timezone
@@ -1076,7 +1075,6 @@ def has_dag(self):
             PrevDagrunDep(),
             TriggerRuleDep(),
             NotPreviouslySkippedDep(),
-            ReadyToRescheduleDep(),
         }
     )
     """
diff --git a/airflow/sensors/base.py b/airflow/sensors/base.py
index f5f4e65c1b360..a3621bf1a49c1 100644
--- a/airflow/sensors/base.py
+++ b/airflow/sensors/base.py
@@ -34,6 +34,7 @@
 from airflow.models.baseoperator import BaseOperator
 from airflow.models.skipmixin import SkipMixin
 from airflow.models.taskreschedule import TaskReschedule
+from airflow.ti_deps.deps.ready_to_reschedule import ReadyToRescheduleDep
 from airflow.utils import timezone
 from airflow.utils.context import Context
 
@@ -100,6 +101,10 @@ class BaseSensorOperator(BaseOperator, SkipMixin):
     ui_color = '#e6f1f2'  # type: str
     valid_modes = ['poke', 'reschedule']  # type: Iterable[str]
 
+    # Adds one additional dependency for all sensor operators that checks if a
+    # sensor task instance can be rescheduled.
+    deps = BaseOperator.deps | {ReadyToRescheduleDep()}
+
     def __init__(
         self,
         *,
diff --git a/airflow/serialization/serialized_objects.py b/airflow/serialization/serialized_objects.py
index 9ddec956c5425..c5bc9567fe188 100644
--- a/airflow/serialization/serialized_objects.py
+++ b/airflow/serialization/serialized_objects.py
@@ -642,8 +642,7 @@ def task_type(self, task_type: str):
 
     @classmethod
     def serialize_mapped_operator(cls, op: MappedOperator) -> Dict[str, Any]:
-        serialized_op = cls._serialize_node(op, include_deps=op.deps is MappedOperator.deps_for(BaseOperator))
-
+        serialized_op = cls._serialize_node(op, include_deps=op.deps != MappedOperator.deps_for(BaseOperator))
         # Handle expand_input and op_kwargs_expand_input.
         expansion_kwargs = op._get_specified_expand_input()
         serialized_op[op._expand_input_attr] = {
diff --git a/tests/serialization/test_dag_serialization.py b/tests/serialization/test_dag_serialization.py
index 6c5e306b07402..e2268e7f45137 100644
--- a/tests/serialization/test_dag_serialization.py
+++ b/tests/serialization/test_dag_serialization.py
@@ -45,6 +45,7 @@
 from airflow.models.xcom import XCOM_RETURN_KEY, XCom
 from airflow.operators.bash import BashOperator
 from airflow.security import permissions
+from airflow.sensors.bash import BashSensor
 from airflow.serialization.json_schema import load_dag_schema_dict
 from airflow.serialization.serialized_objects import (
     DagDependency,
@@ -1325,7 +1326,6 @@ class DummyTask(BaseOperator):
             'airflow.ti_deps.deps.not_in_retry_period_dep.NotInRetryPeriodDep',
             'airflow.ti_deps.deps.not_previously_skipped_dep.NotPreviouslySkippedDep',
             'airflow.ti_deps.deps.prev_dagrun_dep.PrevDagrunDep',
-            'airflow.ti_deps.deps.ready_to_reschedule.ReadyToRescheduleDep',
             'airflow.ti_deps.deps.trigger_rule_dep.TriggerRuleDep',
             'test_plugin.CustomTestTriggerRule',
         ]
@@ -1336,7 +1336,6 @@ class DummyTask(BaseOperator):
             '<TIDep(Not In Retry Period)>',
             '<TIDep(Not Previously Skipped)>',
             '<TIDep(Previous Dagrun State)>',
-            '<TIDep(Ready To Reschedule)>',
             '<TIDep(Trigger Rule)>',
         ]
 
@@ -1608,7 +1607,24 @@ def poke(self, context: Context):
         op = DummySensor(task_id='dummy', mode=mode, poke_interval=23)
 
         blob = SerializedBaseOperator.serialize_operator(op)
-        assert op.deps == BaseOperator.deps
+        assert "deps" in blob
+
+        serialized_op = SerializedBaseOperator.deserialize_operator(blob)
+        assert serialized_op.reschedule == (mode == "reschedule")
+        assert op.deps == serialized_op.deps
+
+    @pytest.mark.parametrize("mode", ["poke", "reschedule"])
+    def test_serialize_mapped_sensor(self, mode):
+        from airflow.sensors.base import BaseSensorOperator
+
+        class DummySensor(BaseSensorOperator):
+            def poke(self, context: Context):
+                return False
+
+        op = DummySensor.partial(task_id='dummy', mode=mode).expand(poke_interval=[23])
+
+        blob = SerializedBaseOperator.serialize_mapped_operator(op)
+        assert "deps" in blob
 
         serialized_op = SerializedBaseOperator.deserialize_operator(blob)
         assert serialized_op.reschedule == (mode == "reschedule")
@@ -1983,6 +1999,17 @@ def test_operator_expand_deserialized_unmap():
     assert deserialize(serialize(mapped)).unmap(None) == deserialize(serialize(normal))
 
 
+def test_sensor_expand_deserialized_unmap():
+    """Unmap a deserialized mapped sensor should be similar to deserializing a non-mapped sensor"""
+    normal = BashSensor(task_id='a', bash_command=[1, 2], mode='reschedule')
+    mapped = BashSensor.partial(task_id='a', mode='reschedule').expand(bash_command=[1, 2])
+
+    serialize = SerializedBaseOperator._serialize
+
+    deserialize = SerializedBaseOperator.deserialize_operator
+    assert deserialize(serialize(mapped)).unmap(None) == deserialize(serialize(normal))
+
+
 def test_task_resources_serde():
     """
     Test task resources serialization/deserialization.

From 243983577fc6decc7ee45bd11f9f8e8b099c0de2 Mon Sep 17 00:00:00 2001
From: Ephraim Anierobi <splendidzigy24@gmail.com>
Date: Tue, 9 Aug 2022 15:26:29 +0100
Subject: [PATCH 4/6] Apply suggestions from code review

Co-authored-by: Tzu-ping Chung <uranusjr@gmail.com>
---
 tests/ti_deps/deps/test_ready_to_reschedule_dep.py | 6 ++----
 1 file changed, 2 insertions(+), 4 deletions(-)

diff --git a/tests/ti_deps/deps/test_ready_to_reschedule_dep.py b/tests/ti_deps/deps/test_ready_to_reschedule_dep.py
index 6448a0ec98423..fa4d2cf2f08fd 100644
--- a/tests/ti_deps/deps/test_ready_to_reschedule_dep.py
+++ b/tests/ti_deps/deps/test_ready_to_reschedule_dep.py
@@ -31,14 +31,12 @@
 class TestNotInReschedulePeriodDep(unittest.TestCase):
     def _get_task_instance(self, state):
         dag = DAG('test_dag')
-        task = Mock(dag=dag, reschedule=True)
-        task.is_mapped = False
+        task = Mock(dag=dag, reschedule=True, is_mapped=False)
         ti = TaskInstance(task=task, state=state, run_id=None)
         return ti
 
     def _get_task_reschedule(self, reschedule_date):
-        task = Mock(dag_id='test_dag', task_id='test_task')
-        task.is_mapped = False
+        task = Mock(dag_id='test_dag', task_id='test_task', is_mapped=False)
         reschedule = TaskReschedule(
             task=task,
             run_id=None,

From 6235c94fecb93e6dfda9e856a67b4c096e0d0e64 Mon Sep 17 00:00:00 2001
From: Ephraim Anierobi <splendidzigy24@gmail.com>
Date: Tue, 9 Aug 2022 15:28:18 +0100
Subject: [PATCH 5/6] Add more tests

---
 airflow/ti_deps/deps/ready_to_reschedule.py   |   3 +
 tests/models/test_taskinstance.py             | 217 ++++++++++++++++++
 tests/serialization/test_dag_serialization.py |   6 +-
 .../deps/test_ready_to_reschedule_dep.py      |  77 +++++++
 4 files changed, 299 insertions(+), 4 deletions(-)

diff --git a/airflow/ti_deps/deps/ready_to_reschedule.py b/airflow/ti_deps/deps/ready_to_reschedule.py
index 224f50ee7eeb6..a3f7ab59b7aa7 100644
--- a/airflow/ti_deps/deps/ready_to_reschedule.py
+++ b/airflow/ti_deps/deps/ready_to_reschedule.py
@@ -42,6 +42,7 @@ def _get_dep_statuses(self, ti, session, dep_context):
         """
         is_mapped = ti.task.is_mapped
         if not is_mapped and not getattr(ti.task, "reschedule", False):
+            # Mapped sensors don't currently have the reschedule property,
             yield self._passing_status(reason="Task is not in reschedule mode.")
             return
 
@@ -63,6 +64,8 @@ def _get_dep_statuses(self, ti, session, dep_context):
             .first()
         )
         if not task_reschedule:
+            # Because mapped sensors don't have the reschedule property, here's the last resort
+            # and we need a slightly different passing reason
             if is_mapped:
                 yield self._passing_status(reason="The task is mapped and not in reschedule mode")
                 return
diff --git a/tests/models/test_taskinstance.py b/tests/models/test_taskinstance.py
index 0ef15d3018b3d..76a9a303f519b 100644
--- a/tests/models/test_taskinstance.py
+++ b/tests/models/test_taskinstance.py
@@ -801,6 +801,174 @@ def run_ti_and_assert(
         done, fail = True, False
         run_ti_and_assert(date4, date3, date4, 60, State.SUCCESS, 3, 0)
 
+    def test_mapped_reschedule_handling(self, dag_maker):
+        """
+        Test that mapped task reschedules are handled properly
+        """
+        # Return values of the python sensor callable, modified during tests
+        done = False
+        fail = False
+
+        def func():
+            if fail:
+                raise AirflowException()
+            return done
+
+        with dag_maker(dag_id='test_reschedule_handling') as dag:
+
+            task = PythonSensor.partial(
+                task_id='test_reschedule_handling_sensor',
+                mode='reschedule',
+                python_callable=func,
+                retries=1,
+                retry_delay=datetime.timedelta(seconds=0),
+            ).expand(poke_interval=[0])
+
+        ti = dag_maker.create_dagrun(execution_date=timezone.utcnow()).task_instances[0]
+
+        ti.task = task
+        assert ti._try_number == 0
+        assert ti.try_number == 1
+
+        def run_ti_and_assert(
+            run_date,
+            expected_start_date,
+            expected_end_date,
+            expected_duration,
+            expected_state,
+            expected_try_number,
+            expected_task_reschedule_count,
+        ):
+            ti.refresh_from_task(task)
+            with freeze_time(run_date):
+                try:
+                    ti.run()
+                except AirflowException:
+                    if not fail:
+                        raise
+            ti.refresh_from_db()
+            assert ti.state == expected_state
+            assert ti._try_number == expected_try_number
+            assert ti.try_number == expected_try_number + 1
+            assert ti.start_date == expected_start_date
+            assert ti.end_date == expected_end_date
+            assert ti.duration == expected_duration
+            trs = TaskReschedule.find_for_task_instance(ti)
+            assert len(trs) == expected_task_reschedule_count
+
+        date1 = timezone.utcnow()
+        date2 = date1 + datetime.timedelta(minutes=1)
+        date3 = date2 + datetime.timedelta(minutes=1)
+        date4 = date3 + datetime.timedelta(minutes=1)
+
+        # Run with multiple reschedules.
+        # During reschedule the try number remains the same, but each reschedule is recorded.
+        # The start date is expected to remain the initial date, hence the duration increases.
+        # When finished the try number is incremented and there is no reschedule expected
+        # for this try.
+
+        done, fail = False, False
+        run_ti_and_assert(date1, date1, date1, 0, State.UP_FOR_RESCHEDULE, 0, 1)
+
+        done, fail = False, False
+        run_ti_and_assert(date2, date1, date2, 60, State.UP_FOR_RESCHEDULE, 0, 2)
+
+        done, fail = False, False
+        run_ti_and_assert(date3, date1, date3, 120, State.UP_FOR_RESCHEDULE, 0, 3)
+
+        done, fail = True, False
+        run_ti_and_assert(date4, date1, date4, 180, State.SUCCESS, 1, 0)
+
+        # Clear the task instance.
+        dag.clear()
+        ti.refresh_from_db()
+        assert ti.state == State.NONE
+        assert ti._try_number == 1
+
+        # Run again after clearing with reschedules and a retry.
+        # The retry increments the try number, and for that try no reschedule is expected.
+        # After the retry the start date is reset, hence the duration is also reset.
+
+        done, fail = False, False
+        run_ti_and_assert(date1, date1, date1, 0, State.UP_FOR_RESCHEDULE, 1, 1)
+
+        done, fail = False, True
+        run_ti_and_assert(date2, date1, date2, 60, State.UP_FOR_RETRY, 2, 0)
+
+        done, fail = False, False
+        run_ti_and_assert(date3, date3, date3, 0, State.UP_FOR_RESCHEDULE, 2, 1)
+
+        done, fail = True, False
+        run_ti_and_assert(date4, date3, date4, 60, State.SUCCESS, 3, 0)
+
+    @pytest.mark.usefixtures('test_pool')
+    def test_mapped_task_reschedule_handling_clear_reschedules(self, dag_maker):
+        """
+        Test that mapped task reschedules clearing are handled properly
+        """
+        # Return values of the python sensor callable, modified during tests
+        done = False
+        fail = False
+
+        def func():
+            if fail:
+                raise AirflowException()
+            return done
+
+        with dag_maker(dag_id='test_reschedule_handling') as dag:
+            task = PythonSensor.partial(
+                task_id='test_reschedule_handling_sensor',
+                mode='reschedule',
+                python_callable=func,
+                retries=1,
+                retry_delay=datetime.timedelta(seconds=0),
+                pool='test_pool',
+            ).expand(poke_interval=[0])
+        ti = dag_maker.create_dagrun(execution_date=timezone.utcnow()).task_instances[0]
+        ti.task = task
+        assert ti._try_number == 0
+        assert ti.try_number == 1
+
+        def run_ti_and_assert(
+            run_date,
+            expected_start_date,
+            expected_end_date,
+            expected_duration,
+            expected_state,
+            expected_try_number,
+            expected_task_reschedule_count,
+        ):
+            ti.refresh_from_task(task)
+            with freeze_time(run_date):
+                try:
+                    ti.run()
+                except AirflowException:
+                    if not fail:
+                        raise
+            ti.refresh_from_db()
+            assert ti.state == expected_state
+            assert ti._try_number == expected_try_number
+            assert ti.try_number == expected_try_number + 1
+            assert ti.start_date == expected_start_date
+            assert ti.end_date == expected_end_date
+            assert ti.duration == expected_duration
+            trs = TaskReschedule.find_for_task_instance(ti)
+            assert len(trs) == expected_task_reschedule_count
+
+        date1 = timezone.utcnow()
+
+        done, fail = False, False
+        run_ti_and_assert(date1, date1, date1, 0, State.UP_FOR_RESCHEDULE, 0, 1)
+
+        # Clear the task instance.
+        dag.clear()
+        ti.refresh_from_db()
+        assert ti.state == State.NONE
+        assert ti._try_number == 0
+        # Check that reschedules for ti have also been cleared.
+        trs = TaskReschedule.find_for_task_instance(ti)
+        assert not trs
+
     @pytest.mark.usefixtures('test_pool')
     def test_reschedule_handling_clear_reschedules(self, dag_maker):
         """
@@ -2541,6 +2709,55 @@ def timeout():
     assert ti.state == State.FAILED
 
 
+@pytest.mark.parametrize("mode", ["poke", "reschedule"])
+@pytest.mark.parametrize("retries", [0, 1])
+def test_mapped_sensor_timeout(mode, retries, dag_maker):
+    """
+    Test that AirflowSensorTimeout does not cause mapped sensor to retry.
+    """
+
+    def timeout():
+        raise AirflowSensorTimeout
+
+    mock_on_failure = mock.MagicMock()
+    with dag_maker(dag_id=f'test_sensor_timeout_{mode}_{retries}'):
+        PythonSensor.partial(
+            task_id='test_raise_sensor_timeout',
+            python_callable=timeout,
+            on_failure_callback=mock_on_failure,
+            retries=retries,
+        ).expand(mode=[mode])
+    ti = dag_maker.create_dagrun(execution_date=timezone.utcnow()).task_instances[0]
+
+    with pytest.raises(AirflowSensorTimeout):
+        ti.run()
+
+    assert mock_on_failure.called
+    assert ti.state == State.FAILED
+
+
+@pytest.mark.parametrize("mode", ["poke", "reschedule"])
+@pytest.mark.parametrize("retries", [0, 1])
+def test_mapped_sensor_works(mode, retries, dag_maker):
+    """
+    Test that mapped sensors reaches success state.
+    """
+
+    def timeout(ti):
+        return 1
+
+    with dag_maker(dag_id=f'test_sensor_timeout_{mode}_{retries}'):
+        PythonSensor.partial(
+            task_id='test_raise_sensor_timeout',
+            python_callable=timeout,
+            retries=retries,
+        ).expand(mode=[mode])
+    ti = dag_maker.create_dagrun().task_instances[0]
+
+    ti.run()
+    assert ti.state == State.SUCCESS
+
+
 class TestTaskInstanceRecordTaskMapXComPush:
     """Test TI.xcom_push() correctly records return values for task-mapping."""
 
diff --git a/tests/serialization/test_dag_serialization.py b/tests/serialization/test_dag_serialization.py
index e2268e7f45137..ec0996bf43fd0 100644
--- a/tests/serialization/test_dag_serialization.py
+++ b/tests/serialization/test_dag_serialization.py
@@ -1614,7 +1614,7 @@ def poke(self, context: Context):
         assert op.deps == serialized_op.deps
 
     @pytest.mark.parametrize("mode", ["poke", "reschedule"])
-    def test_serialize_mapped_sensor(self, mode):
+    def test_serialize_mapped_sensor_has_reschedule_dep(self, mode):
         from airflow.sensors.base import BaseSensorOperator
 
         class DummySensor(BaseSensorOperator):
@@ -1626,9 +1626,7 @@ def poke(self, context: Context):
         blob = SerializedBaseOperator.serialize_mapped_operator(op)
         assert "deps" in blob
 
-        serialized_op = SerializedBaseOperator.deserialize_operator(blob)
-        assert serialized_op.reschedule == (mode == "reschedule")
-        assert op.deps == serialized_op.deps
+        assert 'airflow.ti_deps.deps.ready_to_reschedule.ReadyToRescheduleDep' in blob['deps']
 
     @pytest.mark.parametrize(
         "passed_success_callback, expected_value",
diff --git a/tests/ti_deps/deps/test_ready_to_reschedule_dep.py b/tests/ti_deps/deps/test_ready_to_reschedule_dep.py
index fa4d2cf2f08fd..2ab8c539dc50b 100644
--- a/tests/ti_deps/deps/test_ready_to_reschedule_dep.py
+++ b/tests/ti_deps/deps/test_ready_to_reschedule_dep.py
@@ -47,6 +47,24 @@ def _get_task_reschedule(self, reschedule_date):
         )
         return reschedule
 
+    def _get_mapped_task_instance(self, state):
+        dag = DAG('test_dag')
+        task = Mock(dag=dag, reschedule=True, is_mapped=True)
+        ti = TaskInstance(task=task, state=state, run_id=None)
+        return ti
+
+    def _get_mapped_task_reschedule(self, reschedule_date):
+        task = Mock(dag_id='test_dag', task_id='test_task', is_mapped=True)
+        reschedule = TaskReschedule(
+            task=task,
+            run_id=None,
+            try_number=None,
+            start_date=reschedule_date,
+            end_date=reschedule_date,
+            reschedule_date=reschedule_date,
+        )
+        return reschedule
+
     def test_should_pass_if_ignore_in_reschedule_period_is_set(self):
         ti = self._get_task_instance(State.UP_FOR_RESCHEDULE)
         dep_context = DepContext(ignore_in_reschedule_period=True)
@@ -103,3 +121,62 @@ def test_should_fail_before_reschedule_date_multiple(self, mock_query_for_task_i
         ][-1]
         ti = self._get_task_instance(State.UP_FOR_RESCHEDULE)
         assert not ReadyToRescheduleDep().is_met(ti=ti)
+
+    def test_mapped_task_should_pass_if_ignore_in_reschedule_period_is_set(self):
+        ti = self._get_mapped_task_instance(State.UP_FOR_RESCHEDULE)
+        dep_context = DepContext(ignore_in_reschedule_period=True)
+        assert ReadyToRescheduleDep().is_met(ti=ti, dep_context=dep_context)
+
+    @patch('airflow.models.taskreschedule.TaskReschedule.query_for_task_instance')
+    def test_mapped_task_should_pass_if_not_reschedule_mode(self, mock_query_for_task_instance):
+        mock_query_for_task_instance.return_value.with_entities.return_value.first.return_value = []
+        ti = self._get_mapped_task_instance(State.UP_FOR_RESCHEDULE)
+        del ti.task.reschedule
+        assert ReadyToRescheduleDep().is_met(ti=ti)
+
+    def test_mapped_task_should_pass_if_not_in_none_state(self):
+        ti = self._get_mapped_task_instance(State.UP_FOR_RETRY)
+        assert ReadyToRescheduleDep().is_met(ti=ti)
+
+    @patch('airflow.models.taskreschedule.TaskReschedule.query_for_task_instance')
+    def test_mapped_should_pass_if_no_reschedule_record_exists(self, mock_query_for_task_instance):
+        mock_query_for_task_instance.return_value.with_entities.return_value.first.return_value = []
+        ti = self._get_mapped_task_instance(State.NONE)
+        assert ReadyToRescheduleDep().is_met(ti=ti)
+
+    @patch('airflow.models.taskreschedule.TaskReschedule.query_for_task_instance')
+    def test_mapped_should_pass_after_reschedule_date_one(self, mock_query_for_task_instance):
+        mock_query_for_task_instance.return_value.with_entities.return_value.first.return_value = (
+            self._get_mapped_task_reschedule(utcnow() - timedelta(minutes=1))
+        )
+        ti = self._get_mapped_task_instance(State.UP_FOR_RESCHEDULE)
+        assert ReadyToRescheduleDep().is_met(ti=ti)
+
+    @patch('airflow.models.taskreschedule.TaskReschedule.query_for_task_instance')
+    def test_mapped_task_should_pass_after_reschedule_date_multiple(self, mock_query_for_task_instance):
+        mock_query_for_task_instance.return_value.with_entities.return_value.first.return_value = [
+            self._get_mapped_task_reschedule(utcnow() - timedelta(minutes=21)),
+            self._get_mapped_task_reschedule(utcnow() - timedelta(minutes=11)),
+            self._get_mapped_task_reschedule(utcnow() - timedelta(minutes=1)),
+        ][-1]
+        ti = self._get_mapped_task_instance(State.UP_FOR_RESCHEDULE)
+        assert ReadyToRescheduleDep().is_met(ti=ti)
+
+    @patch('airflow.models.taskreschedule.TaskReschedule.query_for_task_instance')
+    def test_mapped_task_should_fail_before_reschedule_date_one(self, mock_query_for_task_instance):
+        mock_query_for_task_instance.return_value.with_entities.return_value.first.return_value = (
+            self._get_mapped_task_reschedule(utcnow() + timedelta(minutes=1))
+        )
+
+        ti = self._get_mapped_task_instance(State.UP_FOR_RESCHEDULE)
+        assert not ReadyToRescheduleDep().is_met(ti=ti)
+
+    @patch('airflow.models.taskreschedule.TaskReschedule.query_for_task_instance')
+    def test_mapped_task_should_fail_before_reschedule_date_multiple(self, mock_query_for_task_instance):
+        mock_query_for_task_instance.return_value.with_entities.return_value.first.return_value = [
+            self._get_mapped_task_reschedule(utcnow() - timedelta(minutes=19)),
+            self._get_mapped_task_reschedule(utcnow() - timedelta(minutes=9)),
+            self._get_mapped_task_reschedule(utcnow() + timedelta(minutes=1)),
+        ][-1]
+        ti = self._get_mapped_task_instance(State.UP_FOR_RESCHEDULE)
+        assert not ReadyToRescheduleDep().is_met(ti=ti)

From 11aca93a5a702e82c123daf87cd4c8b06c0cb3a7 Mon Sep 17 00:00:00 2001
From: Ephraim Anierobi <splendidzigy24@gmail.com>
Date: Wed, 10 Aug 2022 13:31:32 +0100
Subject: [PATCH 6/6] Update airflow/ti_deps/deps/ready_to_reschedule.py

Co-authored-by: Tzu-ping Chung <uranusjr@gmail.com>
---
 airflow/ti_deps/deps/ready_to_reschedule.py | 4 +++-
 1 file changed, 3 insertions(+), 1 deletion(-)

diff --git a/airflow/ti_deps/deps/ready_to_reschedule.py b/airflow/ti_deps/deps/ready_to_reschedule.py
index a3f7ab59b7aa7..88219bb401dca 100644
--- a/airflow/ti_deps/deps/ready_to_reschedule.py
+++ b/airflow/ti_deps/deps/ready_to_reschedule.py
@@ -42,7 +42,9 @@ def _get_dep_statuses(self, ti, session, dep_context):
         """
         is_mapped = ti.task.is_mapped
         if not is_mapped and not getattr(ti.task, "reschedule", False):
-            # Mapped sensors don't currently have the reschedule property,
+            # Mapped sensors don't have the reschedule property (it can only
+            # be calculated after unmapping), so we don't check them here.
+            # They are handled below by checking TaskReschedule instead.
             yield self._passing_status(reason="Task is not in reschedule mode.")
             return
 
