(airflow) branch main updated: Parallelize per-dag auth checks in KeycloakAuthManager (#69107)

[el...@apache.org]

This is an automated email from the ASF dual-hosted git repository.

eladkal 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 3ead1c8039f Parallelize per-dag auth checks in KeycloakAuthManager (#69107)
3ead1c8039f is described below

commit 3ead1c8039fb44debbdd76ac32bc7f0bccd24ea4
Author: Andrew Chang <69...@users.noreply.github.com>
AuthorDate: Tue Jun 30 06:59:22 2026 +0000

    Parallelize per-dag auth checks in KeycloakAuthManager (#69107)
    
    * Parallelize per-dag auth checks in KeycloakAuthManager
    
    * Fix team_name args not found
    
    * Add comment of thread pool size decision
    
    * Add test for checking thread pool size
---
 .../keycloak/auth_manager/keycloak_auth_manager.py | 30 ++++++++++++++++----
 .../auth_manager/test_keycloak_auth_manager.py     | 32 ++++++++++++++++++++++
 2 files changed, 57 insertions(+), 5 deletions(-)

diff --git a/providers/keycloak/src/airflow/providers/keycloak/auth_manager/keycloak_auth_manager.py b/providers/keycloak/src/airflow/providers/keycloak/auth_manager/keycloak_auth_manager.py
index 64f12603ce6..bef6632fe3d 100644
--- a/providers/keycloak/src/airflow/providers/keycloak/auth_manager/keycloak_auth_manager.py
+++ b/providers/keycloak/src/airflow/providers/keycloak/auth_manager/keycloak_auth_manager.py
@@ -22,6 +22,7 @@ import logging
 import time
 import warnings
 from base64 import urlsafe_b64decode
+from concurrent.futures import ThreadPoolExecutor
 from typing import TYPE_CHECKING, Any
 from urllib.parse import urljoin
 
@@ -34,6 +35,7 @@ from urllib3.util import Retry
 
 from airflow.api_fastapi.app import AUTH_MANAGER_FASTAPI_APP_PREFIX
 from airflow.api_fastapi.auth.managers.base_auth_manager import BaseAuthManager
+from airflow.api_fastapi.auth.managers.models.resource_details import DagDetails
 from airflow.exceptions import AirflowProviderDeprecationWarning
 
 try:
@@ -68,7 +70,6 @@ if TYPE_CHECKING:
         ConfigurationDetails,
         ConnectionDetails,
         DagAccessEntity,
-        DagDetails,
         PoolDetails,
         TeamDetails,
         VariableDetails,
@@ -458,10 +459,29 @@ class KeycloakAuthManager(BaseAuthManager[KeycloakAuthManagerUser]):
         cache_key = (user.get_id(), method, team_name, frozenset(dag_ids))
 
         def query_keycloak() -> set[str]:
-            kwargs: dict = dict(dag_ids=dag_ids, user=user, method=method)
-            if team_name is not None:
-                kwargs["team_name"] = team_name
-            return super(KeycloakAuthManager, self).filter_authorized_dag_ids(**kwargs)
+            if not dag_ids:
+                return set()
+            # Cap workers at the HTTP connection pool size: each is_authorized_dag() call
+            # goes through the shared requests.Session, so extra threads would just block
+            # waiting for a free connection in urllib3's pool.
+            max_workers = min(
+                len(dag_ids), conf.getint(CONF_SECTION_NAME, CONF_REQUESTS_POOL_SIZE_KEY, fallback=10)
+            )
+
+            def check(dag_id: str) -> tuple[str, bool]:
+                details_kwargs: dict[str, Any] = {"id": dag_id}
+                if team_name is not None:
+                    details_kwargs["team_name"] = team_name
+                return dag_id, self.is_authorized_dag(
+                    method=method,
+                    user=user,
+                    details=DagDetails(**details_kwargs),
+                )
+
+            with ThreadPoolExecutor(max_workers=max_workers) as executor:
+                results = executor.map(check, dag_ids)
+
+            return {dag_id for dag_id, authorized in results if authorized}
 
         return single_flight(cache_key, query_keycloak)
 
diff --git a/providers/keycloak/tests/unit/keycloak/auth_manager/test_keycloak_auth_manager.py b/providers/keycloak/tests/unit/keycloak/auth_manager/test_keycloak_auth_manager.py
index 48791d11d38..ca4bda4f9e5 100644
--- a/providers/keycloak/tests/unit/keycloak/auth_manager/test_keycloak_auth_manager.py
+++ b/providers/keycloak/tests/unit/keycloak/auth_manager/test_keycloak_auth_manager.py
@@ -18,6 +18,7 @@ from __future__ import annotations
 
 import base64
 import json
+from concurrent.futures import ThreadPoolExecutor
 from contextlib import ExitStack
 from unittest.mock import Mock, patch
 
@@ -57,6 +58,7 @@ from airflow.providers.keycloak.auth_manager.constants import (
     CONF_CLIENT_ID_KEY,
     CONF_CLIENT_SECRET_KEY,
     CONF_REALM_KEY,
+    CONF_REQUESTS_POOL_SIZE_KEY,
     CONF_SECTION_NAME,
     CONF_SERVER_URL_KEY,
 )
@@ -1147,3 +1149,33 @@ class TestKeycloakAuthManager:
         assert result2 == dag_ids
         # is_authorized_dag should only be called for the first invocation (2 dag_ids × 1 call)
         assert mock_is_authorized.call_count == 2
+
+    @pytest.mark.parametrize(
+        ("dag_count", "pool_size", "expected_max_workers"),
+        [
+            pytest.param(5, 10, 5, id="dag-count-smaller-than-pool"),
+            pytest.param(50, 10, 10, id="dag-count-larger-than-pool"),
+            pytest.param(10, 10, 10, id="dag-count-same-as-pool"),
+        ],
+    )
+    @patch.object(KeycloakAuthManager, "is_authorized_dag", return_value=True)
+    @patch(
+        "airflow.providers.keycloak.auth_manager.keycloak_auth_manager.ThreadPoolExecutor",
+        wraps=ThreadPoolExecutor,
+    )
+    def test_filter_authorized_dag_ids_caps_max_workers(
+        self,
+        mock_executor,
+        _mock_is_authorized,
+        auth_manager,
+        user,
+        dag_count,
+        pool_size,
+        expected_max_workers,
+    ):
+        dag_ids = {f"dag-{index}" for index in range(dag_count)}
+
+        with conf_vars({(CONF_SECTION_NAME, CONF_REQUESTS_POOL_SIZE_KEY): str(pool_size)}):
+            auth_manager.filter_authorized_dag_ids(dag_ids=dag_ids, user=user)
+
+        mock_executor.assert_called_once_with(max_workers=expected_max_workers)