This is an automated email from the ASF dual-hosted git repository.
jedcunningham 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 3dc4cd2e012 Add run_on_latest_version support for backfill and clear
operations (#52177)
3dc4cd2e012 is described below
commit 3dc4cd2e01209efbb555b35debeb115367cbeb6d
Author: Ephraim Anierobi <[email protected]>
AuthorDate: Tue Jul 22 18:12:00 2025 +0100
Add run_on_latest_version support for backfill and clear operations (#52177)
With this option, users are able to choose the dag version they want
to run their dag/task after clearing or when running backfill. This
only applies to versioned bundles as non-versioned bundles run with
the latest dag version.
When the user choose the run with latest version, the bundle_version
associated with the dagrun is updated to the latest and the associated
serialized dag version updated to the latest. Choosing not to run
with latest version which is the default means that the bundle version
and serialized dag version that the dag ran with initially would be used
in running it again.
For backfill, there's now --run-on-latest-version flag that makes it run
with the latest version, otherwise it will run with the original bundle
the dagrun was created with. Note that it's only useful when rerunning
a dagrun using backfill. The default behaviour is using the initial
bundle/version
and this is intentional otherwise running backfill will fail if there was
task rename in the latest version.
Summary of changes:
- Use SchedulerDagBag instead of DagBag for execution API
- Add run_on_latest_version field to DAGRunClearBody and
ClearTaskInstancesBody models
- Add --run-on-latest-version CLI flag for backfill command
- Update backfill.py to support running tasks with latest DAG version
- Add UI checkbox for "Run with latest version" in clear dialogs
- Update SchedulerDagBag to handle latest version parameter
- Update API endpoints to support run_on_latest_version parameter
---
airflow-core/docs/core-concepts/dag-run.rst | 2 ++
airflow-core/src/airflow/api_fastapi/app.py | 4 ++-
.../api_fastapi/core_api/datamodels/dag_run.py | 4 +++
.../core_api/datamodels/task_instances.py | 5 +++
.../core_api/openapi/v2-rest-api-generated.yaml | 12 +++++++
.../api_fastapi/core_api/routes/public/dag_run.py | 2 ++
.../core_api/routes/public/task_instances.py | 5 +++
.../api_fastapi/execution_api/routes/dag_runs.py | 16 +++++++--
.../execution_api/routes/task_instances.py | 25 ++++++++------
airflow-core/src/airflow/cli/cli_config.py | 9 +++++
.../src/airflow/cli/commands/backfill_command.py | 2 ++
.../src/airflow/jobs/scheduler_job_runner.py | 15 +++++----
airflow-core/src/airflow/models/backfill.py | 7 +++-
airflow-core/src/airflow/models/dag.py | 7 ++++
airflow-core/src/airflow/models/taskinstance.py | 10 +++---
.../airflow/ui/openapi-gen/requests/schemas.gen.ts | 12 +++++++
.../airflow/ui/openapi-gen/requests/types.gen.ts | 8 +++++
.../airflow/ui/public/i18n/locales/en/dags.json | 1 +
.../ui/src/components/Clear/Run/ClearRunDialog.tsx | 38 +++++++++++++++++++---
.../Clear/TaskInstance/ClearTaskInstanceDialog.tsx | 34 +++++++++++++++++--
.../src/airflow/ui/src/components/ui/index.ts | 1 +
.../core_api/routes/public/test_task_instances.py | 2 +-
.../versions/v2025_04_28/test_task_instances.py | 7 ++--
.../unit/cli/commands/test_backfill_command.py | 28 ++++++++++++++++
airflow-core/tests/unit/models/test_backfill.py | 17 ++++++++--
airflow-core/tests/unit/models/test_dag.py | 15 +++++----
.../src/airflowctl/api/datamodels/generated.py | 14 ++++++++
.../tests/unit/common/sql/operators/test_sql.py | 8 ++++-
.../tests/unit/standard/decorators/test_bash.py | 8 ++---
.../tests/unit/standard/operators/test_python.py | 13 ++++----
.../src/airflow/sdk/execution_time/supervisor.py | 9 ++---
31 files changed, 276 insertions(+), 64 deletions(-)
diff --git a/airflow-core/docs/core-concepts/dag-run.rst
b/airflow-core/docs/core-concepts/dag-run.rst
index af7e727de2b..1cb121c9b72 100644
--- a/airflow-core/docs/core-concepts/dag-run.rst
+++ b/airflow-core/docs/core-concepts/dag-run.rst
@@ -167,6 +167,8 @@ the errors after going through the logs, you can re-run the
tasks by clearing th
scheduled date. Clearing a task instance creates a record of the task instance.
The ``try_number`` of the current task instance is incremented, the
``max_tries`` set to ``0`` and the state set to ``None``, which causes the task
to re-run.
+An experimental feature in Airflow 3.1.0 allows you to clear the task
instances and re-run with the latest bundle version.
+
Click on the failed task in the Tree or Graph views and then click on
**Clear**.
The executor will re-run it.
diff --git a/airflow-core/src/airflow/api_fastapi/app.py
b/airflow-core/src/airflow/api_fastapi/app.py
index f515be6992c..17de4f677cd 100644
--- a/airflow-core/src/airflow/api_fastapi/app.py
+++ b/airflow-core/src/airflow/api_fastapi/app.py
@@ -84,8 +84,10 @@ def create_app(apps: str = "all") -> FastAPI:
dag_bag = create_dag_bag()
if "execution" in apps_list or "all" in apps_list:
+ from airflow.jobs.scheduler_job_runner import SchedulerDagBag
+
task_exec_api_app = create_task_execution_api_app()
- task_exec_api_app.state.dag_bag = dag_bag
+ task_exec_api_app.state.dag_bag = SchedulerDagBag()
init_error_handlers(task_exec_api_app)
app.mount("/execution", task_exec_api_app)
diff --git
a/airflow-core/src/airflow/api_fastapi/core_api/datamodels/dag_run.py
b/airflow-core/src/airflow/api_fastapi/core_api/datamodels/dag_run.py
index e184a3acfe8..bb79fa671ed 100644
--- a/airflow-core/src/airflow/api_fastapi/core_api/datamodels/dag_run.py
+++ b/airflow-core/src/airflow/api_fastapi/core_api/datamodels/dag_run.py
@@ -55,6 +55,10 @@ class DAGRunClearBody(StrictBaseModel):
dry_run: bool = True
only_failed: bool = False
+ run_on_latest_version: bool = Field(
+ default=False,
+ description="(Experimental) Run on the latest bundle version of the
DAG after clearing the DAG Run.",
+ )
class DAGRunResponse(BaseModel):
diff --git
a/airflow-core/src/airflow/api_fastapi/core_api/datamodels/task_instances.py
b/airflow-core/src/airflow/api_fastapi/core_api/datamodels/task_instances.py
index e2470119a1c..bb04e829652 100644
--- a/airflow-core/src/airflow/api_fastapi/core_api/datamodels/task_instances.py
+++ b/airflow-core/src/airflow/api_fastapi/core_api/datamodels/task_instances.py
@@ -179,6 +179,11 @@ class ClearTaskInstancesBody(StrictBaseModel):
include_downstream: bool = False
include_future: bool = False
include_past: bool = False
+ run_on_latest_version: bool = Field(
+ default=False,
+ description="(Experimental) Run on the latest bundle version of the
DAG after "
+ "clearing the task instances.",
+ )
@model_validator(mode="before")
@classmethod
diff --git
a/airflow-core/src/airflow/api_fastapi/core_api/openapi/v2-rest-api-generated.yaml
b/airflow-core/src/airflow/api_fastapi/core_api/openapi/v2-rest-api-generated.yaml
index 894c306581c..74363c05fe0 100644
---
a/airflow-core/src/airflow/api_fastapi/core_api/openapi/v2-rest-api-generated.yaml
+++
b/airflow-core/src/airflow/api_fastapi/core_api/openapi/v2-rest-api-generated.yaml
@@ -8503,6 +8503,12 @@ components:
type: boolean
title: Include Past
default: false
+ run_on_latest_version:
+ type: boolean
+ title: Run On Latest Version
+ description: (Experimental) Run on the latest bundle version of the
DAG
+ after clearing the task instances.
+ default: false
additionalProperties: false
type: object
title: ClearTaskInstancesBody
@@ -9140,6 +9146,12 @@ components:
type: boolean
title: Only Failed
default: false
+ run_on_latest_version:
+ type: boolean
+ title: Run On Latest Version
+ description: (Experimental) Run on the latest bundle version of the
DAG
+ after clearing the DAG Run.
+ default: false
additionalProperties: false
type: object
title: DAGRunClearBody
diff --git
a/airflow-core/src/airflow/api_fastapi/core_api/routes/public/dag_run.py
b/airflow-core/src/airflow/api_fastapi/core_api/routes/public/dag_run.py
index 5d974ec012c..00f93754baf 100644
--- a/airflow-core/src/airflow/api_fastapi/core_api/routes/public/dag_run.py
+++ b/airflow-core/src/airflow/api_fastapi/core_api/routes/public/dag_run.py
@@ -281,6 +281,7 @@ def clear_dag_run(
run_id=dag_run_id,
task_ids=None,
only_failed=body.only_failed,
+ run_on_latest_version=body.run_on_latest_version,
dry_run=True,
session=session,
)
@@ -293,6 +294,7 @@ def clear_dag_run(
run_id=dag_run_id,
task_ids=None,
only_failed=body.only_failed,
+ run_on_latest_version=body.run_on_latest_version,
session=session,
)
dag_run_cleared = session.scalar(select(DagRun).where(DagRun.id ==
dag_run.id))
diff --git
a/airflow-core/src/airflow/api_fastapi/core_api/routes/public/task_instances.py
b/airflow-core/src/airflow/api_fastapi/core_api/routes/public/task_instances.py
index 915054cdf79..33bafc0506b 100644
---
a/airflow-core/src/airflow/api_fastapi/core_api/routes/public/task_instances.py
+++
b/airflow-core/src/airflow/api_fastapi/core_api/routes/public/task_instances.py
@@ -675,7 +675,11 @@ def post_clear_task_instances(
if dag_run is None:
error_message = f"Dag Run id {dag_run_id} not found in dag
{dag_id}"
raise HTTPException(status.HTTP_404_NOT_FOUND, error_message)
+ # If dag_run_id is provided, we should get the dag from SchedulerDagBag
+ # to ensure we get the right version.
+ from airflow.jobs.scheduler_job_runner import SchedulerDagBag
+ dag = SchedulerDagBag().get_dag(dag_run=dag_run, session=session)
if past or future:
raise HTTPException(
status.HTTP_400_BAD_REQUEST,
@@ -724,6 +728,7 @@ def post_clear_task_instances(
task_instances,
session,
DagRunState.QUEUED if reset_dag_runs else False,
+ run_on_latest_version=body.run_on_latest_version,
)
return TaskInstanceCollectionResponse(
diff --git
a/airflow-core/src/airflow/api_fastapi/execution_api/routes/dag_runs.py
b/airflow-core/src/airflow/api_fastapi/execution_api/routes/dag_runs.py
index db27f7bd93e..04a3e2ce7ec 100644
--- a/airflow-core/src/airflow/api_fastapi/execution_api/routes/dag_runs.py
+++ b/airflow-core/src/airflow/api_fastapi/execution_api/routes/dag_runs.py
@@ -29,7 +29,6 @@ from airflow.api_fastapi.common.types import UtcDateTime
from airflow.api_fastapi.execution_api.datamodels.dagrun import
DagRunStateResponse, TriggerDAGRunPayload
from airflow.exceptions import DagRunAlreadyExists
from airflow.models.dag import DagModel
-from airflow.models.dagbag import DagBag
from airflow.models.dagrun import DagRun
from airflow.utils.types import DagRunTriggeredByType
@@ -122,9 +121,20 @@ def clear_dag_run(
"message": f"DAG with dag_id: '{dag_id}' has import errors and
cannot be triggered",
},
)
+ from airflow.jobs.scheduler_job_runner import SchedulerDagBag
+
+ dag_run = session.scalar(select(DagRun).where(DagRun.dag_id == dag_id,
DagRun.run_id == run_id))
+ dag_bag = SchedulerDagBag()
+ dag = dag_bag.get_dag(dag_run=dag_run, session=session)
+ if not dag:
+ raise HTTPException(
+ status.HTTP_404_NOT_FOUND,
+ detail={
+ "reason": "Not Found",
+ "message": f"DAG with dag_id: '{dag_id}' was not found in the
DagBag",
+ },
+ )
- dag_bag = DagBag(dag_folder=dm.fileloc, read_dags_from_db=True)
- dag = dag_bag.get_dag(dag_id)
dag.clear(run_id=run_id)
diff --git
a/airflow-core/src/airflow/api_fastapi/execution_api/routes/task_instances.py
b/airflow-core/src/airflow/api_fastapi/execution_api/routes/task_instances.py
index eef96323d2b..def4dd155c6 100644
---
a/airflow-core/src/airflow/api_fastapi/execution_api/routes/task_instances.py
+++
b/airflow-core/src/airflow/api_fastapi/execution_api/routes/task_instances.py
@@ -29,7 +29,7 @@ from uuid import UUID
import attrs
import structlog
from cadwyn import VersionedAPIRouter
-from fastapi import Body, HTTPException, Query, status
+from fastapi import Body, Depends, HTTPException, Query, status
from pydantic import JsonValue
from sqlalchemy import func, or_, tuple_, update
from sqlalchemy.exc import NoResultFound, SQLAlchemyError
@@ -37,7 +37,7 @@ from sqlalchemy.orm import joinedload
from sqlalchemy.sql import select
from structlog.contextvars import bind_contextvars
-from airflow.api_fastapi.common.dagbag import DagBagDep
+from airflow.api_fastapi.common.dagbag import dag_bag_from_app
from airflow.api_fastapi.common.db.common import SessionDep
from airflow.api_fastapi.common.types import UtcDateTime
from airflow.api_fastapi.execution_api.datamodels.taskinstance import (
@@ -76,6 +76,9 @@ if TYPE_CHECKING:
from airflow.models.expandinput import SchedulerExpandInput
from airflow.sdk.types import Operator
+from airflow.jobs.scheduler_job_runner import SchedulerDagBag
+
+SchedulerDagBagDep = Annotated[SchedulerDagBag, Depends(dag_bag_from_app)]
router = VersionedAPIRouter()
@@ -104,7 +107,7 @@ def ti_run(
task_instance_id: UUID,
ti_run_payload: Annotated[TIEnterRunningPayload, Body()],
session: SessionDep,
- dag_bag: DagBagDep,
+ dag_bag: SchedulerDagBagDep,
) -> TIRunContext:
"""
Run a TaskInstance.
@@ -255,7 +258,7 @@ def ti_run(
or 0
)
- if dag := dag_bag.get_dag(ti.dag_id):
+ if dag := dag_bag.get_dag(dag_run=dr, session=session):
upstream_map_indexes = dict(
_get_upstream_map_indexes(dag.get_task(ti.task_id),
ti.map_index, ti.run_id, session)
)
@@ -330,7 +333,7 @@ def ti_update_state(
task_instance_id: UUID,
ti_patch_payload: Annotated[TIStateUpdate, Body()],
session: SessionDep,
- dag_bag: DagBagDep,
+ dag_bag: SchedulerDagBagDep,
):
"""
Update the state of a TaskInstance.
@@ -417,8 +420,9 @@ def ti_update_state(
)
-def _handle_fail_fast_for_dag(ti: TI, dag_id: str, session: SessionDep,
dag_bag: DagBagDep) -> None:
- ser_dag = dag_bag.get_dag(dag_id)
+def _handle_fail_fast_for_dag(ti: TI, dag_id: str, session: SessionDep,
dag_bag: SchedulerDagBagDep) -> None:
+ dr = ti.dag_run
+ ser_dag = dag_bag.get_dag(dag_run=dr, session=session)
if ser_dag and getattr(ser_dag, "fail_fast", False):
task_dict = getattr(ser_dag, "task_dict")
task_teardown_map = {k: v.is_teardown for k, v in task_dict.items()}
@@ -432,7 +436,7 @@ def _create_ti_state_update_query_and_update_state(
query: Update,
updated_state,
session: SessionDep,
- dag_bag: DagBagDep,
+ dag_bag: SchedulerDagBagDep,
dag_id: str,
) -> tuple[Update, TaskInstanceState]:
if isinstance(ti_patch_payload, (TITerminalStatePayload,
TIRetryStatePayload, TISuccessStatePayload)):
@@ -893,7 +897,7 @@ def _get_group_tasks(dag_id: str, task_group_id: str,
session: SessionDep, logic
def validate_inlets_and_outlets(
task_instance_id: UUID,
session: SessionDep,
- dag_bag: DagBagDep,
+ dag_bag: SchedulerDagBagDep,
) -> InactiveAssetsResponse:
"""Validate whether there're inactive assets in inlets and outlets of a
given task instance."""
ti_id_str = str(task_instance_id)
@@ -911,7 +915,8 @@ def validate_inlets_and_outlets(
)
if not ti.task:
- dag = dag_bag.get_dag(ti.dag_id)
+ dr = ti.dag_run
+ dag = dag_bag.get_dag(dag_run=dr, session=session)
if dag:
with contextlib.suppress(TaskNotFound):
ti.task = dag.get_task(ti.task_id)
diff --git a/airflow-core/src/airflow/cli/cli_config.py
b/airflow-core/src/airflow/cli/cli_config.py
index 363839a8580..2006b5f42b9 100644
--- a/airflow-core/src/airflow/cli/cli_config.py
+++ b/airflow-core/src/airflow/cli/cli_config.py
@@ -340,6 +340,14 @@ ARG_BACKFILL_REPROCESS_BEHAVIOR = Arg(
),
choices=("none", "completed", "failed"),
)
+ARG_BACKFILL_RUN_ON_LATEST_VERSION = Arg(
+ ("--run-on-latest-version",),
+ help=(
+ "(Experimental) If set, the backfill will run tasks using the latest
bundle version instead of "
+ "the version that was active when the original Dag run was created."
+ ),
+ action="store_true",
+)
# misc
@@ -968,6 +976,7 @@ BACKFILL_COMMANDS = (
ARG_RUN_BACKWARDS,
ARG_MAX_ACTIVE_RUNS,
ARG_BACKFILL_REPROCESS_BEHAVIOR,
+ ARG_BACKFILL_RUN_ON_LATEST_VERSION,
ARG_BACKFILL_DRY_RUN,
),
),
diff --git a/airflow-core/src/airflow/cli/commands/backfill_command.py
b/airflow-core/src/airflow/cli/commands/backfill_command.py
index 19108dc4e54..444bc35aea3 100644
--- a/airflow-core/src/airflow/cli/commands/backfill_command.py
+++ b/airflow-core/src/airflow/cli/commands/backfill_command.py
@@ -57,6 +57,7 @@ def create_backfill(args) -> None:
reverse=args.run_backwards,
dag_run_conf=args.dag_run_conf,
reprocess_behavior=reprocess_behavior,
+ run_on_latest_version=args.run_on_latest_version,
)
for k, v in params.items():
console.print(f" - {k} = {v}")
@@ -88,4 +89,5 @@ def create_backfill(args) -> None:
dag_run_conf=args.dag_run_conf,
triggering_user_name=user,
reprocess_behavior=reprocess_behavior,
+ run_on_latest_version=args.run_on_latest_version,
)
diff --git a/airflow-core/src/airflow/jobs/scheduler_job_runner.py
b/airflow-core/src/airflow/jobs/scheduler_job_runner.py
index 4eb99616133..1157c3b51f8 100644
--- a/airflow-core/src/airflow/jobs/scheduler_job_runner.py
+++ b/airflow-core/src/airflow/jobs/scheduler_job_runner.py
@@ -131,15 +131,16 @@ class SchedulerDagBag:
return dag
@staticmethod
- def _version_from_dag_run(dag_run, session):
- if dag_run.bundle_version:
- dag_version = dag_run.created_dag_version
- else:
+ def _version_from_dag_run(dag_run, latest, session):
+ if latest or not dag_run.bundle_version:
dag_version = DagVersion.get_latest_version(dag_id=dag_run.dag_id,
session=session)
- return dag_version
+ if dag_version:
+ return dag_version
+
+ return dag_run.created_dag_version
- def get_dag(self, dag_run: DagRun, session: Session) -> DAG | None:
- version = self._version_from_dag_run(dag_run=dag_run, session=session)
+ def get_dag(self, dag_run: DagRun, session: Session, latest=False) -> DAG
| None:
+ version = self._version_from_dag_run(dag_run=dag_run, latest=latest,
session=session)
if not version:
return None
return self._get_dag(version_id=version.id, session=session)
diff --git a/airflow-core/src/airflow/models/backfill.py
b/airflow-core/src/airflow/models/backfill.py
index e8eb3f8b5cd..e1fd0e506bf 100644
--- a/airflow-core/src/airflow/models/backfill.py
+++ b/airflow-core/src/airflow/models/backfill.py
@@ -291,6 +291,7 @@ def _create_backfill_dag_run(
dag_run_conf,
backfill_sort_ordinal,
triggering_user_name,
+ run_on_latest_version,
session,
):
from airflow.models.dagrun import DagRun
@@ -328,6 +329,7 @@ def _create_backfill_dag_run(
info=info,
backfill_id=backfill_id,
sort_ordinal=backfill_sort_ordinal,
+ run_on_latest=run_on_latest_version,
)
else:
session.add(
@@ -401,7 +403,7 @@ def _get_info_list(
return dagrun_info_list
-def _handle_clear_run(session, dag, dr, info, backfill_id, sort_ordinal):
+def _handle_clear_run(session, dag, dr, info, backfill_id, sort_ordinal,
run_on_latest=False):
"""Clear the existing DAG run and update backfill metadata."""
from sqlalchemy.sql import update
@@ -415,6 +417,7 @@ def _handle_clear_run(session, dag, dr, info, backfill_id,
sort_ordinal):
session=session,
confirm_prompt=False,
dry_run=False,
+ run_on_latest_version=run_on_latest,
)
# Update backfill_id and run_type in DagRun table
@@ -447,6 +450,7 @@ def _create_backfill(
dag_run_conf: dict | None,
triggering_user_name: str | None,
reprocess_behavior: ReprocessBehavior | None = None,
+ run_on_latest_version: bool = False,
) -> Backfill | None:
from airflow.models import DagModel
from airflow.models.serialized_dag import SerializedDagModel
@@ -510,6 +514,7 @@ def _create_backfill(
reprocess_behavior=br.reprocess_behavior,
backfill_sort_ordinal=backfill_sort_ordinal,
triggering_user_name=br.triggering_user_name,
+ run_on_latest_version=run_on_latest_version,
session=session,
)
log.info(
diff --git a/airflow-core/src/airflow/models/dag.py
b/airflow-core/src/airflow/models/dag.py
index a141d96a1c3..43fc62b098d 100644
--- a/airflow-core/src/airflow/models/dag.py
+++ b/airflow-core/src/airflow/models/dag.py
@@ -1282,6 +1282,7 @@ class DAG(TaskSDKDag, LoggingMixin):
only_running: bool = False,
confirm_prompt: bool = False,
dag_run_state: DagRunState = DagRunState.QUEUED,
+ run_on_latest_version: bool = False,
session: Session = NEW_SESSION,
dag_bag: DagBag | None = None,
exclude_task_ids: frozenset[str] | frozenset[tuple[str, int]] | None =
frozenset(),
@@ -1299,6 +1300,7 @@ class DAG(TaskSDKDag, LoggingMixin):
confirm_prompt: bool = False,
dag_run_state: DagRunState = DagRunState.QUEUED,
dry_run: Literal[False] = False,
+ run_on_latest_version: bool = False,
session: Session = NEW_SESSION,
dag_bag: DagBag | None = None,
exclude_task_ids: frozenset[str] | frozenset[tuple[str, int]] | None =
frozenset(),
@@ -1317,6 +1319,7 @@ class DAG(TaskSDKDag, LoggingMixin):
only_running: bool = False,
confirm_prompt: bool = False,
dag_run_state: DagRunState = DagRunState.QUEUED,
+ run_on_latest_version: bool = False,
session: Session = NEW_SESSION,
dag_bag: DagBag | None = None,
exclude_task_ids: frozenset[str] | frozenset[tuple[str, int]] | None =
frozenset(),
@@ -1335,6 +1338,7 @@ class DAG(TaskSDKDag, LoggingMixin):
confirm_prompt: bool = False,
dag_run_state: DagRunState = DagRunState.QUEUED,
dry_run: Literal[False] = False,
+ run_on_latest_version: bool = False,
session: Session = NEW_SESSION,
dag_bag: DagBag | None = None,
exclude_task_ids: frozenset[str] | frozenset[tuple[str, int]] | None =
frozenset(),
@@ -1354,6 +1358,7 @@ class DAG(TaskSDKDag, LoggingMixin):
confirm_prompt: bool = False,
dag_run_state: DagRunState = DagRunState.QUEUED,
dry_run: bool = False,
+ run_on_latest_version: bool = False,
session: Session = NEW_SESSION,
dag_bag: DagBag | None = None,
exclude_task_ids: frozenset[str] | frozenset[tuple[str, int]] | None =
frozenset(),
@@ -1372,6 +1377,7 @@ class DAG(TaskSDKDag, LoggingMixin):
:param dag_run_state: state to set DagRun to. If set to False, dagrun
state will not
be changed.
:param dry_run: Find the tasks to clear but don't clear them.
+ :param run_on_latest_version: whether to run on latest serialized DAG
and Bundle version
:param session: The sqlalchemy session to use
:param dag_bag: The DagBag used to find the dags (Optional)
:param exclude_task_ids: A set of ``task_id`` or (``task_id``,
``map_index``)
@@ -1417,6 +1423,7 @@ class DAG(TaskSDKDag, LoggingMixin):
list(tis),
session,
dag_run_state=dag_run_state,
+ run_on_latest_version=run_on_latest_version,
)
else:
count = 0
diff --git a/airflow-core/src/airflow/models/taskinstance.py
b/airflow-core/src/airflow/models/taskinstance.py
index f4ed5573e59..ca46b480256 100644
--- a/airflow-core/src/airflow/models/taskinstance.py
+++ b/airflow-core/src/airflow/models/taskinstance.py
@@ -203,6 +203,7 @@ def clear_task_instances(
tis: list[TaskInstance],
session: Session,
dag_run_state: DagRunState | Literal[False] = DagRunState.QUEUED,
+ run_on_latest_version: bool = False,
) -> None:
"""
Clear a set of task instances, but make sure the running ones get killed.
@@ -217,6 +218,7 @@ def clear_task_instances(
:param session: current session
:param dag_run_state: state to set finished DagRuns to.
If set to False, DagRuns state will not be changed.
+ :param run_on_latest_version: whether to run on latest serialized DAG and
Bundle version
:meta private:
"""
@@ -234,7 +236,7 @@ def clear_task_instances(
ti.state = TaskInstanceState.RESTARTING
else:
dr = ti.dag_run
- ti_dag = scheduler_dagbag.get_dag(dag_run=dr, session=session)
+ ti_dag = scheduler_dagbag.get_dag(dag_run=dr, session=session,
latest=run_on_latest_version)
if not ti_dag:
log.warning("No serialized dag found for dag '%s'", dr.dag_id)
task_id = ti.task_id
@@ -277,12 +279,12 @@ def clear_task_instances(
if dr.state in State.finished_dr_states:
dr.state = dag_run_state
dr.start_date = timezone.utcnow()
- dr_dag = scheduler_dagbag.get_dag(dag_run=dr, session=session)
+ dr_dag = scheduler_dagbag.get_dag(dag_run=dr, session=session,
latest=run_on_latest_version)
if not dr_dag:
log.warning("No serialized dag found for dag '%s'",
dr.dag_id)
- if dr_dag and not dr_dag.disable_bundle_versioning:
+ if dr_dag and not dr_dag.disable_bundle_versioning and
run_on_latest_version:
bundle_version = dr.dag_model.bundle_version
- if bundle_version is not None:
+ if bundle_version is not None and run_on_latest_version:
dr.bundle_version = bundle_version
if dag_run_state == DagRunState.QUEUED:
dr.last_scheduling_decision = None
diff --git a/airflow-core/src/airflow/ui/openapi-gen/requests/schemas.gen.ts
b/airflow-core/src/airflow/ui/openapi-gen/requests/schemas.gen.ts
index a91264eb077..2b997947993 100644
--- a/airflow-core/src/airflow/ui/openapi-gen/requests/schemas.gen.ts
+++ b/airflow-core/src/airflow/ui/openapi-gen/requests/schemas.gen.ts
@@ -1205,6 +1205,12 @@ export const $ClearTaskInstancesBody = {
type: 'boolean',
title: 'Include Past',
default: false
+ },
+ run_on_latest_version: {
+ type: 'boolean',
+ title: 'Run On Latest Version',
+ description: '(Experimental) Run on the latest bundle version of
the DAG after clearing the task instances.',
+ default: false
}
},
additionalProperties: false,
@@ -2177,6 +2183,12 @@ export const $DAGRunClearBody = {
type: 'boolean',
title: 'Only Failed',
default: false
+ },
+ run_on_latest_version: {
+ type: 'boolean',
+ title: 'Run On Latest Version',
+ description: '(Experimental) Run on the latest bundle version of
the DAG after clearing the DAG Run.',
+ default: false
}
},
additionalProperties: false,
diff --git a/airflow-core/src/airflow/ui/openapi-gen/requests/types.gen.ts
b/airflow-core/src/airflow/ui/openapi-gen/requests/types.gen.ts
index bed40429143..06887e2caa2 100644
--- a/airflow-core/src/airflow/ui/openapi-gen/requests/types.gen.ts
+++ b/airflow-core/src/airflow/ui/openapi-gen/requests/types.gen.ts
@@ -389,6 +389,10 @@ export type ClearTaskInstancesBody = {
include_downstream?: boolean;
include_future?: boolean;
include_past?: boolean;
+ /**
+ * (Experimental) Run on the latest bundle version of the DAG after
clearing the task instances.
+ */
+ run_on_latest_version?: boolean;
};
/**
@@ -592,6 +596,10 @@ export type DAGResponse = {
export type DAGRunClearBody = {
dry_run?: boolean;
only_failed?: boolean;
+ /**
+ * (Experimental) Run on the latest bundle version of the DAG after
clearing the DAG Run.
+ */
+ run_on_latest_version?: boolean;
};
/**
diff --git a/airflow-core/src/airflow/ui/public/i18n/locales/en/dags.json
b/airflow-core/src/airflow/ui/public/i18n/locales/en/dags.json
index d1ee6a3d0ee..6705e101962 100644
--- a/airflow-core/src/airflow/ui/public/i18n/locales/en/dags.json
+++ b/airflow-core/src/airflow/ui/public/i18n/locales/en/dags.json
@@ -62,6 +62,7 @@
"onlyFailed": "Clear only failed tasks",
"past": "Past",
"queueNew": "Queue up new tasks",
+ "runOnLatestVersion": "Run with latest bundle version",
"upstream": "Upstream"
}
},
diff --git
a/airflow-core/src/airflow/ui/src/components/Clear/Run/ClearRunDialog.tsx
b/airflow-core/src/airflow/ui/src/components/Clear/Run/ClearRunDialog.tsx
index 696ada0ee38..3b3c9fe5f0e 100644
--- a/airflow-core/src/airflow/ui/src/components/Clear/Run/ClearRunDialog.tsx
+++ b/airflow-core/src/airflow/ui/src/components/Clear/Run/ClearRunDialog.tsx
@@ -21,9 +21,10 @@ import { useState } from "react";
import { useTranslation } from "react-i18next";
import { CgRedo } from "react-icons/cg";
+import { useDagServiceGetDagDetails } from "openapi/queries";
import type { DAGRunResponse } from "openapi/requests/types.gen";
import { ActionAccordion } from "src/components/ActionAccordion";
-import { Button, Dialog } from "src/components/ui";
+import { Button, Dialog, Checkbox } from "src/components/ui";
import SegmentedControl from "src/components/ui/SegmentedControl";
import { useClearDagRunDryRun } from "src/queries/useClearDagRunDryRun";
import { useClearDagRun } from "src/queries/useClearRun";
@@ -43,11 +44,17 @@ const ClearRunDialog = ({ dagRun, onClose, open }: Props)
=> {
const [note, setNote] = useState<string | null>(dagRun.note);
const [selectedOptions, setSelectedOptions] =
useState<Array<string>>(["existingTasks"]);
const onlyFailed = selectedOptions.includes("onlyFailed");
+ const [runOnLatestVersion, setRunOnLatestVersion] = useState(false);
+
+ // Get current DAG's bundle version to compare with DAG run's bundle version
+ const { data: dagDetails } = useDagServiceGetDagDetails({
+ dagId,
+ });
const { data: affectedTasks = { task_instances: [], total_entries: 0 } } =
useClearDagRunDryRun({
dagId,
dagRunId,
- requestBody: { only_failed: onlyFailed },
+ requestBody: { only_failed: onlyFailed, run_on_latest_version:
runOnLatestVersion },
});
const { isPending, mutate } = useClearDagRun({
@@ -62,6 +69,13 @@ const ClearRunDialog = ({ dagRun, onClose, open }: Props) =>
{
onSuccess: onClose,
});
+ // Check if bundle versions are different
+ const currentDagBundleVersion = dagDetails?.bundle_version;
+ const dagRunBundleVersion = dagRun.bundle_version;
+ const bundleVersionsDiffer = currentDagBundleVersion !== dagRunBundleVersion;
+ const shouldShowBundleVersionOption =
+ bundleVersionsDiffer && dagRunBundleVersion !== null &&
dagRunBundleVersion !== "";
+
return (
<Dialog.Root lazyMount onOpenChange={onClose} open={open} size="xl">
<Dialog.Content backdrop>
@@ -101,7 +115,19 @@ const ClearRunDialog = ({ dagRun, onClose, open }: Props)
=> {
/>
</Flex>
<ActionAccordion affectedTasks={affectedTasks} note={note}
setNote={setNote} />
- <Flex justifyContent="end" mt={3}>
+ <Flex
+ {...(shouldShowBundleVersionOption ? { alignItems: "center" } :
{})}
+ justifyContent={shouldShowBundleVersionOption ? "space-between" :
"end"}
+ mt={3}
+ >
+ {shouldShowBundleVersionOption ? (
+ <Checkbox
+ checked={runOnLatestVersion}
+ onCheckedChange={(event) =>
setRunOnLatestVersion(Boolean(event.checked))}
+ >
+
{translate("dags:runAndTaskActions.options.runOnLatestVersion")}
+ </Checkbox>
+ ) : undefined}
<Button
colorPalette="blue"
disabled={affectedTasks.total_entries === 0}
@@ -110,7 +136,11 @@ const ClearRunDialog = ({ dagRun, onClose, open }: Props)
=> {
mutate({
dagId,
dagRunId,
- requestBody: { dry_run: false, only_failed: onlyFailed },
+ requestBody: {
+ dry_run: false,
+ only_failed: onlyFailed,
+ run_on_latest_version: runOnLatestVersion,
+ },
});
if (note !== dagRun.note) {
mutatePatchDagRun({
diff --git
a/airflow-core/src/airflow/ui/src/components/Clear/TaskInstance/ClearTaskInstanceDialog.tsx
b/airflow-core/src/airflow/ui/src/components/Clear/TaskInstance/ClearTaskInstanceDialog.tsx
index 7458b204898..1f4e1886ea7 100644
---
a/airflow-core/src/airflow/ui/src/components/Clear/TaskInstance/ClearTaskInstanceDialog.tsx
+++
b/airflow-core/src/airflow/ui/src/components/Clear/TaskInstance/ClearTaskInstanceDialog.tsx
@@ -21,10 +21,11 @@ import { useState } from "react";
import { useTranslation } from "react-i18next";
import { CgRedo } from "react-icons/cg";
+import { useDagServiceGetDagDetails } from "openapi/queries";
import type { TaskInstanceResponse } from "openapi/requests/types.gen";
import { ActionAccordion } from "src/components/ActionAccordion";
import Time from "src/components/Time";
-import { Button, Dialog } from "src/components/ui";
+import { Button, Dialog, Checkbox } from "src/components/ui";
import SegmentedControl from "src/components/ui/SegmentedControl";
import { useClearTaskInstances } from "src/queries/useClearTaskInstances";
import { useClearTaskInstancesDryRun } from
"src/queries/useClearTaskInstancesDryRun";
@@ -57,6 +58,7 @@ const ClearTaskInstanceDialog = ({ onClose, open,
taskInstance }: Props) => {
const future = selectedOptions.includes("future");
const upstream = selectedOptions.includes("upstream");
const downstream = selectedOptions.includes("downstream");
+ const [runOnLatestVersion, setRunOnLatestVersion] = useState(false);
const [note, setNote] = useState<string | null>(taskInstance.note);
const { isPending: isPendingPatchDagRun, mutate: mutatePatchTaskInstance } =
usePatchTaskInstance({
@@ -66,6 +68,11 @@ const ClearTaskInstanceDialog = ({ onClose, open,
taskInstance }: Props) => {
taskId,
});
+ // Get current DAG's bundle version to compare with task instance's DAG
version bundle version
+ const { data: dagDetails } = useDagServiceGetDagDetails({
+ dagId,
+ });
+
const { data } = useClearTaskInstancesDryRun({
dagId,
options: {
@@ -79,6 +86,7 @@ const ClearTaskInstanceDialog = ({ onClose, open,
taskInstance }: Props) => {
include_past: past,
include_upstream: upstream,
only_failed: onlyFailed,
+ run_on_latest_version: runOnLatestVersion,
task_ids: [[taskId, mapIndex]],
},
});
@@ -88,6 +96,15 @@ const ClearTaskInstanceDialog = ({ onClose, open,
taskInstance }: Props) => {
total_entries: 0,
};
+ // Check if bundle versions are different
+ const currentDagBundleVersion = dagDetails?.bundle_version;
+ const taskInstanceDagVersionBundleVersion =
taskInstance.dag_version.bundle_version;
+ const bundleVersionsDiffer = currentDagBundleVersion !==
taskInstanceDagVersionBundleVersion;
+ const shouldShowBundleVersionOption =
+ bundleVersionsDiffer &&
+ taskInstanceDagVersionBundleVersion !== null &&
+ taskInstanceDagVersionBundleVersion !== "";
+
return (
<Dialog.Root lazyMount onOpenChange={onClose} open={open} size="xl">
<Dialog.Content backdrop>
@@ -140,7 +157,19 @@ const ClearTaskInstanceDialog = ({ onClose, open,
taskInstance }: Props) => {
/>
</Flex>
<ActionAccordion affectedTasks={affectedTasks} note={note}
setNote={setNote} />
- <Flex justifyContent="end" mt={3}>
+ <Flex
+ {...(shouldShowBundleVersionOption ? { alignItems: "center" } :
{})}
+ justifyContent={shouldShowBundleVersionOption ? "space-between" :
"end"}
+ mt={3}
+ >
+ {shouldShowBundleVersionOption ? (
+ <Checkbox
+ checked={runOnLatestVersion}
+ onCheckedChange={(event) =>
setRunOnLatestVersion(Boolean(event.checked))}
+ >
+
{translate("dags:runAndTaskActions.options.runOnLatestVersion")}
+ </Checkbox>
+ ) : undefined}
<Button
colorPalette="blue"
disabled={affectedTasks.total_entries === 0}
@@ -156,6 +185,7 @@ const ClearTaskInstanceDialog = ({ onClose, open,
taskInstance }: Props) => {
include_past: past,
include_upstream: upstream,
only_failed: onlyFailed,
+ run_on_latest_version: runOnLatestVersion,
task_ids: [[taskId, mapIndex]],
},
});
diff --git a/airflow-core/src/airflow/ui/src/components/ui/index.ts
b/airflow-core/src/airflow/ui/src/components/ui/index.ts
index 9339646d415..2146daedb7a 100644
--- a/airflow-core/src/airflow/ui/src/components/ui/index.ts
+++ b/airflow-core/src/airflow/ui/src/components/ui/index.ts
@@ -33,3 +33,4 @@ export * from "./Toaster";
export * from "./Breadcrumb";
export * from "./Clipboard";
export * from "./Popover";
+export * from "./Checkbox";
diff --git
a/airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_task_instances.py
b/airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_task_instances.py
index 16ee780f076..b2e40ef5d5f 100644
---
a/airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_task_instances.py
+++
b/airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_task_instances.py
@@ -2339,7 +2339,7 @@ class
TestPostClearTaskInstances(TestTaskInstanceEndpoint):
# dag (3rd argument) is a different session object. Manually asserting
that the dag_id
# is the same.
- mock_clearti.assert_called_once_with([], mock.ANY, DagRunState.QUEUED)
+ mock_clearti.assert_called_once_with([], mock.ANY, DagRunState.QUEUED,
run_on_latest_version=False)
def test_clear_taskinstance_is_called_with_invalid_task_ids(self,
test_client, session):
"""Test that dagrun is running when invalid task_ids are passed to
clearTaskInstances API."""
diff --git
a/airflow-core/tests/unit/api_fastapi/execution_api/versions/v2025_04_28/test_task_instances.py
b/airflow-core/tests/unit/api_fastapi/execution_api/versions/v2025_04_28/test_task_instances.py
index b5ef9a3ce5e..5757a3f6465 100644
---
a/airflow-core/tests/unit/api_fastapi/execution_api/versions/v2025_04_28/test_task_instances.py
+++
b/airflow-core/tests/unit/api_fastapi/execution_api/versions/v2025_04_28/test_task_instances.py
@@ -21,7 +21,8 @@ from unittest.mock import patch
import pytest
-from airflow.api_fastapi.common.dagbag import create_dag_bag, dag_bag_from_app
+from airflow.api_fastapi.common.dagbag import dag_bag_from_app
+from airflow.jobs.scheduler_job_runner import SchedulerDagBag
from airflow.utils import timezone
from airflow.utils.state import State
@@ -106,9 +107,7 @@ class TestTIUpdateState:
start_date=instant,
)
- dag = ti.task.dag
- dagbag = create_dag_bag()
- dagbag.dags = {dag.dag_id: dag}
+ dagbag = SchedulerDagBag()
execution_app = get_execution_app(ver_client)
execution_app.dependency_overrides[dag_bag_from_app] = lambda: dagbag
session.commit()
diff --git a/airflow-core/tests/unit/cli/commands/test_backfill_command.py
b/airflow-core/tests/unit/cli/commands/test_backfill_command.py
index ff23765fdd4..1c87d8c2ef8 100644
--- a/airflow-core/tests/unit/cli/commands/test_backfill_command.py
+++ b/airflow-core/tests/unit/cli/commands/test_backfill_command.py
@@ -101,6 +101,34 @@ class TestCliBackfill:
dag_run_conf=None,
reprocess_behavior=expected_repro,
triggering_user_name="root",
+ run_on_latest_version=False,
+ )
+
+ @mock.patch("airflow.cli.commands.backfill_command._create_backfill")
+ def test_backfill_with_run_on_latest_version(self, mock_create):
+ args = [
+ "backfill",
+ "create",
+ "--dag-id",
+ "example_bash_operator",
+ "--from-date",
+ DEFAULT_DATE.isoformat(),
+ "--to-date",
+ DEFAULT_DATE.isoformat(),
+ "--run-on-latest-version",
+ ]
+
airflow.cli.commands.backfill_command.create_backfill(self.parser.parse_args(args))
+
+ mock_create.assert_called_once_with(
+ dag_id="example_bash_operator",
+ from_date=DEFAULT_DATE,
+ to_date=DEFAULT_DATE,
+ max_active_runs=None,
+ reverse=False,
+ dag_run_conf=None,
+ reprocess_behavior=None,
+ run_on_latest_version=True,
+ triggering_user_name="root",
)
@mock.patch("airflow.cli.commands.backfill_command._do_dry_run")
diff --git a/airflow-core/tests/unit/models/test_backfill.py
b/airflow-core/tests/unit/models/test_backfill.py
index e69ff85aa04..109d2145941 100644
--- a/airflow-core/tests/unit/models/test_backfill.py
+++ b/airflow-core/tests/unit/models/test_backfill.py
@@ -154,7 +154,8 @@ def test_create_backfill_simple(reverse, existing,
dag_maker, session):
assert all(x.conf == expected_run_conf for x in dag_runs)
-def test_create_backfill_clear_existing_bundle_version(dag_maker, session):
[email protected]("run_on_latest_version", [True, False])
+def test_create_backfill_clear_existing_bundle_version(dag_maker, session,
run_on_latest_version):
"""
Verify that when backfill clears an existing dag run, bundle version is
cleared.
"""
@@ -193,11 +194,13 @@ def
test_create_backfill_clear_existing_bundle_version(dag_maker, session):
dag_run_conf=None,
triggering_user_name="pytest",
reprocess_behavior=ReprocessBehavior.FAILED,
+ run_on_latest_version=run_on_latest_version,
)
session.commit()
# verify that the old dag run (not included in backfill) still has first
bundle version
- # but the latter 5, which are included in the backfill, have the latest
bundle version
+ # but the latter 5, which are included in the backfill, have the latest
bundle version if run_on_latest_version
+ # is True, otherwise they have the first bundle version
dag_runs = sorted(
session.scalars(
select(DagRun).where(
@@ -206,7 +209,15 @@ def
test_create_backfill_clear_existing_bundle_version(dag_maker, session):
),
key=lambda x: x.logical_date,
)
- expected = [first_bundle_version] + 5 * [new_bundle_version]
+ if run_on_latest_version:
+ expected = [first_bundle_version] + 5 * [new_bundle_version]
+ else:
+ expected = (
+ [first_bundle_version]
+ + [new_bundle_version]
+ + 2 * [first_bundle_version]
+ + 2 * [new_bundle_version]
+ )
assert [x.bundle_version for x in dag_runs] == expected
diff --git a/airflow-core/tests/unit/models/test_dag.py
b/airflow-core/tests/unit/models/test_dag.py
index 9a99eba5974..dab3d4d9916 100644
--- a/airflow-core/tests/unit/models/test_dag.py
+++ b/airflow-core/tests/unit/models/test_dag.py
@@ -1480,8 +1480,7 @@ class TestDag:
def test_dag_test_basic(self):
dag = DAG(dag_id="test_local_testing_conn_file", schedule=None,
start_date=DEFAULT_DATE)
- dag.sync_to_db()
- SerializedDagModel.write_dag(dag, bundle_name="testing")
+
mock_object = mock.MagicMock()
@task_decorator
@@ -1491,14 +1490,14 @@ class TestDag:
with dag:
check_task()
+ dag.sync_to_db()
+ SerializedDagModel.write_dag(dag, bundle_name="testing")
dag.test()
mock_object.assert_called_once()
def test_dag_test_with_dependencies(self):
dag = DAG(dag_id="test_local_testing_conn_file", schedule=None,
start_date=DEFAULT_DATE)
- dag.sync_to_db()
- SerializedDagModel.write_dag(dag, bundle_name="testing")
mock_object = mock.MagicMock()
@task_decorator
@@ -1513,6 +1512,9 @@ class TestDag:
with dag:
check_task_2(check_task())
+ dag.sync_to_db()
+ SerializedDagModel.write_dag(dag, bundle_name="testing")
+
dag.test()
mock_object.assert_called_with("output of first task")
@@ -1538,8 +1540,6 @@ class TestDag:
mock_task_object_1 = mock.MagicMock()
mock_task_object_2 = mock.MagicMock()
- dag.sync_to_db()
- SerializedDagModel.write_dag(dag, bundle_name="testing")
@task_decorator
def check_task():
@@ -1553,7 +1553,8 @@ class TestDag:
with dag:
check_task_2(check_task())
-
+ dag.sync_to_db()
+ SerializedDagModel.write_dag(dag, bundle_name="testing")
dr = dag.test()
ti1 = dr.get_task_instance("check_task")
diff --git a/airflow-ctl/src/airflowctl/api/datamodels/generated.py
b/airflow-ctl/src/airflowctl/api/datamodels/generated.py
index 0fe1d48ae40..0910c25a074 100644
--- a/airflow-ctl/src/airflowctl/api/datamodels/generated.py
+++ b/airflow-ctl/src/airflowctl/api/datamodels/generated.py
@@ -201,6 +201,13 @@ class ClearTaskInstancesBody(BaseModel):
include_downstream: Annotated[bool | None, Field(title="Include
Downstream")] = False
include_future: Annotated[bool | None, Field(title="Include Future")] =
False
include_past: Annotated[bool | None, Field(title="Include Past")] = False
+ run_on_latest_version: Annotated[
+ bool | None,
+ Field(
+ description="(Experimental) Run on the latest bundle version of
the DAG after clearing the task instances.",
+ title="Run On Latest Version",
+ ),
+ ] = False
class Value(RootModel[list]):
@@ -308,6 +315,13 @@ class DAGRunClearBody(BaseModel):
)
dry_run: Annotated[bool | None, Field(title="Dry Run")] = True
only_failed: Annotated[bool | None, Field(title="Only Failed")] = False
+ run_on_latest_version: Annotated[
+ bool | None,
+ Field(
+ description="(Experimental) Run on the latest bundle version of
the DAG after clearing the DAG Run.",
+ title="Run On Latest Version",
+ ),
+ ] = False
class DAGRunPatchStates(str, Enum):
diff --git a/providers/common/sql/tests/unit/common/sql/operators/test_sql.py
b/providers/common/sql/tests/unit/common/sql/operators/test_sql.py
index 489e228dd80..23065cd3a71 100644
--- a/providers/common/sql/tests/unit/common/sql/operators/test_sql.py
+++ b/providers/common/sql/tests/unit/common/sql/operators/test_sql.py
@@ -1140,7 +1140,7 @@ class TestSqlBranch:
@pytest.fixture
def branch_op(self):
- return BranchSQLOperator(
+ branch_op = BranchSQLOperator(
task_id="make_choice",
conn_id="mysql_default",
sql="SELECT 1",
@@ -1148,6 +1148,10 @@ class TestSqlBranch:
follow_task_ids_if_false=["branch_2"],
dag=self.dag,
)
+ if AIRFLOW_V_3_0_PLUS:
+ self.dag.sync_to_db()
+ SerializedDagModel.write_dag(self.dag, bundle_name="testing")
+ return branch_op
def test_unsupported_conn_type(self):
"""Check if BranchSQLOperator throws an exception for unsupported
connection type"""
@@ -1387,6 +1391,8 @@ class TestSqlBranch:
self.dag.clear()
if AIRFLOW_V_3_0_PLUS:
+ self.dag.sync_to_db()
+ SerializedDagModel.write_dag(self.dag, bundle_name="testing")
dagrun_kwargs = {
"logical_date": DEFAULT_DATE,
"run_after": DEFAULT_DATE,
diff --git a/providers/standard/tests/unit/standard/decorators/test_bash.py
b/providers/standard/tests/unit/standard/decorators/test_bash.py
index f4303b345eb..e53e8e75f6b 100644
--- a/providers/standard/tests/unit/standard/decorators/test_bash.py
+++ b/providers/standard/tests/unit/standard/decorators/test_bash.py
@@ -342,7 +342,7 @@ class TestBashDecorator:
"""Verify task failure for non-existent, user-defined working
directory."""
cwd_path = tmp_path / "test_cwd"
- with self.dag:
+ with self.dag_maker():
@task.bash(cwd=os.fspath(cwd_path))
def bash():
@@ -363,7 +363,7 @@ class TestBashDecorator:
cwd_file = tmp_path / "testfile.var.env"
cwd_file.touch()
- with self.dag:
+ with self.dag_maker():
@task.bash(cwd=os.fspath(cwd_file))
def bash():
@@ -381,7 +381,7 @@ class TestBashDecorator:
def test_command_not_found(self):
"""Fail task if executed command is not found on path."""
- with self.dag:
+ with self.dag_maker():
@task.bash
def bash():
@@ -481,7 +481,7 @@ class TestBashDecorator:
def test_rtif_updates_upon_failure(self):
"""Veriy RenderedTaskInstanceField data should contain the rendered
command even if the task fails."""
- with self.dag:
+ with self.dag_maker():
@task.bash
def bash():
diff --git a/providers/standard/tests/unit/standard/operators/test_python.py
b/providers/standard/tests/unit/standard/operators/test_python.py
index 7a3ad2167a9..95c503aed70 100644
--- a/providers/standard/tests/unit/standard/operators/test_python.py
+++ b/providers/standard/tests/unit/standard/operators/test_python.py
@@ -551,11 +551,11 @@ class TestBranchOperator(BasePythonTest):
("join", [State.SUCCESS, State.SKIPPED, State.SUCCESS]),
],
)
- def test_empty_branch(self, choice, expected_states):
+ def test_empty_branch(self, choice, expected_states, session):
"""
Tests that BranchPythonOperator handles empty branches properly.
"""
- with self.dag_non_serialized:
+ with self.dag_maker(self.dag_id,
template_searchpath=TEMPLATE_SEARCHPATH, serialized=True):
def f():
return choice
@@ -567,13 +567,13 @@ class TestBranchOperator(BasePythonTest):
branch >> [task1, join]
task1 >> join
- dr = self.create_dag_run()
+ dr = self.dag_maker.create_dagrun()
task_ids = [self.task_id, "task1", "join"]
tis = {ti.task_id: ti for ti in dr.task_instances}
for task_id in task_ids: # Mimic the specific order the scheduling
would run the tests.
task_instance = tis[task_id]
-
task_instance.refresh_from_task(self.dag_non_serialized.get_task(task_id))
+
task_instance.refresh_from_task(self.dag_maker.dag.get_task(task_id))
if AIRFLOW_V_3_0_1:
from airflow.exceptions import DownstreamTasksSkipped
@@ -758,10 +758,11 @@ class TestShortCircuitOperator(BasePythonTest):
After a downstream task is skipped by ShortCircuitOperator, clearing
the skipped task
should not cause it to be executed.
"""
- with self.dag_non_serialized:
+ with self.dag_maker(self.dag_id,
template_searchpath=TEMPLATE_SEARCHPATH, serialized=True):
short_circuit = ShortCircuitOperator(task_id=self.task_id,
python_callable=lambda: False)
short_circuit >> self.op1 >> self.op2
- dr = self.create_dag_run()
+
+ dr = self.dag_maker.create_dagrun()
if AIRFLOW_V_3_0_1:
from airflow.exceptions import DownstreamTasksSkipped
diff --git a/task-sdk/src/airflow/sdk/execution_time/supervisor.py
b/task-sdk/src/airflow/sdk/execution_time/supervisor.py
index e36b3667e2e..2e8d46e5e74 100644
--- a/task-sdk/src/airflow/sdk/execution_time/supervisor.py
+++ b/task-sdk/src/airflow/sdk/execution_time/supervisor.py
@@ -1402,20 +1402,17 @@ class InProcessTestSupervisor(ActivitySubprocess):
@staticmethod
def _api_client(dag=None):
- from airflow.models.dagbag import DagBag
from airflow.sdk.api.client import Client
api = in_process_api_server()
if dag is not None:
from airflow.api_fastapi.common.dagbag import dag_bag_from_app
- from airflow.serialization.serialized_objects import SerializedDAG
+ from airflow.jobs.scheduler_job_runner import SchedulerDagBag
- # This is needed since the Execution API server uses the DagBag in
its "state".
+ # This is needed since the Execution API server uses the
SchedulerDagBag in its "state".
# This `app.state.dag_bag` is used to get some DAG properties like
`fail_fast`.
- dag_bag = DagBag(include_examples=False, collect_dags=False,
load_op_links=False)
+ dag_bag = SchedulerDagBag()
- # Mimic the behavior of the DagBag in the API server by converting
the DAG to a SerializedDAG
- dag_bag.dags[dag.dag_id] =
SerializedDAG.from_dict(SerializedDAG.to_dict(dag))
api.app.dependency_overrides[dag_bag_from_app] = lambda: dag_bag
client = Client(base_url=None, token="", dry_run=True,
transport=api.transport)
