This is an automated email from the ASF dual-hosted git repository.
pierrejeambrun 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 40b6796fac5 Avoid N+1 team-name queries in bulk Dag run authorization
(#68286)
40b6796fac5 is described below
commit 40b6796fac5de56df6025bd06c3b9c5e37db107e
Author: Yuseok Jo <[email protected]>
AuthorDate: Thu Jun 11 00:30:57 2026 +0900
Avoid N+1 team-name queries in bulk Dag run authorization (#68286)
* Avoid N+1 team-name queries in bulk Dag run authorization
* Update airflow-core/src/airflow/api_fastapi/core_api/security.py
Co-authored-by: Henry Chen <[email protected]>
* Update airflow-core/tests/unit/api_fastapi/core_api/test_security.py
Co-authored-by: Henry Chen <[email protected]>
* Assert resolved team name in bulk dag-run access requests test
---------
Co-authored-by: Henry Chen <[email protected]>
---
.../src/airflow/api_fastapi/core_api/security.py | 9 +++--
.../unit/api_fastapi/core_api/test_security.py | 41 ++++++++++++++++++++++
2 files changed, 47 insertions(+), 3 deletions(-)
diff --git a/airflow-core/src/airflow/api_fastapi/core_api/security.py
b/airflow-core/src/airflow/api_fastapi/core_api/security.py
index 154929af064..8a009231f12 100644
--- a/airflow-core/src/airflow/api_fastapi/core_api/security.py
+++ b/airflow-core/src/airflow/api_fastapi/core_api/security.py
@@ -738,10 +738,13 @@ def _build_dag_run_access_requests(
``entity_methods`` is a list of ``(dag_id, method)`` pairs with
unresolvable
entries (no dag_id or the ``~`` wildcard) already filtered out by the
caller.
- The team for each dag is resolved once and shared across that dag's
requests.
+ Teams for all Dags are resolved in a single batched query and shared
across each
+ Dag's requests.
"""
- resolved_dag_ids = {dag_id for dag_id, _ in entity_methods}
- dag_id_to_team = {dag_id: DagModel.get_team_name(dag_id) for dag_id in
resolved_dag_ids}
+ if not entity_methods:
+ return []
+ resolved_dag_ids = list({dag_id for dag_id, _ in entity_methods})
+ dag_id_to_team = DagModel.get_dag_id_to_team_name_mapping(resolved_dag_ids)
return [
{
"method": method,
diff --git a/airflow-core/tests/unit/api_fastapi/core_api/test_security.py
b/airflow-core/tests/unit/api_fastapi/core_api/test_security.py
index c6f684854af..e916bc19e5d 100644
--- a/airflow-core/tests/unit/api_fastapi/core_api/test_security.py
+++ b/airflow-core/tests/unit/api_fastapi/core_api/test_security.py
@@ -23,6 +23,7 @@ import pytest
from fastapi import HTTPException
from jwt import ExpiredSignatureError, InvalidTokenError
+from airflow import settings
from airflow.api_fastapi.app import create_app
from airflow.api_fastapi.auth.managers.base_auth_manager import
COOKIE_NAME_JWT_TOKEN
from airflow.api_fastapi.auth.managers.models.resource_details import (
@@ -38,6 +39,7 @@ from airflow.api_fastapi.core_api.datamodels.connections
import ConnectionBody
from airflow.api_fastapi.core_api.datamodels.pools import PoolBody
from airflow.api_fastapi.core_api.datamodels.variables import VariableBody
from airflow.api_fastapi.core_api.security import (
+ _build_dag_run_access_requests,
get_user,
is_safe_url,
requires_access_backfill,
@@ -53,9 +55,48 @@ from airflow.api_fastapi.core_api.security import (
)
from airflow.models import Connection, Pool, Variable
from airflow.models.dag import DagModel
+from airflow.models.dagbundle import DagBundleModel
from airflow.models.team import Team
+from tests_common.test_utils.asserts import assert_queries_count
from tests_common.test_utils.config import conf_vars
+from tests_common.test_utils.db import clear_db_dag_bundles, clear_db_dags
+
+
[email protected]_test
+def test_build_dag_run_access_requests_batches_team_lookup():
+ """The bulk dag-run team resolution must be a single query regardless of
Dag count (no N+1)."""
+ entity_methods = [(f"dag_{i}", "GET") for i in range(25)]
+ with assert_queries_count(1):
+ requests = _build_dag_run_access_requests(entity_methods)
+ assert len(requests) == 25
+
+
[email protected]_test
+def test_build_dag_run_access_requests_empty_skips_query():
+ with assert_queries_count(0):
+ assert _build_dag_run_access_requests([]) == []
+
+
[email protected]_test
+def test_build_dag_run_access_requests_includes_team_name(testing_team):
+ """A team-owned Dag must surface its team name in the generated access
request."""
+ session = settings.Session()
+ bundle = DagBundleModel(name="team-owned-bundle")
+ bundle.teams.append(testing_team)
+ session.add(bundle)
+ session.flush()
+ session.add(DagModel(dag_id="team_owned_dag",
bundle_name="team-owned-bundle", is_stale=False))
+ session.flush()
+ try:
+ requests = _build_dag_run_access_requests([("team_owned_dag", "GET")])
+ finally:
+ clear_db_dags()
+ clear_db_dag_bundles()
+
+ assert len(requests) == 1
+ assert requests[0]["details"].id == "team_owned_dag"
+ assert requests[0]["details"].team_name == "testing"
@pytest.mark.asyncio
