diff --git a/airflow-core/src/airflow/api/common/mark_tasks.py b/airflow-core/src/airflow/api/common/mark_tasks.py index 7004156ef2423..5274d8c706e04 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_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, ) 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 4504ddceda020..e1f5695ada8ce 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]):