From fce529479a72041850e5870950b2c7262e4886a9 Mon Sep 17 00:00:00 2001 From: Niko Oliveira Date: Mon, 29 Jun 2026 09:51:27 -0700 Subject: [PATCH 1/4] Thread version_data to callbacks PR #67217 wired version_data into the task execution path but not callbacks, so callbacks for a pinned run initialized their bundle without the manifest tasks used. Populate version_data on the callback producer paths (ExecuteCallback.make and the CallbackRequest creation sites) under the same pin guard as tasks, carry it on BaseCallbackRequest, and forward it through prepare_callback_bundle to get_bundle. --- .../airflow/callbacks/callback_requests.py | 6 ++ .../src/airflow/dag_processing/manager.py | 6 +- .../airflow/executors/workloads/callback.py | 4 ++ .../src/airflow/executors/workloads/task.py | 7 +-- .../src/airflow/jobs/scheduler_job_runner.py | 14 ++++- .../src/airflow/models/dag_version.py | 20 ++++++- airflow-core/src/airflow/models/dagrun.py | 6 ++ airflow-core/src/airflow/models/trigger.py | 6 ++ .../unit/callbacks/test_callback_requests.py | 27 +++++++++ .../tests/unit/dag_processing/test_manager.py | 23 ++++++- .../tests/unit/executors/test_workloads.py | 60 ++++++++++++++++++- .../tests/unit/jobs/test_scheduler_job.py | 1 + .../tests/unit/models/test_dag_version.py | 29 +++++++++ 13 files changed, 200 insertions(+), 9 deletions(-) diff --git a/airflow-core/src/airflow/callbacks/callback_requests.py b/airflow-core/src/airflow/callbacks/callback_requests.py index ce48438f0e70a..2573c0d89bbde 100644 --- a/airflow-core/src/airflow/callbacks/callback_requests.py +++ b/airflow-core/src/airflow/callbacks/callback_requests.py @@ -46,6 +46,12 @@ class BaseCallbackRequest(BaseModel): """File Path to use to run the callback""" bundle_name: str bundle_version: str | None + version_data: dict[str, Any] | None = None + """Optional structured metadata for the pinned bundle version (e.g. an S3 object manifest). + + Populated only for pinned runs so the callback initializes the bundle against the same + version the task ran with; ``None`` for unpinned runs. + """ msg: str | None = None """Additional Message that can be used for logging to determine failure/task heartbeat timeout""" diff --git a/airflow-core/src/airflow/dag_processing/manager.py b/airflow-core/src/airflow/dag_processing/manager.py index b2428985a6506..f41883edf31fd 100644 --- a/airflow-core/src/airflow/dag_processing/manager.py +++ b/airflow-core/src/airflow/dag_processing/manager.py @@ -705,7 +705,11 @@ def prepare_callback_bundle(self, request: CallbackRequest) -> BaseDagBundle | N Override to source the bundle from an API. """ try: - bundle = DagBundlesManager().get_bundle(name=request.bundle_name, version=request.bundle_version) + bundle = DagBundlesManager().get_bundle( + name=request.bundle_name, + version=request.bundle_version, + version_data=request.version_data, + ) except ValueError: self.log.error("Bundle %s no longer configured, skipping callback", request.bundle_name) return None diff --git a/airflow-core/src/airflow/executors/workloads/callback.py b/airflow-core/src/airflow/executors/workloads/callback.py index 04f26b8e787c1..a83deebc101e8 100644 --- a/airflow-core/src/airflow/executors/workloads/callback.py +++ b/airflow-core/src/airflow/executors/workloads/callback.py @@ -114,9 +114,13 @@ def make( ) -> ExecuteCallback: """Create an ExecuteCallback workload from a Callback ORM model.""" if not bundle_info: + from airflow.models.dag_version import resolve_pinned_version_data + + version_data = resolve_pinned_version_data(dag_run.created_dag_version, dag_run.bundle_version) bundle_info = BundleInfo( name=dag_run.dag_model.bundle_name, version=dag_run.bundle_version, + version_data=version_data, ) fname = f"executor_callbacks/{dag_run.dag_id}/{dag_run.run_id}/{callback.id}" diff --git a/airflow-core/src/airflow/executors/workloads/task.py b/airflow-core/src/airflow/executors/workloads/task.py index 611457f88a957..6c4e0ffa7e466 100644 --- a/airflow-core/src/airflow/executors/workloads/task.py +++ b/airflow-core/src/airflow/executors/workloads/task.py @@ -102,13 +102,12 @@ def make( ser_ti = TaskInstanceDTO.model_validate(ti, from_attributes=True) if not bundle_info: - version_data = None - if ti.dag_version is not None and ti.dag_run.bundle_version is not None: - version_data = ti.dag_version.version_data + from airflow.models.dag_version import resolve_pinned_version_data + bundle_info = BundleInfo( name=ti.dag_model.bundle_name, version=ti.dag_run.bundle_version, - version_data=version_data, + version_data=resolve_pinned_version_data(ti.dag_version, ti.dag_run.bundle_version), ) fname = log_filename_template_renderer()(ti=ti) diff --git a/airflow-core/src/airflow/jobs/scheduler_job_runner.py b/airflow-core/src/airflow/jobs/scheduler_job_runner.py index 1718987cff683..5b0f901fc9249 100644 --- a/airflow-core/src/airflow/jobs/scheduler_job_runner.py +++ b/airflow-core/src/airflow/jobs/scheduler_job_runner.py @@ -97,7 +97,7 @@ ConnectionTestState, ) from airflow.models.dag import DagModel -from airflow.models.dag_version import DagVersion +from airflow.models.dag_version import DagVersion, resolve_pinned_version_data from airflow.models.dagbag import DBDagBag from airflow.models.dagbundle import DagBundleModel from airflow.models.dagrun import DagRun @@ -1514,6 +1514,7 @@ def process_executor_events( if ti.dag_version and ti.dag_run.bundle_version is not None else ti.dag_run.bundle_version ) + _version_data = resolve_pinned_version_data(ti.dag_version, ti.dag_run.bundle_version) # Backfill dag_version_id for legacy tasks (Pydantic requires uuid.UUID). if not _ensure_ti_has_dag_version_id(ti, session, cls.logger()): continue @@ -1521,6 +1522,7 @@ def process_executor_events( filepath=ti.dag_model.relative_fileloc or "", bundle_name=_bundle_name, bundle_version=_bundle_version, + version_data=_version_data, ti=ti, msg=msg, task_callback_type=( @@ -1563,6 +1565,9 @@ def process_executor_events( _email_bundle_version = ( ti.dag_version.bundle_version if ti.dag_version else ti.dag_run.bundle_version ) + _email_version_data = resolve_pinned_version_data( + ti.dag_version, ti.dag_run.bundle_version + ) # Backfill dag_version_id for legacy tasks (Pydantic requires uuid.UUID). if not _ensure_ti_has_dag_version_id(ti, session, cls.logger()): continue @@ -1570,6 +1575,7 @@ def process_executor_events( filepath=ti.dag_model.relative_fileloc or "", bundle_name=_email_bundle_name, bundle_version=_email_bundle_version, + version_data=_email_version_data, ti=ti, msg=msg, email_type="retry" if ti.is_eligible_to_retry() else "failure", @@ -3082,6 +3088,9 @@ def _maybe_requeue_stuck_ti(self, *, ti, session, executor): if ti.dag_version and ti.dag_run.bundle_version is not None else ti.dag_run.bundle_version ) + _stuck_version_data = resolve_pinned_version_data( + ti.dag_version, ti.dag_run.bundle_version + ) # Backfill dag_version_id for legacy tasks (Pydantic requires uuid.UUID). # Note: we cannot use `continue` here because this method is not # inside a loop. If backfilling fails we simply skip the callback. @@ -3090,6 +3099,7 @@ def _maybe_requeue_stuck_ti(self, *, ti, session, executor): filepath=ti.dag_model.relative_fileloc or "", bundle_name=_stuck_bundle_name, bundle_version=_stuck_bundle_version, + version_data=_stuck_version_data, ti=ti, msg=msg, context_from_server=TIRunContext( @@ -3553,6 +3563,7 @@ def _purge_task_instances_without_heartbeats( if ti.dag_version and ti.dag_run.bundle_version is not None else ti.dag_run.bundle_version ) + _hb_version_data = resolve_pinned_version_data(ti.dag_version, ti.dag_run.bundle_version) # Backfill dag_version_id for legacy tasks (Pydantic requires uuid.UUID). if not _ensure_ti_has_dag_version_id(ti, session, self.log): continue @@ -3560,6 +3571,7 @@ def _purge_task_instances_without_heartbeats( filepath=ti.dag_model.relative_fileloc or "", bundle_name=_hb_bundle_name, bundle_version=_hb_bundle_version, + version_data=_hb_version_data, ti=ti, msg=str(task_instance_heartbeat_timeout_message_details), context_from_server=TIRunContext( diff --git a/airflow-core/src/airflow/models/dag_version.py b/airflow-core/src/airflow/models/dag_version.py index e6564a6da0668..678857a47728d 100644 --- a/airflow-core/src/airflow/models/dag_version.py +++ b/airflow-core/src/airflow/models/dag_version.py @@ -19,7 +19,7 @@ import logging from datetime import datetime -from typing import TYPE_CHECKING +from typing import TYPE_CHECKING, Any from uuid import UUID import sqlalchemy as sa @@ -235,3 +235,21 @@ def get_version( def version(self) -> str: """A human-friendly representation of the version.""" return f"{self.dag_id}-{self.version_number}" + + +def resolve_pinned_version_data( + dag_version: DagVersion | None, bundle_version: str | None +) -> dict[str, Any] | None: + """ + Return a bundle version's ``version_data`` manifest, but only for pinned runs. + + Mirrors the bundle-version pinning rule used when building task and callback + workloads: ``version_data`` is exposed only when the run is pinned + (``bundle_version`` is set) and a ``DagVersion`` is available, so the worker + initializes the bundle against the exact version the run used. Returns ``None`` + for unpinned runs (which should follow the latest bundle state) and for legacy + rows without a ``DagVersion``. + """ + if dag_version is not None and bundle_version is not None: + return dag_version.version_data + return None diff --git a/airflow-core/src/airflow/models/dagrun.py b/airflow-core/src/airflow/models/dagrun.py index 056151cefa1c1..d143daf9e5f9f 100644 --- a/airflow-core/src/airflow/models/dagrun.py +++ b/airflow-core/src/airflow/models/dagrun.py @@ -1449,12 +1449,18 @@ def produce_dag_callback( ) relevant_ti = None if not execute: + from airflow.models.dag_version import resolve_pinned_version_data + + # Only carry version_data for pinned runs so the callback initializes the bundle + # against the same version the run used. + version_data = resolve_pinned_version_data(self.created_dag_version, self.bundle_version) return DagCallbackRequest( filepath=self.dag_model.relative_fileloc, dag_id=self.dag_id, run_id=self.run_id, bundle_name=self.dag_model.bundle_name, bundle_version=self.bundle_version, + version_data=version_data, context_from_server=DagRunContext( dag_run=self, last_ti=relevant_ti, diff --git a/airflow-core/src/airflow/models/trigger.py b/airflow-core/src/airflow/models/trigger.py index e457b2ec93691..3dffea1a2d4c4 100644 --- a/airflow-core/src/airflow/models/trigger.py +++ b/airflow-core/src/airflow/models/trigger.py @@ -570,12 +570,18 @@ def _submit_callback_if_necessary() -> None: if event.task_instance_state in (TaskInstanceState.SUCCESS, TaskInstanceState.FAILED): if task_instance.dag_model.relative_fileloc is None: raise RuntimeError("relative_fileloc should not be None for a finished task") + from airflow.models.dag_version import resolve_pinned_version_data + + version_data = resolve_pinned_version_data( + task_instance.dag_version, task_instance.dag_run.bundle_version + ) request = TaskCallbackRequest( filepath=task_instance.dag_model.relative_fileloc, ti=task_instance, task_callback_type=event.task_instance_state, bundle_name=task_instance.dag_model.bundle_name, bundle_version=task_instance.dag_run.bundle_version, + version_data=version_data, ) log.info("Sending callback: %s", request) try: diff --git a/airflow-core/tests/unit/callbacks/test_callback_requests.py b/airflow-core/tests/unit/callbacks/test_callback_requests.py index 434223e37475f..9b85d61e5dba0 100644 --- a/airflow-core/tests/unit/callbacks/test_callback_requests.py +++ b/airflow-core/tests/unit/callbacks/test_callback_requests.py @@ -118,6 +118,33 @@ def test_is_failure_callback_property( assert request.is_failure_callback == expected_is_failure + def test_version_data_round_trips_and_defaults_none(self): + """version_data survives JSON serialization and defaults to None when omitted.""" + version_data = {"schema_version": 1, "files": {"dags/my_dag.py": "ver123"}} + request = DagCallbackRequest( + filepath="filepath", + dag_id="fake_dag", + run_id="fake_run", + is_failure_callback=False, + bundle_name="testing", + bundle_version="abc123", + version_data=version_data, + ) + result = DagCallbackRequest.from_json(request.to_json()) + assert result.version_data == version_data + + # Omitted -> defaults to None and round-trips as None. + unpinned = DagCallbackRequest( + filepath="filepath", + dag_id="fake_dag", + run_id="fake_run", + is_failure_callback=False, + bundle_name="testing", + bundle_version=None, + ) + assert unpinned.version_data is None + assert DagCallbackRequest.from_json(unpinned.to_json()).version_data is None + class TestDagRunContext: def test_dagrun_context_creation(self): diff --git a/airflow-core/tests/unit/dag_processing/test_manager.py b/airflow-core/tests/unit/dag_processing/test_manager.py index 4b55b6b2f1ee9..746d7ce076bcb 100644 --- a/airflow-core/tests/unit/dag_processing/test_manager.py +++ b/airflow-core/tests/unit/dag_processing/test_manager.py @@ -1967,7 +1967,28 @@ def test_prepare_callback_bundle_initializes_versioned_bundle(self, mock_bundle_ bundle.initialize.assert_called_once() @mock.patch("airflow.dag_processing.manager.DagBundlesManager") - def test_prepare_callback_bundle_skips_initialize_for_unversioned_request(self, mock_bundle_manager): + def test_prepare_callback_bundle_forwards_version_data(self, mock_bundle_manager): + manager = DagFileProcessorManager(max_runs=1) + bundle = MagicMock(spec=BaseDagBundle) + bundle.supports_versioning = True + mock_bundle_manager.return_value.get_bundle.return_value = bundle + + version_data = {"schema_version": 1, "files": {"dags/my_dag.py": "ver123"}} + request = DagCallbackRequest( + filepath="file1.py", + dag_id="dag1", + run_id="run1", + is_failure_callback=False, + bundle_name="testing", + bundle_version="some_commit_hash", + version_data=version_data, + msg=None, + ) + + manager.prepare_callback_bundle(request) + mock_bundle_manager.return_value.get_bundle.assert_called_once_with( + name="testing", version="some_commit_hash", version_data=version_data + ) manager = DagFileProcessorManager(max_runs=1) bundle = MagicMock(spec=BaseDagBundle) bundle.supports_versioning = True diff --git a/airflow-core/tests/unit/executors/test_workloads.py b/airflow-core/tests/unit/executors/test_workloads.py index 63a249a36fd1b..1ef04477d081c 100644 --- a/airflow-core/tests/unit/executors/test_workloads.py +++ b/airflow-core/tests/unit/executors/test_workloads.py @@ -28,7 +28,7 @@ from airflow.executors import workloads from airflow.executors.workloads import TaskInstance, TaskInstanceDTO, base as workloads_base from airflow.executors.workloads.base import BaseWorkloadSchema, BundleInfo -from airflow.executors.workloads.callback import CallbackDTO, CallbackFetchMethod +from airflow.executors.workloads.callback import CallbackDTO, CallbackFetchMethod, ExecuteCallback from airflow.executors.workloads.task import ExecuteTask from airflow.executors.workloads.types import state_class_for_key from airflow.models.callback import CallbackKey @@ -250,3 +250,61 @@ def test_missing_dag_version_yields_none(self): assert workload.bundle_info.version == "abc123" assert workload.bundle_info.version_data is None + + +class TestExecuteCallbackMakeVersionData: + """Tests for ExecuteCallback.make() threading version_data through BundleInfo.""" + + @staticmethod + def _make_mocks(bundle_version, version_data, *, has_created_dag_version=True): + """Build mock Callback + DagRun with the attributes ExecuteCallback.make() reads.""" + from unittest.mock import Mock + + callback = Mock() + callback.id = uuid4() + callback.fetch_method = CallbackFetchMethod.IMPORT_PATH + callback.data = {"path": "my_module.my_callback"} + + dag_run = Mock() + dag_run.dag_id = "test_dag" + dag_run.run_id = "test_run" + dag_run.bundle_version = bundle_version + dag_run.dag_model.bundle_name = "test-bundle" + dag_run.dag_model.relative_fileloc = "dags/test_dag.py" + if has_created_dag_version: + dag_run.created_dag_version.version_data = version_data + else: + dag_run.created_dag_version = None + + return callback, dag_run + + def test_pinned_run_populates_version_data(self): + """When the run is pinned, version_data from created_dag_version flows to BundleInfo.""" + version_data = {"schema_version": 1, "files": {"dags/my_dag.py": "ver123"}} + callback, dag_run = self._make_mocks(bundle_version="abc123", version_data=version_data) + + workload = ExecuteCallback.make(callback=callback, dag_run=dag_run) + + assert workload.bundle_info.version == "abc123" + assert workload.bundle_info.version_data == version_data + + def test_unpinned_run_suppresses_present_version_data(self): + """An unpinned run must not expose version_data even when created_dag_version carries it.""" + version_data = {"schema_version": 1, "files": {"dags/my_dag.py": "ver123"}} + callback, dag_run = self._make_mocks(bundle_version=None, version_data=version_data) + + workload = ExecuteCallback.make(callback=callback, dag_run=dag_run) + + assert workload.bundle_info.version is None + assert workload.bundle_info.version_data is None + + def test_missing_created_dag_version_yields_none(self): + """A pinned run without a created_dag_version yields no version_data.""" + callback, dag_run = self._make_mocks( + bundle_version="abc123", version_data=None, has_created_dag_version=False + ) + + workload = ExecuteCallback.make(callback=callback, dag_run=dag_run) + + assert workload.bundle_info.version == "abc123" + assert workload.bundle_info.version_data is None diff --git a/airflow-core/tests/unit/jobs/test_scheduler_job.py b/airflow-core/tests/unit/jobs/test_scheduler_job.py index 603fcf4bf66cf..e5bc8827b4776 100644 --- a/airflow-core/tests/unit/jobs/test_scheduler_job.py +++ b/airflow-core/tests/unit/jobs/test_scheduler_job.py @@ -730,6 +730,7 @@ def test_process_executor_events_with_callback( ti=mock.ANY, bundle_name="dag_maker", bundle_version=None, + version_data=None, msg=f"Executor {executor} reported that the task instance " f" " "finished with state failed, but the task instance's state attribute is queued. " diff --git a/airflow-core/tests/unit/models/test_dag_version.py b/airflow-core/tests/unit/models/test_dag_version.py index ec600c7be66e3..def8a0a35127c 100644 --- a/airflow-core/tests/unit/models/test_dag_version.py +++ b/airflow-core/tests/unit/models/test_dag_version.py @@ -17,6 +17,7 @@ from __future__ import annotations from datetime import timedelta +from unittest import mock import pytest from sqlalchemy import func, select @@ -181,3 +182,31 @@ def test_write_dag_without_version_data(self, dag_maker, session): retrieved = DagVersion.get_latest_version("test_no_version_data", session=session) assert retrieved.version_data is None assert retrieved.bundle_version == "abc123" + + +class TestResolvePinnedVersionData: + """Unit tests for the resolve_pinned_version_data pin-guard helper.""" + + @pytest.mark.parametrize( + ("dag_version", "bundle_version", "expected"), + [ + pytest.param( + mock.Mock(version_data={"schema_version": 1}), + "abc123", + {"schema_version": 1}, + id="pinned-with-data", + ), + pytest.param( + mock.Mock(version_data={"schema_version": 1}), + None, + None, + id="unpinned-suppresses-present-data", + ), + pytest.param(None, "abc123", None, id="missing-dag-version"), + pytest.param(None, None, None, id="unpinned-and-missing"), + ], + ) + def test_resolve_pinned_version_data(self, dag_version, bundle_version, expected): + from airflow.models.dag_version import resolve_pinned_version_data + + assert resolve_pinned_version_data(dag_version, bundle_version) == expected From 504af36f1d4ba1c75ed4b3d04bd1ef9ef9a28fc3 Mon Sep 17 00:00:00 2001 From: Niko Oliveira Date: Tue, 30 Jun 2026 13:25:12 -0700 Subject: [PATCH 2/4] Fix test failures --- .../tests/unit/dag_processing/test_manager.py | 1 + .../sdk/execution_time/schema/schema.json | 39 +++++++++++++++++++ 2 files changed, 40 insertions(+) diff --git a/airflow-core/tests/unit/dag_processing/test_manager.py b/airflow-core/tests/unit/dag_processing/test_manager.py index 746d7ce076bcb..8ff1e46d52801 100644 --- a/airflow-core/tests/unit/dag_processing/test_manager.py +++ b/airflow-core/tests/unit/dag_processing/test_manager.py @@ -1375,6 +1375,7 @@ def test_collect_results_processes_remaining_files_when_one_persist_fails(self, "filepath": "dag_callback_dag.py", "bundle_name": "testing", "bundle_version": None, + "version_data": None, "msg": None, "dag_id": "dag_id", "run_id": "run_id", diff --git a/task-sdk/src/airflow/sdk/execution_time/schema/schema.json b/task-sdk/src/airflow/sdk/execution_time/schema/schema.json index e6ce8aa3d066e..ba9be425a2506 100644 --- a/task-sdk/src/airflow/sdk/execution_time/schema/schema.json +++ b/task-sdk/src/airflow/sdk/execution_time/schema/schema.json @@ -753,6 +753,19 @@ ], "title": "Bundle Version" }, + "version_data": { + "anyOf": [ + { + "additionalProperties": true, + "type": "object" + }, + { + "type": "null" + } + ], + "default": null, + "title": "Version Data" + }, "msg": { "anyOf": [ { @@ -1651,6 +1664,19 @@ ], "title": "Bundle Version" }, + "version_data": { + "anyOf": [ + { + "additionalProperties": true, + "type": "object" + }, + { + "type": "null" + } + ], + "default": null, + "title": "Version Data" + }, "msg": { "anyOf": [ { @@ -3859,6 +3885,19 @@ ], "title": "Bundle Version" }, + "version_data": { + "anyOf": [ + { + "additionalProperties": true, + "type": "object" + }, + { + "type": "null" + } + ], + "default": null, + "title": "Version Data" + }, "msg": { "anyOf": [ { From ccb2d0eccb27467274d6759e25f901449b89b1bd Mon Sep 17 00:00:00 2001 From: Niko Oliveira Date: Mon, 6 Jul 2026 14:59:00 -0700 Subject: [PATCH 3/4] Address feedback from PR --- .../airflow/callbacks/callback_requests.py | 6 +---- .../airflow/executors/workloads/callback.py | 4 +-- .../src/airflow/executors/workloads/task.py | 4 +-- .../src/airflow/jobs/scheduler_job_runner.py | 14 ++++------ .../src/airflow/models/dag_version.py | 16 ++++------- airflow-core/src/airflow/models/dagrun.py | 4 +-- airflow-core/src/airflow/models/trigger.py | 23 ++++++++++++---- .../tests/unit/dag_processing/test_manager.py | 3 +++ .../tests/unit/models/test_dag_version.py | 10 +++---- .../tests/unit/models/test_trigger.py | 27 +++++++++++++++++++ 10 files changed, 70 insertions(+), 41 deletions(-) diff --git a/airflow-core/src/airflow/callbacks/callback_requests.py b/airflow-core/src/airflow/callbacks/callback_requests.py index 2573c0d89bbde..6609dc30c6d9d 100644 --- a/airflow-core/src/airflow/callbacks/callback_requests.py +++ b/airflow-core/src/airflow/callbacks/callback_requests.py @@ -47,11 +47,7 @@ class BaseCallbackRequest(BaseModel): bundle_name: str bundle_version: str | None version_data: dict[str, Any] | None = None - """Optional structured metadata for the pinned bundle version (e.g. an S3 object manifest). - - Populated only for pinned runs so the callback initializes the bundle against the same - version the task ran with; ``None`` for unpinned runs. - """ + """Optional structured metadata for the pinned bundle version (e.g. an S3 object manifest).""" msg: str | None = None """Additional Message that can be used for logging to determine failure/task heartbeat timeout""" diff --git a/airflow-core/src/airflow/executors/workloads/callback.py b/airflow-core/src/airflow/executors/workloads/callback.py index a83deebc101e8..c1842a5900685 100644 --- a/airflow-core/src/airflow/executors/workloads/callback.py +++ b/airflow-core/src/airflow/executors/workloads/callback.py @@ -114,9 +114,9 @@ def make( ) -> ExecuteCallback: """Create an ExecuteCallback workload from a Callback ORM model.""" if not bundle_info: - from airflow.models.dag_version import resolve_pinned_version_data + from airflow.models.dag_version import _resolve_version_data - version_data = resolve_pinned_version_data(dag_run.created_dag_version, dag_run.bundle_version) + version_data = _resolve_version_data(dag_run.created_dag_version, dag_run.bundle_version) bundle_info = BundleInfo( name=dag_run.dag_model.bundle_name, version=dag_run.bundle_version, diff --git a/airflow-core/src/airflow/executors/workloads/task.py b/airflow-core/src/airflow/executors/workloads/task.py index 6c4e0ffa7e466..68a917118e39c 100644 --- a/airflow-core/src/airflow/executors/workloads/task.py +++ b/airflow-core/src/airflow/executors/workloads/task.py @@ -102,12 +102,12 @@ def make( ser_ti = TaskInstanceDTO.model_validate(ti, from_attributes=True) if not bundle_info: - from airflow.models.dag_version import resolve_pinned_version_data + from airflow.models.dag_version import _resolve_version_data bundle_info = BundleInfo( name=ti.dag_model.bundle_name, version=ti.dag_run.bundle_version, - version_data=resolve_pinned_version_data(ti.dag_version, ti.dag_run.bundle_version), + version_data=_resolve_version_data(ti.dag_version, ti.dag_run.bundle_version), ) fname = log_filename_template_renderer()(ti=ti) diff --git a/airflow-core/src/airflow/jobs/scheduler_job_runner.py b/airflow-core/src/airflow/jobs/scheduler_job_runner.py index 5b0f901fc9249..1ca61b8ddc861 100644 --- a/airflow-core/src/airflow/jobs/scheduler_job_runner.py +++ b/airflow-core/src/airflow/jobs/scheduler_job_runner.py @@ -97,7 +97,7 @@ ConnectionTestState, ) from airflow.models.dag import DagModel -from airflow.models.dag_version import DagVersion, resolve_pinned_version_data +from airflow.models.dag_version import DagVersion, _resolve_version_data from airflow.models.dagbag import DBDagBag from airflow.models.dagbundle import DagBundleModel from airflow.models.dagrun import DagRun @@ -1514,7 +1514,7 @@ def process_executor_events( if ti.dag_version and ti.dag_run.bundle_version is not None else ti.dag_run.bundle_version ) - _version_data = resolve_pinned_version_data(ti.dag_version, ti.dag_run.bundle_version) + _version_data = _resolve_version_data(ti.dag_version, ti.dag_run.bundle_version) # Backfill dag_version_id for legacy tasks (Pydantic requires uuid.UUID). if not _ensure_ti_has_dag_version_id(ti, session, cls.logger()): continue @@ -1565,9 +1565,7 @@ def process_executor_events( _email_bundle_version = ( ti.dag_version.bundle_version if ti.dag_version else ti.dag_run.bundle_version ) - _email_version_data = resolve_pinned_version_data( - ti.dag_version, ti.dag_run.bundle_version - ) + _email_version_data = _resolve_version_data(ti.dag_version, ti.dag_run.bundle_version) # Backfill dag_version_id for legacy tasks (Pydantic requires uuid.UUID). if not _ensure_ti_has_dag_version_id(ti, session, cls.logger()): continue @@ -3088,9 +3086,7 @@ def _maybe_requeue_stuck_ti(self, *, ti, session, executor): if ti.dag_version and ti.dag_run.bundle_version is not None else ti.dag_run.bundle_version ) - _stuck_version_data = resolve_pinned_version_data( - ti.dag_version, ti.dag_run.bundle_version - ) + _stuck_version_data = _resolve_version_data(ti.dag_version, ti.dag_run.bundle_version) # Backfill dag_version_id for legacy tasks (Pydantic requires uuid.UUID). # Note: we cannot use `continue` here because this method is not # inside a loop. If backfilling fails we simply skip the callback. @@ -3563,7 +3559,7 @@ def _purge_task_instances_without_heartbeats( if ti.dag_version and ti.dag_run.bundle_version is not None else ti.dag_run.bundle_version ) - _hb_version_data = resolve_pinned_version_data(ti.dag_version, ti.dag_run.bundle_version) + _hb_version_data = _resolve_version_data(ti.dag_version, ti.dag_run.bundle_version) # Backfill dag_version_id for legacy tasks (Pydantic requires uuid.UUID). if not _ensure_ti_has_dag_version_id(ti, session, self.log): continue diff --git a/airflow-core/src/airflow/models/dag_version.py b/airflow-core/src/airflow/models/dag_version.py index 678857a47728d..eb0a5204956be 100644 --- a/airflow-core/src/airflow/models/dag_version.py +++ b/airflow-core/src/airflow/models/dag_version.py @@ -237,19 +237,13 @@ def version(self) -> str: return f"{self.dag_id}-{self.version_number}" -def resolve_pinned_version_data( +def _resolve_version_data( dag_version: DagVersion | None, bundle_version: str | None ) -> dict[str, Any] | None: - """ - Return a bundle version's ``version_data`` manifest, but only for pinned runs. - - Mirrors the bundle-version pinning rule used when building task and callback - workloads: ``version_data`` is exposed only when the run is pinned - (``bundle_version`` is set) and a ``DagVersion`` is available, so the worker - initializes the bundle against the exact version the run used. Returns ``None`` - for unpinned runs (which should follow the latest bundle state) and for legacy - rows without a ``DagVersion``. - """ + """Return a bundle version's ``version_data`` manifest, but only for pinned runs.""" + # Expose version_data only when the run is pinned (bundle_version set) and a DagVersion is + # present, so the bundle initializes against the exact version the run used. Unpinned runs + # follow the latest bundle state, and legacy rows have no DagVersion. if dag_version is not None and bundle_version is not None: return dag_version.version_data return None diff --git a/airflow-core/src/airflow/models/dagrun.py b/airflow-core/src/airflow/models/dagrun.py index d143daf9e5f9f..27751ed8b565b 100644 --- a/airflow-core/src/airflow/models/dagrun.py +++ b/airflow-core/src/airflow/models/dagrun.py @@ -1449,11 +1449,11 @@ def produce_dag_callback( ) relevant_ti = None if not execute: - from airflow.models.dag_version import resolve_pinned_version_data + from airflow.models.dag_version import _resolve_version_data # Only carry version_data for pinned runs so the callback initializes the bundle # against the same version the run used. - version_data = resolve_pinned_version_data(self.created_dag_version, self.bundle_version) + version_data = _resolve_version_data(self.created_dag_version, self.bundle_version) return DagCallbackRequest( filepath=self.dag_model.relative_fileloc, dag_id=self.dag_id, diff --git a/airflow-core/src/airflow/models/trigger.py b/airflow-core/src/airflow/models/trigger.py index 3dffea1a2d4c4..d52daa649bc63 100644 --- a/airflow-core/src/airflow/models/trigger.py +++ b/airflow-core/src/airflow/models/trigger.py @@ -570,17 +570,30 @@ def _submit_callback_if_necessary() -> None: if event.task_instance_state in (TaskInstanceState.SUCCESS, TaskInstanceState.FAILED): if task_instance.dag_model.relative_fileloc is None: raise RuntimeError("relative_fileloc should not be None for a finished task") - from airflow.models.dag_version import resolve_pinned_version_data - - version_data = resolve_pinned_version_data( + from airflow.models.dag_version import _resolve_version_data + + # Derive bundle identity from the TI's dag_version (falling back to dag_run/dag_model + # for legacy/unpinned runs), mirroring the other callback sites so bundle_version and + # version_data always describe the same version. + bundle_name = ( + task_instance.dag_version.bundle_name + if task_instance.dag_version + else task_instance.dag_model.bundle_name + ) + bundle_version = ( + task_instance.dag_version.bundle_version + if task_instance.dag_version and task_instance.dag_run.bundle_version is not None + else task_instance.dag_run.bundle_version + ) + version_data = _resolve_version_data( task_instance.dag_version, task_instance.dag_run.bundle_version ) request = TaskCallbackRequest( filepath=task_instance.dag_model.relative_fileloc, ti=task_instance, task_callback_type=event.task_instance_state, - bundle_name=task_instance.dag_model.bundle_name, - bundle_version=task_instance.dag_run.bundle_version, + bundle_name=bundle_name, + bundle_version=bundle_version, version_data=version_data, ) log.info("Sending callback: %s", request) diff --git a/airflow-core/tests/unit/dag_processing/test_manager.py b/airflow-core/tests/unit/dag_processing/test_manager.py index 8ff1e46d52801..ee2d7cd76d4d9 100644 --- a/airflow-core/tests/unit/dag_processing/test_manager.py +++ b/airflow-core/tests/unit/dag_processing/test_manager.py @@ -1990,6 +1990,9 @@ def test_prepare_callback_bundle_forwards_version_data(self, mock_bundle_manager mock_bundle_manager.return_value.get_bundle.assert_called_once_with( name="testing", version="some_commit_hash", version_data=version_data ) + + @mock.patch("airflow.dag_processing.manager.DagBundlesManager") + def test_prepare_callback_bundle_skips_initialize_for_unversioned_request(self, mock_bundle_manager): manager = DagFileProcessorManager(max_runs=1) bundle = MagicMock(spec=BaseDagBundle) bundle.supports_versioning = True diff --git a/airflow-core/tests/unit/models/test_dag_version.py b/airflow-core/tests/unit/models/test_dag_version.py index def8a0a35127c..da4923d8771e4 100644 --- a/airflow-core/tests/unit/models/test_dag_version.py +++ b/airflow-core/tests/unit/models/test_dag_version.py @@ -184,8 +184,8 @@ def test_write_dag_without_version_data(self, dag_maker, session): assert retrieved.bundle_version == "abc123" -class TestResolvePinnedVersionData: - """Unit tests for the resolve_pinned_version_data pin-guard helper.""" +class TestResolveVersionData: + """Unit tests for the _resolve_version_data pin-guard helper.""" @pytest.mark.parametrize( ("dag_version", "bundle_version", "expected"), @@ -206,7 +206,7 @@ class TestResolvePinnedVersionData: pytest.param(None, None, None, id="unpinned-and-missing"), ], ) - def test_resolve_pinned_version_data(self, dag_version, bundle_version, expected): - from airflow.models.dag_version import resolve_pinned_version_data + def test_resolve_version_data(self, dag_version, bundle_version, expected): + from airflow.models.dag_version import _resolve_version_data - assert resolve_pinned_version_data(dag_version, bundle_version) == expected + assert _resolve_version_data(dag_version, bundle_version) == expected diff --git a/airflow-core/tests/unit/models/test_trigger.py b/airflow-core/tests/unit/models/test_trigger.py index 669d98a77045d..91a6a92e27cb3 100644 --- a/airflow-core/tests/unit/models/test_trigger.py +++ b/airflow-core/tests/unit/models/test_trigger.py @@ -335,6 +335,33 @@ def get_xcoms(ti): assert actual_xcoms == expected_xcoms +@patch("airflow.callbacks.database_callback_sink.DatabaseCallbackSink.send") +def test_submit_event_task_end_callback_includes_version_data(mock_send, session, create_task_instance): + """A finished deferred task's callback carries version_data for a pinned run, and its + bundle_version is derived from the same dag_version so the two cannot diverge.""" + version_data = {"schema_version": 1, "files": {"dags/my_dag.py": "ver123"}} + + trigger = Trigger(classpath="does.not.matter", kwargs={}) + session.add(trigger) + task_instance = create_task_instance( + session=session, logical_date=timezone.utcnow(), state=State.DEFERRED + ) + task_instance.trigger_id = trigger.id + # Pin the run and attach a manifest to the TI's dag_version. + task_instance.dag_run.bundle_version = "some_hash" + task_instance.dag_version.bundle_version = "some_hash" + task_instance.dag_version.version_data = version_data + session.commit() + + Trigger.submit_event(trigger.id, TaskSuccessEvent(), session=session) + session.flush() + + mock_send.assert_called_once() + request = mock_send.call_args.kwargs["callback"] + assert request.bundle_version == "some_hash" + assert request.version_data == version_data + + @pytest.fixture def create_triggerer(): """Fixture factory which creates individual test Triggerer instances.""" From bb04f5e18cbf72d56b4d0cbeca9a1ef3f229ad2d Mon Sep 17 00:00:00 2001 From: Niko Oliveira Date: Mon, 6 Jul 2026 17:47:22 -0700 Subject: [PATCH 4/4] Regenerate ts-sdk supervisor.ts for version_data on callback requests --- ts-sdk/src/generated/supervisor.ts | 12 ++++++++++++ 1 file changed, 12 insertions(+) diff --git a/ts-sdk/src/generated/supervisor.ts b/ts-sdk/src/generated/supervisor.ts index fe61155f39e4c..004ad49e7baba 100644 --- a/ts-sdk/src/generated/supervisor.ts +++ b/ts-sdk/src/generated/supervisor.ts @@ -120,6 +120,9 @@ export type Type11 = "DRCount"; export type Filepath = string; export type BundleName = string; export type BundleVersion = string | null; +export type VersionData1 = { + [k: string]: unknown; +} | null; export type Msg = string | null; export type DagId1 = string; export type RunId1 = string; @@ -192,6 +195,9 @@ export type BundleName1 = string; export type Filepath1 = string; export type BundleName2 = string; export type BundleVersion1 = string | null; +export type VersionData2 = { + [k: string]: unknown; +} | null; export type Msg1 = string | null; /** * All possible states that a Task Instance can be in. @@ -243,6 +249,9 @@ export type Type13 = "TaskCallbackRequest"; export type Filepath2 = string; export type BundleName3 = string; export type BundleVersion2 = string | null; +export type VersionData3 = { + [k: string]: unknown; +} | null; export type Msg2 = string | null; export type EmailType = "failure" | "retry"; export type Type14 = "EmailRequest"; @@ -849,6 +858,7 @@ export interface DagCallbackRequest { filepath: Filepath; bundle_name: BundleName; bundle_version: BundleVersion; + version_data?: VersionData1; msg?: Msg; dag_id: DagId1; run_id: RunId1; @@ -959,6 +969,7 @@ export interface TaskCallbackRequest { filepath: Filepath1; bundle_name: BundleName2; bundle_version: BundleVersion1; + version_data?: VersionData2; msg?: Msg1; ti: TaskInstance; task_callback_type?: TaskInstanceState | null; @@ -1019,6 +1030,7 @@ export interface EmailRequest { filepath: Filepath2; bundle_name: BundleName3; bundle_version: BundleVersion2; + version_data?: VersionData3; msg?: Msg2; ti: TaskInstance; email_type?: EmailType;