From 91076843d6e3de9326acc9207495a106a4c41968 Mon Sep 17 00:00:00 2001
From: ReadytoRocc <75824493+ReadytoRocc@users.noreply.github.com>
Date: Fri, 22 Oct 2021 23:21:53 -0400
Subject: [PATCH 1/4] Added ti.clear_next_method_args() to clear ti.next_method
 and ti.next_kwargs when a task finishes running.

---
 airflow/models/taskinstance.py    |  18 ++++-
 tests/models/test_taskinstance.py | 125 ++++++++++++++++++++++++++++--
 2 files changed, 134 insertions(+), 9 deletions(-)

diff --git a/airflow/models/taskinstance.py b/airflow/models/taskinstance.py
index 6289ecaf06dd5..f9690371757c6 100644
--- a/airflow/models/taskinstance.py
+++ b/airflow/models/taskinstance.py
@@ -1278,6 +1278,14 @@ def _log_state(self, lead_msg: str = ''):
             self._date_or_empty('end_date'),
         )
 
+    # Ensure we unset next_method and next_kwargs to ensure that any
+    # retries don't re-use them.
+    def clear_next_method_args(self):
+        self.log.debug("Clearing next_method and next_kwargs.")
+
+        self.next_method = None
+        self.next_kwargs = None
+
     @provide_session
     @Sentry.enrich_errors
     def _run_raw_task(
@@ -1367,6 +1375,8 @@ def _run_raw_task(
             # or dagrun timed out and task is marked as skipped
             # current behavior doesn't hit the callbacks
             if self.state in State.finished:
+                self.clear_next_method_args()
+                session.commit()
                 return
             else:
                 self.handle_failure(e, test_mode, error_file=error_file, session=session)
@@ -1380,6 +1390,7 @@ def _run_raw_task(
             Stats.incr(f'ti.finish.{self.task.dag_id}.{self.task.task_id}.{self.state}')
 
         # Recording SKIPPED or SUCCESS
+        self.clear_next_method_args()
         self.end_date = timezone.utcnow()
         self._log_state()
         self.set_duration()
@@ -1665,6 +1676,8 @@ def _handle_reschedule(self, actual_start_date, reschedule_exception, test_mode=
         # to same log file.
         self._try_number -= 1
 
+        self.clear_next_method_args()
+
         session.merge(self)
         session.commit()
         self.log.info('Rescheduling task, marking task as UP_FOR_RESCHEDULE')
@@ -1706,10 +1719,7 @@ def handle_failure(
             dag_run = self.get_dagrun(session=session)  # self.dag_run not populated by refresh_from_db
             session.add(TaskFail(task, dag_run.execution_date, self.start_date, self.end_date))
 
-        # Ensure we unset next_method and next_kwargs to ensure that any
-        # retries don't re-use them.
-        self.next_method = None
-        self.next_kwargs = None
+        self.clear_next_method_args()
 
         # Set state correctly and figure out how to log it and decide whether
         # to email
diff --git a/tests/models/test_taskinstance.py b/tests/models/test_taskinstance.py
index d22aef8c1c7b1..af60770eb4c55 100644
--- a/tests/models/test_taskinstance.py
+++ b/tests/models/test_taskinstance.py
@@ -33,6 +33,7 @@
 from airflow.exceptions import (
     AirflowException,
     AirflowFailException,
+    AirflowRescheduleException,
     AirflowSensorTimeout,
     AirflowSkipException,
 )
@@ -481,16 +482,71 @@ def task_function(ti):
         ti.refresh_from_db()
         ti.state == state
 
-    def test_task_retry_wipes_next_fields(self, session, dag_maker):
+    def test_task_wipes_next_fields_success(self, session, dag_maker):
         """
         Test that ensures that tasks wipe their next_method and next_kwargs
-        fields when they are queued for retry after a failure.
+        when they go into a state of SUCCESS.
         """
 
-        with dag_maker('test_mark_failure_2'):
+        with dag_maker("test_deferred_method_clear"):
             task = BashOperator(
-                task_id='test_retry_handling_op',
-                bash_command='exit 1',
+                task_id="test_deferred_method_clear_task",
+                bash_command=""
+            )
+
+        dr = dag_maker.create_dagrun()
+        ti = dr.task_instances[0]
+        ti.next_method = "execute"
+        ti.next_kwargs = {}
+        session.merge(ti)
+        session.commit()
+
+        ti.task = task
+        ti.run()
+        ti.refresh_from_db()
+
+        assert ti.next_method is None
+        assert ti.next_kwargs is None
+        assert ti.state == State.SUCCESS
+
+    def test_task_wipes_next_fields_failure(self, session, dag_maker):
+        """
+        Test that ensures that tasks wipe their next_method and next_kwargs
+        when they go into a state of FAILED.
+        """
+
+        with dag_maker("test_deferred_method_clear"):
+            task = BashOperator(
+                task_id="test_deferred_method_clear_task",
+                bash_command="exit 1"
+            )
+
+        dr = dag_maker.create_dagrun()
+        ti = dr.task_instances[0]
+        ti.next_method = "execute"
+        ti.next_kwargs = {}
+        session.merge(ti)
+        session.commit()
+
+        ti.task = task
+        with pytest.raises(AirflowException):
+            ti.run()
+        ti.refresh_from_db()
+
+        assert ti.next_method is None
+        assert ti.next_kwargs is None
+        assert ti.state == State.FAILED
+
+    def test_task_wipes_next_fields_up_for_retry(self, session, dag_maker):
+        """
+        Test that ensures that tasks wipe their next_method and next_kwargs
+        when they go into a state of UP_FOR_RETRY.
+        """
+
+        with dag_maker("test_deferred_method_clear"):
+            task = BashOperator(
+                task_id="test_deferred_method_clear_task",
+                bash_command="exit 1",
                 retries=1,
                 retry_delay=datetime.timedelta(seconds=2),
             )
@@ -511,6 +567,65 @@ def test_task_retry_wipes_next_fields(self, session, dag_maker):
         assert ti.next_kwargs is None
         assert ti.state == State.UP_FOR_RETRY
 
+    def test_task_wipes_next_fields_skipped(self, session, dag_maker):
+        """
+        Test that ensures that tasks wipe their next_method and next_kwargs
+        when they go into a state of SKIPPED.
+        """
+        def skip():
+            raise AirflowSkipException
+
+        with dag_maker("test_deferred_method_clear"):
+            task = PythonOperator(
+                task_id="test_deferred_method_clear_task",
+                python_callable=skip,
+            )
+
+        dr = dag_maker.create_dagrun()
+        ti = dr.task_instances[0]
+        ti.next_method = "execute"
+        ti.next_kwargs = {}
+        session.merge(ti)
+        session.commit()
+
+        ti.task = task
+        ti.run()
+        ti.refresh_from_db()
+
+        assert ti.next_method is None
+        assert ti.next_kwargs is None
+        assert ti.state == State.SKIPPED
+
+    def test_task_wipes_next_fields_reschedule(self, session, dag_maker):
+        """
+        Test that ensures that tasks wipe their next_method and next_kwargs
+        when they go into a state of SKIPPED.
+        """
+        def reschedule():
+            reschedule_date = timezone.utcnow()
+            raise AirflowRescheduleException(reschedule_date)
+
+        with dag_maker("test_deferred_method_clear"):
+            task = PythonOperator(
+                task_id="test_deferred_method_clear_task",
+                python_callable=reschedule,
+            )
+
+        dr = dag_maker.create_dagrun()
+        ti = dr.task_instances[0]
+        ti.next_method = "execute"
+        ti.next_kwargs = {}
+        session.merge(ti)
+        session.commit()
+
+        ti.task = task
+        ti.run()
+        ti.refresh_from_db()
+
+        assert ti.next_method is None
+        assert ti.next_kwargs is None
+        assert ti.state == State.UP_FOR_RESCHEDULE
+
     @freeze_time('2021-09-19 04:56:35', as_kwarg='frozen_time')
     def test_retry_delay(self, dag_maker, frozen_time=None):
         """

From 2a3bbe79c414f447f242c9a9766a2bd840d0487b Mon Sep 17 00:00:00 2001
From: ReadytoRocc <75824493+ReadytoRocc@users.noreply.github.com>
Date: Sat, 23 Oct 2021 10:25:40 -0400
Subject: [PATCH 2/4] Updated formatting.

---
 tests/models/test_taskinstance.py | 12 ++++--------
 1 file changed, 4 insertions(+), 8 deletions(-)

diff --git a/tests/models/test_taskinstance.py b/tests/models/test_taskinstance.py
index af60770eb4c55..20ac1a92571c8 100644
--- a/tests/models/test_taskinstance.py
+++ b/tests/models/test_taskinstance.py
@@ -489,10 +489,7 @@ def test_task_wipes_next_fields_success(self, session, dag_maker):
         """
 
         with dag_maker("test_deferred_method_clear"):
-            task = BashOperator(
-                task_id="test_deferred_method_clear_task",
-                bash_command=""
-            )
+            task = BashOperator(task_id="test_deferred_method_clear_task", bash_command="")
 
         dr = dag_maker.create_dagrun()
         ti = dr.task_instances[0]
@@ -516,10 +513,7 @@ def test_task_wipes_next_fields_failure(self, session, dag_maker):
         """
 
         with dag_maker("test_deferred_method_clear"):
-            task = BashOperator(
-                task_id="test_deferred_method_clear_task",
-                bash_command="exit 1"
-            )
+            task = BashOperator(task_id="test_deferred_method_clear_task", bash_command="exit 1")
 
         dr = dag_maker.create_dagrun()
         ti = dr.task_instances[0]
@@ -572,6 +566,7 @@ def test_task_wipes_next_fields_skipped(self, session, dag_maker):
         Test that ensures that tasks wipe their next_method and next_kwargs
         when they go into a state of SKIPPED.
         """
+
         def skip():
             raise AirflowSkipException
 
@@ -601,6 +596,7 @@ def test_task_wipes_next_fields_reschedule(self, session, dag_maker):
         Test that ensures that tasks wipe their next_method and next_kwargs
         when they go into a state of SKIPPED.
         """
+
         def reschedule():
             reschedule_date = timezone.utcnow()
             raise AirflowRescheduleException(reschedule_date)

From e338f6cdd297c41b79b479828ef641b152a7b3d9 Mon Sep 17 00:00:00 2001
From: ReadytoRocc <75824493+ReadytoRocc@users.noreply.github.com>
Date: Sat, 23 Oct 2021 11:20:34 -0400
Subject: [PATCH 3/4] Consolidated into a Parameterized test.

---
 tests/models/test_taskinstance.py | 144 ++++++++----------------------
 1 file changed, 35 insertions(+), 109 deletions(-)

diff --git a/tests/models/test_taskinstance.py b/tests/models/test_taskinstance.py
index 20ac1a92571c8..211e67f662417 100644
--- a/tests/models/test_taskinstance.py
+++ b/tests/models/test_taskinstance.py
@@ -482,129 +482,51 @@ def task_function(ti):
         ti.refresh_from_db()
         ti.state == state
 
-    def test_task_wipes_next_fields_success(self, session, dag_maker):
-        """
-        Test that ensures that tasks wipe their next_method and next_kwargs
-        when they go into a state of SUCCESS.
-        """
-
-        with dag_maker("test_deferred_method_clear"):
-            task = BashOperator(task_id="test_deferred_method_clear_task", bash_command="")
-
-        dr = dag_maker.create_dagrun()
-        ti = dr.task_instances[0]
-        ti.next_method = "execute"
-        ti.next_kwargs = {}
-        session.merge(ti)
-        session.commit()
-
-        ti.task = task
-        ti.run()
-        ti.refresh_from_db()
-
-        assert ti.next_method is None
-        assert ti.next_kwargs is None
-        assert ti.state == State.SUCCESS
-
-    def test_task_wipes_next_fields_failure(self, session, dag_maker):
+    @pytest.mark.parametrize(
+        "state",
+        [State.FAILED, State.SKIPPED, State.SUCCESS, State.UP_FOR_RESCHEDULE, State.UP_FOR_RETRY],
+    )
+    def test_task_wipes_next_fields(self, session, state, dag_maker):
         """
         Test that ensures that tasks wipe their next_method and next_kwargs
-        when they go into a state of FAILED.
+        when they go into a state of FAILED, SKIPPED, SUCCESS, UP_FOR_RESCHEDULE, or UP_FOR_RETRY.
         """
 
-        with dag_maker("test_deferred_method_clear"):
-            task = BashOperator(task_id="test_deferred_method_clear_task", bash_command="exit 1")
-
-        dr = dag_maker.create_dagrun()
-        ti = dr.task_instances[0]
-        ti.next_method = "execute"
-        ti.next_kwargs = {}
-        session.merge(ti)
-        session.commit()
-
-        ti.task = task
-        with pytest.raises(AirflowException):
-            ti.run()
-        ti.refresh_from_db()
-
-        assert ti.next_method is None
-        assert ti.next_kwargs is None
-        assert ti.state == State.FAILED
-
-    def test_task_wipes_next_fields_up_for_retry(self, session, dag_maker):
-        """
-        Test that ensures that tasks wipe their next_method and next_kwargs
-        when they go into a state of UP_FOR_RETRY.
-        """
-
-        with dag_maker("test_deferred_method_clear"):
-            task = BashOperator(
-                task_id="test_deferred_method_clear_task",
-                bash_command="exit 1",
-                retries=1,
-                retry_delay=datetime.timedelta(seconds=2),
-            )
-
-        dr = dag_maker.create_dagrun()
-        ti = dr.task_instances[0]
-        ti.next_method = "execute"
-        ti.next_kwargs = {}
-        session.merge(ti)
-        session.commit()
-
-        ti.task = task
-        with pytest.raises(AirflowException):
-            ti.run()
-        ti.refresh_from_db()
-
-        assert ti.next_method is None
-        assert ti.next_kwargs is None
-        assert ti.state == State.UP_FOR_RETRY
-
-    def test_task_wipes_next_fields_skipped(self, session, dag_maker):
-        """
-        Test that ensures that tasks wipe their next_method and next_kwargs
-        when they go into a state of SKIPPED.
-        """
+        def failure():
+            raise AirflowException
 
         def skip():
             raise AirflowSkipException
 
-        with dag_maker("test_deferred_method_clear"):
-            task = PythonOperator(
-                task_id="test_deferred_method_clear_task",
-                python_callable=skip,
-            )
-
-        dr = dag_maker.create_dagrun()
-        ti = dr.task_instances[0]
-        ti.next_method = "execute"
-        ti.next_kwargs = {}
-        session.merge(ti)
-        session.commit()
-
-        ti.task = task
-        ti.run()
-        ti.refresh_from_db()
-
-        assert ti.next_method is None
-        assert ti.next_kwargs is None
-        assert ti.state == State.SKIPPED
-
-    def test_task_wipes_next_fields_reschedule(self, session, dag_maker):
-        """
-        Test that ensures that tasks wipe their next_method and next_kwargs
-        when they go into a state of SKIPPED.
-        """
+        def success():
+            return None
 
         def reschedule():
             reschedule_date = timezone.utcnow()
             raise AirflowRescheduleException(reschedule_date)
 
+        _retries = 0
+        _retry_delay = datetime.timedelta(seconds=0)
+
+        if state == State.FAILED:
+            _python_callable = failure
+        elif state == State.SKIPPED:
+            _python_callable = skip
+        elif state == State.SUCCESS:
+            _python_callable = success
+        elif state == State.UP_FOR_RESCHEDULE:
+            _python_callable = reschedule
+        elif state in [State.FAILED, State.UP_FOR_RETRY]:
+            _python_callable = failure
+            _retries = 1
+            _retry_delay = datetime.timedelta(seconds=2)
+
         with dag_maker("test_deferred_method_clear"):
             task = PythonOperator(
                 task_id="test_deferred_method_clear_task",
-                python_callable=reschedule,
+                python_callable=_python_callable,
+                retries=_retries,
+                retry_delay=_retry_delay,
             )
 
         dr = dag_maker.create_dagrun()
@@ -615,12 +537,16 @@ def reschedule():
         session.commit()
 
         ti.task = task
-        ti.run()
+        if state in [State.FAILED, State.UP_FOR_RETRY]:
+            with pytest.raises(AirflowException):
+                ti.run()
+        elif state in [State.SKIPPED, State.SUCCESS, State.UP_FOR_RESCHEDULE]:
+            ti.run()
         ti.refresh_from_db()
 
         assert ti.next_method is None
         assert ti.next_kwargs is None
-        assert ti.state == State.UP_FOR_RESCHEDULE
+        assert ti.state == state
 
     @freeze_time('2021-09-19 04:56:35', as_kwarg='frozen_time')
     def test_retry_delay(self, dag_maker, frozen_time=None):

From 296a350abf0b5b531d218564fd718b0cb745d625 Mon Sep 17 00:00:00 2001
From: ReadytoRocc <75824493+ReadytoRocc@users.noreply.github.com>
Date: Sat, 23 Oct 2021 11:24:45 -0400
Subject: [PATCH 4/4] Added session.merge to match other cases.

---
 airflow/models/taskinstance.py | 1 +
 1 file changed, 1 insertion(+)

diff --git a/airflow/models/taskinstance.py b/airflow/models/taskinstance.py
index f9690371757c6..38ae39727eb01 100644
--- a/airflow/models/taskinstance.py
+++ b/airflow/models/taskinstance.py
@@ -1376,6 +1376,7 @@ def _run_raw_task(
             # current behavior doesn't hit the callbacks
             if self.state in State.finished:
                 self.clear_next_method_args()
+                session.merge(self)
                 session.commit()
                 return
             else:
