Parallelize per-dag auth checks in KeycloakAuthManager

providers/keycloak/src/airflow/providers/keycloak/auth_manager/keycloak_auth_manager.py

@@ -22,6 +22,7 @@
 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 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 @@
         ConfigurationDetails,
         ConnectionDetails,
         DagAccessEntity,
-        DagDetails,
         PoolDetails,
         TeamDetails,
         VariableDetails,
@@ -458,10 +459,29 @@ def filter_authorized_dag_ids(
         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)
 

providers/keycloak/tests/unit/keycloak/auth_manager/test_keycloak_auth_manager.py

@@ -18,6 +18,7 @@
 
 import base64
 import json
+from concurrent.futures import ThreadPoolExecutor
 from contextlib import ExitStack
 from unittest.mock import Mock, patch
 
@@ -57,6 +58,7 @@
     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 @@ def test_filter_authorized_dag_ids_cache_hit(self, mock_is_authorized, auth_mana
         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)