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 197907e1e46 KeycloakAuthManager - Create CLI to create resources in
KeyCloak (#51691)
197907e1e46 is described below
commit 197907e1e46ec2434695b05d12661e0a90cc0a2a
Author: Vincent <[email protected]>
AuthorDate: Fri Jun 13 15:21:29 2025 -0400
KeycloakAuthManager - Create CLI to create resources in KeyCloak (#51691)
---
.pre-commit-config.yaml | 1 +
.../keycloak/auth_manager/cli/__init__.py | 16 +
.../keycloak/auth_manager/cli/commands.py | 224 ++++++++++++++
.../keycloak/auth_manager/cli/definition.py | 77 +++++
.../providers/keycloak/auth_manager/constants.py | 25 ++
.../keycloak/auth_manager/keycloak_auth_manager.py | 25 +-
.../keycloak/auth_manager/routes/login.py | 15 +-
.../unit/keycloak/auth_manager/cli/__init__.py | 16 +
.../keycloak/auth_manager/cli/test_commands.py | 331 +++++++++++++++++++++
.../keycloak/auth_manager/cli/test_definition.py | 24 ++
.../keycloak/auth_manager/routes/test_login.py | 14 +-
.../unit/keycloak/auth_manager/test_constants.py | 42 +++
.../auth_manager/test_keycloak_auth_manager.py | 15 +-
13 files changed, 811 insertions(+), 14 deletions(-)
diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml
index c6d9954346b..2f22b1d9e41 100644
--- a/.pre-commit-config.yaml
+++ b/.pre-commit-config.yaml
@@ -648,6 +648,7 @@ repos:
^providers/google/src/airflow/providers/google/cloud/operators/cloud_build\.py$|
^providers/google/src/airflow/providers/google/cloud/operators/dataproc\.py$|
^providers/google/src/airflow/providers/google/cloud/operators/mlengine\.py$|
+
^providers/keycloak/src/airflow/providers/keycloak/auth_manager/cli/definition.py|
^providers/microsoft/azure/src/airflow/providers/microsoft/azure/hooks/cosmos\.py$|
^providers/microsoft/winrm/src/airflow/providers/microsoft/winrm/hooks/winrm\.py$|
^airflow-core/docs/.*commits\.rst$|
diff --git
a/providers/keycloak/src/airflow/providers/keycloak/auth_manager/cli/__init__.py
b/providers/keycloak/src/airflow/providers/keycloak/auth_manager/cli/__init__.py
new file mode 100644
index 00000000000..13a83393a91
--- /dev/null
+++
b/providers/keycloak/src/airflow/providers/keycloak/auth_manager/cli/__init__.py
@@ -0,0 +1,16 @@
+# 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.
diff --git
a/providers/keycloak/src/airflow/providers/keycloak/auth_manager/cli/commands.py
b/providers/keycloak/src/airflow/providers/keycloak/auth_manager/cli/commands.py
new file mode 100644
index 00000000000..1dd3ed822d7
--- /dev/null
+++
b/providers/keycloak/src/airflow/providers/keycloak/auth_manager/cli/commands.py
@@ -0,0 +1,224 @@
+# 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
+
+import json
+import logging
+from typing import get_args
+
+from keycloak import KeycloakAdmin, KeycloakError
+
+from airflow.api_fastapi.auth.managers.base_auth_manager import ResourceMethod
+from airflow.configuration import conf
+from airflow.providers.keycloak.auth_manager.constants import (
+ CONF_CLIENT_ID_KEY,
+ CONF_REALM_KEY,
+ CONF_SECTION_NAME,
+ CONF_SERVER_URL_KEY,
+)
+from airflow.providers.keycloak.auth_manager.resources import KeycloakResource
+from airflow.utils import cli as cli_utils
+from airflow.utils.providers_configuration_loader import
providers_configuration_loaded
+
+log = logging.getLogger(__name__)
+
+
+@cli_utils.action_cli
+@providers_configuration_loaded
+def create_scopes_command(args):
+ """Create Keycloak auth manager scopes in Keycloak."""
+ client = _get_client(args)
+ client_uuid = _get_client_uuid(args)
+
+ _create_scopes(client, client_uuid)
+
+
+@cli_utils.action_cli
+@providers_configuration_loaded
+def create_resources_command(args):
+ """Create Keycloak auth manager resources in Keycloak."""
+ client = _get_client(args)
+ client_uuid = _get_client_uuid(args)
+
+ _create_resources(client, client_uuid)
+
+
+@cli_utils.action_cli
+@providers_configuration_loaded
+def create_permissions_command(args):
+ """Create Keycloak auth manager permissions in Keycloak."""
+ client = _get_client(args)
+ client_uuid = _get_client_uuid(args)
+
+ _create_permissions(client, client_uuid)
+
+
+@cli_utils.action_cli
+@providers_configuration_loaded
+def create_all_command(args):
+ """Create Keycloak auth manager scopes in Keycloak."""
+ client = _get_client(args)
+ client_uuid = _get_client_uuid(args)
+
+ _create_scopes(client, client_uuid)
+ _create_resources(client, client_uuid)
+ _create_permissions(client, client_uuid)
+
+
+def _get_client(args):
+ server_url = conf.get(CONF_SECTION_NAME, CONF_SERVER_URL_KEY)
+ realm = conf.get(CONF_SECTION_NAME, CONF_REALM_KEY)
+
+ return KeycloakAdmin(
+ server_url=server_url,
+ username=args.username,
+ password=args.password,
+ realm_name=realm,
+ user_realm_name=args.user_realm,
+ client_id=args.client_id,
+ verify=True,
+ )
+
+
+def _get_client_uuid(args):
+ client = _get_client(args)
+ clients = client.get_clients()
+ client_id = conf.get(CONF_SECTION_NAME, CONF_CLIENT_ID_KEY)
+
+ matches = [client for client in clients if client["clientId"] == client_id]
+ if not matches:
+ raise ValueError(f"Client with ID='{client_id}' not found in realm
'{client.realm_name}'")
+
+ return matches[0]["id"]
+
+
+def _create_scopes(client: KeycloakAdmin, client_uuid: str):
+ scopes = [{"name": method} for method in get_args(ResourceMethod)]
+ for scope in scopes:
+ client.create_client_authz_scopes(client_id=client_uuid, payload=scope)
+
+ print("Scopes created successfully.")
+
+
+def _create_resources(client: KeycloakAdmin, client_uuid: str):
+ # Fetch existing scopes
+ all_scopes = client.get_client_authz_scopes(client_uuid)
+ scopes = [
+ {"id": scope["id"], "name": scope["name"]}
+ for scope in all_scopes
+ if scope["name"] in get_args(ResourceMethod)
+ ]
+
+ for resource in KeycloakResource:
+ client.create_client_authz_resource(
+ client_id=client_uuid,
+ payload={
+ "name": resource.value,
+ "scopes": scopes,
+ },
+ skip_exists=True,
+ )
+
+ print("Resources created successfully.")
+
+
+def _create_permissions(client: KeycloakAdmin, client_uuid: str):
+ _create_read_only_permission(client, client_uuid)
+ _create_admin_permission(client, client_uuid)
+ _create_user_permission(client, client_uuid)
+ _create_op_permission(client, client_uuid)
+
+ print("Permissions created successfully.")
+
+
+def _create_read_only_permission(client: KeycloakAdmin, client_uuid: str):
+ all_scopes = client.get_client_authz_scopes(client_uuid)
+ scopes = [scope["id"] for scope in all_scopes if scope["name"] in ["GET",
"MENU"]]
+ payload = {
+ "name": "ReadOnly",
+ "type": "scope",
+ "logic": "POSITIVE",
+ "decisionStrategy": "UNANIMOUS",
+ "scopes": scopes,
+ }
+ _create_permission(client, client_uuid, payload)
+
+
+def _create_admin_permission(client: KeycloakAdmin, client_uuid: str):
+ all_scopes = client.get_client_authz_scopes(client_uuid)
+ scopes = [scope["id"] for scope in all_scopes if scope["name"] in
get_args(ResourceMethod)]
+ payload = {
+ "name": "Admin",
+ "type": "scope",
+ "logic": "POSITIVE",
+ "decisionStrategy": "UNANIMOUS",
+ "scopes": scopes,
+ }
+ _create_permission(client, client_uuid, payload)
+
+
+def _create_user_permission(client: KeycloakAdmin, client_uuid: str):
+ _create_resource_based_permission(
+ client, client_uuid, "User", [KeycloakResource.DAG.value,
KeycloakResource.ASSET.value]
+ )
+
+
+def _create_op_permission(client: KeycloakAdmin, client_uuid: str):
+ _create_resource_based_permission(
+ client,
+ client_uuid,
+ "Op",
+ [
+ KeycloakResource.CONNECTION.value,
+ KeycloakResource.POOL.value,
+ KeycloakResource.VARIABLE.value,
+ KeycloakResource.BACKFILL.value,
+ ],
+ )
+
+
+def _create_permission(client: KeycloakAdmin, client_uuid: str, payload: dict):
+ try:
+ client.create_client_authz_scope_permission(
+ client_id=client_uuid,
+ payload=payload,
+ )
+ except KeycloakError as e:
+ if e.response_body:
+ error = json.loads(e.response_body.decode("utf-8"))
+ if error.get("error_description") == "Conflicting policy":
+ print(f"Policy creation skipped. {error.get('error')}")
+
+
+def _create_resource_based_permission(
+ client: KeycloakAdmin, client_uuid: str, name: str, allowed_resources:
list[str]
+):
+ all_resources = client.get_client_authz_resources(client_uuid)
+ resources = [resource["_id"] for resource in all_resources if
resource["name"] in allowed_resources]
+
+ payload = {
+ "name": name,
+ "type": "scope",
+ "logic": "POSITIVE",
+ "decisionStrategy": "UNANIMOUS",
+ "resources": resources,
+ }
+ client.create_client_authz_resource_based_permission(
+ client_id=client_uuid,
+ payload=payload,
+ skip_exists=True,
+ )
diff --git
a/providers/keycloak/src/airflow/providers/keycloak/auth_manager/cli/definition.py
b/providers/keycloak/src/airflow/providers/keycloak/auth_manager/cli/definition.py
new file mode 100644
index 00000000000..4d06c412ce5
--- /dev/null
+++
b/providers/keycloak/src/airflow/providers/keycloak/auth_manager/cli/definition.py
@@ -0,0 +1,77 @@
+# 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.cli.cli_config import (
+ ActionCommand,
+ Arg,
+ lazy_load_command,
+)
+
+############
+# # ARGS # #
+############
+
+ARG_USERNAME = Arg(
+ ("--username",),
+ help="Username associated to the user used to create resources",
+)
+ARG_PASSWORD = Arg(
+ ("--password",),
+ help="Password associated to the user used to create resources",
+)
+ARG_USER_REALM = Arg(
+ ("--user-realm",), help="Realm name where the user used to create
resources is", default="master"
+)
+ARG_CLIENT_ID = Arg(("--client-id",), help="ID of the client used to create
resources", default="admin-cli")
+
+
+################
+# # COMMANDS # #
+################
+
+KEYCLOAK_AUTH_MANAGER_COMMANDS = (
+ ActionCommand(
+ name="create-scopes",
+ help="Create scopes in Keycloak",
+
func=lazy_load_command("airflow.providers.keycloak.auth_manager.cli.commands.create_scopes_command"),
+ args=(ARG_USERNAME, ARG_PASSWORD, ARG_USER_REALM, ARG_CLIENT_ID),
+ ),
+ ActionCommand(
+ name="create-resources",
+ help="Create resources in Keycloak",
+ func=lazy_load_command(
+
"airflow.providers.keycloak.auth_manager.cli.commands.create_resources_command"
+ ),
+ args=(ARG_USERNAME, ARG_PASSWORD, ARG_USER_REALM, ARG_CLIENT_ID),
+ ),
+ ActionCommand(
+ name="create-permissions",
+ help="Create permissions in Keycloak",
+ func=lazy_load_command(
+
"airflow.providers.keycloak.auth_manager.cli.commands.create_permissions_command"
+ ),
+ args=(ARG_USERNAME, ARG_PASSWORD, ARG_USER_REALM, ARG_CLIENT_ID),
+ ),
+ ActionCommand(
+ name="create-all",
+ help="Create all entities (scopes, resources and permissions) in
Keycloak",
+
func=lazy_load_command("airflow.providers.keycloak.auth_manager.cli.commands.create_all_command"),
+ args=(ARG_USERNAME, ARG_PASSWORD, ARG_USER_REALM, ARG_CLIENT_ID),
+ ),
+)
diff --git
a/providers/keycloak/src/airflow/providers/keycloak/auth_manager/constants.py
b/providers/keycloak/src/airflow/providers/keycloak/auth_manager/constants.py
new file mode 100644
index 00000000000..aa62dd821ce
--- /dev/null
+++
b/providers/keycloak/src/airflow/providers/keycloak/auth_manager/constants.py
@@ -0,0 +1,25 @@
+# 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.
+
+# Configuration keys
+from __future__ import annotations
+
+CONF_SECTION_NAME = "keycloak_auth_manager"
+CONF_CLIENT_ID_KEY = "client_id"
+CONF_CLIENT_SECRET_KEY = "client_secret"
+CONF_REALM_KEY = "realm"
+CONF_SERVER_URL_KEY = "server_url"
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 0d7a415453e..87c2fed8219 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
@@ -26,8 +26,16 @@ 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
+from airflow.cli.cli_config import CLICommand, GroupCommand
from airflow.configuration import conf
from airflow.exceptions import AirflowException
+from airflow.providers.keycloak.auth_manager.cli.definition import
KEYCLOAK_AUTH_MANAGER_COMMANDS
+from airflow.providers.keycloak.auth_manager.constants import (
+ CONF_CLIENT_ID_KEY,
+ CONF_REALM_KEY,
+ CONF_SECTION_NAME,
+ CONF_SERVER_URL_KEY,
+)
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
@@ -207,6 +215,17 @@ class
KeycloakAuthManager(BaseAuthManager[KeycloakAuthManagerUser]):
return app
+ @staticmethod
+ def get_cli_commands() -> list[CLICommand]:
+ """Vends CLI commands to be included in Airflow CLI."""
+ return [
+ GroupCommand(
+ name="keycloak-auth-manager",
+ help="Manage resources used by Keycloak auth manager",
+ subcommands=KEYCLOAK_AUTH_MANAGER_COMMANDS,
+ ),
+ ]
+
def _is_authorized(
self,
*,
@@ -216,9 +235,9 @@ class
KeycloakAuthManager(BaseAuthManager[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")
+ client_id = conf.get(CONF_SECTION_NAME, CONF_CLIENT_ID_KEY)
+ realm = conf.get(CONF_SECTION_NAME, CONF_REALM_KEY)
+ server_url = conf.get(CONF_SECTION_NAME, CONF_SERVER_URL_KEY)
context_attributes = prune_dict(attributes or {})
if resource_id:
diff --git
a/providers/keycloak/src/airflow/providers/keycloak/auth_manager/routes/login.py
b/providers/keycloak/src/airflow/providers/keycloak/auth_manager/routes/login.py
index b8e2abdcba8..ef5b6ffb0e3 100644
---
a/providers/keycloak/src/airflow/providers/keycloak/auth_manager/routes/login.py
+++
b/providers/keycloak/src/airflow/providers/keycloak/auth_manager/routes/login.py
@@ -27,6 +27,13 @@ from airflow.api_fastapi.app import get_auth_manager
from airflow.api_fastapi.auth.managers.base_auth_manager import
COOKIE_NAME_JWT_TOKEN
from airflow.api_fastapi.common.router import AirflowRouter
from airflow.configuration import conf
+from airflow.providers.keycloak.auth_manager.constants import (
+ CONF_CLIENT_ID_KEY,
+ CONF_CLIENT_SECRET_KEY,
+ CONF_REALM_KEY,
+ CONF_SECTION_NAME,
+ CONF_SERVER_URL_KEY,
+)
from airflow.providers.keycloak.auth_manager.user import
KeycloakAuthManagerUser
log = logging.getLogger(__name__)
@@ -73,10 +80,10 @@ def login_callback(request: Request):
def _get_keycloak_client() -> KeycloakOpenID:
- client_id = conf.get("keycloak_auth_manager", "client_id")
- client_secret = conf.get("keycloak_auth_manager", "client_secret")
- realm = conf.get("keycloak_auth_manager", "realm")
- server_url = conf.get("keycloak_auth_manager", "server_url")
+ client_id = conf.get(CONF_SECTION_NAME, CONF_CLIENT_ID_KEY)
+ client_secret = conf.get(CONF_SECTION_NAME, CONF_CLIENT_SECRET_KEY)
+ realm = conf.get(CONF_SECTION_NAME, CONF_REALM_KEY)
+ server_url = conf.get(CONF_SECTION_NAME, CONF_SERVER_URL_KEY)
return KeycloakOpenID(
server_url=server_url,
diff --git
a/providers/keycloak/tests/unit/keycloak/auth_manager/cli/__init__.py
b/providers/keycloak/tests/unit/keycloak/auth_manager/cli/__init__.py
new file mode 100644
index 00000000000..13a83393a91
--- /dev/null
+++ b/providers/keycloak/tests/unit/keycloak/auth_manager/cli/__init__.py
@@ -0,0 +1,16 @@
+# 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.
diff --git
a/providers/keycloak/tests/unit/keycloak/auth_manager/cli/test_commands.py
b/providers/keycloak/tests/unit/keycloak/auth_manager/cli/test_commands.py
new file mode 100644
index 00000000000..81e9a723b23
--- /dev/null
+++ b/providers/keycloak/tests/unit/keycloak/auth_manager/cli/test_commands.py
@@ -0,0 +1,331 @@
+# 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
+
+import importlib
+from typing import get_args
+from unittest.mock import Mock, call, patch
+
+import pytest
+
+from airflow.api_fastapi.auth.managers.base_auth_manager import ResourceMethod
+from airflow.cli import cli_parser
+from airflow.providers.keycloak.auth_manager.cli.commands import (
+ _create_admin_permission,
+ _create_op_permission,
+ _create_read_only_permission,
+ _create_user_permission,
+ create_all_command,
+ create_permissions_command,
+ create_resources_command,
+ create_scopes_command,
+)
+from airflow.providers.keycloak.auth_manager.resources import KeycloakResource
+
+from tests_common.test_utils.config import conf_vars
+
+
[email protected]_test
+class TestCommands:
+ @classmethod
+ def setup_class(cls):
+ with conf_vars(
+ {
+ (
+ "core",
+ "auth_manager",
+ ):
"airflow.providers.keycloak.auth_manager.keycloak_auth_manager.KeycloakAuthManager",
+ }
+ ):
+ importlib.reload(cli_parser)
+ cls.arg_parser = cli_parser.get_parser()
+
+ @patch("airflow.providers.keycloak.auth_manager.cli.commands._get_client")
+ def test_create_scopes(self, mock_get_client):
+ client = Mock()
+ mock_get_client.return_value = client
+
+ client.get_clients.return_value = [
+ {"id": "dummy-id", "clientId": "dummy-client"},
+ {"id": "test-id", "clientId": "test_client_id"},
+ ]
+
+ params = [
+ "keycloak-auth-manager",
+ "create-scopes",
+ "--username",
+ "test",
+ "--password",
+ "test",
+ ]
+ with conf_vars(
+ {
+ ("keycloak_auth_manager", "client_id"): "test_client_id",
+ }
+ ):
+ create_scopes_command(self.arg_parser.parse_args(params))
+
+ client.get_clients.assert_called_once_with()
+ scopes = [{"name": method} for method in get_args(ResourceMethod)]
+ calls = [call(client_id="test-id", payload=scope) for scope in scopes]
+ client.create_client_authz_scopes.assert_has_calls(calls)
+
+ @patch("airflow.providers.keycloak.auth_manager.cli.commands._get_client")
+ def test_create_scopes_with_client_not_found(self, mock_get_client):
+ client = Mock()
+ mock_get_client.return_value = client
+
+ client.get_clients.return_value = [
+ {"id": "dummy-id", "clientId": "dummy-client"},
+ ]
+
+ params = [
+ "keycloak-auth-manager",
+ "create-scopes",
+ "--username",
+ "test",
+ "--password",
+ "test",
+ ]
+ with conf_vars(
+ {
+ ("keycloak_auth_manager", "client_id"): "test_client_id",
+ }
+ ):
+ with pytest.raises(ValueError):
+ create_scopes_command(self.arg_parser.parse_args(params))
+
+ client.get_clients.assert_called_once_with()
+ client.create_client_authz_scopes.assert_not_called()
+
+ @patch("airflow.providers.keycloak.auth_manager.cli.commands._get_client")
+ def test_create_resources(self, mock_get_client):
+ client = Mock()
+ mock_get_client.return_value = client
+ scopes = [{"id": "1", "name": "GET"}]
+
+ client.get_clients.return_value = [
+ {"id": "dummy-id", "clientId": "dummy-client"},
+ {"id": "test-id", "clientId": "test_client_id"},
+ ]
+ client.get_client_authz_scopes.return_value = scopes
+
+ params = [
+ "keycloak-auth-manager",
+ "create-resources",
+ "--username",
+ "test",
+ "--password",
+ "test",
+ ]
+ with conf_vars(
+ {
+ ("keycloak_auth_manager", "client_id"): "test_client_id",
+ }
+ ):
+ create_resources_command(self.arg_parser.parse_args(params))
+
+ client.get_clients.assert_called_once_with()
+ client.get_client_authz_scopes.assert_called_once_with("test-id")
+ calls = []
+ for resource in KeycloakResource:
+ calls.append(
+ call(
+ client_id="test-id",
+ payload={
+ "name": resource.value,
+ "scopes": scopes,
+ },
+ skip_exists=True,
+ )
+ )
+ client.create_client_authz_resource.assert_has_calls(calls)
+
+
@patch("airflow.providers.keycloak.auth_manager.cli.commands._create_op_permission")
+
@patch("airflow.providers.keycloak.auth_manager.cli.commands._create_user_permission")
+
@patch("airflow.providers.keycloak.auth_manager.cli.commands._create_admin_permission")
+
@patch("airflow.providers.keycloak.auth_manager.cli.commands._create_read_only_permission")
+ @patch("airflow.providers.keycloak.auth_manager.cli.commands._get_client")
+ def test_create_permissions(
+ self,
+ mock_get_client,
+ mock_create_read_only_permission,
+ mock_create_admin_permission,
+ mock__create_user_permission,
+ mock_create_op_permission,
+ ):
+ client = Mock()
+ mock_get_client.return_value = client
+ scopes = [{"id": "1", "name": "GET"}]
+
+ client.get_clients.return_value = [
+ {"id": "dummy-id", "clientId": "dummy-client"},
+ {"id": "test-id", "clientId": "test_client_id"},
+ ]
+ client.get_client_authz_scopes.return_value = scopes
+
+ params = [
+ "keycloak-auth-manager",
+ "create-permissions",
+ "--username",
+ "test",
+ "--password",
+ "test",
+ ]
+ with conf_vars(
+ {
+ ("keycloak_auth_manager", "client_id"): "test_client_id",
+ }
+ ):
+ create_permissions_command(self.arg_parser.parse_args(params))
+
+ client.get_clients.assert_called_once_with()
+ mock_create_read_only_permission.assert_called_once_with(client,
"test-id")
+ mock_create_admin_permission.assert_called_once_with(client, "test-id")
+ mock__create_user_permission.assert_called_once_with(client, "test-id")
+ mock_create_op_permission.assert_called_once_with(client, "test-id")
+
+
@patch("airflow.providers.keycloak.auth_manager.cli.commands._create_permissions")
+
@patch("airflow.providers.keycloak.auth_manager.cli.commands._create_resources")
+
@patch("airflow.providers.keycloak.auth_manager.cli.commands._create_scopes")
+ @patch("airflow.providers.keycloak.auth_manager.cli.commands._get_client")
+ def test_create_all(
+ self,
+ mock_get_client,
+ mock_create_scopes,
+ mock_create_resources,
+ mock_create_permissions,
+ ):
+ client = Mock()
+ mock_get_client.return_value = client
+ scopes = [{"id": "1", "name": "GET"}]
+
+ client.get_clients.return_value = [
+ {"id": "dummy-id", "clientId": "dummy-client"},
+ {"id": "test-id", "clientId": "test_client_id"},
+ ]
+ client.get_client_authz_scopes.return_value = scopes
+
+ params = [
+ "keycloak-auth-manager",
+ "create-all",
+ "--username",
+ "test",
+ "--password",
+ "test",
+ ]
+ with conf_vars(
+ {
+ ("keycloak_auth_manager", "client_id"): "test_client_id",
+ }
+ ):
+ create_all_command(self.arg_parser.parse_args(params))
+
+ client.get_clients.assert_called_once_with()
+ mock_create_scopes.assert_called_once_with(client, "test-id")
+ mock_create_resources.assert_called_once_with(client, "test-id")
+ mock_create_permissions.assert_called_once_with(client, "test-id")
+
+ def test_create_permissions_read_only(self):
+ client = Mock()
+ scopes = [{"id": "1", "name": "GET"}, {"id": "2", "name": "MENU"},
{"id": "3", "name": "PUT"}]
+
+ client.get_client_authz_scopes.return_value = scopes
+
+ _create_read_only_permission(client, "test-id")
+
+ client.get_client_authz_scopes.assert_called_once_with("test-id")
+ client.create_client_authz_scope_permission.assert_called_once_with(
+ client_id="test-id",
+ payload={
+ "name": "ReadOnly",
+ "type": "scope",
+ "logic": "POSITIVE",
+ "decisionStrategy": "UNANIMOUS",
+ "scopes": ["1", "2"],
+ },
+ )
+
+ def test_create_permissions_admin(self):
+ client = Mock()
+ scopes = [{"id": "1", "name": "GET"}, {"id": "2", "name": "MENU"},
{"id": "3", "name": "PUT"}]
+
+ client.get_client_authz_scopes.return_value = scopes
+
+ _create_admin_permission(client, "test-id")
+
+ client.get_client_authz_scopes.assert_called_once_with("test-id")
+ client.create_client_authz_scope_permission.assert_called_once_with(
+ client_id="test-id",
+ payload={
+ "name": "Admin",
+ "type": "scope",
+ "logic": "POSITIVE",
+ "decisionStrategy": "UNANIMOUS",
+ "scopes": ["1", "2", "3"],
+ },
+ )
+
+ def test_create_permissions_user(self):
+ client = Mock()
+ resources = [
+ {"_id": "1", "name": "Dag"},
+ {"_id": "2", "name": "Asset"},
+ {"_id": "3", "name": "Variable"},
+ ]
+
+ client.get_client_authz_resources.return_value = resources
+
+ _create_user_permission(client, "test-id")
+
+ client.get_client_authz_resources.assert_called_once_with("test-id")
+
client.create_client_authz_resource_based_permission.assert_called_once_with(
+ client_id="test-id",
+ payload={
+ "name": "User",
+ "type": "scope",
+ "logic": "POSITIVE",
+ "decisionStrategy": "UNANIMOUS",
+ "resources": ["1", "2"],
+ },
+ skip_exists=True,
+ )
+
+ def test_create_permissions_op(self):
+ client = Mock()
+ resources = [
+ {"_id": "1", "name": "Dag"},
+ {"_id": "2", "name": "Connection"},
+ {"_id": "3", "name": "Variable"},
+ ]
+
+ client.get_client_authz_resources.return_value = resources
+
+ _create_op_permission(client, "test-id")
+
+ client.get_client_authz_resources.assert_called_once_with("test-id")
+
client.create_client_authz_resource_based_permission.assert_called_once_with(
+ client_id="test-id",
+ payload={
+ "name": "Op",
+ "type": "scope",
+ "logic": "POSITIVE",
+ "decisionStrategy": "UNANIMOUS",
+ "resources": ["2", "3"],
+ },
+ skip_exists=True,
+ )
diff --git
a/providers/keycloak/tests/unit/keycloak/auth_manager/cli/test_definition.py
b/providers/keycloak/tests/unit/keycloak/auth_manager/cli/test_definition.py
new file mode 100644
index 00000000000..2b7b0087c0e
--- /dev/null
+++ b/providers/keycloak/tests/unit/keycloak/auth_manager/cli/test_definition.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.cli.definition import
KEYCLOAK_AUTH_MANAGER_COMMANDS
+
+
+class TestKeycloakCliDefinition:
+ def test_aws_auth_manager_cli_commands(self):
+ assert len(KEYCLOAK_AUTH_MANAGER_COMMANDS) == 4
diff --git
a/providers/keycloak/tests/unit/keycloak/auth_manager/routes/test_login.py
b/providers/keycloak/tests/unit/keycloak/auth_manager/routes/test_login.py
index 0e3d2714b2f..d74040467b0 100644
--- a/providers/keycloak/tests/unit/keycloak/auth_manager/routes/test_login.py
+++ b/providers/keycloak/tests/unit/keycloak/auth_manager/routes/test_login.py
@@ -22,6 +22,12 @@ import pytest
from fastapi.testclient import TestClient
from airflow.api_fastapi.app import AUTH_MANAGER_FASTAPI_APP_PREFIX, create_app
+from airflow.providers.keycloak.auth_manager.constants import (
+ CONF_CLIENT_ID_KEY,
+ CONF_CLIENT_SECRET_KEY,
+ CONF_REALM_KEY,
+ CONF_SECTION_NAME,
+)
from tests_common.test_utils.config import conf_vars
@@ -34,10 +40,10 @@ def client():
"core",
"auth_manager",
):
"airflow.providers.keycloak.auth_manager.keycloak_auth_manager.KeycloakAuthManager",
- ("keycloak_auth_manager", "client_id"): "test",
- ("keycloak_auth_manager", "client_secret"): "test",
- ("keycloak_auth_manager", "realm"): "test",
- ("keycloak_auth_manager", "base_url"):
"http://host.docker.internal:48080";,
+ (CONF_SECTION_NAME, CONF_CLIENT_ID_KEY): "test",
+ (CONF_SECTION_NAME, CONF_CLIENT_SECRET_KEY): "test",
+ (CONF_SECTION_NAME, CONF_REALM_KEY): "test",
+ (CONF_SECTION_NAME, "base_url"):
"http://host.docker.internal:48080";,
}
):
yield TestClient(create_app())
diff --git
a/providers/keycloak/tests/unit/keycloak/auth_manager/test_constants.py
b/providers/keycloak/tests/unit/keycloak/auth_manager/test_constants.py
new file mode 100644
index 00000000000..923e6cf256b
--- /dev/null
+++ b/providers/keycloak/tests/unit/keycloak/auth_manager/test_constants.py
@@ -0,0 +1,42 @@
+# 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.constants import (
+ CONF_CLIENT_ID_KEY,
+ CONF_CLIENT_SECRET_KEY,
+ CONF_REALM_KEY,
+ CONF_SECTION_NAME,
+ CONF_SERVER_URL_KEY,
+)
+
+
+class TestKeycloakAuthManagerConstants:
+ def test_conf_section_name(self):
+ assert CONF_SECTION_NAME == "keycloak_auth_manager"
+
+ def test_conf_client_id_key(self):
+ assert CONF_CLIENT_ID_KEY == "client_id"
+
+ def test_conf_client_secret_key(self):
+ assert CONF_CLIENT_SECRET_KEY == "client_secret"
+
+ def test_conf_realm_key(self):
+ assert CONF_REALM_KEY == "realm"
+
+ def test_conf_server_url_key(self):
+ assert CONF_SERVER_URL_KEY == "server_url"
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 1cf2c3343d7..7bd580de2dc 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
@@ -35,6 +35,12 @@ from
airflow.api_fastapi.auth.managers.models.resource_details import (
VariableDetails,
)
from airflow.exceptions import AirflowException
+from airflow.providers.keycloak.auth_manager.constants import (
+ CONF_CLIENT_ID_KEY,
+ CONF_REALM_KEY,
+ CONF_SECTION_NAME,
+ CONF_SERVER_URL_KEY,
+)
from airflow.providers.keycloak.auth_manager.keycloak_auth_manager import (
RESOURCE_ID_ATTRIBUTE_NAME,
KeycloakAuthManager,
@@ -48,9 +54,9 @@ from tests_common.test_utils.config import conf_vars
def auth_manager():
with conf_vars(
{
- ("keycloak_auth_manager", "client_id"): "client_id",
- ("keycloak_auth_manager", "realm"): "realm",
- ("keycloak_auth_manager", "server_url"): "server_url",
+ (CONF_SECTION_NAME, CONF_CLIENT_ID_KEY): "client_id",
+ (CONF_SECTION_NAME, CONF_REALM_KEY): "realm",
+ (CONF_SECTION_NAME, CONF_SERVER_URL_KEY): "server_url",
}
):
yield KeycloakAuthManager()
@@ -362,3 +368,6 @@ class TestKeycloakAuthManager:
headers = auth_manager._get_headers("access_token")
mock_requests.post.assert_called_once_with(token_url, data=payload,
headers=headers)
assert result == expected
+
+ def test_get_cli_commands_return_cli_commands(self, auth_manager):
+ assert len(auth_manager.get_cli_commands()) == 1
