Update TI.dag_version after dag/taskinstance clearing (#55180)

ephraimbuddy · web-flow · commit bcf512543802 · 2025-09-02T15:16:36.000+01:00
If run with latest bundle is selected during clear, then we should
also update the TI.dag_version same way we updated dagrun. This is
primarily for consistency in the UI
diff --git a/airflow-core/src/airflow/models/taskinstance.py b/airflow-core/src/airflow/models/taskinstance.py
@@ -288,6 +288,8 @@ def clear_task_instances(
                         dr.created_dag_version_id = dag_version.id
                         dr.dag = dr_dag
                         dr.verify_integrity(session=session, dag_version_id=dag_version.id)
+                        for ti in dr.task_instances:
+                            ti.dag_version_id = dag_version.id
                 else:
                     dr_dag = scheduler_dagbag.get_dag_for_run(dag_run=dr, session=session)
                 if not dr_dag:
diff --git a/airflow-core/tests/unit/models/test_cleartasks.py b/airflow-core/tests/unit/models/test_cleartasks.py
@@ -717,7 +717,11 @@ def test_clear_task_instances_with_run_on_latest_version(self, run_on_latest_ver
             assert dr.created_dag_version_id == new_dag_version.id
             assert dr.bundle_version == new_dag_version.bundle_version
             assert TaskInstanceState.REMOVED in [ti.state for ti in dr.task_instances]
+            for ti in dr.task_instances:
+                assert ti.dag_version_id == new_dag_version.id
         else:
             assert dr.created_dag_version_id == old_dag_version.id
             assert dr.bundle_version == old_dag_version.bundle_version
             assert TaskInstanceState.REMOVED not in [ti.state for ti in dr.task_instances]
+            for ti in dr.task_instances:
+                assert ti.dag_version_id == old_dag_version.id
