Skip to content

Use attribute based access control in KeycloakAuthManager #51486

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
vincbeck merged 1 commit into from
Jun 9, 2025
+179 −175
Merged

Use attribute based access control in KeycloakAuthManager #51486

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
105 changes: 45 additions & 60 deletions providers/keycloak/src/airflow/providers/keycloak/auth_manager/keycloak_auth_manager.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -25,11 +25,12 @@
from fastapi import FastAPI

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.api_fastapi.auth.managers.base_auth_manager import BaseAuthManager
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
from airflow.utils.helpers import prune_dict

if TYPE_CHECKING:
from airflow.api_fastapi.auth.managers.base_auth_manager import ResourceMethod
Expand All @@ -49,6 +50,8 @@

log = logging.getLogger(__name__)

RESOURCE_ID_ATTRIBUTE_NAME = "resource_id"


class KeycloakAuthManager(BaseAuthManager[KeycloakAuthManagerUser]):
"""
Expand Down Expand Up @@ -86,15 +89,10 @@ def is_authorized_configuration(
) -> bool:
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,
)
method=method,
resource_type=KeycloakResource.CONFIGURATION,
user=user,
resource_id=config_section,
)

def is_authorized_connection(
Expand All @@ -105,43 +103,42 @@ def is_authorized_connection(
details: ConnectionDetails | None = None,
) -> bool:
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
)
return self._is_authorized(
method=method, resource_type=KeycloakResource.CONNECTION, user=user, resource_id=connection_id
)

def is_authorized_dag(
self,
*,
method: ResourceMethod,
user: T,
user: KeycloakAuthManagerUser,
access_entity: DagAccessEntity | None = None,
details: DagDetails | None = None,
) -> bool:
return True
dag_id = details.id if details else None
access_entity_str = access_entity.value if access_entity else None
return self._is_authorized(
method=method,
resource_type=KeycloakResource.DAG,
user=user,
resource_id=dag_id,
attributes={"dag_entity": access_entity_str},
)

def is_authorized_backfill(
self, *, method: ResourceMethod, user: KeycloakAuthManagerUser, details: BackfillDetails | None = None
) -> bool:
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
)
return 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:
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
)
return self._is_authorized(
method=method, resource_type=KeycloakResource.ASSET, user=user, resource_id=asset_id
)

def is_authorized_asset_alias(
Expand All @@ -152,42 +149,31 @@ def is_authorized_asset_alias(
details: AssetAliasDetails | None = None,
) -> bool:
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,
)
return 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:
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
)
return 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:
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
)
return 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 self._is_authorized(
method="GET", resource_type=KeycloakResource.VIEW, user=user
) or self._is_authorized(
method="GET",
resource_type=KeycloakResource.VIEW,
user=user,
Expand All @@ -198,8 +184,6 @@ def is_authorized_custom_view(
self, *, method: ResourceMethod | str, resource_name: str, user: KeycloakAuthManagerUser
) -> bool:
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
)

Expand Down Expand Up @@ -230,19 +214,19 @@ def _is_authorized(
resource_type: KeycloakResource,
user: KeycloakAuthManagerUser,
resource_id: str | None = None,
attributes: dict[str, str | None] | 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}"
)
context_attributes = prune_dict(attributes or {})
if resource_id:
context_attributes[RESOURCE_ID_ATTRIBUTE_NAME] = resource_id

resp = requests.post(
self._get_token_url(server_url, realm),
data=self._get_payload(client_id, permission),
data=self._get_payload(client_id, f"{resource_type.value}#{method}", context_attributes),
headers=self._get_headers(user.access_token),
)

Expand All @@ -252,9 +236,6 @@ def _is_authorized(
return False
if resp.status_code == 400:
error = json.loads(resp.text)
if error.get("error") == "invalid_resource":
log.debug(error["error_description"])
return False
raise AirflowException(
f"Request not recognized by Keycloak. {error.get('error')}. {error.get('error_description')}"
)
Expand All @@ -265,12 +246,16 @@ 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 {
def _get_payload(client_id: str, permission: str, attributes: dict[str, str] | None = None):
payload: dict[str, Any] = {
"grant_type": "urn:ietf:params:oauth:grant-type:uma-ticket",
"audience": client_id,
"permission": permission,
}
if attributes:
payload["context"] = {"attributes": attributes}

return payload

@staticmethod
def _get_headers(access_token):
Expand Down
Loading