Skip to content

[v3-3-test] Allow keeping finished task states on Dag run patch (#68666) #69662

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
pierrejeambrun merged 1 commit into from
Jul 10, 2026
+87 −83
Merged

[v3-3-test] Allow keeping finished task states on Dag run patch (#68666) #69662

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
145 changes: 76 additions & 69 deletions airflow-core/src/airflow/api/common/mark_tasks.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -215,80 +215,26 @@ def _set_dag_run_state(dag_id: str, run_id: str, state: DagRunState, session: SA
session.merge(dag_run)


@provide_session
def set_dag_run_state_to_success(
def _set_dag_run_terminal_state(
*,
dag: SerializedDAG,
run_id: str | None = None,
commit: bool = False,
session: SASession = NEW_SESSION,
) -> list[TaskInstance]:
"""
Set the dag run's state to success.

Set for a specific logical date and its task instances to success.

:param dag: the Dag of which to alter state
:param run_id: the run_id to start looking from
:param commit: commit Dag and tasks to be altered to the database
:param session: database session
:return: If commit is true, list of tasks that have been updated,
otherwise list of tasks that will be updated
:raises: ValueError if dag or logical_date is invalid
"""
if not dag:
return []
if not run_id:
raise ValueError(f"Invalid dag_run_id: {run_id}")

tasks = dag.tasks

# Mark all task instances of the dag run to success - except for unfinished teardown as they need to complete work.
teardown_tasks = [task for task in tasks if task.is_teardown]
unfinished_teardown_task_ids = set(
session.scalars(
select(TaskInstance.task_id).where(
TaskInstance.dag_id == dag.dag_id,
TaskInstance.run_id == run_id,
TaskInstance.task_id.in_(task.task_id for task in teardown_tasks),
or_(TaskInstance.state.is_(None), TaskInstance.state.in_(State.unfinished)),
)
)
)

# Mark the dag run to success if there are no unfinished teardown tasks.
if commit and len(unfinished_teardown_task_ids) == 0:
_set_dag_run_state(dag.dag_id, run_id, DagRunState.SUCCESS, session)

tasks_to_mark_success = [task for task in tasks if not task.is_teardown] + [
task for task in teardown_tasks if task.task_id not in unfinished_teardown_task_ids
]
for task in tasks_to_mark_success:
task.dag = dag
return set_state(
tasks=tasks_to_mark_success,
run_id=run_id,
state=TaskInstanceState.SUCCESS,
commit=commit,
session=session,
)


@provide_session
def set_dag_run_state_to_failed(
*,
dag: SerializedDAG,
run_id: str | None = None,
commit: bool = False,
session: SASession = NEW_SESSION,
run_id: str | None,
run_state: DagRunState,
ti_state: TaskInstanceState,
commit: bool,
session: SASession,
) -> list[TaskInstance]:
"""
Set the dag run's state to failed.
Set the dag run's state to the given terminal state.

Set for a specific logical date and its task instances to failed.
Running task instances are set to ``ti_state`` and non-finished ones to
SKIPPED; finished task instances keep their state. Teardown tasks are left
untouched so they can complete their work.

:param dag: the Dag of which to alter state
:param run_id: the Dag run_id to start looking from
:param run_state: terminal state to set on the dag run
:param ti_state: state to set on running task instances
:param commit: commit Dag and tasks to be altered to the database
:param session: database session
:return: If commit is true, list of tasks that have been updated,
Expand Down Expand Up @@ -352,14 +298,75 @@ def _set_running_task(task: Operator) -> Operator:
for ti in pending_normal_tis:
ti.set_state(TaskInstanceState.SKIPPED)

# Mark the dag run to failed if there is no pending teardown (else this would not be scheduled later).
# Set the dag run state only if there is no pending teardown (else this would not be scheduled later).
if not any(dag.task_dict[ti.task_id].is_teardown for ti in (running_tis + pending_tis)):
_set_dag_run_state(dag.dag_id, run_id, DagRunState.FAILED, session)
_set_dag_run_state(dag.dag_id, run_id, run_state, session)

return pending_normal_tis + set_state(
tasks=running_tasks,
run_id=run_id,
state=TaskInstanceState.FAILED,
state=ti_state,
commit=commit,
session=session,
)


@provide_session
def set_dag_run_state_to_success(
*,
dag: SerializedDAG,
run_id: str | None = None,
commit: bool = False,
session: SASession = NEW_SESSION,
) -> list[TaskInstance]:
"""
Set the dag run's state to success.

Set for a specific logical date and its task instances to success.

:param dag: the Dag of which to alter state
:param run_id: the run_id to start looking from
:param commit: commit Dag and tasks to be altered to the database
:param session: database session
:return: If commit is true, list of tasks that have been updated,
otherwise list of tasks that will be updated
:raises: ValueError if dag or logical_date is invalid
"""
return _set_dag_run_terminal_state(
dag=dag,
run_id=run_id,
run_state=DagRunState.SUCCESS,
ti_state=TaskInstanceState.SUCCESS,
commit=commit,
session=session,
)


@provide_session
def set_dag_run_state_to_failed(
*,
dag: SerializedDAG,
run_id: str | None = None,
commit: bool = False,
session: SASession = NEW_SESSION,
) -> list[TaskInstance]:
"""
Set the dag run's state to failed.

Set for a specific logical date and its task instances to failed.

:param dag: the Dag of which to alter state
:param run_id: the Dag run_id to start looking from
:param commit: commit Dag and tasks to be altered to the database
:param session: database session
:return: If commit is true, list of tasks that have been updated,
otherwise list of tasks that will be updated
"""
return _set_dag_run_terminal_state(
dag=dag,
run_id=run_id,
run_state=DagRunState.FAILED,
ti_state=TaskInstanceState.FAILED,
commit=commit,
session=session,
)
Expand Down
25 changes: 11 additions & 14 deletions airflow-core/tests/unit/api/common/test_mark_tasks.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -91,21 +91,23 @@ def test_set_dag_run_state_to_success_unfinished_teardown(
assert len(updated_tis) == 2
task_dict = {ti.task_id: ti for ti in updated_tis}
assert task_dict["running"].state == TaskInstanceState.SUCCESS
assert task_dict["pending"].state == TaskInstanceState.SUCCESS
assert task_dict["pending"].state == TaskInstanceState.SKIPPED
assert "teardown" not in task_dict


@pytest.mark.parametrize("finished_state", sorted(list(State.finished)))
def test_set_dag_run_state_to_success_finished_teardown(dag_maker: DagMaker[SerializedDAG], finished_state):
def test_set_dag_run_state_to_success_keeps_finished_task_states(
dag_maker: DagMaker[SerializedDAG], finished_state
):
with dag_maker("TEST_DAG_1") as dag:
with EmptyOperator(task_id="teardown").as_teardown():
EmptyOperator(task_id="failed")
EmptyOperator(task_id="finished")
dr = dag_maker.create_dagrun()
for ti in dr.get_task_instances():
if ti.task_id == "failed":
ti.set_state(TaskInstanceState.FAILED)
if ti.task_id == "teardown":
if ti.task_id == "finished":
ti.set_state(finished_state)
if ti.task_id == "teardown":
ti.set_state(TaskInstanceState.SUCCESS)
dag_maker.session.flush()
dr.set_state(DagRunState.FAILED)

Expand All @@ -115,14 +117,9 @@ def test_set_dag_run_state_to_success_finished_teardown(dag_maker: DagMaker[Seri
run = dag_maker.session.scalar(select(DagRun).filter_by(dag_id=dr.dag_id, run_id=dr.run_id))
assert run is not None
assert run.state == DagRunState.SUCCESS
if finished_state == TaskInstanceState.SUCCESS:
assert len(updated_tis) == 1
else:
assert len(updated_tis) == 2
task_dict = {ti.task_id: ti for ti in updated_tis}
assert task_dict["failed"].state == TaskInstanceState.SUCCESS
if finished_state != TaskInstanceState.SUCCESS:
assert task_dict["teardown"].state == TaskInstanceState.SUCCESS
assert updated_tis == []
states = {ti.task_id: ti.state for ti in dr.get_task_instances(session=dag_maker.session)}
assert states == {"finished": finished_state, "teardown": TaskInstanceState.SUCCESS}


def test_find_task_relatives_downstream_skips_teardowns(dag_maker: DagMaker[SerializedDAG]):
Expand Down