From 946ce2bc67ce55a00b6aecca9a13e87b558d09c0 Mon Sep 17 00:00:00 2001 From: Ephraim Anierobi 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 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): '', '', '', + '', '', ] @@ -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 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): '', '', '', - '', '', ] @@ -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 Date: Tue, 9 Aug 2022 15:26:29 +0100 Subject: [PATCH 4/6] Apply suggestions from code review Co-authored-by: Tzu-ping Chung --- 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 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 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 --- 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