From 1969a74db349c9d0f9da3b848b4d14d3e3788bee Mon Sep 17 00:00:00 2001 From: PoAn Yang Date: Thu, 27 Aug 2026 22:47:44 +0900 Subject: [PATCH] Fix unscheduled tasks left on a stale Dag version Signed-off-by: PoAn Yang --- .../src/airflow/jobs/scheduler_job_runner.py | 2 +- .../tests/unit/jobs/test_scheduler_job.py | 49 +++++++++++++++++++ 2 files changed, 50 insertions(+), 1 deletion(-) diff --git a/airflow-core/src/airflow/jobs/scheduler_job_runner.py b/airflow-core/src/airflow/jobs/scheduler_job_runner.py index f8b2df4aac2ca..ccc273d759a90 100644 --- a/airflow-core/src/airflow/jobs/scheduler_job_runner.py +++ b/airflow-core/src/airflow/jobs/scheduler_job_runner.py @@ -3094,7 +3094,7 @@ def _verify_integrity_if_dag_changed(self, dag_run: DagRun, session: Session) -> .where( TI.dag_id == dag_run.dag_id, TI.run_id == dag_run.run_id, - TI.state.in_(State.unfinished), + or_(TI.state.is_(None), TI.state.in_(State.unfinished)), ) .values(dag_version_id=latest_dag_version.id), execution_options={"synchronize_session": False}, diff --git a/airflow-core/tests/unit/jobs/test_scheduler_job.py b/airflow-core/tests/unit/jobs/test_scheduler_job.py index a7b332e763f1c..46846106d0de1 100644 --- a/airflow-core/tests/unit/jobs/test_scheduler_job.py +++ b/airflow-core/tests/unit/jobs/test_scheduler_job.py @@ -5204,6 +5204,55 @@ def test_verify_integrity_if_dag_changed(self, dag_maker): session.rollback() session.close() + def test_verify_integrity_if_dag_changed_moves_unscheduled_tis(self, dag_maker, session): + """Task instances that have not been scheduled yet must also be moved to the new Dag version.""" + with dag_maker(dag_id="test_verify_integrity_unscheduled_tis", serialized=False) as dag: + BashOperator(task_id="upstream", bash_command="echo hi") >> BashOperator( + task_id="downstream", bash_command="echo hi" + ) + + orm_dag = dag_maker.dag_model + SerializedDagModel.write_dag( + LazyDeserializedDAG.from_dag(dag), bundle_name="testing", session=session + ) + session.commit() + assert orm_dag.bundle_version is None + + self.job_runner = SchedulerJobRunner(job=Job()) + self.job_runner._create_dag_runs([orm_dag], session) + dr = DagRun.find(dag_id=dag.dag_id, session=session)[0] + + self.job_runner._schedule_dag_run(dag_run=dr, session=session) + session.commit() + dag_version_1 = DagVersion.get_latest_version(dr.dag_id, session=session) + + states = {ti.task_id: ti.state for ti in dr.task_instances} + assert states == {"upstream": State.SCHEDULED, "downstream": None} + + BashOperator(task_id="extra", dag=dag, bash_command="echo hi") + SerializedDagModel.write_dag( + LazyDeserializedDAG.from_dag(dag), bundle_name="testing", session=session + ) + session.commit() + dag_version_2 = DagVersion.get_latest_version(dr.dag_id, session=session) + assert dag_version_2.id != dag_version_1.id + + self.job_runner._schedule_dag_run(dr, session) + session.commit() + + versions = dict( + session.execute( + select(TaskInstance.task_id, TaskInstance.dag_version_id).where( + TaskInstance.dag_id == dr.dag_id, TaskInstance.run_id == dr.run_id + ) + ).all() + ) + assert versions == { + "upstream": dag_version_2.id, + "downstream": dag_version_2.id, + "extra": dag_version_2.id, + } + def test_verify_integrity_not_called_for_versioned_bundles(self, dag_maker, session): with dag_maker("test_verify_integrity_if_dag_not_changed") as dag: BashOperator(task_id="dummy", bash_command="echo hi")