From 1a767723ec0b9009fdd78fda65335cfbf2850fbf Mon Sep 17 00:00:00 2001 From: Anmol Mishra Date: Sat, 16 May 2026 09:01:50 +0000 Subject: [PATCH] fix(keycloak): attach default role policies --- .../keycloak/auth_manager/cli/commands.py | 61 ++++++++++ .../auth_manager/cli/test_commands.py | 108 +++++++++++++++++- 2 files changed, 168 insertions(+), 1 deletion(-) 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 index 22bb187996d9f..705fc4eb0bfac 100644 --- a/providers/keycloak/src/airflow/providers/keycloak/auth_manager/cli/commands.py +++ b/providers/keycloak/src/airflow/providers/keycloak/auth_manager/cli/commands.py @@ -137,7 +137,11 @@ def create_permissions_command(args): for role_name in TEAM_ROLE_NAMES: _ensure_role_policy(client, client_uuid, role_name, _dry_run=args.dry_run) _ensure_role_policy(client, client_uuid, SUPER_ADMIN_ROLE_NAME, _dry_run=args.dry_run) + else: + _ensure_default_role_policies(client, client_uuid, _dry_run=args.dry_run) _create_permissions(client, client_uuid, teams=teams, _dry_run=args.dry_run) + if not teams: + _attach_default_role_permissions(client, client_uuid, _dry_run=args.dry_run) @cli_utils.action_cli @@ -158,7 +162,11 @@ def create_all_command(args): for role_name in TEAM_ROLE_NAMES: _ensure_role_policy(client, client_uuid, role_name, _dry_run=args.dry_run) _ensure_role_policy(client, client_uuid, SUPER_ADMIN_ROLE_NAME, _dry_run=args.dry_run) + else: + _ensure_default_role_policies(client, client_uuid, _dry_run=args.dry_run) _create_permissions(client, client_uuid, teams=teams, _dry_run=args.dry_run) + if not teams: + _attach_default_role_permissions(client, client_uuid, _dry_run=args.dry_run) def _get_client(args): @@ -244,6 +252,59 @@ def _ensure_multi_team_enabled(*, teams: list[str], command_name: str) -> None: raise SystemExit(f"{command_name} requires core.multi_team=True when --teams is used.") +def _ensure_default_role_policies(client: KeycloakAdmin, client_uuid: str, *, _dry_run: bool = False) -> None: + for role_name in (*TEAM_ROLE_NAMES, SUPER_ADMIN_ROLE_NAME): + _ensure_role_policy(client, client_uuid, role_name, _dry_run=_dry_run) + + +def _attach_default_role_permissions( + client: KeycloakAdmin, client_uuid: str, *, _dry_run: bool = False +) -> None: + for role_name in TEAM_ROLE_NAMES: + _attach_policy_to_scope_permission( + client, + client_uuid, + permission_name="ReadOnly", + policy_name=_role_policy_name(role_name), + scope_names=["GET", "MENU", "LIST"], + resource_names=[], + decision_strategy="AFFIRMATIVE", + _dry_run=_dry_run, + ) + + _attach_policy_to_resource_permission( + client, + client_uuid, + permission_name="User", + policy_name=_role_policy_name("User"), + resource_names=[KeycloakResource.DAG.value, KeycloakResource.ASSET.value], + _dry_run=_dry_run, + ) + _attach_policy_to_resource_permission( + client, + client_uuid, + permission_name="Op", + policy_name=_role_policy_name("Op"), + resource_names=[ + KeycloakResource.CONNECTION.value, + KeycloakResource.POOL.value, + KeycloakResource.VARIABLE.value, + KeycloakResource.BACKFILL.value, + ], + _dry_run=_dry_run, + ) + for role_name in ("Admin", SUPER_ADMIN_ROLE_NAME): + _attach_policy_to_scope_permission( + client, + client_uuid, + permission_name="Admin", + policy_name=_role_policy_name(role_name), + scope_names=_get_extended_resource_methods() + ["LIST"], + resource_names=[], + _dry_run=_dry_run, + ) + + def _preview_scopes(*args, **kwargs): """Preview scopes that would be created.""" scopes = _get_scopes_to_create() 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 index 07bb80c944897..62c61f95f7249 100644 --- a/providers/keycloak/tests/unit/keycloak/auth_manager/cli/test_commands.py +++ b/providers/keycloak/tests/unit/keycloak/auth_manager/cli/test_commands.py @@ -24,6 +24,8 @@ from airflow.api_fastapi.common.types import MenuItem from airflow.cli import cli_parser from airflow.providers.keycloak.auth_manager.cli.commands import ( + SUPER_ADMIN_ROLE_NAME, + TEAM_ROLE_NAMES, TEAM_SCOPED_RESOURCE_NAMES, _get_extended_resource_methods, _get_resource_methods, @@ -213,10 +215,14 @@ def test_create_resources_with_teams(self, mock_get_client): } assert expected_team_resources.issubset(created_resource_names) + @patch("airflow.providers.keycloak.auth_manager.cli.commands._attach_default_role_permissions") + @patch("airflow.providers.keycloak.auth_manager.cli.commands._ensure_default_role_policies") @patch("airflow.providers.keycloak.auth_manager.cli.commands._get_client") def test_create_permissions( self, mock_get_client, + mock_ensure_default_role_policies, + mock_attach_default_role_permissions, ): client = Mock() mock_get_client.return_value = client @@ -309,6 +315,8 @@ def test_create_permissions( ), ] client.create_client_authz_resource_based_permission.assert_has_calls(resource_calls, any_order=True) + mock_ensure_default_role_policies.assert_called_once_with(client, "test-id", _dry_run=False) + mock_attach_default_role_permissions.assert_called_once_with(client, "test-id", _dry_run=False) @patch("airflow.providers.keycloak.auth_manager.cli.commands._get_client") def test_create_permissions_with_teams(self, mock_get_client): @@ -419,6 +427,86 @@ def test_create_permissions_with_teams(self, mock_get_client): skip_exists=True, ) + @patch("airflow.providers.keycloak.auth_manager.cli.commands._attach_policy_to_resource_permission") + @patch("airflow.providers.keycloak.auth_manager.cli.commands._attach_policy_to_scope_permission") + @patch("airflow.providers.keycloak.auth_manager.cli.commands._ensure_role_policy") + @patch("airflow.providers.keycloak.auth_manager.cli.commands._create_permissions") + @patch("airflow.providers.keycloak.auth_manager.cli.commands._get_client") + def test_create_permissions_attaches_default_role_policies( + self, + mock_get_client, + mock_create_permissions, + mock_ensure_role_policy, + mock_attach_scope_policy, + mock_attach_resource_policy, + ): + client = Mock() + mock_get_client.return_value = client + client.get_clients.return_value = [ + {"id": "test-id", "clientId": "test_client_id"}, + ] + + 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)) + + for role_name in (*TEAM_ROLE_NAMES, SUPER_ADMIN_ROLE_NAME): + mock_ensure_role_policy.assert_any_call(client, "test-id", role_name, _dry_run=False) + + mock_create_permissions.assert_called_once_with(client, "test-id", teams=[], _dry_run=False) + for role_name in TEAM_ROLE_NAMES: + mock_attach_scope_policy.assert_any_call( + client, + "test-id", + permission_name="ReadOnly", + policy_name=f"Allow-{role_name}", + scope_names=["GET", "MENU", "LIST"], + resource_names=[], + decision_strategy="AFFIRMATIVE", + _dry_run=False, + ) + mock_attach_scope_policy.assert_any_call( + client, + "test-id", + permission_name="Admin", + policy_name="Allow-Admin", + scope_names=_get_extended_resource_methods() + ["LIST"], + resource_names=[], + _dry_run=False, + ) + mock_attach_scope_policy.assert_any_call( + client, + "test-id", + permission_name="Admin", + policy_name="Allow-SuperAdmin", + scope_names=_get_extended_resource_methods() + ["LIST"], + resource_names=[], + _dry_run=False, + ) + mock_attach_resource_policy.assert_any_call( + client, + "test-id", + permission_name="User", + policy_name="Allow-User", + resource_names=["Dag", "Asset"], + _dry_run=False, + ) + mock_attach_resource_policy.assert_any_call( + client, + "test-id", + permission_name="Op", + policy_name="Allow-Op", + resource_names=["Connection", "Pool", "Variable", "Backfill"], + _dry_run=False, + ) + @patch("airflow.providers.keycloak.auth_manager.cli.commands._update_admin_permission_resources") @patch("airflow.providers.keycloak.auth_manager.cli.commands._ensure_scope_permission") @patch("airflow.providers.keycloak.auth_manager.cli.commands._attach_policy_to_resource_permission") @@ -628,6 +716,8 @@ def test_add_user_to_team_command(self, mock_get_client, mock_ensure_group, mock mock_ensure_group.assert_called_once_with(client, "team-a", _dry_run=False) mock_add_user.assert_called_once_with(client, username="user-a", team="team-a", _dry_run=False) + @patch("airflow.providers.keycloak.auth_manager.cli.commands._attach_default_role_permissions") + @patch("airflow.providers.keycloak.auth_manager.cli.commands._ensure_default_role_policies") @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") @@ -638,6 +728,8 @@ def test_create_all( mock_create_scopes, mock_create_resources, mock_create_permissions, + mock_ensure_default_role_policies, + mock_attach_default_role_permissions, ): client = Mock() mock_get_client.return_value = client @@ -674,6 +766,8 @@ def test_create_all( mock_create_scopes.assert_called_once_with(client, "test-id", _dry_run=False) mock_create_resources.assert_called_once_with(client, "test-id", teams=[], _dry_run=False) mock_create_permissions.assert_called_once_with(client, "test-id", teams=[], _dry_run=False) + mock_ensure_default_role_policies.assert_called_once_with(client, "test-id", _dry_run=False) + mock_attach_default_role_permissions.assert_called_once_with(client, "test-id", _dry_run=False) @patch("airflow.providers.keycloak.auth_manager.cli.commands._get_client") def test_create_scopes_dry_run(self, mock_get_client): @@ -738,8 +832,12 @@ def test_create_resources_dry_run(self, mock_get_client): # In dry-run mode, no resources should be created client.create_client_authz_resource.assert_not_called() + @patch("airflow.providers.keycloak.auth_manager.cli.commands._attach_default_role_permissions") + @patch("airflow.providers.keycloak.auth_manager.cli.commands._ensure_default_role_policies") @patch("airflow.providers.keycloak.auth_manager.cli.commands._get_client") - def test_create_permissions_dry_run(self, mock_get_client): + def test_create_permissions_dry_run( + self, mock_get_client, mock_ensure_default_role_policies, mock_attach_default_role_permissions + ): client = Mock() mock_get_client.return_value = client scopes = [{"id": "1", "name": "GET"}, {"id": "2", "name": "MENU"}, {"id": "3", "name": "LIST"}] @@ -777,7 +875,11 @@ def test_create_permissions_dry_run(self, mock_get_client): # In dry-run mode, no permissions should be created client.create_client_authz_scope_permission.assert_not_called() client.create_client_authz_resource_based_permission.assert_not_called() + mock_ensure_default_role_policies.assert_called_once_with(client, "test-id", _dry_run=True) + mock_attach_default_role_permissions.assert_called_once_with(client, "test-id", _dry_run=True) + @patch("airflow.providers.keycloak.auth_manager.cli.commands._attach_default_role_permissions") + @patch("airflow.providers.keycloak.auth_manager.cli.commands._ensure_default_role_policies") @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") @@ -788,6 +890,8 @@ def test_create_all_dry_run( mock_create_scopes, mock_create_resources, mock_create_permissions, + mock_ensure_default_role_policies, + mock_attach_default_role_permissions, ): client = Mock() mock_get_client.return_value = client @@ -824,3 +928,5 @@ def test_create_all_dry_run( mock_create_scopes.assert_called_once_with(client, "test-id", _dry_run=True) mock_create_resources.assert_called_once_with(client, "test-id", teams=[], _dry_run=True) mock_create_permissions.assert_called_once_with(client, "test-id", teams=[], _dry_run=True) + mock_ensure_default_role_policies.assert_called_once_with(client, "test-id", _dry_run=True) + mock_attach_default_role_permissions.assert_called_once_with(client, "test-id", _dry_run=True)