diff --git a/airflow/providers/fab/auth_manager/cli_commands/role_command.py b/airflow/providers/fab/auth_manager/cli_commands/role_command.py index cc30bc8fafec4..9288474ab0400 100644 --- a/airflow/providers/fab/auth_manager/cli_commands/role_command.py +++ b/airflow/providers/fab/auth_manager/cli_commands/role_command.py @@ -21,6 +21,7 @@ import itertools import json import os +from argparse import Namespace from collections import defaultdict from typing import TYPE_CHECKING @@ -155,30 +156,36 @@ def roles_del_perms(args): @suppress_logs_and_warning @providers_configuration_loaded def roles_export(args): - """ - Export all the roles from the database to a file. - - Note, this function does not export the permissions associated for each role. - Strictly, it exports the role names into the passed role json file. - """ + """Export all the roles from the database to a file including permissions.""" with get_application_builder() as appbuilder: roles = appbuilder.sm.get_all_roles() - exporting_roles = [role.name for role in roles if role.name not in EXISTING_ROLES] + exporting_roles = [role for role in roles if role.name not in EXISTING_ROLES] filename = os.path.expanduser(args.file) - kwargs = {} if not args.pretty else {"sort_keys": True, "indent": 4} + + permission_map: dict[tuple[str, str], list[str]] = defaultdict(list) + for role in exporting_roles: + if role.permissions: + for permission in role.permissions: + permission_map[(role.name, permission.resource.name)].append(permission.action.name) + else: + permission_map[(role.name, "")].append("") + export_data = [ + {"name": role, "resource": resource, "action": ",".join(sorted(permissions))} + for (role, resource), permissions in permission_map.items() + ] + kwargs = {} if not args.pretty else {"sort_keys": False, "indent": 4} with open(filename, "w", encoding="utf-8") as f: - json.dump(exporting_roles, f, **kwargs) - print(f"{len(exporting_roles)} roles successfully exported to {filename}") + json.dump(export_data, f, **kwargs) + print(f"{len(export_data)} roles with permissions successfully exported to {filename}") @cli_utils.action_cli @suppress_logs_and_warning def roles_import(args): """ - Import all the roles into the db from the given json file. + Import all the roles into the db from the given json file including their permissions. - Note, this function does not import the permissions for different roles and import them as well. - Strictly, it imports the role names in the role json file passed. + Note, if a role already exists in the db, it is not overwritten, even when the permissions change. """ json_file = args.file try: @@ -193,7 +200,30 @@ def roles_import(args): with get_application_builder() as appbuilder: existing_roles = [role.name for role in appbuilder.sm.get_all_roles()] - roles_to_import = [role for role in role_list if role not in existing_roles] - for role_name in roles_to_import: - appbuilder.sm.add_role(role_name) - print(f"roles '{roles_to_import}' successfully imported") + roles_to_import = [role_dict for role_dict in role_list if role_dict["name"] not in existing_roles] + for role_dict in roles_to_import: + if role_dict["name"] not in appbuilder.sm.get_all_roles(): + if role_dict["action"] == "" or role_dict["resource"] == "": + appbuilder.sm.add_role(role_dict["name"]) + else: + appbuilder.sm.add_role(role_dict["name"]) + role_args = Namespace( + subcommand="add-perms", + role=[role_dict["name"]], + resource=[role_dict["resource"]], + action=role_dict["action"].split(","), + ) + __roles_add_or_remove_permissions(role_args) + + if role_dict["name"] in appbuilder.sm.get_all_roles(): + if role_dict["action"] == "" or role_dict["resource"] == "": + pass + else: + role_args = Namespace( + subcommand="add-perms", + role=[role_dict["name"]], + resource=[role_dict["resource"]], + action=role_dict["action"].split(","), + ) + __roles_add_or_remove_permissions(role_args) + print("roles and permissions successfully imported") diff --git a/tests/providers/fab/auth_manager/cli_commands/test_role_command.py b/tests/providers/fab/auth_manager/cli_commands/test_role_command.py index 9a82a6c893a63..a6bf870d6d443 100644 --- a/tests/providers/fab/auth_manager/cli_commands/test_role_command.py +++ b/tests/providers/fab/auth_manager/cli_commands/test_role_command.py @@ -151,27 +151,52 @@ def test_cli_roles_add_and_del_perms(self): role: Role = self.appbuilder.sm.find_role("FakeTeamC") assert len(role.permissions) == 0 - def test_cli_import_roles(self, tmp_path): - fn = tmp_path / "import_roles.json" - fn.touch() - roles_list = ["FakeTeamA", "FakeTeamB"] - with open(fn, "w") as outfile: - json.dump(roles_list, outfile) - role_command.roles_import(self.parser.parse_args(["roles", "import", str(fn)])) - assert self.appbuilder.sm.find_role("FakeTeamA") is not None - assert self.appbuilder.sm.find_role("FakeTeamB") is not None - def test_cli_export_roles(self, tmp_path): fn = tmp_path / "export_roles.json" fn.touch() args = self.parser.parse_args(["roles", "create", "FakeTeamA", "FakeTeamB"]) role_command.roles_create(args) - - assert self.appbuilder.sm.find_role("FakeTeamA") is not None - assert self.appbuilder.sm.find_role("FakeTeamB") is not None - + role_command.roles_add_perms( + self.parser.parse_args( + [ + "roles", + "add-perms", + "FakeTeamA", + "-r", + permissions.RESOURCE_POOL, + "-a", + permissions.ACTION_CAN_EDIT, + permissions.ACTION_CAN_READ, + ] + ) + ) role_command.roles_export(self.parser.parse_args(["roles", "export", str(fn)])) with open(fn) as outfile: roles_exported = json.load(outfile) - assert "FakeTeamA" in roles_exported - assert "FakeTeamB" in roles_exported + assert {"name": "FakeTeamA", "resource": "Pools", "action": "can_edit,can_read"} in roles_exported + assert {"name": "FakeTeamB", "resource": "", "action": ""} in roles_exported + + def test_cli_import_roles(self, tmp_path): + fn = tmp_path / "import_roles.json" + fn.touch() + roles_list = [ + {"name": "FakeTeamA", "resource": "Pools", "action": "can_edit,can_read"}, + {"name": "FakeTeamA", "resource": "Admin", "action": "menu_access"}, + {"name": "FakeTeamB", "resource": "", "action": ""}, + ] + with open(fn, "w") as outfile: + json.dump(roles_list, outfile) + role_command.roles_import(self.parser.parse_args(["roles", "import", str(fn)])) + fakeTeamA: Role = self.appbuilder.sm.find_role("FakeTeamA") + fakeTeamB: Role = self.appbuilder.sm.find_role("FakeTeamB") + + assert fakeTeamA is not None + assert fakeTeamB is not None + assert len(fakeTeamB.permissions) == 0 + assert len(fakeTeamA.permissions) == 3 + assert fakeTeamA.permissions[0].resource.name == permissions.RESOURCE_POOL + assert fakeTeamA.permissions[0].action.name == permissions.ACTION_CAN_EDIT + assert fakeTeamA.permissions[1].resource.name == permissions.RESOURCE_POOL + assert fakeTeamA.permissions[1].action.name == permissions.ACTION_CAN_READ + assert fakeTeamA.permissions[2].resource.name == permissions.RESOURCE_ADMIN_MENU + assert fakeTeamA.permissions[2].action.name == permissions.ACTION_CAN_ACCESS_MENU