This is an automated email from the ASF dual-hosted git repository.
vincbeck 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 1fb77d1d104 Cache bundle-to-team name lookups in dag processor manager
(#68730)
1fb77d1d104 is described below
commit 1fb77d1d104b584405a6f88c69dde418d7fd4513
Author: D. Ferruzzi <[email protected]>
AuthorDate: Mon Jun 22 07:17:59 2026 -0700
Cache bundle-to-team name lookups in dag processor manager (#68730)
---
airflow-core/src/airflow/dag_processing/manager.py | 69 ++++++-----------
airflow-core/src/airflow/models/dagbundle.py | 17 +++++
.../tests/unit/dag_processing/test_manager.py | 89 +++++++++++++++++-----
airflow-core/tests/unit/models/test_dagbundle.py | 17 +++++
4 files changed, 131 insertions(+), 61 deletions(-)
diff --git a/airflow-core/src/airflow/dag_processing/manager.py
b/airflow-core/src/airflow/dag_processing/manager.py
index e43920f2618..b2428985a65 100644
--- a/airflow-core/src/airflow/dag_processing/manager.py
+++ b/airflow-core/src/airflow/dag_processing/manager.py
@@ -85,7 +85,7 @@ from airflow.utils.sqlalchemy import (
)
if TYPE_CHECKING:
- from collections.abc import Callable, Iterable, Iterator, Sequence
+ from collections.abc import Callable, Collection, Iterable, Iterator,
Sequence
from socket import socket
from sqlalchemy.orm import Session
@@ -266,6 +266,8 @@ class DagFileProcessorManager(LoggingMixin):
_dag_bundles: list[BaseDagBundle] = attrs.field(factory=list, init=False)
_bundle_versions: dict[str, str | None] = attrs.field(factory=dict,
init=False)
_bundle_version_data: dict[str, dict | None] = attrs.field(factory=dict,
init=False)
+ _multi_team: bool = attrs.field(factory=lambda: conf.getboolean("core",
"multi_team"), init=False)
+ _bundle_name_to_team_name: dict[str, str | None] =
attrs.field(factory=dict, init=False)
_processors: dict[DagFileInfo, DagFileProcessorProcess] =
attrs.field(factory=dict, init=False)
@@ -307,6 +309,19 @@ class DagFileProcessorManager(LoggingMixin):
# So that we ignore the debug dump signal, making it easier to send
signal.signal(signal.SIGUSR2, signal.SIG_IGN)
+ def _get_team_names(self, bundle_names: Collection[str]) -> dict[str, str
| None]:
+ if not self._multi_team or not bundle_names:
+ return {}
+ missing = [name for name in bundle_names if name not in
self._bundle_name_to_team_name]
+ if missing:
+ queried = DagBundleModel.get_team_names(missing)
+ for name in missing:
+ self._bundle_name_to_team_name[name] = queried.get(name)
+ return {name: self._bundle_name_to_team_name.get(name) for name in
bundle_names}
+
+ def _get_team_name(self, bundle_name: str) -> str | None:
+ return self._get_team_names({bundle_name}).get(bundle_name)
+
def _exit_gracefully(self, signum, frame):
"""Clean up DAG file processors to avoid leaving orphan processes."""
self.log.info("Exiting gracefully upon receiving signal %s", signum)
@@ -720,11 +735,7 @@ class DagFileProcessorManager(LoggingMixin):
)
self._callback_to_execute[file_info].append(request)
self._add_files_to_queue([file_info], mode="front")
- team_name = (
- DagBundleModel.get_team_name(file_info.bundle_name)
- if conf.getboolean("core", "multi_team")
- else None
- )
+ team_name = self._get_team_name(file_info.bundle_name)
stats.incr("dag_processing.other_callback_count",
tags=prune_dict({"team_name": team_name}))
@provide_session
@@ -891,6 +902,8 @@ class DagFileProcessorManager(LoggingMixin):
)
if any_refreshed:
+ # Bundle-to-team assignments can only change on bundle refresh, so
clear the cache.
+ self._bundle_name_to_team_name = {}
self.handle_removed_files(known_files=known_files)
self._resort_file_queue()
self._add_new_files_to_queue(known_files=known_files)
@@ -1026,13 +1039,7 @@ class DagFileProcessorManager(LoggingMixin):
utcnow = timezone.utcnow()
now = time.monotonic()
- if conf.getboolean("core", "multi_team"):
- bundle_names = {bundle_name for bundle_name in known_files}
- bundle_to_team = {
- bundle_name: DagBundleModel.get_team_name(bundle_name) for
bundle_name in bundle_names
- }
- else:
- bundle_to_team = {}
+ bundle_to_team = self._get_team_names({bundle_name for bundle_name in
known_files})
for files in known_files.values():
for file in files:
@@ -1154,13 +1161,7 @@ class DagFileProcessorManager(LoggingMixin):
"""Stop processors that are working on deleted files."""
present_keys = {file.presence_key for file in present}
- if conf.getboolean("core", "multi_team"):
- bundle_names = {file.bundle_name for file in self._processors}
- bundle_to_team = {
- bundle_name: DagBundleModel.get_team_name(bundle_name) for
bundle_name in bundle_names
- }
- else:
- bundle_to_team = {}
+ bundle_to_team = self._get_team_names({file.bundle_name for file in
self._processors})
for file in list(self._processors.keys()):
if file.presence_key not in present_keys:
@@ -1215,9 +1216,7 @@ class DagFileProcessorManager(LoggingMixin):
run_duration = time.monotonic() - proc.start_time
finish_time = timezone.utcnow()
- team_name = (
- DagBundleModel.get_team_name(file.bundle_name) if
conf.getboolean("core", "multi_team") else None
- )
+ team_name = self._get_team_name(file.bundle_name)
next_stat = process_parse_results(
run_duration=run_duration,
finish_time=finish_time,
@@ -1394,13 +1393,7 @@ class DagFileProcessorManager(LoggingMixin):
def _start_new_processes(self):
"""Start more processors if we have enough slots and files to
process."""
- if conf.getboolean("core", "multi_team"):
- bundle_names = {file.bundle_name for file in self._file_queue}
- bundle_to_team = {
- bundle_name: DagBundleModel.get_team_name(bundle_name) for
bundle_name in bundle_names
- }
- else:
- bundle_to_team = {}
+ bundle_to_team = self._get_team_names({file.bundle_name for file in
self._file_queue})
while self._parallelism > len(self._processors) and self._file_queue:
file, _ = self._file_queue.popitem(last=False)
@@ -1584,13 +1577,7 @@ class DagFileProcessorManager(LoggingMixin):
now = time.monotonic()
processors_to_remove = []
- if conf.getboolean("core", "multi_team"):
- bundle_names = {file.bundle_name for file in self._processors}
- bundle_to_team = {
- bundle_name: DagBundleModel.get_team_name(bundle_name) for
bundle_name in bundle_names
- }
- else:
- bundle_to_team = {}
+ bundle_to_team = self._get_team_names({file.bundle_name for file in
self._processors})
for file, processor in self._processors.items():
duration = now - processor.start_time
@@ -1674,13 +1661,7 @@ class DagFileProcessorManager(LoggingMixin):
def terminate(self):
"""Stop all running processors."""
- if conf.getboolean("core", "multi_team"):
- bundle_names = {file.bundle_name for file in self._processors}
- bundle_to_team = {
- bundle_name: DagBundleModel.get_team_name(bundle_name) for
bundle_name in bundle_names
- }
- else:
- bundle_to_team = {}
+ bundle_to_team = self._get_team_names({file.bundle_name for file in
self._processors})
for file, processor in self._processors.items():
stats.decr(
diff --git a/airflow-core/src/airflow/models/dagbundle.py
b/airflow-core/src/airflow/models/dagbundle.py
index abd723a4e67..f18c548a019 100644
--- a/airflow-core/src/airflow/models/dagbundle.py
+++ b/airflow-core/src/airflow/models/dagbundle.py
@@ -30,6 +30,8 @@ from airflow.utils.session import NEW_SESSION, provide_session
from airflow.utils.sqlalchemy import UtcDateTime
if TYPE_CHECKING:
+ from collections.abc import Collection
+
from sqlalchemy.orm import Session
@@ -121,3 +123,18 @@ class DagBundleModel(Base, LoggingMixin):
return session.scalar(
select(Team.name).join(DagBundleModel.teams).where(DagBundleModel.name ==
bundle_name)
)
+
+ @staticmethod
+ @provide_session
+ def get_team_names(
+ bundle_names: Collection[str], *, session: Session = NEW_SESSION
+ ) -> dict[str, str | None]:
+ """Return a mapping of bundle name to team name (None for bundles not
mapped to a team)."""
+ if not bundle_names:
+ return {}
+ rows = session.execute(
+ select(DagBundleModel.name, Team.name)
+ .join(DagBundleModel.teams)
+ .where(DagBundleModel.name.in_(bundle_names))
+ ).all()
+ return {name: team_name for name, team_name in rows}
diff --git a/airflow-core/tests/unit/dag_processing/test_manager.py
b/airflow-core/tests/unit/dag_processing/test_manager.py
index 6d10e5c1a06..4b55b6b2f1e 100644
--- a/airflow-core/tests/unit/dag_processing/test_manager.py
+++ b/airflow-core/tests/unit/dag_processing/test_manager.py
@@ -2732,6 +2732,15 @@ class TestDagFileProcessorManager:
mock_update.assert_called_once_with("mock_bundle",
last_refreshed=mock.ANY, version=None)
assert manager._bundle_versions["mock_bundle"] is None
+ def test_refresh_dag_bundles_clears_team_name_cache(self):
+ manager = DagFileProcessorManager(max_runs=1)
+ manager._bundle_name_to_team_name = {"stale_bundle": "old_team"}
+ bundle = self._make_refresh_bundle(supports_versioning=False)
+
+ self._refresh_with_mocked_state(manager, bundle,
BundleState(last_refreshed=None, version=None))
+
+ assert manager._bundle_name_to_team_name == {}
+
def
test_refresh_dag_bundles_versioned_version_changed_calls_update_bundle_state(self):
"""Versioned bundle with new version: update_bundle_state called with
the new version."""
manager = DagFileProcessorManager(max_runs=1)
@@ -2948,7 +2957,9 @@ class TestMultiTeamMetrics:
return ret
@conf_vars({("core", "multi_team"): "true"})
- @mock.patch("airflow.dag_processing.manager.DagBundleModel.get_team_name",
return_value="team_alpha")
+ @mock.patch(
+ "airflow.dag_processing.manager.DagBundleModel.get_team_names",
return_value={"testing": "team_alpha"}
+ )
@mock.patch("airflow.dag_processing.manager.stats.gauge")
def test_log_file_processing_stats_includes_team_name(self, mock_gauge,
mock_get_team_name):
manager = DagFileProcessorManager(max_runs=1)
@@ -2975,7 +2986,7 @@ class TestMultiTeamMetrics:
)
@conf_vars({("core", "multi_team"): "false"})
- @mock.patch("airflow.dag_processing.manager.DagBundleModel.get_team_name")
+ @mock.patch("airflow.dag_processing.manager.DagBundleModel.get_team_names")
@mock.patch("airflow.dag_processing.manager.stats.gauge")
def test_log_file_processing_stats_omits_team_name_when_not_multi_team(
self, mock_gauge, mock_get_team_name
@@ -3004,7 +3015,9 @@ class TestMultiTeamMetrics:
)
@conf_vars({("core", "multi_team"): "true"})
- @mock.patch("airflow.dag_processing.manager.DagBundleModel.get_team_name",
return_value="team_alpha")
+ @mock.patch(
+ "airflow.dag_processing.manager.DagBundleModel.get_team_names",
return_value={"testing": "team_alpha"}
+ )
@mock.patch("airflow.dag_processing.manager.stats.incr")
def test_start_new_processes_includes_team_name(self, mock_incr,
mock_get_team_name):
manager = DagFileProcessorManager(max_runs=1)
@@ -3024,7 +3037,9 @@ class TestMultiTeamMetrics:
)
@conf_vars({("core", "multi_team"): "true"})
- @mock.patch("airflow.dag_processing.manager.DagBundleModel.get_team_name",
return_value="team_alpha")
+ @mock.patch(
+ "airflow.dag_processing.manager.DagBundleModel.get_team_names",
return_value={"testing": "team_alpha"}
+ )
@mock.patch("airflow.dag_processing.manager.stats.incr")
@mock.patch("airflow.dag_processing.manager.stats.decr")
def test_kill_timed_out_processors_includes_team_name(self, mock_decr,
mock_incr, mock_get_team_name):
@@ -3049,7 +3064,9 @@ class TestMultiTeamMetrics:
)
@conf_vars({("core", "multi_team"): "true"})
- @mock.patch("airflow.dag_processing.manager.DagBundleModel.get_team_name",
return_value="team_alpha")
+ @mock.patch(
+ "airflow.dag_processing.manager.DagBundleModel.get_team_names",
return_value={"testing": "team_alpha"}
+ )
@mock.patch("airflow.dag_processing.manager.stats.timing")
def test_process_parse_results_includes_team_name(self, mock_timing,
mock_get_team_name):
from airflow.dag_processing.manager import process_parse_results
@@ -3097,13 +3114,13 @@ class TestMultiTeamMetrics:
("multi_team", "team_name", "expected_tags"),
[
pytest.param(
- "true",
+ True,
"team_alpha",
{"file_path": "dag_file.py", "action": "stop", "team_name":
"team_alpha"},
id="with_team",
),
pytest.param(
- "false",
+ False,
None,
{"file_path": "dag_file.py", "action": "stop"},
id="without_team",
@@ -3115,6 +3132,7 @@ class TestMultiTeamMetrics:
self, mock_decr, multi_team, team_name, expected_tags
):
manager = DagFileProcessorManager(max_runs=1)
+ manager._multi_team = multi_team
dag_file = DagFileInfo(
bundle_name="testing", rel_path=Path("dag_file.py"),
bundle_path=TEST_DAGS_FOLDER
)
@@ -3122,8 +3140,10 @@ class TestMultiTeamMetrics:
manager._processors = {dag_file: processor}
with (
- conf_vars({("core", "multi_team"): multi_team}),
-
mock.patch("airflow.dag_processing.manager.DagBundleModel.get_team_name",
return_value=team_name),
+ mock.patch(
+ "airflow.dag_processing.manager.DagBundleModel.get_team_names",
+ return_value={"testing": team_name},
+ ),
mock.patch.object(type(processor), "kill"),
):
# Empty "present" set means the file is orphaned, so its processor
is stopped.
@@ -3135,13 +3155,13 @@ class TestMultiTeamMetrics:
("multi_team", "team_name", "expected_tags"),
[
pytest.param(
- "true",
+ True,
"team_alpha",
{"file_path": "dag_file.py", "action": "terminate",
"team_name": "team_alpha"},
id="with_team",
),
pytest.param(
- "false",
+ False,
None,
{"file_path": "dag_file.py", "action": "terminate"},
id="without_team",
@@ -3151,6 +3171,7 @@ class TestMultiTeamMetrics:
@mock.patch("airflow.dag_processing.manager.stats.decr")
def test_terminate_includes_team_name(self, mock_decr, multi_team,
team_name, expected_tags):
manager = DagFileProcessorManager(max_runs=1)
+ manager._multi_team = multi_team
dag_file = DagFileInfo(
bundle_name="testing", rel_path=Path("dag_file.py"),
bundle_path=TEST_DAGS_FOLDER
)
@@ -3158,8 +3179,10 @@ class TestMultiTeamMetrics:
manager._processors = {dag_file: processor}
with (
- conf_vars({("core", "multi_team"): multi_team}),
-
mock.patch("airflow.dag_processing.manager.DagBundleModel.get_team_name",
return_value=team_name),
+ mock.patch(
+ "airflow.dag_processing.manager.DagBundleModel.get_team_names",
+ return_value={"testing": team_name},
+ ),
mock.patch.object(type(processor), "kill"),
):
manager.terminate()
@@ -3195,22 +3218,54 @@ class TestMultiTeamMetrics:
@pytest.mark.parametrize(
("multi_team", "team_name", "expected_tags"),
[
- pytest.param("true", "team_alpha", {"team_name": "team_alpha"},
id="with_team"),
- pytest.param("false", None, {}, id="without_team"),
+ pytest.param(True, "team_alpha", {"team_name": "team_alpha"},
id="with_team"),
+ pytest.param(False, None, {}, id="without_team"),
],
)
@mock.patch("airflow.dag_processing.manager.stats.incr")
def test_add_callback_to_queue_includes_team_name(self, mock_incr,
multi_team, team_name, expected_tags):
manager = DagFileProcessorManager(max_runs=1)
+ manager._multi_team = multi_team
request = MagicMock(filepath="test_dag.py", bundle_name="testing",
bundle_version=None)
bundle = MagicMock(path=TEST_DAGS_FOLDER)
with (
- conf_vars({("core", "multi_team"): multi_team}),
mock.patch.object(manager, "prepare_callback_bundle",
return_value=bundle),
mock.patch.object(manager, "_add_files_to_queue"),
-
mock.patch("airflow.dag_processing.manager.DagBundleModel.get_team_name",
return_value=team_name),
+ mock.patch(
+ "airflow.dag_processing.manager.DagBundleModel.get_team_names",
+ return_value={"testing": team_name},
+ ),
):
manager._add_callback_to_queue(request)
mock_incr.assert_called_once_with("dag_processing.other_callback_count",
tags=expected_tags)
+
+ @mock.patch(
+ "airflow.dag_processing.manager.DagBundleModel.get_team_names",
+ return_value={"bundle_a": "team_alpha"},
+ )
+ def test_get_team_name_caches_lookup(self, mock_get_team_names):
+ manager = DagFileProcessorManager(max_runs=1)
+ manager._multi_team = True
+
+ assert manager._get_team_name("bundle_a") == "team_alpha"
+ assert manager._get_team_name("bundle_a") == "team_alpha"
+ assert manager._get_team_name("bundle_a") == "team_alpha"
+
+ mock_get_team_names.assert_called_once()
+
+ @mock.patch(
+ "airflow.dag_processing.manager.DagBundleModel.get_team_names",
+ return_value={"bundle_a": "team_alpha", "bundle_b": "team_alpha"},
+ )
+ def test_get_team_names_batches_and_caches(self, mock_get_team_names):
+ manager = DagFileProcessorManager(max_runs=1)
+ manager._multi_team = True
+
+ manager._get_team_names({"bundle_a", "bundle_b"})
+ manager._get_team_names({"bundle_a", "bundle_b"})
+
+ # Two bundles resolved in a single batched query; the repeat call is
served from cache.
+ mock_get_team_names.assert_called_once()
+ assert manager._bundle_name_to_team_name == {"bundle_a": "team_alpha",
"bundle_b": "team_alpha"}
diff --git a/airflow-core/tests/unit/models/test_dagbundle.py
b/airflow-core/tests/unit/models/test_dagbundle.py
index 84223e34a9b..0c5c5def9d7 100644
--- a/airflow-core/tests/unit/models/test_dagbundle.py
+++ b/airflow-core/tests/unit/models/test_dagbundle.py
@@ -58,3 +58,20 @@ class TestDagBundleModel:
def test_get_team_name_unknown_bundle(self, session: Session):
assert DagBundleModel.get_team_name("does_not_exist", session=session)
is None
+
+ def test_get_team_names(self, testing_team: Team, session: Session):
+ mapped = DagBundleModel(name="mapped_bundle")
+ mapped.teams.append(testing_team)
+ unmapped = DagBundleModel(name="unmapped_bundle")
+ session.add_all([mapped, unmapped])
+ session.flush()
+
+ result = DagBundleModel.get_team_names(
+ ["mapped_bundle", "unmapped_bundle", "does_not_exist"],
session=session
+ )
+
+ # Only bundles actually mapped to a team are returned; callers treat
absent keys as None.
+ assert result == {"mapped_bundle": "testing"}
+
+ def test_get_team_names_empty(self, session: Session):
+ assert DagBundleModel.get_team_names([], session=session) == {}
