This is an automated email from the ASF dual-hosted git repository.
ephraimanierobi 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 ac1c1beb948 Run verify_integrity when cleared to run on latest (#54984)
ac1c1beb948 is described below
commit ac1c1beb94866bf18835aa679969f5380f0c333f
Author: Ephraim Anierobi <[email protected]>
AuthorDate: Thu Aug 28 08:39:29 2025 +0100
Run verify_integrity when cleared to run on latest (#54984)
* Run verify_integrity when cleared to run on latest
This helps recreating or removing tasks depending on the new serialized
dag version as run_on_latest_version is selected
* fixup! Run verify_integrity when cleared to run on latest
---
airflow-core/src/airflow/models/taskinstance.py | 9 +++-
airflow-core/tests/unit/models/test_cleartasks.py | 55 +++++++++++++++++++++++
devel-common/src/tests_common/pytest_plugin.py | 1 -
3 files changed, 63 insertions(+), 2 deletions(-)
diff --git a/airflow-core/src/airflow/models/taskinstance.py
b/airflow-core/src/airflow/models/taskinstance.py
index 608aadd1d78..e81b1691f70 100644
--- a/airflow-core/src/airflow/models/taskinstance.py
+++ b/airflow-core/src/airflow/models/taskinstance.py
@@ -79,6 +79,7 @@ from airflow.exceptions import (
from airflow.listeners.listener import get_listener_manager
from airflow.models.asset import AssetEvent, AssetModel
from airflow.models.base import Base, StringID, TaskInstanceDependencies
+from airflow.models.dag_version import DagVersion
# Import HITLDetail at runtime so SQLAlchemy can resolve the relationship
from airflow.models.hitl import HITLDetail # noqa: F401
@@ -222,7 +223,6 @@ def clear_task_instances(
from airflow.models.dagbag import DBDagBag
scheduler_dagbag = DBDagBag(load_op_links=False)
-
for ti in tis:
task_instance_ids.append(ti.id)
ti.prepare_db_for_next_try(session)
@@ -281,6 +281,13 @@ def clear_task_instances(
dr.start_date = timezone.utcnow()
if run_on_latest_version:
dr_dag =
scheduler_dagbag.get_latest_version_of_dag(dr.dag_id, session=session)
+ dag_version = DagVersion.get_latest_version(dr.dag_id,
session=session)
+ if dag_version:
+ # Change the dr.created_dag_version_id so the
scheduler doesn't reject this
+ # version when it sets the dag_run.dag
+ dr.created_dag_version_id = dag_version.id
+ dr.dag = dr_dag
+ dr.verify_integrity(session=session,
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
index 801862e8ad8..8e4e23a7bba 100644
--- a/airflow-core/tests/unit/models/test_cleartasks.py
+++ b/airflow-core/tests/unit/models/test_cleartasks.py
@@ -23,7 +23,9 @@ import random
import pytest
from sqlalchemy import select
+from airflow.models import DagRun
from airflow.models.dag import DAG
+from airflow.models.dag_version import DagVersion
from airflow.models.serialized_dag import SerializedDagModel
from airflow.models.taskinstance import TaskInstance, TaskInstance as TI,
clear_task_instances
from airflow.models.taskinstancehistory import TaskInstanceHistory
@@ -674,3 +676,56 @@ class TestClearTasks:
assert ti.state == State.SUCCESS
assert ti.try_number == 2
assert ti.max_tries == 1
+
+ @pytest.mark.parametrize("run_on_latest_version", [True, False])
+ def test_clear_task_instances_with_run_on_latest_version(self,
run_on_latest_version, dag_maker, session):
+ # Explicitly needs catchup as True as test is creating history runs
+ with dag_maker(
+ "test_clear_task_instances",
+ start_date=DEFAULT_DATE,
+ end_date=DEFAULT_DATE + datetime.timedelta(days=10),
+ catchup=True,
+ bundle_version="v1",
+ ):
+ task0 = EmptyOperator(task_id="0")
+ task1 = EmptyOperator(task_id="1", retries=2)
+ dr = dag_maker.create_dagrun(
+ state=State.RUNNING,
+ run_type=DagRunType.SCHEDULED,
+ )
+
+ old_dag_version = DagVersion.get_latest_version(dr.dag_id)
+ ti0, ti1 = sorted(dr.task_instances, key=lambda ti: ti.task_id)
+ ti0.refresh_from_task(task0)
+ ti1.refresh_from_task(task1)
+
+ ti0.run()
+ ti1.run()
+ dr.state = DagRunState.SUCCESS
+ session.merge(dr)
+ session.flush()
+
+ with dag_maker(
+ "test_clear_task_instances",
+ start_date=DEFAULT_DATE,
+ end_date=DEFAULT_DATE + datetime.timedelta(days=10),
+ catchup=True,
+ bundle_version="v2",
+ ) as dag:
+ EmptyOperator(task_id="0")
+ SerializedDagModel.write_dag(dag=dag, bundle_name="dag_maker",
bundle_version="v2")
+ new_dag_version = DagVersion.get_latest_version(dag.dag_id)
+
+ assert old_dag_version.id != new_dag_version.id
+ qry = session.query(TI).filter(TI.dag_id ==
dag.dag_id).order_by(TI.task_id).all()
+ clear_task_instances(qry, session,
run_on_latest_version=run_on_latest_version)
+ session.commit()
+ dr = session.query(DagRun).filter(DagRun.dag_id == dag.dag_id).one()
+ if run_on_latest_version:
+ 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]
+ 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]
diff --git a/devel-common/src/tests_common/pytest_plugin.py
b/devel-common/src/tests_common/pytest_plugin.py
index 20b39ff7172..4e6263491b3 100644
--- a/devel-common/src/tests_common/pytest_plugin.py
+++ b/devel-common/src/tests_common/pytest_plugin.py
@@ -949,7 +949,6 @@ def dag_maker(request) -> Generator[DagMaker, None, None]:
if type is not None:
return
- dag.clear(session=self.session)
if AIRFLOW_V_3_0_PLUS:
dag.bulk_write_to_db(self.bundle_name, self.bundle_version,
[dag], session=self.session)
else:
