This is an automated email from the ASF dual-hosted git repository.
Lee-W 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 1b8b0ff40e5 Extract seams for Dag processor callback fetching to allow
DB-to-API swap (#65543)
1b8b0ff40e5 is described below
commit 1b8b0ff40e5db84b3e96e968be5f27c590921158
Author: Ephraim Anierobi <[email protected]>
AuthorDate: Tue Apr 28 09:41:57 2026 +0100
Extract seams for Dag processor callback fetching to allow DB-to-API swap
(#65543)
---
airflow-core/src/airflow/dag_processing/manager.py | 30 ++++-
.../tests/unit/dag_processing/test_manager.py | 147 +++++++++++++++++++--
2 files changed, 163 insertions(+), 14 deletions(-)
diff --git a/airflow-core/src/airflow/dag_processing/manager.py
b/airflow-core/src/airflow/dag_processing/manager.py
index 8b1f6f7d546..4065baf88b5 100644
--- a/airflow-core/src/airflow/dag_processing/manager.py
+++ b/airflow-core/src/airflow/dag_processing/manager.py
@@ -471,7 +471,7 @@ class DagFileProcessorManager(LoggingMixin):
self._collect_results()
- for callback in self._fetch_callbacks():
+ for callback in self.fetch_callbacks():
self._add_callback_to_queue(callback)
self._scan_stale_dags()
self._cleanup_stale_bundle_versions()
@@ -568,9 +568,17 @@ class DagFileProcessorManager(LoggingMixin):
session.delete(request)
return files
+ def fetch_callbacks(self) -> list[CallbackRequest]:
+ """
+ Fetch and claim callbacks for this manager's bundles.
+
+ Default implementation reads from the metadata DB; override to source
callbacks from an API.
+ """
+ return self._fetch_callbacks_from_db()
+
@provide_session
@retry_db_transaction
- def _fetch_callbacks(
+ def _fetch_callbacks_from_db(
self,
session: Session = NEW_SESSION,
) -> list[CallbackRequest]:
@@ -602,12 +610,17 @@ class DagFileProcessorManager(LoggingMixin):
guard.commit()
return callback_queue
- def _add_callback_to_queue(self, request: CallbackRequest):
- self.log.debug("Queuing %s CallbackRequest: %s",
type(request).__name__, request)
+ def prepare_callback_bundle(self, request: CallbackRequest) ->
BaseDagBundle | None:
+ """
+ Return the bundle to run the callback against, or ``None`` to skip the
callback.
+
+ Default implementation looks the bundle up via
:class:`DagBundlesManager` and, for
+ versioned requests on bundles that support versioning, calls
``bundle.initialize()``.
+ Override to source the bundle from an API.
+ """
try:
bundle = DagBundlesManager().get_bundle(name=request.bundle_name,
version=request.bundle_version)
except ValueError:
- # Bundle no longer configured
self.log.error("Bundle %s no longer configured, skipping
callback", request.bundle_name)
return None
if bundle.supports_versioning and request.bundle_version:
@@ -620,6 +633,13 @@ class DagFileProcessorManager(LoggingMixin):
request.bundle_version,
)
return None
+ return bundle
+
+ def _add_callback_to_queue(self, request: CallbackRequest) -> None:
+ self.log.debug("Queuing %s CallbackRequest: %s",
type(request).__name__, request)
+ bundle = self.prepare_callback_bundle(request)
+ if bundle is None:
+ return
file_info = DagFileInfo(
rel_path=Path(request.filepath),
diff --git a/airflow-core/tests/unit/dag_processing/test_manager.py
b/airflow-core/tests/unit/dag_processing/test_manager.py
index e44cf43c9da..5b6ff6af014 100644
--- a/airflow-core/tests/unit/dag_processing/test_manager.py
+++ b/airflow-core/tests/unit/dag_processing/test_manager.py
@@ -1281,7 +1281,7 @@ class TestDagFileProcessorManager:
@conf_vars({("core", "load_examples"): "False"})
def test_fetch_callbacks_from_database(self, configure_testing_dag_bundle):
- """Test _fetch_callbacks returns callbacks ordered by priority_weight
desc."""
+ """Test _fetch_callbacks_from_db returns callbacks ordered by
priority_weight desc."""
dag_filepath = TEST_DAG_FOLDER / "test_on_failure_callback_dag.py"
@@ -1311,7 +1311,7 @@ class TestDagFileProcessorManager:
manager._dag_bundles =
list(DagBundlesManager().get_all_dag_bundles())
with create_session() as session:
- callbacks = manager._fetch_callbacks(session=session)
+ callbacks = manager._fetch_callbacks_from_db(session=session)
# Should return callbacks ordered by priority_weight desc
(highest first)
assert callbacks[0].run_id == "123"
@@ -1326,7 +1326,7 @@ class TestDagFileProcessorManager:
}
)
def test_fetch_callbacks_from_database_max_per_loop(self, tmp_path,
configure_testing_dag_bundle):
- """Test DagFileProcessorManager._fetch_callbacks method"""
+ """Test DagFileProcessorManager.fetch_callbacks method."""
dag_filepath = TEST_DAG_FOLDER / "test_on_failure_callback_dag.py"
with create_session() as session:
@@ -1385,7 +1385,7 @@ class TestDagFileProcessorManager:
manager._dag_bundles =
list(DagBundlesManager().get_all_dag_bundles())
with create_session() as session:
- callbacks = manager._fetch_callbacks(session=session)
+ callbacks = manager._fetch_callbacks_from_db(session=session)
# Only the matching callback should be returned
assert [c.run_id for c in callbacks] == ["match"]
@@ -1441,7 +1441,7 @@ class TestDagFileProcessorManager:
manager._dag_bundles =
list(DagBundlesManager().get_all_dag_bundles())
with create_session() as session:
- callbacks = manager._fetch_callbacks(session=session)
+ callbacks = manager._fetch_callbacks_from_db(session=session)
assert [c.run_id for c in callbacks] == ["match"]
@@ -1564,9 +1564,107 @@ class TestDagFileProcessorManager:
assert dag2_path not in manager._callback_to_execute
@mock.patch("airflow.dag_processing.manager.DagBundlesManager")
- def test_add_callback_initializes_versioned_bundle(self,
mock_bundle_manager):
+ def test_prepare_callback_bundle_initializes_versioned_bundle(self,
mock_bundle_manager):
manager = DagFileProcessorManager(max_runs=1)
- bundle = MagicMock()
+ bundle = MagicMock(spec=BaseDagBundle)
+ bundle.supports_versioning = True
+ mock_bundle_manager.return_value.get_bundle.return_value = bundle
+
+ request = DagCallbackRequest(
+ filepath="file1.py",
+ dag_id="dag1",
+ run_id="run1",
+ is_failure_callback=False,
+ bundle_name="testing",
+ bundle_version="some_commit_hash",
+ msg=None,
+ )
+
+ assert manager.prepare_callback_bundle(request) is 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):
+ manager = DagFileProcessorManager(max_runs=1)
+ bundle = MagicMock(spec=BaseDagBundle)
+ bundle.supports_versioning = True
+ mock_bundle_manager.return_value.get_bundle.return_value = bundle
+
+ request = DagCallbackRequest(
+ filepath="file1.py",
+ dag_id="dag1",
+ run_id="run1",
+ is_failure_callback=False,
+ bundle_name="testing",
+ bundle_version=None,
+ msg=None,
+ )
+
+ assert manager.prepare_callback_bundle(request) is bundle
+ bundle.initialize.assert_not_called()
+
+ @mock.patch("airflow.dag_processing.manager.DagBundlesManager")
+ def
test_prepare_callback_bundle_skips_initialize_for_non_versioning_bundle(self,
mock_bundle_manager):
+ manager = DagFileProcessorManager(max_runs=1)
+ bundle = MagicMock(spec=BaseDagBundle)
+ bundle.supports_versioning = False
+ mock_bundle_manager.return_value.get_bundle.return_value = bundle
+
+ request = DagCallbackRequest(
+ filepath="file1.py",
+ dag_id="dag1",
+ run_id="run1",
+ is_failure_callback=False,
+ bundle_name="testing",
+ bundle_version="some_commit_hash",
+ msg=None,
+ )
+
+ assert manager.prepare_callback_bundle(request) is bundle
+ bundle.initialize.assert_not_called()
+
+ @mock.patch("airflow.dag_processing.manager.DagBundlesManager")
+ def test_prepare_callback_bundle_returns_none_when_unconfigured(self,
mock_bundle_manager):
+ manager = DagFileProcessorManager(max_runs=1)
+ mock_bundle_manager.return_value.get_bundle.side_effect =
ValueError("missing bundle")
+
+ request = DagCallbackRequest(
+ filepath="file1.py",
+ dag_id="dag1",
+ run_id="run1",
+ is_failure_callback=False,
+ bundle_name="testing",
+ bundle_version="some_commit_hash",
+ msg=None,
+ )
+
+ assert manager.prepare_callback_bundle(request) is None
+
+ @mock.patch("airflow.dag_processing.manager.DagBundlesManager")
+ def test_prepare_callback_bundle_returns_none_when_initialize_fails(self,
mock_bundle_manager):
+ manager = DagFileProcessorManager(max_runs=1)
+ bundle = MagicMock(spec=BaseDagBundle)
+ bundle.supports_versioning = True
+ bundle.initialize.side_effect = RuntimeError("clone failed")
+ mock_bundle_manager.return_value.get_bundle.return_value = bundle
+
+ request = DagCallbackRequest(
+ filepath="file1.py",
+ dag_id="dag1",
+ run_id="run1",
+ is_failure_callback=False,
+ bundle_name="testing",
+ bundle_version="some_commit_hash",
+ msg=None,
+ )
+
+ assert manager.prepare_callback_bundle(request) is None
+ bundle.initialize.assert_called_once()
+
+ @mock.patch("airflow.dag_processing.manager.DagBundlesManager")
+ def test_add_callback_queues_file_info_on_success(self,
mock_bundle_manager):
+ manager = DagFileProcessorManager(max_runs=1)
+ bundle = MagicMock(spec=BaseDagBundle)
bundle.supports_versioning = True
bundle.path = Path("/tmp/bundle")
mock_bundle_manager.return_value.get_bundle.return_value = bundle
@@ -1584,11 +1682,16 @@ class TestDagFileProcessorManager:
manager._add_callback_to_queue(request)
bundle.initialize.assert_called_once()
+ assert manager._callback_to_execute
+ [(file_info, requests)] = manager._callback_to_execute.items()
+ assert file_info.rel_path == Path("file1.py")
+ assert file_info.bundle_name == "testing"
+ assert requests == [request]
@mock.patch("airflow.dag_processing.manager.DagBundlesManager")
def test_add_callback_skips_when_bundle_init_fails(self,
mock_bundle_manager):
manager = DagFileProcessorManager(max_runs=1)
- bundle = MagicMock()
+ bundle = MagicMock(spec=BaseDagBundle)
bundle.supports_versioning = True
bundle.initialize.side_effect = Exception("clone failed")
mock_bundle_manager.return_value.get_bundle.return_value = bundle
@@ -1606,7 +1709,33 @@ class TestDagFileProcessorManager:
manager._add_callback_to_queue(request)
bundle.initialize.assert_called_once()
- assert len(manager._callback_to_execute) == 0
+ assert not manager._callback_to_execute
+
+ @mock.patch("airflow.dag_processing.manager.DagBundlesManager")
+ def test_add_callback_skips_when_bundle_unconfigured(self,
mock_bundle_manager):
+ manager = DagFileProcessorManager(max_runs=1)
+ mock_bundle_manager.return_value.get_bundle.side_effect =
ValueError("missing bundle")
+
+ request = DagCallbackRequest(
+ filepath="file1.py",
+ dag_id="dag1",
+ run_id="run1",
+ is_failure_callback=False,
+ bundle_name="testing",
+ bundle_version="some_commit_hash",
+ msg=None,
+ )
+
+ manager._add_callback_to_queue(request)
+
+ assert not manager._callback_to_execute
+
+ def test_fetch_callbacks_delegates_to_private_method(self):
+ manager = DagFileProcessorManager(max_runs=1)
+ expected: list = [mock.sentinel.callback]
+ with mock.patch.object(manager, "_fetch_callbacks_from_db",
return_value=expected) as private:
+ assert manager.fetch_callbacks() is expected
+ private.assert_called_once_with()
def test_dag_with_assets(self, session, configure_testing_dag_bundle):
"""'Integration' test to ensure that the assets get parsed and stored
correctly for parsed dags."""
