This is an automated email from the ASF dual-hosted git repository.
pierrejeambrun pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/airflow.git
The following commit(s) were added to refs/heads/main by this push:
new bb29517c618 Allow keeping finished task states on Dag run patch
(#68666)
bb29517c618 is described below
commit bb29517c618f43f5150df66cab910fb1b3d23ad1
Author: Guan-Ming Chiu <[email protected]>
AuthorDate: Thu Jul 9 23:31:38 2026 +0800
Allow keeping finished task states on Dag run patch (#68666)
---
airflow-core/src/airflow/api/common/mark_tasks.py | 145 +++++++++++----------
.../tests/unit/api/common/test_mark_tasks.py | 25 ++--
2 files changed, 87 insertions(+), 83 deletions(-)
diff --git a/airflow-core/src/airflow/api/common/mark_tasks.py
b/airflow-core/src/airflow/api/common/mark_tasks.py
index 7004156ef24..5274d8c706e 100644
--- a/airflow-core/src/airflow/api/common/mark_tasks.py
+++ b/airflow-core/src/airflow/api/common/mark_tasks.py
@@ -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,
@@ -352,14 +298,75 @@ def set_dag_run_state_to_failed(
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,
)
diff --git a/airflow-core/tests/unit/api/common/test_mark_tasks.py
b/airflow-core/tests/unit/api/common/test_mark_tasks.py
index 4504ddceda0..e1f5695ada8 100644
--- a/airflow-core/tests/unit/api/common/test_mark_tasks.py
+++ b/airflow-core/tests/unit/api/common/test_mark_tasks.py
@@ -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)
@@ -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]):
