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 f2ccc0fe7db Implement bulk authorization methods in
KeycloakAuthManager (#70647)
f2ccc0fe7db is described below
commit f2ccc0fe7dbe2739e8dc2e0a6b1ea2079fd40a5b
Author: abhishekmauryaKsolves <[email protected]>
AuthorDate: Fri Jul 31 19:54:11 2026 +0530
Implement bulk authorization methods in KeycloakAuthManager (#70647)
---
.../keycloak/auth_manager/keycloak_auth_manager.py | 209 ++++++++++++++++++++-
.../auth_manager/test_keycloak_auth_manager.py | 120 ++++++++++++
2 files changed, 325 insertions(+), 4 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 9755d574279..a680be488f7 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 collections.abc import Sequence
from concurrent.futures import ThreadPoolExecutor
from typing import TYPE_CHECKING, Annotated, Any
from urllib.parse import urljoin
@@ -35,7 +36,12 @@ 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.api_fastapi.auth.managers.models.resource_details import (
+ ConnectionDetails,
+ DagDetails,
+ PoolDetails,
+ VariableDetails,
+)
from airflow.exceptions import AirflowProviderDeprecationWarning
try:
@@ -62,17 +68,20 @@ from airflow.utils.helpers import prune_dict
if TYPE_CHECKING:
from airflow.api_fastapi.auth.managers.base_auth_manager import
ResourceMethod
+ from airflow.api_fastapi.auth.managers.models.batch_apis import (
+ IsAuthorizedConnectionRequest,
+ IsAuthorizedDagRequest,
+ IsAuthorizedPoolRequest,
+ IsAuthorizedVariableRequest,
+ )
from airflow.api_fastapi.auth.managers.models.resource_details import (
AccessView,
AssetAliasDetails,
AssetDetails,
BackfillDetails,
ConfigurationDetails,
- ConnectionDetails,
DagAccessEntity,
- PoolDetails,
TeamDetails,
- VariableDetails,
)
from airflow.cli.cli_config import CLICommand
@@ -492,6 +501,198 @@ class
KeycloakAuthManager(BaseAuthManager[KeycloakAuthManagerUser]):
return single_flight(cache_key, query_keycloak)
+ def batch_is_authorized_connection(
+ self,
+ requests: Sequence[IsAuthorizedConnectionRequest],
+ *,
+ user: KeycloakAuthManagerUser,
+ ) -> bool:
+ if not requests:
+ return True
+ max_workers = min(
+ len(requests), conf.getint(CONF_SECTION_NAME,
CONF_REQUESTS_POOL_SIZE_KEY, fallback=10)
+ )
+
+ def check(request: IsAuthorizedConnectionRequest) -> bool:
+ return self.is_authorized_connection(
+ method=request["method"],
+ details=request.get("details"),
+ user=user,
+ )
+
+ with ThreadPoolExecutor(max_workers=max_workers) as executor:
+ results = executor.map(check, requests)
+ return all(results)
+
+ def batch_is_authorized_dag(
+ self,
+ requests: Sequence[IsAuthorizedDagRequest],
+ *,
+ user: KeycloakAuthManagerUser,
+ ) -> bool:
+ if not requests:
+ return True
+ max_workers = min(
+ len(requests), conf.getint(CONF_SECTION_NAME,
CONF_REQUESTS_POOL_SIZE_KEY, fallback=10)
+ )
+
+ def check(request: IsAuthorizedDagRequest) -> bool:
+ return self.is_authorized_dag(
+ method=request["method"],
+ access_entity=request.get("access_entity"),
+ details=request.get("details"),
+ user=user,
+ )
+
+ with ThreadPoolExecutor(max_workers=max_workers) as executor:
+ results = executor.map(check, requests)
+ return all(results)
+
+ def batch_is_authorized_pool(
+ self,
+ requests: Sequence[IsAuthorizedPoolRequest],
+ *,
+ user: KeycloakAuthManagerUser,
+ ) -> bool:
+ if not requests:
+ return True
+ max_workers = min(
+ len(requests), conf.getint(CONF_SECTION_NAME,
CONF_REQUESTS_POOL_SIZE_KEY, fallback=10)
+ )
+
+ def check(request: IsAuthorizedPoolRequest) -> bool:
+ return self.is_authorized_pool(
+ method=request["method"],
+ details=request.get("details"),
+ user=user,
+ )
+
+ with ThreadPoolExecutor(max_workers=max_workers) as executor:
+ results = executor.map(check, requests)
+ return all(results)
+
+ def batch_is_authorized_variable(
+ self,
+ requests: Sequence[IsAuthorizedVariableRequest],
+ *,
+ user: KeycloakAuthManagerUser,
+ ) -> bool:
+ if not requests:
+ return True
+ max_workers = min(
+ len(requests), conf.getint(CONF_SECTION_NAME,
CONF_REQUESTS_POOL_SIZE_KEY, fallback=10)
+ )
+
+ def check(request: IsAuthorizedVariableRequest) -> bool:
+ return self.is_authorized_variable(
+ method=request["method"],
+ details=request.get("details"),
+ user=user,
+ )
+
+ with ThreadPoolExecutor(max_workers=max_workers) as executor:
+ results = executor.map(check, requests)
+ return all(results)
+
+ def filter_authorized_connections(
+ self,
+ *,
+ conn_ids: set[str],
+ user: KeycloakAuthManagerUser,
+ method: ResourceMethod = "GET",
+ team_name: str | None = None,
+ ) -> set[str]:
+ cache_key = (user.get_id(), method, team_name, frozenset(conn_ids))
+
+ def query_keycloak() -> set[str]:
+ if not conn_ids:
+ return set()
+ max_workers = min(
+ len(conn_ids), conf.getint(CONF_SECTION_NAME,
CONF_REQUESTS_POOL_SIZE_KEY, fallback=10)
+ )
+
+ def check(conn_id: str) -> tuple[str, bool]:
+ details_kwargs: dict[str, Any] = {"conn_id": conn_id}
+ if team_name is not None:
+ details_kwargs["team_name"] = team_name
+ return conn_id, self.is_authorized_connection(
+ method=method,
+ user=user,
+ details=ConnectionDetails(**details_kwargs),
+ )
+
+ with ThreadPoolExecutor(max_workers=max_workers) as executor:
+ results = executor.map(check, conn_ids)
+ return {conn_id for conn_id, authorized in results if authorized}
+
+ return single_flight(cache_key, query_keycloak)
+
+ def filter_authorized_pools(
+ self,
+ *,
+ pool_names: set[str],
+ user: KeycloakAuthManagerUser,
+ method: ResourceMethod = "GET",
+ team_name: str | None = None,
+ ) -> set[str]:
+ cache_key = (user.get_id(), method, team_name, frozenset(pool_names))
+
+ def query_keycloak() -> set[str]:
+ if not pool_names:
+ return set()
+ max_workers = min(
+ len(pool_names), conf.getint(CONF_SECTION_NAME,
CONF_REQUESTS_POOL_SIZE_KEY, fallback=10)
+ )
+
+ def check(pool_name: str) -> tuple[str, bool]:
+ details_kwargs: dict[str, Any] = {"name": pool_name}
+ if team_name is not None:
+ details_kwargs["team_name"] = team_name
+ return pool_name, self.is_authorized_pool(
+ method=method,
+ user=user,
+ details=PoolDetails(**details_kwargs),
+ )
+
+ with ThreadPoolExecutor(max_workers=max_workers) as executor:
+ results = executor.map(check, pool_names)
+ return {pool_name for pool_name, authorized in results if
authorized}
+
+ return single_flight(cache_key, query_keycloak)
+
+ def filter_authorized_variables(
+ self,
+ *,
+ variable_keys: set[str],
+ user: KeycloakAuthManagerUser,
+ method: ResourceMethod = "GET",
+ team_name: str | None = None,
+ ) -> set[str]:
+ cache_key = (user.get_id(), method, team_name,
frozenset(variable_keys))
+
+ def query_keycloak() -> set[str]:
+ if not variable_keys:
+ return set()
+ max_workers = min(
+ len(variable_keys), conf.getint(CONF_SECTION_NAME,
CONF_REQUESTS_POOL_SIZE_KEY, fallback=10)
+ )
+
+ def check(variable_key: str) -> tuple[str, bool]:
+ details_kwargs: dict[str, Any] = {"key": variable_key}
+ if team_name is not None:
+ details_kwargs["team_name"] = team_name
+ return variable_key, self.is_authorized_variable(
+ method=method,
+ user=user,
+ details=VariableDetails(**details_kwargs),
+ )
+
+ with ThreadPoolExecutor(max_workers=max_workers) as executor:
+ results = executor.map(check, variable_keys)
+ return {variable_key for variable_key, authorized in results if
authorized}
+
+ return single_flight(cache_key, query_keycloak)
+
def _is_batch_authorized(
self,
*,
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 58698292bc9..d3f075d2a46 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
@@ -1162,6 +1162,126 @@ class TestKeycloakAuthManager:
# is_authorized_dag should only be called for the first invocation (2
dag_ids × 1 call)
assert mock_is_authorized.call_count == 2
+ @patch.object(
+ KeycloakAuthManager,
+ "is_authorized_connection",
+ side_effect=lambda *, details, **kw: {"conn_0": True, "conn_1": False,
"conn_2": True}[
+ details.conn_id
+ ],
+ )
+ def test_filter_authorized_connections(self, mock_is_authorized,
auth_manager, user):
+ result = auth_manager.filter_authorized_connections(
+ conn_ids={"conn_0", "conn_1", "conn_2"}, user=user, method="GET"
+ )
+
+ assert result == {"conn_0", "conn_2"}
+ assert mock_is_authorized.call_count == 3
+
+ def test_filter_authorized_connections_empty(self, auth_manager, user):
+ result = auth_manager.filter_authorized_connections(conn_ids=set(),
user=user, method="GET")
+ assert result == set()
+
+ @patch.object(KeycloakAuthManager, "is_authorized_connection",
return_value=False)
+ def test_filter_authorized_connections_all_denied(self,
mock_is_authorized, auth_manager, user):
+ result = auth_manager.filter_authorized_connections(
+ conn_ids={"conn_0", "conn_1"}, user=user, method="GET"
+ )
+
+ assert result == set()
+ assert mock_is_authorized.call_count == 2
+
+ @patch.object(KeycloakAuthManager, "is_authorized_connection",
return_value=True)
+ def test_filter_authorized_connections_cache_hit(self, mock_is_authorized,
auth_manager, user):
+ """Second call with same args should return cached result without
hitting Keycloak."""
+ conn_ids = {"conn_0", "conn_1"}
+
+ result1 =
auth_manager.filter_authorized_connections(conn_ids=conn_ids, user=user,
method="GET")
+ result2 =
auth_manager.filter_authorized_connections(conn_ids=conn_ids, user=user,
method="GET")
+
+ assert result1 == conn_ids
+ assert result2 == conn_ids
+ assert mock_is_authorized.call_count == 2
+
+ @patch.object(
+ KeycloakAuthManager,
+ "is_authorized_pool",
+ side_effect=lambda *, details, **kw: {"pool_0": True, "pool_1": False,
"pool_2": True}[details.name],
+ )
+ def test_filter_authorized_pools(self, mock_is_authorized, auth_manager,
user):
+ result = auth_manager.filter_authorized_pools(
+ pool_names={"pool_0", "pool_1", "pool_2"}, user=user, method="GET"
+ )
+
+ assert result == {"pool_0", "pool_2"}
+ assert mock_is_authorized.call_count == 3
+
+ def test_filter_authorized_pools_empty(self, auth_manager, user):
+ result = auth_manager.filter_authorized_pools(pool_names=set(),
user=user, method="GET")
+ assert result == set()
+
+ @patch.object(KeycloakAuthManager, "is_authorized_pool",
return_value=False)
+ def test_filter_authorized_pools_all_denied(self, mock_is_authorized,
auth_manager, user):
+ result = auth_manager.filter_authorized_pools(
+ pool_names={"pool_0", "pool_1"}, user=user, method="GET"
+ )
+
+ assert result == set()
+ assert mock_is_authorized.call_count == 2
+
+ @patch.object(KeycloakAuthManager, "is_authorized_pool", return_value=True)
+ def test_filter_authorized_pools_cache_hit(self, mock_is_authorized,
auth_manager, user):
+ """Second call with same args should return cached result without
hitting Keycloak."""
+ pool_names = {"pool_0", "pool_1"}
+
+ result1 = auth_manager.filter_authorized_pools(pool_names=pool_names,
user=user, method="GET")
+ result2 = auth_manager.filter_authorized_pools(pool_names=pool_names,
user=user, method="GET")
+
+ assert result1 == pool_names
+ assert result2 == pool_names
+ assert mock_is_authorized.call_count == 2
+
+ @patch.object(
+ KeycloakAuthManager,
+ "is_authorized_variable",
+ side_effect=lambda *, details, **kw: {"var_0": True, "var_1": False,
"var_2": True}[details.key],
+ )
+ def test_filter_authorized_variables(self, mock_is_authorized,
auth_manager, user):
+ result = auth_manager.filter_authorized_variables(
+ variable_keys={"var_0", "var_1", "var_2"}, user=user, method="GET"
+ )
+
+ assert result == {"var_0", "var_2"}
+ assert mock_is_authorized.call_count == 3
+
+ def test_filter_authorized_variables_empty(self, auth_manager, user):
+ result = auth_manager.filter_authorized_variables(variable_keys=set(),
user=user, method="GET")
+ assert result == set()
+
+ @patch.object(KeycloakAuthManager, "is_authorized_variable",
return_value=False)
+ def test_filter_authorized_variables_all_denied(self, mock_is_authorized,
auth_manager, user):
+ result = auth_manager.filter_authorized_variables(
+ variable_keys={"var_0", "var_1"}, user=user, method="GET"
+ )
+
+ assert result == set()
+ assert mock_is_authorized.call_count == 2
+
+ @patch.object(KeycloakAuthManager, "is_authorized_variable",
return_value=True)
+ def test_filter_authorized_variables_cache_hit(self, mock_is_authorized,
auth_manager, user):
+ """Second call with same args should return cached result without
hitting Keycloak."""
+ variable_keys = {"var_0", "var_1"}
+
+ result1 = auth_manager.filter_authorized_variables(
+ variable_keys=variable_keys, user=user, method="GET"
+ )
+ result2 = auth_manager.filter_authorized_variables(
+ variable_keys=variable_keys, user=user, method="GET"
+ )
+
+ assert result1 == variable_keys
+ assert result2 == variable_keys
+ assert mock_is_authorized.call_count == 2
+
@pytest.mark.parametrize(
("dag_count", "pool_size", "expected_max_workers"),
[
