From c8f0902d7439a4e405b26bc19ea54ce9852a8dbe Mon Sep 17 00:00:00 2001 From: Karthikeyan Singaravelan Date: Sat, 13 Apr 2024 14:50:50 +0530 Subject: [PATCH 1/2] Increment try_number while clearing deferred tasks. --- airflow/models/taskinstance.py | 6 +++++ tests/www/views/test_views_tasks.py | 36 +++++++++++++++++++++++++++++ 2 files changed, 42 insertions(+) diff --git a/airflow/models/taskinstance.py b/airflow/models/taskinstance.py index 2fd5df910afd0..a0a58af684e11 100644 --- a/airflow/models/taskinstance.py +++ b/airflow/models/taskinstance.py @@ -273,6 +273,12 @@ def clear_task_instances( ti.state = TaskInstanceState.RESTARTING job_ids.append(ti.job_id) else: + # When the task is deferred the try_number is decremented so that the same try + # number is used when the task resumes execution to process the event. But in case of + # clearing the try number should be incremented so that the next run doesn't reuse the same try_number + if ti.state == TaskInstanceState.DEFERRED: + ti._try_number += 1 + ti_dag = dag if dag and dag.dag_id == ti.dag_id else dag_bag.get_dag(ti.dag_id, session=session) task_id = ti.task_id if ti_dag and ti_dag.has_task(task_id): diff --git a/tests/www/views/test_views_tasks.py b/tests/www/views/test_views_tasks.py index 2b893c40f671d..226e6afab48cc 100644 --- a/tests/www/views/test_views_tasks.py +++ b/tests/www/views/test_views_tasks.py @@ -856,6 +856,42 @@ def test_task_instance_clear(session, request, client_fixture, should_succeed): assert state == (State.NONE if should_succeed else initial_state) +def test_task_instance_clear_deferred(session, admin_client, create_task_instance): + """Ensures clearing a task instance in deferred state increments _try_number for next execution.""" + task_instance = create_task_instance( + dag_id="example_bash_operator", + task_id="run_this_last", + execution_date=timezone.utcnow(), + state=State.DEFERRED, + ) + task_id = task_instance.task_id + run_id = task_instance.run_id + old_try_number = task_instance._try_number + + data = { + "dag_id": task_instance.dag_id, + "dag_run_id": task_instance.dag_run.run_id, + "execution_date": DEFAULT_DATE, + "task_id": task_instance.task_id, + "upstream": "false", + "downstream": "false", + "future": "true", + "past": "true", + "only_failed": "false", + "confirmed": "true", + } + + resp = admin_client.post("clear", data=data, follow_redirects=True) + assert resp.status_code == 200 + + try_number = ( + session.query(TaskInstance._try_number) + .filter(TaskInstance.task_id == task_id, TaskInstance.run_id == run_id) + .scalar() + ) + assert try_number == old_try_number + 1 + + def test_task_instance_clear_downstream(session, admin_client, dag_maker): """Ensures clearing a task instance clears its downstream dependencies exclusively""" with dag_maker( From f45054267b01773f3299b6ff411968ea60cf3a88 Mon Sep 17 00:00:00 2001 From: Karthikeyan Singaravelan Date: Sat, 13 Apr 2024 15:47:34 +0530 Subject: [PATCH 2/2] Fix test since using example_bash_operator fails for some reason. --- tests/www/views/test_views_tasks.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tests/www/views/test_views_tasks.py b/tests/www/views/test_views_tasks.py index 226e6afab48cc..0813e638e7c8e 100644 --- a/tests/www/views/test_views_tasks.py +++ b/tests/www/views/test_views_tasks.py @@ -859,8 +859,8 @@ def test_task_instance_clear(session, request, client_fixture, should_succeed): def test_task_instance_clear_deferred(session, admin_client, create_task_instance): """Ensures clearing a task instance in deferred state increments _try_number for next execution.""" task_instance = create_task_instance( - dag_id="example_bash_operator", - task_id="run_this_last", + dag_id="example_bash_decorator", + task_id="runme_0", execution_date=timezone.utcnow(), state=State.DEFERRED, )