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 fffbb5d07b8 Implement authorization methods of `KeycloakAuthManager`
(#51403)
fffbb5d07b8 is described below
commit fffbb5d07b8b23eb5aebf2fd8747cbcbfd5e7fdc
Author: Vincent <[email protected]>
AuthorDate: Thu Jun 5 11:53:52 2025 -0400
Implement authorization methods of `KeycloakAuthManager` (#51403)
---
.../keycloak/auth_manager/keycloak_auth_manager.py | 145 +++++++++++--
.../providers/keycloak/auth_manager/resources.py | 35 +++
.../auth_manager/test_keycloak_auth_manager.py | 235 ++++++++++++++++++++-
.../unit/keycloak/auth_manager/test_resources.py | 24 +++
4 files changed, 416 insertions(+), 23 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 dee92d23304..9857e8d1fd2 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
@@ -16,16 +16,17 @@
# under the License.
from __future__ import annotations
-from functools import cached_property
from typing import TYPE_CHECKING, Any
from urllib.parse import urljoin
+import requests
from fastapi import FastAPI
-from starlette.middleware.sessions import SessionMiddleware
from airflow.api_fastapi.app import AUTH_MANAGER_FASTAPI_APP_PREFIX
from airflow.api_fastapi.auth.managers.base_auth_manager import
BaseAuthManager, T
from airflow.configuration import conf
+from airflow.exceptions import AirflowException
+from airflow.providers.keycloak.auth_manager.resources import KeycloakResource
from airflow.providers.keycloak.auth_manager.user import
KeycloakAuthManagerUser
if TYPE_CHECKING:
@@ -52,10 +53,6 @@ class
KeycloakAuthManager(BaseAuthManager[KeycloakAuthManagerUser]):
Leverages Keycloak to perform authentication and authorization in Airflow.
"""
- @cached_property
- def api_server_endpoint(self) -> str:
- return conf.get("api", "base_url", fallback="/")
-
def deserialize_user(self, token: dict[str, Any]) ->
KeycloakAuthManagerUser:
return KeycloakAuthManagerUser(
user_id=token.pop("user_id"),
@@ -73,7 +70,8 @@ class
KeycloakAuthManager(BaseAuthManager[KeycloakAuthManagerUser]):
}
def get_url_login(self, **kwargs) -> str:
- return urljoin(self.api_server_endpoint,
f"{AUTH_MANAGER_FASTAPI_APP_PREFIX}/login")
+ base_url = conf.get("api", "base_url", fallback="/")
+ return urljoin(base_url, f"{AUTH_MANAGER_FASTAPI_APP_PREFIX}/login")
def is_authorized_configuration(
self,
@@ -82,7 +80,18 @@ class
KeycloakAuthManager(BaseAuthManager[KeycloakAuthManagerUser]):
user: KeycloakAuthManagerUser,
details: ConfigurationDetails | None = None,
) -> bool:
- return True
+ config_section = details.section if details else None
+ return self._is_authorized(
+ method=method, resource_type=KeycloakResource.CONFIGURATION,
user=user
+ ) or (
+ config_section is not None
+ and self._is_authorized(
+ method=method,
+ resource_type=KeycloakResource.CONFIGURATION,
+ user=user,
+ resource_id=config_section,
+ )
+ )
def is_authorized_connection(
self,
@@ -91,7 +100,13 @@ class
KeycloakAuthManager(BaseAuthManager[KeycloakAuthManagerUser]):
user: KeycloakAuthManagerUser,
details: ConnectionDetails | None = None,
) -> bool:
- return True
+ connection_id = details.conn_id if details else None
+ return self._is_authorized(method=method,
resource_type=KeycloakResource.CONNECTION, user=user) or (
+ connection_id is not None
+ and self._is_authorized(
+ method=method, resource_type=KeycloakResource.CONNECTION,
user=user, resource_id=connection_id
+ )
+ )
def is_authorized_dag(
self,
@@ -106,12 +121,24 @@ class
KeycloakAuthManager(BaseAuthManager[KeycloakAuthManagerUser]):
def is_authorized_backfill(
self, *, method: ResourceMethod, user: KeycloakAuthManagerUser,
details: BackfillDetails | None = None
) -> bool:
- return True
+ backfill_id = str(details.id) if details else None
+ return self._is_authorized(method=method,
resource_type=KeycloakResource.BACKFILL, user=user) or (
+ backfill_id is not None
+ and self._is_authorized(
+ method=method, resource_type=KeycloakResource.BACKFILL,
user=user, resource_id=backfill_id
+ )
+ )
def is_authorized_asset(
self, *, method: ResourceMethod, user: KeycloakAuthManagerUser,
details: AssetDetails | None = None
) -> bool:
- return True
+ asset_id = details.id if details else None
+ return self._is_authorized(method=method,
resource_type=KeycloakResource.ASSET, user=user) or (
+ asset_id is not None
+ and self._is_authorized(
+ method=method, resource_type=KeycloakResource.ASSET,
user=user, resource_id=asset_id
+ )
+ )
def is_authorized_asset_alias(
self,
@@ -120,25 +147,57 @@ class
KeycloakAuthManager(BaseAuthManager[KeycloakAuthManagerUser]):
user: KeycloakAuthManagerUser,
details: AssetAliasDetails | None = None,
) -> bool:
- return True
+ asset_alias_id = details.id if details else None
+ return self._is_authorized(method=method,
resource_type=KeycloakResource.ASSET_ALIAS, user=user) or (
+ asset_alias_id is not None
+ and self._is_authorized(
+ method=method,
+ resource_type=KeycloakResource.ASSET_ALIAS,
+ user=user,
+ resource_id=asset_alias_id,
+ )
+ )
def is_authorized_variable(
self, *, method: ResourceMethod, user: KeycloakAuthManagerUser,
details: VariableDetails | None = None
) -> bool:
- return True
+ variable_key = details.key if details else None
+ return self._is_authorized(method=method,
resource_type=KeycloakResource.VARIABLE, user=user) or (
+ variable_key is not None
+ and self._is_authorized(
+ method=method, resource_type=KeycloakResource.VARIABLE,
user=user, resource_id=variable_key
+ )
+ )
def is_authorized_pool(
self, *, method: ResourceMethod, user: KeycloakAuthManagerUser,
details: PoolDetails | None = None
) -> bool:
- return True
+ pool_name = details.name if details else None
+ return self._is_authorized(method=method,
resource_type=KeycloakResource.POOL, user=user) or (
+ pool_name is not None
+ and self._is_authorized(
+ method=method, resource_type=KeycloakResource.POOL, user=user,
resource_id=pool_name
+ )
+ )
def is_authorized_view(self, *, access_view: AccessView, user:
KeycloakAuthManagerUser) -> bool:
- return True
+ return self._is_authorized(
+ method="GET", resource_type=KeycloakResource.VIEW, user=user
+ ) or self._is_authorized(
+ method="GET",
+ resource_type=KeycloakResource.VIEW,
+ user=user,
+ resource_id=access_view.value,
+ )
def is_authorized_custom_view(
self, *, method: ResourceMethod | str, resource_name: str, user:
KeycloakAuthManagerUser
) -> bool:
- return True
+ return self._is_authorized(
+ method=method, resource_type=KeycloakResource.CUSTOM, user=user
+ ) or self._is_authorized(
+ method=method, resource_type=KeycloakResource.CUSTOM, user=user,
resource_id=resource_name
+ )
def filter_authorized_menu_items(
self, menu_items: list[MenuItem], *, user: KeycloakAuthManagerUser
@@ -156,12 +215,54 @@ class
KeycloakAuthManager(BaseAuthManager[KeycloakAuthManagerUser]):
"This sub application provides login routes."
),
)
- # authlib requires ``SessionMiddleware``
- app.add_middleware(
- SessionMiddleware,
- secret_key=conf.get("api_auth", "jwt_secret"),
- https_only=False,
- )
app.include_router(login_router)
return app
+
+ def _is_authorized(
+ self,
+ *,
+ method: ResourceMethod | str,
+ resource_type: KeycloakResource,
+ user: KeycloakAuthManagerUser,
+ resource_id: str | None = None,
+ ) -> bool:
+ client_id = conf.get("keycloak_auth_manager", "client_id")
+ realm = conf.get("keycloak_auth_manager", "realm")
+ server_url = conf.get("keycloak_auth_manager", "server_url")
+
+ permission = (
+ f"{resource_type.value}:{resource_id}#{method}"
+ if resource_id
+ else f"{resource_type.value}#{method}"
+ )
+ resp = requests.post(
+ self._get_token_url(server_url, realm),
+ data=self._get_payload(client_id, permission),
+ headers=self._get_headers(user.access_token),
+ )
+
+ if resp.status_code == 200:
+ return True
+ if resp.status_code == 403:
+ return False
+ raise AirflowException(f"Unexpected error: {resp.status_code} -
{resp.text}")
+
+ @staticmethod
+ def _get_token_url(server_url, realm):
+ return f"{server_url}/realms/{realm}/protocol/openid-connect/token"
+
+ @staticmethod
+ def _get_payload(client_id, permission):
+ return {
+ "grant_type": "urn:ietf:params:oauth:grant-type:uma-ticket",
+ "audience": client_id,
+ "permission": permission,
+ }
+
+ @staticmethod
+ def _get_headers(access_token):
+ return {
+ "Authorization": f"Bearer {access_token}",
+ "Content-Type": "application/x-www-form-urlencoded",
+ }
diff --git
a/providers/keycloak/src/airflow/providers/keycloak/auth_manager/resources.py
b/providers/keycloak/src/airflow/providers/keycloak/auth_manager/resources.py
new file mode 100644
index 00000000000..a2d1b1a5bbf
--- /dev/null
+++
b/providers/keycloak/src/airflow/providers/keycloak/auth_manager/resources.py
@@ -0,0 +1,35 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from __future__ import annotations
+
+from enum import Enum
+
+
+class KeycloakResource(Enum):
+ """Enum of Keycloak resources."""
+
+ ASSET = "Asset"
+ ASSET_ALIAS = "AssetAlias"
+ BACKFILL = "Backfill"
+ CONFIGURATION = "Configuration"
+ CONNECTION = "Connection"
+ CUSTOM = "Custom"
+ DAG = "Dag"
+ MENU = "Menu"
+ POOL = "Pool"
+ VARIABLE = "Variable"
+ VIEW = "View"
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 ab63ef3d596..f9de61de79a 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
@@ -16,16 +16,45 @@
# under the License.
from __future__ import annotations
+from unittest.mock import Mock, call, patch
+
import pytest
from airflow.api_fastapi.app import AUTH_MANAGER_FASTAPI_APP_PREFIX
+from airflow.api_fastapi.auth.managers.models.resource_details import (
+ AccessView,
+ AssetAliasDetails,
+ AssetDetails,
+ BackfillDetails,
+ ConfigurationDetails,
+ ConnectionDetails,
+ PoolDetails,
+ VariableDetails,
+)
+from airflow.exceptions import AirflowException
from airflow.providers.keycloak.auth_manager.keycloak_auth_manager import
KeycloakAuthManager
from airflow.providers.keycloak.auth_manager.user import
KeycloakAuthManagerUser
+from tests_common.test_utils.config import conf_vars
+
@pytest.fixture
def auth_manager():
- return KeycloakAuthManager()
+ with conf_vars(
+ {
+ ("keycloak_auth_manager", "client_id"): "client_id",
+ ("keycloak_auth_manager", "realm"): "realm",
+ ("keycloak_auth_manager", "server_url"): "server_url",
+ }
+ ):
+ yield KeycloakAuthManager()
+
+
[email protected]
+def user():
+ user = Mock()
+ user.access_token = "access_token"
+ return user
class TestKeycloakAuthManager:
@@ -59,3 +88,207 @@ class TestKeycloakAuthManager:
def test_get_url_login(self, auth_manager):
result = auth_manager.get_url_login()
assert result == f"{AUTH_MANAGER_FASTAPI_APP_PREFIX}/login"
+
+ @pytest.mark.parametrize(
+ "function, method, details, permission_type, permission_resource",
+ [
+ [
+ "is_authorized_configuration",
+ "GET",
+ ConfigurationDetails(section="test"),
+ "Configuration#GET",
+ "Configuration:test#GET",
+ ],
+ ["is_authorized_configuration", "GET", None, "Configuration#GET",
None],
+ [
+ "is_authorized_configuration",
+ "PUT",
+ ConfigurationDetails(section="test"),
+ "Configuration#PUT",
+ "Configuration:test#PUT",
+ ],
+ [
+ "is_authorized_connection",
+ "DELETE",
+ ConnectionDetails(conn_id="test"),
+ "Connection#DELETE",
+ "Connection:test#DELETE",
+ ],
+ ["is_authorized_connection", "GET", None, "Connection#GET", None],
+ ["is_authorized_backfill", "POST", BackfillDetails(id=1),
"Backfill#POST", "Backfill:1#POST"],
+ ["is_authorized_backfill", "GET", None, "Backfill#GET", None],
+ ["is_authorized_asset", "GET", AssetDetails(id="test"),
"Asset#GET", "Asset:test#GET"],
+ ["is_authorized_asset", "GET", None, "Asset#GET", None],
+ [
+ "is_authorized_asset_alias",
+ "GET",
+ AssetAliasDetails(id="test"),
+ "AssetAlias#GET",
+ "AssetAlias:test#GET",
+ ],
+ ["is_authorized_asset_alias", "GET", None, "AssetAlias#GET", None],
+ [
+ "is_authorized_variable",
+ "PUT",
+ VariableDetails(key="test"),
+ "Variable#PUT",
+ "Variable:test#PUT",
+ ],
+ ["is_authorized_variable", "GET", None, "Variable#GET", None],
+ ["is_authorized_pool", "POST", PoolDetails(name="test"),
"Pool#POST", "Pool:test#POST"],
+ ["is_authorized_pool", "GET", None, "Pool#GET", None],
+ ],
+ )
+ @pytest.mark.parametrize(
+ "status_code_type, status_code_resource, expected_one_call,
expected_two_calls",
+ [
+ [200, 200, True, True],
+ [200, 403, True, True],
+ [403, 200, False, True],
+ [403, 403, False, False],
+ ],
+ )
+
@patch("airflow.providers.keycloak.auth_manager.keycloak_auth_manager.requests")
+ def test_is_authorized(
+ self,
+ mock_requests,
+ function,
+ method,
+ details,
+ permission_type,
+ permission_resource,
+ status_code_type,
+ status_code_resource,
+ expected_one_call,
+ expected_two_calls,
+ auth_manager,
+ user,
+ ):
+ expected_num_calls = 1 if status_code_type == 200 or not
permission_resource else 2
+ resp1 = Mock()
+ resp1.status_code = status_code_type
+ resp2 = Mock()
+ resp2.status_code = status_code_resource
+
+ mock_requests.post.side_effect = [resp1, resp2]
+
+ result = getattr(auth_manager, function)(method=method, user=user,
details=details)
+
+ token_url = auth_manager._get_token_url("server_url", "realm")
+ payload = auth_manager._get_payload("client_id", permission_type)
+ headers = auth_manager._get_headers("access_token")
+ expected_calls = [call(token_url, data=payload, headers=headers)]
+ if expected_num_calls == 2:
+ expected_calls.append(
+ call(token_url, data={**payload, "permission":
permission_resource}, headers=headers)
+ )
+ mock_requests.post.assert_has_calls(expected_calls)
+ expected = expected_two_calls if expected_num_calls == 2 else
expected_one_call
+ assert result == expected
+
+ @pytest.mark.parametrize(
+ "function",
+ [
+ "is_authorized_configuration",
+ "is_authorized_connection",
+ "is_authorized_backfill",
+ "is_authorized_asset",
+ "is_authorized_asset_alias",
+ "is_authorized_variable",
+ "is_authorized_pool",
+ ],
+ )
+
@patch("airflow.providers.keycloak.auth_manager.keycloak_auth_manager.requests")
+ def test_is_authorized_failure(self, mock_requests, function,
auth_manager, user):
+ resp = Mock()
+ resp.status_code = 500
+ mock_requests.post.return_value = resp
+
+ with pytest.raises(AirflowException):
+ getattr(auth_manager, function)(method="GET", user=user)
+
+ @pytest.mark.parametrize(
+ "status_code_type, status_code_resource, expected_one_call,
expected_two_calls",
+ [
+ [200, 200, True, True],
+ [200, 403, True, True],
+ [403, 200, False, True],
+ [403, 403, False, False],
+ ],
+ )
+
@patch("airflow.providers.keycloak.auth_manager.keycloak_auth_manager.requests")
+ def test_is_authorized_view(
+ self,
+ mock_requests,
+ status_code_type,
+ status_code_resource,
+ expected_one_call,
+ expected_two_calls,
+ auth_manager,
+ user,
+ ):
+ expected_num_calls = 1 if status_code_type == 200 else 2
+ resp1 = Mock()
+ resp1.status_code = status_code_type
+ resp2 = Mock()
+ resp2.status_code = status_code_resource
+
+ mock_requests.post.side_effect = [resp1, resp2]
+
+ result =
auth_manager.is_authorized_view(access_view=AccessView.CLUSTER_ACTIVITY,
user=user)
+
+ token_url = auth_manager._get_token_url("server_url", "realm")
+ payload = auth_manager._get_payload("client_id", "View#GET")
+ headers = auth_manager._get_headers("access_token")
+ expected_calls = [call(token_url, data=payload, headers=headers)]
+ if expected_num_calls == 2:
+ expected_calls.append(
+ call(token_url, data={**payload, "permission":
"View:CLUSTER_ACTIVITY#GET"}, headers=headers)
+ )
+
+ mock_requests.post.assert_has_calls(expected_calls)
+ expected = expected_two_calls if expected_num_calls == 2 else
expected_one_call
+ assert result == expected
+
+ @pytest.mark.parametrize(
+ "status_code_type, status_code_resource, expected_one_call,
expected_two_calls",
+ [
+ [200, 200, True, True],
+ [200, 403, True, True],
+ [403, 200, False, True],
+ [403, 403, False, False],
+ ],
+ )
+
@patch("airflow.providers.keycloak.auth_manager.keycloak_auth_manager.requests")
+ def test_is_authorized_custom_view(
+ self,
+ mock_requests,
+ status_code_type,
+ status_code_resource,
+ expected_one_call,
+ expected_two_calls,
+ auth_manager,
+ user,
+ ):
+ expected_num_calls = 1 if status_code_type == 200 else 2
+ resp1 = Mock()
+ resp1.status_code = status_code_type
+ resp2 = Mock()
+ resp2.status_code = status_code_resource
+
+ mock_requests.post.side_effect = [resp1, resp2]
+
+ result = auth_manager.is_authorized_custom_view(method="GET",
resource_name="test", user=user)
+
+ token_url = auth_manager._get_token_url("server_url", "realm")
+ payload = auth_manager._get_payload("client_id", "Custom#GET")
+ headers = auth_manager._get_headers("access_token")
+ expected_calls = [call(token_url, data=payload, headers=headers)]
+ if expected_num_calls == 2:
+ expected_calls.append(
+ call(token_url, data={**payload, "permission":
"Custom:test#GET"}, headers=headers)
+ )
+
+ mock_requests.post.assert_has_calls(expected_calls)
+ expected = expected_two_calls if expected_num_calls == 2 else
expected_one_call
+ assert result == expected
diff --git
a/providers/keycloak/tests/unit/keycloak/auth_manager/test_resources.py
b/providers/keycloak/tests/unit/keycloak/auth_manager/test_resources.py
new file mode 100644
index 00000000000..a05a3182b4d
--- /dev/null
+++ b/providers/keycloak/tests/unit/keycloak/auth_manager/test_resources.py
@@ -0,0 +1,24 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from __future__ import annotations
+
+from airflow.providers.keycloak.auth_manager.resources import KeycloakResource
+
+
+class TestKeycloakResource:
+ def test_asset(self):
+ assert KeycloakResource.ASSET.value == "Asset"
