From 80ee6b785fda07bc6759ba71fb61d2e185c276b3 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Rapha=C3=ABl=20Vandon?= Date: Mon, 14 Aug 2023 16:05:03 -0700 Subject: [PATCH 01/11] add a way for auth managers to define their own CLI commands --- airflow/auth/managers/base_auth_manager.py | 9 + airflow/auth/managers/fab/fab_auth_manager.py | 198 ++++++++++++++++++ airflow/cli/cli_config.py | 195 +---------------- airflow/cli/cli_parser.py | 8 + 4 files changed, 221 insertions(+), 189 deletions(-) diff --git a/airflow/auth/managers/base_auth_manager.py b/airflow/auth/managers/base_auth_manager.py index bcc5e13892e2b..445b5b0245327 100644 --- a/airflow/auth/managers/base_auth_manager.py +++ b/airflow/auth/managers/base_auth_manager.py @@ -21,6 +21,7 @@ from typing import TYPE_CHECKING from airflow.auth.managers.models.base_user import BaseUser +from airflow.cli.cli_config import GroupCommand from airflow.exceptions import AirflowException from airflow.utils.log.logging_mixin import LoggingMixin @@ -38,6 +39,14 @@ class BaseAuthManager(LoggingMixin): def __init__(self): self._security_manager: AirflowSecurityManager | None = None + @staticmethod + def get_cli_commands() -> list[GroupCommand]: + """Vends CLI commands to be included in Airflow CLI. + + Override this method to expose commands via Airflow CLI to manage this auth manager. + """ + return [] + @abstractmethod def get_user_name(self) -> str: """Return the username associated to the user in session.""" diff --git a/airflow/auth/managers/fab/fab_auth_manager.py b/airflow/auth/managers/fab/fab_auth_manager.py index 848f5ff188d6c..f5e6444c5cfe4 100644 --- a/airflow/auth/managers/fab/fab_auth_manager.py +++ b/airflow/auth/managers/fab/fab_auth_manager.py @@ -17,6 +17,8 @@ # under the License. from __future__ import annotations +import textwrap + from flask import url_for from flask_login import current_user @@ -24,6 +26,187 @@ from airflow.auth.managers.base_auth_manager import BaseAuthManager from airflow.auth.managers.fab.models import User from airflow.auth.managers.fab.security_manager.override import FabAirflowSecurityManagerOverride +from airflow.cli.cli_config import ( + ARG_OUTPUT, + ARG_VERBOSE, + ActionCommand, + Arg, + GroupCommand, + lazy_load_command, +) + +# users +ARG_USERNAME = Arg(("-u", "--username"), help="Username of the user", required=True, type=str) +ARG_USERNAME_OPTIONAL = Arg(("-u", "--username"), help="Username of the user", type=str) +ARG_FIRSTNAME = Arg(("-f", "--firstname"), help="First name of the user", required=True, type=str) +ARG_LASTNAME = Arg(("-l", "--lastname"), help="Last name of the user", required=True, type=str) +ARG_ROLE = Arg( + ("-r", "--role"), + help="Role of the user. Existing roles include Admin, User, Op, Viewer, and Public", + required=True, + type=str, +) +ARG_EMAIL = Arg(("-e", "--email"), help="Email of the user", required=True, type=str) +ARG_EMAIL_OPTIONAL = Arg(("-e", "--email"), help="Email of the user", type=str) +ARG_PASSWORD = Arg( + ("-p", "--password"), + help="Password of the user, required to create a user without --use-random-password", + type=str, +) +ARG_USE_RANDOM_PASSWORD = Arg( + ("--use-random-password",), + help="Do not prompt for password. Use random string instead." + " Required to create a user without --password ", + default=False, + action="store_true", +) +ARG_USER_IMPORT = Arg( + ("import",), + metavar="FILEPATH", + help="Import users from JSON file. Example format::\n" + + textwrap.indent( + textwrap.dedent( + """ + [ + { + "email": "foo@bar.org", + "firstname": "Jon", + "lastname": "Doe", + "roles": ["Public"], + "username": "jondoe" + } + ]""" + ), + " " * 4, + ), +) +ARG_USER_EXPORT = Arg(("export",), metavar="FILEPATH", help="Export all users to JSON file") + +# roles +ARG_CREATE_ROLE = Arg(("-c", "--create"), help="Create a new role", action="store_true") +ARG_LIST_ROLES = Arg(("-l", "--list"), help="List roles", action="store_true") +ARG_ROLES = Arg(("role",), help="The name of a role", nargs="*") +ARG_PERMISSIONS = Arg(("-p", "--permission"), help="Show role permissions", action="store_true") +ARG_ROLE_RESOURCE = Arg(("-r", "--resource"), help="The name of permissions", nargs="*", required=True) +ARG_ROLE_ACTION = Arg(("-a", "--action"), help="The action of permissions", nargs="*") +ARG_ROLE_ACTION_REQUIRED = Arg(("-a", "--action"), help="The action of permissions", nargs="*", required=True) + +ARG_ROLE_IMPORT = Arg(("file",), help="Import roles from JSON file", nargs=None) +ARG_ROLE_EXPORT = Arg(("file",), help="Export all roles to JSON file", nargs=None) +ARG_ROLE_EXPORT_FMT = Arg( + ("-p", "--pretty"), + help="Format output JSON file by sorting role names and indenting by 4 spaces", + action="store_true", +) + +USERS_COMMANDS = ( + ActionCommand( + name="list", + help="List users", + func=lazy_load_command("airflow.cli.commands.user_command.users_list"), + args=(ARG_OUTPUT, ARG_VERBOSE), + ), + ActionCommand( + name="create", + help="Create a user", + func=lazy_load_command("airflow.cli.commands.user_command.users_create"), + args=( + ARG_ROLE, + ARG_USERNAME, + ARG_EMAIL, + ARG_FIRSTNAME, + ARG_LASTNAME, + ARG_PASSWORD, + ARG_USE_RANDOM_PASSWORD, + ARG_VERBOSE, + ), + epilog=( + "examples:\n" + 'To create an user with "Admin" role and username equals to "admin", run:\n' + "\n" + " $ airflow users create \\\n" + " --username admin \\\n" + " --firstname FIRST_NAME \\\n" + " --lastname LAST_NAME \\\n" + " --role Admin \\\n" + " --email admin@example.org" + ), + ), + ActionCommand( + name="delete", + help="Delete a user", + func=lazy_load_command("airflow.cli.commands.user_command.users_delete"), + args=(ARG_USERNAME_OPTIONAL, ARG_EMAIL_OPTIONAL, ARG_VERBOSE), + ), + ActionCommand( + name="add-role", + help="Add role to a user", + func=lazy_load_command("airflow.cli.commands.user_command.add_role"), + args=(ARG_USERNAME_OPTIONAL, ARG_EMAIL_OPTIONAL, ARG_ROLE, ARG_VERBOSE), + ), + ActionCommand( + name="remove-role", + help="Remove role from a user", + func=lazy_load_command("airflow.cli.commands.user_command.remove_role"), + args=(ARG_USERNAME_OPTIONAL, ARG_EMAIL_OPTIONAL, ARG_ROLE, ARG_VERBOSE), + ), + ActionCommand( + name="import", + help="Import users", + func=lazy_load_command("airflow.cli.commands.user_command.users_import"), + args=(ARG_USER_IMPORT, ARG_VERBOSE), + ), + ActionCommand( + name="export", + help="Export all users", + func=lazy_load_command("airflow.cli.commands.user_command.users_export"), + args=(ARG_USER_EXPORT, ARG_VERBOSE), + ), +) +ROLES_COMMANDS = ( + ActionCommand( + name="list", + help="List roles", + func=lazy_load_command("airflow.cli.commands.role_command.roles_list"), + args=(ARG_PERMISSIONS, ARG_OUTPUT, ARG_VERBOSE), + ), + ActionCommand( + name="create", + help="Create role", + func=lazy_load_command("airflow.cli.commands.role_command.roles_create"), + args=(ARG_ROLES, ARG_VERBOSE), + ), + ActionCommand( + name="delete", + help="Delete role", + func=lazy_load_command("airflow.cli.commands.role_command.roles_delete"), + args=(ARG_ROLES, ARG_VERBOSE), + ), + ActionCommand( + name="add-perms", + help="Add roles permissions", + func=lazy_load_command("airflow.cli.commands.role_command.roles_add_perms"), + args=(ARG_ROLES, ARG_ROLE_RESOURCE, ARG_ROLE_ACTION_REQUIRED, ARG_VERBOSE), + ), + ActionCommand( + name="del-perms", + help="Delete roles permissions", + func=lazy_load_command("airflow.cli.commands.role_command.roles_del_perms"), + args=(ARG_ROLES, ARG_ROLE_RESOURCE, ARG_ROLE_ACTION, ARG_VERBOSE), + ), + ActionCommand( + name="export", + help="Export roles (without permissions) from db to JSON file", + func=lazy_load_command("airflow.cli.commands.role_command.roles_export"), + args=(ARG_ROLE_EXPORT, ARG_ROLE_EXPORT_FMT, ARG_VERBOSE), + ), + ActionCommand( + name="import", + help="Import roles (without permissions) from JSON file to db", + func=lazy_load_command("airflow.cli.commands.role_command.roles_import"), + args=(ARG_ROLE_IMPORT, ARG_VERBOSE), + ), +) class FabAuthManager(BaseAuthManager): @@ -33,6 +216,21 @@ class FabAuthManager(BaseAuthManager): This auth manager is responsible for providing a backward compatible user management experience to users. """ + @staticmethod + def get_cli_commands() -> list[GroupCommand]: + return [ + GroupCommand( + name="users", + help="Manage users", + subcommands=USERS_COMMANDS, + ), + GroupCommand( + name="roles", + help="Manage roles", + subcommands=ROLES_COMMANDS, + ), + ] + def get_user_name(self) -> str: """ Return the username associated to the user in session. diff --git a/airflow/cli/cli_config.py b/airflow/cli/cli_config.py index 2ec55ff79498f..13cf8f018ed90 100644 --- a/airflow/cli/cli_config.py +++ b/airflow/cli/cli_config.py @@ -229,6 +229,12 @@ def string_lower_type(val): ), default=None, ) +ARG_SKIP_SERVE_LOGS = Arg( + ("-s", "--skip-serve-logs"), + default=False, + help="Don't start the serve logs process along with the workers", + action="store_true", +) # list_dag_runs ARG_DAG_ID_REQ_FLAG = Arg( @@ -878,76 +884,6 @@ def string_lower_type(val): action="store_true", ) -# users -ARG_USERNAME = Arg(("-u", "--username"), help="Username of the user", required=True, type=str) -ARG_USERNAME_OPTIONAL = Arg(("-u", "--username"), help="Username of the user", type=str) -ARG_FIRSTNAME = Arg(("-f", "--firstname"), help="First name of the user", required=True, type=str) -ARG_LASTNAME = Arg(("-l", "--lastname"), help="Last name of the user", required=True, type=str) -ARG_ROLE = Arg( - ("-r", "--role"), - help="Role of the user. Existing roles include Admin, User, Op, Viewer, and Public", - required=True, - type=str, -) -ARG_EMAIL = Arg(("-e", "--email"), help="Email of the user", required=True, type=str) -ARG_EMAIL_OPTIONAL = Arg(("-e", "--email"), help="Email of the user", type=str) -ARG_PASSWORD = Arg( - ("-p", "--password"), - help="Password of the user, required to create a user without --use-random-password", - type=str, -) -ARG_USE_RANDOM_PASSWORD = Arg( - ("--use-random-password",), - help="Do not prompt for password. Use random string instead." - " Required to create a user without --password ", - default=False, - action="store_true", -) -ARG_USER_IMPORT = Arg( - ("import",), - metavar="FILEPATH", - help="Import users from JSON file. Example format::\n" - + textwrap.indent( - textwrap.dedent( - """ - [ - { - "email": "foo@bar.org", - "firstname": "Jon", - "lastname": "Doe", - "roles": ["Public"], - "username": "jondoe" - } - ]""" - ), - " " * 4, - ), -) -ARG_USER_EXPORT = Arg(("export",), metavar="FILEPATH", help="Export all users to JSON file") - -# roles -ARG_CREATE_ROLE = Arg(("-c", "--create"), help="Create a new role", action="store_true") -ARG_LIST_ROLES = Arg(("-l", "--list"), help="List roles", action="store_true") -ARG_ROLES = Arg(("role",), help="The name of a role", nargs="*") -ARG_PERMISSIONS = Arg(("-p", "--permission"), help="Show role permissions", action="store_true") -ARG_ROLE_RESOURCE = Arg(("-r", "--resource"), help="The name of permissions", nargs="*", required=True) -ARG_ROLE_ACTION = Arg(("-a", "--action"), help="The action of permissions", nargs="*") -ARG_ROLE_ACTION_REQUIRED = Arg(("-a", "--action"), help="The action of permissions", nargs="*", required=True) -ARG_AUTOSCALE = Arg(("-a", "--autoscale"), help="Minimum and Maximum number of worker to autoscale") -ARG_SKIP_SERVE_LOGS = Arg( - ("-s", "--skip-serve-logs"), - default=False, - help="Don't start the serve logs process along with the workers", - action="store_true", -) -ARG_ROLE_IMPORT = Arg(("file",), help="Import roles from JSON file", nargs=None) -ARG_ROLE_EXPORT = Arg(("file",), help="Export all roles to JSON file", nargs=None) -ARG_ROLE_EXPORT_FMT = Arg( - ("-p", "--pretty"), - help="Format output JSON file by sorting role names and indenting by 4 spaces", - action="store_true", -) - # info ARG_ANONYMIZE = Arg( ("--anonymize",), @@ -1839,115 +1775,6 @@ class GroupCommand(NamedTuple): ) -USERS_COMMANDS = ( - ActionCommand( - name="list", - help="List users", - func=lazy_load_command("airflow.cli.commands.user_command.users_list"), - args=(ARG_OUTPUT, ARG_VERBOSE), - ), - ActionCommand( - name="create", - help="Create a user", - func=lazy_load_command("airflow.cli.commands.user_command.users_create"), - args=( - ARG_ROLE, - ARG_USERNAME, - ARG_EMAIL, - ARG_FIRSTNAME, - ARG_LASTNAME, - ARG_PASSWORD, - ARG_USE_RANDOM_PASSWORD, - ARG_VERBOSE, - ), - epilog=( - "examples:\n" - 'To create an user with "Admin" role and username equals to "admin", run:\n' - "\n" - " $ airflow users create \\\n" - " --username admin \\\n" - " --firstname FIRST_NAME \\\n" - " --lastname LAST_NAME \\\n" - " --role Admin \\\n" - " --email admin@example.org" - ), - ), - ActionCommand( - name="delete", - help="Delete a user", - func=lazy_load_command("airflow.cli.commands.user_command.users_delete"), - args=(ARG_USERNAME_OPTIONAL, ARG_EMAIL_OPTIONAL, ARG_VERBOSE), - ), - ActionCommand( - name="add-role", - help="Add role to a user", - func=lazy_load_command("airflow.cli.commands.user_command.add_role"), - args=(ARG_USERNAME_OPTIONAL, ARG_EMAIL_OPTIONAL, ARG_ROLE, ARG_VERBOSE), - ), - ActionCommand( - name="remove-role", - help="Remove role from a user", - func=lazy_load_command("airflow.cli.commands.user_command.remove_role"), - args=(ARG_USERNAME_OPTIONAL, ARG_EMAIL_OPTIONAL, ARG_ROLE, ARG_VERBOSE), - ), - ActionCommand( - name="import", - help="Import users", - func=lazy_load_command("airflow.cli.commands.user_command.users_import"), - args=(ARG_USER_IMPORT, ARG_VERBOSE), - ), - ActionCommand( - name="export", - help="Export all users", - func=lazy_load_command("airflow.cli.commands.user_command.users_export"), - args=(ARG_USER_EXPORT, ARG_VERBOSE), - ), -) -ROLES_COMMANDS = ( - ActionCommand( - name="list", - help="List roles", - func=lazy_load_command("airflow.cli.commands.role_command.roles_list"), - args=(ARG_PERMISSIONS, ARG_OUTPUT, ARG_VERBOSE), - ), - ActionCommand( - name="create", - help="Create role", - func=lazy_load_command("airflow.cli.commands.role_command.roles_create"), - args=(ARG_ROLES, ARG_VERBOSE), - ), - ActionCommand( - name="delete", - help="Delete role", - func=lazy_load_command("airflow.cli.commands.role_command.roles_delete"), - args=(ARG_ROLES, ARG_VERBOSE), - ), - ActionCommand( - name="add-perms", - help="Add roles permissions", - func=lazy_load_command("airflow.cli.commands.role_command.roles_add_perms"), - args=(ARG_ROLES, ARG_ROLE_RESOURCE, ARG_ROLE_ACTION_REQUIRED, ARG_VERBOSE), - ), - ActionCommand( - name="del-perms", - help="Delete roles permissions", - func=lazy_load_command("airflow.cli.commands.role_command.roles_del_perms"), - args=(ARG_ROLES, ARG_ROLE_RESOURCE, ARG_ROLE_ACTION, ARG_VERBOSE), - ), - ActionCommand( - name="export", - help="Export roles (without permissions) from db to JSON file", - func=lazy_load_command("airflow.cli.commands.role_command.roles_export"), - args=(ARG_ROLE_EXPORT, ARG_ROLE_EXPORT_FMT, ARG_VERBOSE), - ), - ActionCommand( - name="import", - help="Import roles (without permissions) from JSON file to db", - func=lazy_load_command("airflow.cli.commands.role_command.roles_import"), - args=(ARG_ROLE_IMPORT, ARG_VERBOSE), - ), -) - CONFIG_COMMANDS = ( ActionCommand( name="get-value", @@ -2171,16 +1998,6 @@ class GroupCommand(NamedTuple): help="Display providers", subcommands=PROVIDERS_COMMANDS, ), - GroupCommand( - name="users", - help="Manage users", - subcommands=USERS_COMMANDS, - ), - GroupCommand( - name="roles", - help="Manage roles", - subcommands=ROLES_COMMANDS, - ), ActionCommand( name="sync-perm", help="Update permissions for existing roles and optionally DAGs", diff --git a/airflow/cli/cli_parser.py b/airflow/cli/cli_parser.py index 8e4d819098c5a..6ffbd9dbd674f 100644 --- a/airflow/cli/cli_parser.py +++ b/airflow/cli/cli_parser.py @@ -44,6 +44,7 @@ from airflow.exceptions import AirflowException from airflow.executors.executor_loader import ExecutorLoader from airflow.utils.helpers import partition +from airflow.www.extensions.init_auth_manager import get_auth_manager airflow_commands = core_commands @@ -62,6 +63,13 @@ # Do no re-raise the exception since we want the CLI to still function for # other commands. +try: + auth_mgr = get_auth_manager() + airflow_commands.extend(auth_mgr.get_cli_commands()) +except Exception: + log.exception("cannot load CLI commands from auth manager") + # do not re-raise for the same reason as above + ALL_COMMANDS_DICT: dict[str, CLICommand] = {sp.name: sp for sp in airflow_commands} From 8a45b1e02a6e2e998c9f9e01eaaeaaffa1f77e2d Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Rapha=C3=ABl=20Vandon?= Date: Mon, 14 Aug 2023 16:49:46 -0700 Subject: [PATCH 02/11] move commands --- .../fab/cli_commands}/role_command.py | 0 .../fab/cli_commands}/user_command.py | 0 airflow/auth/managers/fab/fab_auth_manager.py | 28 +++++++++---------- 3 files changed, 14 insertions(+), 14 deletions(-) rename airflow/{cli/commands => auth/managers/fab/cli_commands}/role_command.py (100%) rename airflow/{cli/commands => auth/managers/fab/cli_commands}/user_command.py (100%) diff --git a/airflow/cli/commands/role_command.py b/airflow/auth/managers/fab/cli_commands/role_command.py similarity index 100% rename from airflow/cli/commands/role_command.py rename to airflow/auth/managers/fab/cli_commands/role_command.py diff --git a/airflow/cli/commands/user_command.py b/airflow/auth/managers/fab/cli_commands/user_command.py similarity index 100% rename from airflow/cli/commands/user_command.py rename to airflow/auth/managers/fab/cli_commands/user_command.py diff --git a/airflow/auth/managers/fab/fab_auth_manager.py b/airflow/auth/managers/fab/fab_auth_manager.py index f5e6444c5cfe4..97255530c563b 100644 --- a/airflow/auth/managers/fab/fab_auth_manager.py +++ b/airflow/auth/managers/fab/fab_auth_manager.py @@ -103,13 +103,13 @@ ActionCommand( name="list", help="List users", - func=lazy_load_command("airflow.cli.commands.user_command.users_list"), + func=lazy_load_command("airflow.auth.managers.fab.cli_commands.user_command.users_list"), args=(ARG_OUTPUT, ARG_VERBOSE), ), ActionCommand( name="create", help="Create a user", - func=lazy_load_command("airflow.cli.commands.user_command.users_create"), + func=lazy_load_command("airflow.auth.managers.fab.cli_commands.user_command.users_create"), args=( ARG_ROLE, ARG_USERNAME, @@ -135,31 +135,31 @@ ActionCommand( name="delete", help="Delete a user", - func=lazy_load_command("airflow.cli.commands.user_command.users_delete"), + func=lazy_load_command("airflow.auth.managers.fab.cli_commands.user_command.users_delete"), args=(ARG_USERNAME_OPTIONAL, ARG_EMAIL_OPTIONAL, ARG_VERBOSE), ), ActionCommand( name="add-role", help="Add role to a user", - func=lazy_load_command("airflow.cli.commands.user_command.add_role"), + func=lazy_load_command("airflow.auth.managers.fab.cli_commands.user_command.add_role"), args=(ARG_USERNAME_OPTIONAL, ARG_EMAIL_OPTIONAL, ARG_ROLE, ARG_VERBOSE), ), ActionCommand( name="remove-role", help="Remove role from a user", - func=lazy_load_command("airflow.cli.commands.user_command.remove_role"), + func=lazy_load_command("airflow.auth.managers.fab.cli_commands.user_command.remove_role"), args=(ARG_USERNAME_OPTIONAL, ARG_EMAIL_OPTIONAL, ARG_ROLE, ARG_VERBOSE), ), ActionCommand( name="import", help="Import users", - func=lazy_load_command("airflow.cli.commands.user_command.users_import"), + func=lazy_load_command("airflow.auth.managers.fab.cli_commands.user_command.users_import"), args=(ARG_USER_IMPORT, ARG_VERBOSE), ), ActionCommand( name="export", help="Export all users", - func=lazy_load_command("airflow.cli.commands.user_command.users_export"), + func=lazy_load_command("airflow.auth.managers.fab.cli_commands.user_command.users_export"), args=(ARG_USER_EXPORT, ARG_VERBOSE), ), ) @@ -167,43 +167,43 @@ ActionCommand( name="list", help="List roles", - func=lazy_load_command("airflow.cli.commands.role_command.roles_list"), + func=lazy_load_command("airflow.auth.managers.fab.cli_commands.role_command.roles_list"), args=(ARG_PERMISSIONS, ARG_OUTPUT, ARG_VERBOSE), ), ActionCommand( name="create", help="Create role", - func=lazy_load_command("airflow.cli.commands.role_command.roles_create"), + func=lazy_load_command("airflow.auth.managers.fab.cli_commands.role_command.roles_create"), args=(ARG_ROLES, ARG_VERBOSE), ), ActionCommand( name="delete", help="Delete role", - func=lazy_load_command("airflow.cli.commands.role_command.roles_delete"), + func=lazy_load_command("airflow.auth.managers.fab.cli_commands.role_command.roles_delete"), args=(ARG_ROLES, ARG_VERBOSE), ), ActionCommand( name="add-perms", help="Add roles permissions", - func=lazy_load_command("airflow.cli.commands.role_command.roles_add_perms"), + func=lazy_load_command("airflow.auth.managers.fab.cli_commands.role_command.roles_add_perms"), args=(ARG_ROLES, ARG_ROLE_RESOURCE, ARG_ROLE_ACTION_REQUIRED, ARG_VERBOSE), ), ActionCommand( name="del-perms", help="Delete roles permissions", - func=lazy_load_command("airflow.cli.commands.role_command.roles_del_perms"), + func=lazy_load_command("airflow.auth.managers.fab.cli_commands.role_command.roles_del_perms"), args=(ARG_ROLES, ARG_ROLE_RESOURCE, ARG_ROLE_ACTION, ARG_VERBOSE), ), ActionCommand( name="export", help="Export roles (without permissions) from db to JSON file", - func=lazy_load_command("airflow.cli.commands.role_command.roles_export"), + func=lazy_load_command("airflow.auth.managers.fab.cli_commands.role_command.roles_export"), args=(ARG_ROLE_EXPORT, ARG_ROLE_EXPORT_FMT, ARG_VERBOSE), ), ActionCommand( name="import", help="Import roles (without permissions) from JSON file to db", - func=lazy_load_command("airflow.cli.commands.role_command.roles_import"), + func=lazy_load_command("airflow.auth.managers.fab.cli_commands.role_command.roles_import"), args=(ARG_ROLE_IMPORT, ARG_VERBOSE), ), ) From 8379ffecb7400794abfbf964dd512f6262675cf4 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Rapha=C3=ABl=20Vandon?= Date: Tue, 15 Aug 2023 08:36:37 -0700 Subject: [PATCH 03/11] move helper --- .../managers/fab/cli_commands/role_command.py | 12 +----------- .../managers/fab/cli_commands/user_command.py | 15 +-------------- .../managers/fab/cli_commands/utils.py} | 0 3 files changed, 2 insertions(+), 25 deletions(-) rename airflow/{utils/cli_app_builder.py => auth/managers/fab/cli_commands/utils.py} (100%) diff --git a/airflow/auth/managers/fab/cli_commands/role_command.py b/airflow/auth/managers/fab/cli_commands/role_command.py index a582b33195320..34ea8fb9d30e0 100644 --- a/airflow/auth/managers/fab/cli_commands/role_command.py +++ b/airflow/auth/managers/fab/cli_commands/role_command.py @@ -23,6 +23,7 @@ import json import os +from airflow.auth.managers.fab.cli_commands.utils import get_application_builder from airflow.auth.managers.fab.models import Action, Permission, Resource, Role from airflow.cli.simple_table import AirflowConsole from airflow.utils import cli as cli_utils @@ -35,8 +36,6 @@ @providers_configuration_loaded def roles_list(args): """List all existing roles.""" - from airflow.utils.cli_app_builder import get_application_builder - with get_application_builder() as appbuilder: roles = appbuilder.sm.get_all_roles() @@ -63,8 +62,6 @@ def roles_list(args): @providers_configuration_loaded def roles_create(args): """Create new empty role in DB.""" - from airflow.utils.cli_app_builder import get_application_builder - with get_application_builder() as appbuilder: for role_name in args.role: appbuilder.sm.add_role(role_name) @@ -76,8 +73,6 @@ def roles_create(args): @providers_configuration_loaded def roles_delete(args): """Delete role in DB.""" - from airflow.utils.cli_app_builder import get_application_builder - with get_application_builder() as appbuilder: for role_name in args.role: role = appbuilder.sm.find_role(role_name) @@ -90,8 +85,6 @@ def roles_delete(args): def __roles_add_or_remove_permissions(args): - from airflow.utils.cli_app_builder import get_application_builder - with get_application_builder() as appbuilder: is_add: bool = args.subcommand.startswith("add") @@ -165,8 +158,6 @@ def roles_export(args): Note, this function does not export the permissions associated for each role. Strictly, it exports the role names into the passed role json file. """ - from airflow.utils.cli_app_builder import get_application_builder - 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] @@ -196,7 +187,6 @@ def roles_import(args): except ValueError as e: print(f"File '{json_file}' is not a valid JSON file. Error: {e}") exit(1) - from airflow.utils.cli_app_builder import get_application_builder with get_application_builder() as appbuilder: existing_roles = [role.name for role in appbuilder.sm.get_all_roles()] diff --git a/airflow/auth/managers/fab/cli_commands/user_command.py b/airflow/auth/managers/fab/cli_commands/user_command.py index bc982719c94df..84e6318e40537 100644 --- a/airflow/auth/managers/fab/cli_commands/user_command.py +++ b/airflow/auth/managers/fab/cli_commands/user_command.py @@ -29,6 +29,7 @@ from marshmallow import Schema, fields, validate from marshmallow.exceptions import ValidationError +from airflow.auth.managers.fab.cli_commands.utils import get_application_builder from airflow.cli.simple_table import AirflowConsole from airflow.utils import cli as cli_utils from airflow.utils.cli import suppress_logs_and_warning @@ -50,8 +51,6 @@ class UserSchema(Schema): @providers_configuration_loaded def users_list(args): """List users at the command line.""" - from airflow.utils.cli_app_builder import get_application_builder - with get_application_builder() as appbuilder: users = appbuilder.sm.get_all_users() fields = ["id", "username", "email", "first_name", "last_name", "roles"] @@ -65,8 +64,6 @@ def users_list(args): @providers_configuration_loaded def users_create(args): """Create new user in the DB.""" - from airflow.utils.cli_app_builder import get_application_builder - with get_application_builder() as appbuilder: role = appbuilder.sm.find_role(args.role) if not role: @@ -101,8 +98,6 @@ def _find_user(args): if args.username and args.email: raise SystemExit("Conflicting args: must supply either --username or --email, but not both") - from airflow.utils.cli_app_builder import get_application_builder - with get_application_builder() as appbuilder: user = appbuilder.sm.find_user(username=args.username, email=args.email) if not user: @@ -119,8 +114,6 @@ def users_delete(args): # Clear the associated user roles first. user.roles.clear() - from airflow.utils.cli_app_builder import get_application_builder - with get_application_builder() as appbuilder: if appbuilder.sm.del_register_user(user): print(f'User "{user.username}" deleted') @@ -134,8 +127,6 @@ def users_manage_role(args, remove=False): """Delete or appends user roles.""" user = _find_user(args) - from airflow.utils.cli_app_builder import get_application_builder - with get_application_builder() as appbuilder: role = appbuilder.sm.find_role(args.role) if not role: @@ -161,8 +152,6 @@ def users_manage_role(args, remove=False): @providers_configuration_loaded def users_export(args): """Export all users to the json file.""" - from airflow.utils.cli_app_builder import get_application_builder - with get_application_builder() as appbuilder: users = appbuilder.sm.get_all_users() fields = ["id", "username", "email", "first_name", "last_name", "roles"] @@ -211,8 +200,6 @@ def users_import(args): def _import_users(users_list: list[dict[str, Any]]): - from airflow.utils.cli_app_builder import get_application_builder - with get_application_builder() as appbuilder: users_created = [] users_updated = [] diff --git a/airflow/utils/cli_app_builder.py b/airflow/auth/managers/fab/cli_commands/utils.py similarity index 100% rename from airflow/utils/cli_app_builder.py rename to airflow/auth/managers/fab/cli_commands/utils.py From 2724f38b274025f7906310f7b60e618920427882 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Rapha=C3=ABl=20Vandon?= Date: Thu, 17 Aug 2023 10:42:31 -0700 Subject: [PATCH 04/11] move sync perm command + some minor changes --- airflow/auth/managers/base_auth_manager.py | 5 ++--- .../fab/cli_commands}/sync_perm_command.py | 2 +- airflow/auth/managers/fab/fab_auth_manager.py | 15 +++++++++++++-- airflow/cli/cli_config.py | 11 ----------- airflow/cli/commands/standalone_command.py | 2 +- .../providers/celery/executors/celery_executor.py | 2 +- 6 files changed, 18 insertions(+), 19 deletions(-) rename airflow/{cli/commands => auth/managers/fab/cli_commands}/sync_perm_command.py (94%) diff --git a/airflow/auth/managers/base_auth_manager.py b/airflow/auth/managers/base_auth_manager.py index 445b5b0245327..53805a5fb1e86 100644 --- a/airflow/auth/managers/base_auth_manager.py +++ b/airflow/auth/managers/base_auth_manager.py @@ -21,11 +21,11 @@ from typing import TYPE_CHECKING from airflow.auth.managers.models.base_user import BaseUser -from airflow.cli.cli_config import GroupCommand from airflow.exceptions import AirflowException from airflow.utils.log.logging_mixin import LoggingMixin if TYPE_CHECKING: + from airflow.cli.cli_config import CLICommand from airflow.www.security import AirflowSecurityManager @@ -39,8 +39,7 @@ class BaseAuthManager(LoggingMixin): def __init__(self): self._security_manager: AirflowSecurityManager | None = None - @staticmethod - def get_cli_commands() -> list[GroupCommand]: + def get_cli_commands(self) -> list[CLICommand]: """Vends CLI commands to be included in Airflow CLI. Override this method to expose commands via Airflow CLI to manage this auth manager. diff --git a/airflow/cli/commands/sync_perm_command.py b/airflow/auth/managers/fab/cli_commands/sync_perm_command.py similarity index 94% rename from airflow/cli/commands/sync_perm_command.py rename to airflow/auth/managers/fab/cli_commands/sync_perm_command.py index 4d4e280637f9c..14b6e58bbb08a 100644 --- a/airflow/cli/commands/sync_perm_command.py +++ b/airflow/auth/managers/fab/cli_commands/sync_perm_command.py @@ -26,7 +26,7 @@ @providers_configuration_loaded def sync_perm(args): """Update permissions for existing roles and DAGs.""" - from airflow.utils.cli_app_builder import get_application_builder + from airflow.auth.managers.fab.cli_commands.utils import get_application_builder with get_application_builder() as appbuilder: print("Updating actions and resources for all existing roles") diff --git a/airflow/auth/managers/fab/fab_auth_manager.py b/airflow/auth/managers/fab/fab_auth_manager.py index 97255530c563b..fc21b76941724 100644 --- a/airflow/auth/managers/fab/fab_auth_manager.py +++ b/airflow/auth/managers/fab/fab_auth_manager.py @@ -31,6 +31,7 @@ ARG_VERBOSE, ActionCommand, Arg, + CLICommand, GroupCommand, lazy_load_command, ) @@ -99,6 +100,11 @@ action="store_true", ) +# sync-perm +ARG_INCLUDE_DAGS = Arg( + ("--include-dags",), help="If passed, DAG specific permissions will also be synced.", action="store_true" +) + USERS_COMMANDS = ( ActionCommand( name="list", @@ -216,8 +222,7 @@ class FabAuthManager(BaseAuthManager): This auth manager is responsible for providing a backward compatible user management experience to users. """ - @staticmethod - def get_cli_commands() -> list[GroupCommand]: + def get_cli_commands(self) -> list[CLICommand]: return [ GroupCommand( name="users", @@ -229,6 +234,12 @@ def get_cli_commands() -> list[GroupCommand]: help="Manage roles", subcommands=ROLES_COMMANDS, ), + ActionCommand( + name="sync-perm", + help="Update permissions for existing roles and optionally DAGs", + func=lazy_load_command("airflow.auth.managers.fab.cli_commands.sync_perm_command.sync_perm"), + args=(ARG_INCLUDE_DAGS, ARG_VERBOSE), + ), ] def get_user_name(self) -> str: diff --git a/airflow/cli/cli_config.py b/airflow/cli/cli_config.py index 13cf8f018ed90..a2b2eaaae8e7b 100644 --- a/airflow/cli/cli_config.py +++ b/airflow/cli/cli_config.py @@ -960,11 +960,6 @@ def string_lower_type(val): help="If passed, this command will be successful even if multiple matching alive jobs are found.", ) -# sync-perm -ARG_INCLUDE_DAGS = Arg( - ("--include-dags",), help="If passed, DAG specific permissions will also be synced.", action="store_true" -) - # triggerer ARG_CAPACITY = Arg( ("--capacity",), @@ -1998,12 +1993,6 @@ class GroupCommand(NamedTuple): help="Display providers", subcommands=PROVIDERS_COMMANDS, ), - ActionCommand( - name="sync-perm", - help="Update permissions for existing roles and optionally DAGs", - func=lazy_load_command("airflow.cli.commands.sync_perm_command.sync_perm"), - args=(ARG_INCLUDE_DAGS, ARG_VERBOSE), - ), ActionCommand( name="rotate-fernet-key", func=lazy_load_command("airflow.cli.commands.rotate_fernet_key_command.rotate_fernet_key"), diff --git a/airflow/cli/commands/standalone_command.py b/airflow/cli/commands/standalone_command.py index ceae1c6dcaec6..0beacb71d159d 100644 --- a/airflow/cli/commands/standalone_command.py +++ b/airflow/cli/commands/standalone_command.py @@ -182,7 +182,7 @@ def initialize_database(self): # server. Thus, we make a random password and store it in AIRFLOW_HOME, # with the reasoning that if you can read that directory, you can see # the database credentials anyway. - from airflow.utils.cli_app_builder import get_application_builder + from airflow.auth.managers.fab.cli_commands.utils import get_application_builder with get_application_builder() as appbuilder: user_exists = appbuilder.sm.find_user("admin") diff --git a/airflow/providers/celery/executors/celery_executor.py b/airflow/providers/celery/executors/celery_executor.py index 8bdff2a25e727..cc1b6e8122744 100644 --- a/airflow/providers/celery/executors/celery_executor.py +++ b/airflow/providers/celery/executors/celery_executor.py @@ -37,7 +37,6 @@ try: from airflow.cli.cli_config import ( - ARG_AUTOSCALE, ARG_DAEMON, ARG_LOG_FILE, ARG_PID, @@ -143,6 +142,7 @@ def __getattr__(name): ) # worker cli args +ARG_AUTOSCALE = Arg(("-a", "--autoscale"), help="Minimum and Maximum number of worker to autoscale") ARG_QUEUES = Arg( ("-q", "--queues"), help="Comma delimited list of queues to serve", From a40803024e7115f580a60f72c304f1b1ce69e379 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Rapha=C3=ABl=20Vandon?= Date: Thu, 17 Aug 2023 10:55:05 -0700 Subject: [PATCH 05/11] extract command definition to a different file for readability --- .../managers/fab/cli_commands/definition.py | 220 ++++++++++++++++++ airflow/auth/managers/fab/fab_auth_manager.py | 198 +--------------- 2 files changed, 227 insertions(+), 191 deletions(-) create mode 100644 airflow/auth/managers/fab/cli_commands/definition.py diff --git a/airflow/auth/managers/fab/cli_commands/definition.py b/airflow/auth/managers/fab/cli_commands/definition.py new file mode 100644 index 0000000000000..478f6d8d309e0 --- /dev/null +++ b/airflow/auth/managers/fab/cli_commands/definition.py @@ -0,0 +1,220 @@ +# 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 textwrap + +from airflow.cli.cli_config import ( + ARG_OUTPUT, + ARG_VERBOSE, + ActionCommand, + Arg, + lazy_load_command, +) + +############ +# # ARGS # # +############ + +# users +ARG_USERNAME = Arg(("-u", "--username"), help="Username of the user", required=True, type=str) +ARG_USERNAME_OPTIONAL = Arg(("-u", "--username"), help="Username of the user", type=str) +ARG_FIRSTNAME = Arg(("-f", "--firstname"), help="First name of the user", required=True, type=str) +ARG_LASTNAME = Arg(("-l", "--lastname"), help="Last name of the user", required=True, type=str) +ARG_ROLE = Arg( + ("-r", "--role"), + help="Role of the user. Existing roles include Admin, User, Op, Viewer, and Public", + required=True, + type=str, +) +ARG_EMAIL = Arg(("-e", "--email"), help="Email of the user", required=True, type=str) +ARG_EMAIL_OPTIONAL = Arg(("-e", "--email"), help="Email of the user", type=str) +ARG_PASSWORD = Arg( + ("-p", "--password"), + help="Password of the user, required to create a user without --use-random-password", + type=str, +) +ARG_USE_RANDOM_PASSWORD = Arg( + ("--use-random-password",), + help="Do not prompt for password. Use random string instead." + " Required to create a user without --password ", + default=False, + action="store_true", +) +ARG_USER_IMPORT = Arg( + ("import",), + metavar="FILEPATH", + help="Import users from JSON file. Example format::\n" + + textwrap.indent( + textwrap.dedent( + """ + [ + { + "email": "foo@bar.org", + "firstname": "Jon", + "lastname": "Doe", + "roles": ["Public"], + "username": "jondoe" + } + ]""" + ), + " " * 4, + ), +) +ARG_USER_EXPORT = Arg(("export",), metavar="FILEPATH", help="Export all users to JSON file") + +# roles +ARG_CREATE_ROLE = Arg(("-c", "--create"), help="Create a new role", action="store_true") +ARG_LIST_ROLES = Arg(("-l", "--list"), help="List roles", action="store_true") +ARG_ROLES = Arg(("role",), help="The name of a role", nargs="*") +ARG_PERMISSIONS = Arg(("-p", "--permission"), help="Show role permissions", action="store_true") +ARG_ROLE_RESOURCE = Arg(("-r", "--resource"), help="The name of permissions", nargs="*", required=True) +ARG_ROLE_ACTION = Arg(("-a", "--action"), help="The action of permissions", nargs="*") +ARG_ROLE_ACTION_REQUIRED = Arg(("-a", "--action"), help="The action of permissions", nargs="*", required=True) + +ARG_ROLE_IMPORT = Arg(("file",), help="Import roles from JSON file", nargs=None) +ARG_ROLE_EXPORT = Arg(("file",), help="Export all roles to JSON file", nargs=None) +ARG_ROLE_EXPORT_FMT = Arg( + ("-p", "--pretty"), + help="Format output JSON file by sorting role names and indenting by 4 spaces", + action="store_true", +) + +# sync-perm +ARG_INCLUDE_DAGS = Arg( + ("--include-dags",), help="If passed, DAG specific permissions will also be synced.", action="store_true" +) + +################ +# # COMMANDS # # +################ + +USERS_COMMANDS = ( + ActionCommand( + name="list", + help="List users", + func=lazy_load_command("airflow.auth.managers.fab.cli_commands.user_command.users_list"), + args=(ARG_OUTPUT, ARG_VERBOSE), + ), + ActionCommand( + name="create", + help="Create a user", + func=lazy_load_command("airflow.auth.managers.fab.cli_commands.user_command.users_create"), + args=( + ARG_ROLE, + ARG_USERNAME, + ARG_EMAIL, + ARG_FIRSTNAME, + ARG_LASTNAME, + ARG_PASSWORD, + ARG_USE_RANDOM_PASSWORD, + ARG_VERBOSE, + ), + epilog=( + "examples:\n" + 'To create an user with "Admin" role and username equals to "admin", run:\n' + "\n" + " $ airflow users create \\\n" + " --username admin \\\n" + " --firstname FIRST_NAME \\\n" + " --lastname LAST_NAME \\\n" + " --role Admin \\\n" + " --email admin@example.org" + ), + ), + ActionCommand( + name="delete", + help="Delete a user", + func=lazy_load_command("airflow.auth.managers.fab.cli_commands.user_command.users_delete"), + args=(ARG_USERNAME_OPTIONAL, ARG_EMAIL_OPTIONAL, ARG_VERBOSE), + ), + ActionCommand( + name="add-role", + help="Add role to a user", + func=lazy_load_command("airflow.auth.managers.fab.cli_commands.user_command.add_role"), + args=(ARG_USERNAME_OPTIONAL, ARG_EMAIL_OPTIONAL, ARG_ROLE, ARG_VERBOSE), + ), + ActionCommand( + name="remove-role", + help="Remove role from a user", + func=lazy_load_command("airflow.auth.managers.fab.cli_commands.user_command.remove_role"), + args=(ARG_USERNAME_OPTIONAL, ARG_EMAIL_OPTIONAL, ARG_ROLE, ARG_VERBOSE), + ), + ActionCommand( + name="import", + help="Import users", + func=lazy_load_command("airflow.auth.managers.fab.cli_commands.user_command.users_import"), + args=(ARG_USER_IMPORT, ARG_VERBOSE), + ), + ActionCommand( + name="export", + help="Export all users", + func=lazy_load_command("airflow.auth.managers.fab.cli_commands.user_command.users_export"), + args=(ARG_USER_EXPORT, ARG_VERBOSE), + ), +) +ROLES_COMMANDS = ( + ActionCommand( + name="list", + help="List roles", + func=lazy_load_command("airflow.auth.managers.fab.cli_commands.role_command.roles_list"), + args=(ARG_PERMISSIONS, ARG_OUTPUT, ARG_VERBOSE), + ), + ActionCommand( + name="create", + help="Create role", + func=lazy_load_command("airflow.auth.managers.fab.cli_commands.role_command.roles_create"), + args=(ARG_ROLES, ARG_VERBOSE), + ), + ActionCommand( + name="delete", + help="Delete role", + func=lazy_load_command("airflow.auth.managers.fab.cli_commands.role_command.roles_delete"), + args=(ARG_ROLES, ARG_VERBOSE), + ), + ActionCommand( + name="add-perms", + help="Add roles permissions", + func=lazy_load_command("airflow.auth.managers.fab.cli_commands.role_command.roles_add_perms"), + args=(ARG_ROLES, ARG_ROLE_RESOURCE, ARG_ROLE_ACTION_REQUIRED, ARG_VERBOSE), + ), + ActionCommand( + name="del-perms", + help="Delete roles permissions", + func=lazy_load_command("airflow.auth.managers.fab.cli_commands.role_command.roles_del_perms"), + args=(ARG_ROLES, ARG_ROLE_RESOURCE, ARG_ROLE_ACTION, ARG_VERBOSE), + ), + ActionCommand( + name="export", + help="Export roles (without permissions) from db to JSON file", + func=lazy_load_command("airflow.auth.managers.fab.cli_commands.role_command.roles_export"), + args=(ARG_ROLE_EXPORT, ARG_ROLE_EXPORT_FMT, ARG_VERBOSE), + ), + ActionCommand( + name="import", + help="Import roles (without permissions) from JSON file to db", + func=lazy_load_command("airflow.auth.managers.fab.cli_commands.role_command.roles_import"), + args=(ARG_ROLE_IMPORT, ARG_VERBOSE), + ), +) + +SYNC_PERM_COMMAND = ActionCommand( + name="sync-perm", + help="Update permissions for existing roles and optionally DAGs", + func=lazy_load_command("airflow.auth.managers.fab.cli_commands.sync_perm_command.sync_perm"), + args=(ARG_INCLUDE_DAGS, ARG_VERBOSE), +) diff --git a/airflow/auth/managers/fab/fab_auth_manager.py b/airflow/auth/managers/fab/fab_auth_manager.py index fc21b76941724..c1b1958f78b05 100644 --- a/airflow/auth/managers/fab/fab_auth_manager.py +++ b/airflow/auth/managers/fab/fab_auth_manager.py @@ -17,201 +17,21 @@ # under the License. from __future__ import annotations -import textwrap - from flask import url_for from flask_login import current_user from airflow import AirflowException from airflow.auth.managers.base_auth_manager import BaseAuthManager +from airflow.auth.managers.fab.cli_commands.definition import ( + ROLES_COMMANDS, + SYNC_PERM_COMMAND, + USERS_COMMANDS, +) from airflow.auth.managers.fab.models import User from airflow.auth.managers.fab.security_manager.override import FabAirflowSecurityManagerOverride from airflow.cli.cli_config import ( - ARG_OUTPUT, - ARG_VERBOSE, - ActionCommand, - Arg, CLICommand, GroupCommand, - lazy_load_command, -) - -# users -ARG_USERNAME = Arg(("-u", "--username"), help="Username of the user", required=True, type=str) -ARG_USERNAME_OPTIONAL = Arg(("-u", "--username"), help="Username of the user", type=str) -ARG_FIRSTNAME = Arg(("-f", "--firstname"), help="First name of the user", required=True, type=str) -ARG_LASTNAME = Arg(("-l", "--lastname"), help="Last name of the user", required=True, type=str) -ARG_ROLE = Arg( - ("-r", "--role"), - help="Role of the user. Existing roles include Admin, User, Op, Viewer, and Public", - required=True, - type=str, -) -ARG_EMAIL = Arg(("-e", "--email"), help="Email of the user", required=True, type=str) -ARG_EMAIL_OPTIONAL = Arg(("-e", "--email"), help="Email of the user", type=str) -ARG_PASSWORD = Arg( - ("-p", "--password"), - help="Password of the user, required to create a user without --use-random-password", - type=str, -) -ARG_USE_RANDOM_PASSWORD = Arg( - ("--use-random-password",), - help="Do not prompt for password. Use random string instead." - " Required to create a user without --password ", - default=False, - action="store_true", -) -ARG_USER_IMPORT = Arg( - ("import",), - metavar="FILEPATH", - help="Import users from JSON file. Example format::\n" - + textwrap.indent( - textwrap.dedent( - """ - [ - { - "email": "foo@bar.org", - "firstname": "Jon", - "lastname": "Doe", - "roles": ["Public"], - "username": "jondoe" - } - ]""" - ), - " " * 4, - ), -) -ARG_USER_EXPORT = Arg(("export",), metavar="FILEPATH", help="Export all users to JSON file") - -# roles -ARG_CREATE_ROLE = Arg(("-c", "--create"), help="Create a new role", action="store_true") -ARG_LIST_ROLES = Arg(("-l", "--list"), help="List roles", action="store_true") -ARG_ROLES = Arg(("role",), help="The name of a role", nargs="*") -ARG_PERMISSIONS = Arg(("-p", "--permission"), help="Show role permissions", action="store_true") -ARG_ROLE_RESOURCE = Arg(("-r", "--resource"), help="The name of permissions", nargs="*", required=True) -ARG_ROLE_ACTION = Arg(("-a", "--action"), help="The action of permissions", nargs="*") -ARG_ROLE_ACTION_REQUIRED = Arg(("-a", "--action"), help="The action of permissions", nargs="*", required=True) - -ARG_ROLE_IMPORT = Arg(("file",), help="Import roles from JSON file", nargs=None) -ARG_ROLE_EXPORT = Arg(("file",), help="Export all roles to JSON file", nargs=None) -ARG_ROLE_EXPORT_FMT = Arg( - ("-p", "--pretty"), - help="Format output JSON file by sorting role names and indenting by 4 spaces", - action="store_true", -) - -# sync-perm -ARG_INCLUDE_DAGS = Arg( - ("--include-dags",), help="If passed, DAG specific permissions will also be synced.", action="store_true" -) - -USERS_COMMANDS = ( - ActionCommand( - name="list", - help="List users", - func=lazy_load_command("airflow.auth.managers.fab.cli_commands.user_command.users_list"), - args=(ARG_OUTPUT, ARG_VERBOSE), - ), - ActionCommand( - name="create", - help="Create a user", - func=lazy_load_command("airflow.auth.managers.fab.cli_commands.user_command.users_create"), - args=( - ARG_ROLE, - ARG_USERNAME, - ARG_EMAIL, - ARG_FIRSTNAME, - ARG_LASTNAME, - ARG_PASSWORD, - ARG_USE_RANDOM_PASSWORD, - ARG_VERBOSE, - ), - epilog=( - "examples:\n" - 'To create an user with "Admin" role and username equals to "admin", run:\n' - "\n" - " $ airflow users create \\\n" - " --username admin \\\n" - " --firstname FIRST_NAME \\\n" - " --lastname LAST_NAME \\\n" - " --role Admin \\\n" - " --email admin@example.org" - ), - ), - ActionCommand( - name="delete", - help="Delete a user", - func=lazy_load_command("airflow.auth.managers.fab.cli_commands.user_command.users_delete"), - args=(ARG_USERNAME_OPTIONAL, ARG_EMAIL_OPTIONAL, ARG_VERBOSE), - ), - ActionCommand( - name="add-role", - help="Add role to a user", - func=lazy_load_command("airflow.auth.managers.fab.cli_commands.user_command.add_role"), - args=(ARG_USERNAME_OPTIONAL, ARG_EMAIL_OPTIONAL, ARG_ROLE, ARG_VERBOSE), - ), - ActionCommand( - name="remove-role", - help="Remove role from a user", - func=lazy_load_command("airflow.auth.managers.fab.cli_commands.user_command.remove_role"), - args=(ARG_USERNAME_OPTIONAL, ARG_EMAIL_OPTIONAL, ARG_ROLE, ARG_VERBOSE), - ), - ActionCommand( - name="import", - help="Import users", - func=lazy_load_command("airflow.auth.managers.fab.cli_commands.user_command.users_import"), - args=(ARG_USER_IMPORT, ARG_VERBOSE), - ), - ActionCommand( - name="export", - help="Export all users", - func=lazy_load_command("airflow.auth.managers.fab.cli_commands.user_command.users_export"), - args=(ARG_USER_EXPORT, ARG_VERBOSE), - ), -) -ROLES_COMMANDS = ( - ActionCommand( - name="list", - help="List roles", - func=lazy_load_command("airflow.auth.managers.fab.cli_commands.role_command.roles_list"), - args=(ARG_PERMISSIONS, ARG_OUTPUT, ARG_VERBOSE), - ), - ActionCommand( - name="create", - help="Create role", - func=lazy_load_command("airflow.auth.managers.fab.cli_commands.role_command.roles_create"), - args=(ARG_ROLES, ARG_VERBOSE), - ), - ActionCommand( - name="delete", - help="Delete role", - func=lazy_load_command("airflow.auth.managers.fab.cli_commands.role_command.roles_delete"), - args=(ARG_ROLES, ARG_VERBOSE), - ), - ActionCommand( - name="add-perms", - help="Add roles permissions", - func=lazy_load_command("airflow.auth.managers.fab.cli_commands.role_command.roles_add_perms"), - args=(ARG_ROLES, ARG_ROLE_RESOURCE, ARG_ROLE_ACTION_REQUIRED, ARG_VERBOSE), - ), - ActionCommand( - name="del-perms", - help="Delete roles permissions", - func=lazy_load_command("airflow.auth.managers.fab.cli_commands.role_command.roles_del_perms"), - args=(ARG_ROLES, ARG_ROLE_RESOURCE, ARG_ROLE_ACTION, ARG_VERBOSE), - ), - ActionCommand( - name="export", - help="Export roles (without permissions) from db to JSON file", - func=lazy_load_command("airflow.auth.managers.fab.cli_commands.role_command.roles_export"), - args=(ARG_ROLE_EXPORT, ARG_ROLE_EXPORT_FMT, ARG_VERBOSE), - ), - ActionCommand( - name="import", - help="Import roles (without permissions) from JSON file to db", - func=lazy_load_command("airflow.auth.managers.fab.cli_commands.role_command.roles_import"), - args=(ARG_ROLE_IMPORT, ARG_VERBOSE), - ), ) @@ -223,6 +43,7 @@ class FabAuthManager(BaseAuthManager): """ def get_cli_commands(self) -> list[CLICommand]: + """Vends CLI commands to be included in Airflow CLI.""" return [ GroupCommand( name="users", @@ -234,12 +55,7 @@ def get_cli_commands(self) -> list[CLICommand]: help="Manage roles", subcommands=ROLES_COMMANDS, ), - ActionCommand( - name="sync-perm", - help="Update permissions for existing roles and optionally DAGs", - func=lazy_load_command("airflow.auth.managers.fab.cli_commands.sync_perm_command.sync_perm"), - args=(ARG_INCLUDE_DAGS, ARG_VERBOSE), - ), + SYNC_PERM_COMMAND, # not in a command group ] def get_user_name(self) -> str: From 3207a6e0046ad0bebd1115bc51ecb02eadedc7b6 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Rapha=C3=ABl=20Vandon?= Date: Thu, 17 Aug 2023 11:40:13 -0700 Subject: [PATCH 06/11] move tests --- .../managers/fab/cli_commands}/test_role_command.py | 10 +++++----- .../fab/cli_commands}/test_sync_perm_command.py | 6 +++--- .../managers/fab/cli_commands}/test_user_command.py | 8 ++++---- 3 files changed, 12 insertions(+), 12 deletions(-) rename tests/{cli/commands => auth/managers/fab/cli_commands}/test_role_command.py (96%) rename tests/{cli/commands => auth/managers/fab/cli_commands}/test_sync_perm_command.py (89%) rename tests/{cli/commands => auth/managers/fab/cli_commands}/test_user_command.py (98%) diff --git a/tests/cli/commands/test_role_command.py b/tests/auth/managers/fab/cli_commands/test_role_command.py similarity index 96% rename from tests/cli/commands/test_role_command.py rename to tests/auth/managers/fab/cli_commands/test_role_command.py index 544d8e9560627..5259ca8482fdb 100644 --- a/tests/cli/commands/test_role_command.py +++ b/tests/auth/managers/fab/cli_commands/test_role_command.py @@ -23,10 +23,11 @@ import pytest +from airflow.auth.managers.fab.cli_commands import role_command +from airflow.auth.managers.fab.cli_commands.utils import get_application_builder from airflow.auth.managers.fab.models import Role -from airflow.cli.commands import role_command +from airflow.cli import cli_parser from airflow.security import permissions -from airflow.utils.cli_app_builder import get_application_builder TEST_USER1_EMAIL = "test-user1@example.com" TEST_USER2_EMAIL = "test-user2@example.com" @@ -34,9 +35,8 @@ class TestCliRoles: @pytest.fixture(autouse=True) - def _set_attrs(self, dagbag, parser): - self.dagbag = dagbag - self.parser = parser + def _set_attrs(self): + self.parser = cli_parser.get_parser() with get_application_builder() as appbuilder: self.appbuilder = appbuilder self.clear_roles_and_roles() diff --git a/tests/cli/commands/test_sync_perm_command.py b/tests/auth/managers/fab/cli_commands/test_sync_perm_command.py similarity index 89% rename from tests/cli/commands/test_sync_perm_command.py rename to tests/auth/managers/fab/cli_commands/test_sync_perm_command.py index d55d4e372498a..d59fb34a8ced6 100644 --- a/tests/cli/commands/test_sync_perm_command.py +++ b/tests/auth/managers/fab/cli_commands/test_sync_perm_command.py @@ -19,8 +19,8 @@ from unittest import mock +from airflow.auth.managers.fab.cli_commands import sync_perm_command from airflow.cli import cli_parser -from airflow.cli.commands import sync_perm_command class TestCliSyncPerm: @@ -28,7 +28,7 @@ class TestCliSyncPerm: def setup_class(cls): cls.parser = cli_parser.get_parser() - @mock.patch("airflow.utils.cli_app_builder.get_application_builder") + @mock.patch("airflow.auth.managers.fab.cli_commands.utils.get_application_builder") def test_cli_sync_perm(self, mock_get_application_builder): mock_appbuilder = mock.MagicMock() mock_get_application_builder.return_value.__enter__.return_value = mock_appbuilder @@ -40,7 +40,7 @@ def test_cli_sync_perm(self, mock_get_application_builder): mock_appbuilder.sm.sync_roles.assert_called_once_with() mock_appbuilder.sm.create_dag_specific_permissions.assert_not_called() - @mock.patch("airflow.utils.cli_app_builder.get_application_builder") + @mock.patch("airflow.auth.managers.fab.cli_commands.utils.get_application_builder") def test_cli_sync_perm_include_dags(self, mock_get_application_builder): mock_appbuilder = mock.MagicMock() mock_get_application_builder.return_value.__enter__.return_value = mock_appbuilder diff --git a/tests/cli/commands/test_user_command.py b/tests/auth/managers/fab/cli_commands/test_user_command.py similarity index 98% rename from tests/cli/commands/test_user_command.py rename to tests/auth/managers/fab/cli_commands/test_user_command.py index 32f2e5db8840c..4daf0e6e36d05 100644 --- a/tests/cli/commands/test_user_command.py +++ b/tests/auth/managers/fab/cli_commands/test_user_command.py @@ -25,7 +25,8 @@ import pytest -from airflow.cli.commands import user_command +from airflow.auth.managers.fab.cli_commands import user_command +from airflow.cli import cli_parser from tests.test_utils.api_connexion_utils import delete_users TEST_USER1_EMAIL = "test-user1@example.com" @@ -44,10 +45,9 @@ def _does_user_belong_to_role(appbuilder, email, rolename): class TestCliUsers: @pytest.fixture(autouse=True) - def _set_attrs(self, app, dagbag, parser): + def _set_attrs(self, app): self.app = app - self.dagbag = dagbag - self.parser = parser + self.parser = cli_parser.get_parser() self.appbuilder = self.app.appbuilder delete_users(app) yield From 572c3ce7a9d0dca9287f72f1df9df6c3f8a42379 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Rapha=C3=ABl=20Vandon?= Date: Thu, 17 Aug 2023 11:43:48 -0700 Subject: [PATCH 07/11] forgot the init file --- tests/auth/managers/fab/cli_commands/__init__.py | 16 ++++++++++++++++ 1 file changed, 16 insertions(+) create mode 100644 tests/auth/managers/fab/cli_commands/__init__.py diff --git a/tests/auth/managers/fab/cli_commands/__init__.py b/tests/auth/managers/fab/cli_commands/__init__.py new file mode 100644 index 0000000000000..13a83393a9124 --- /dev/null +++ b/tests/auth/managers/fab/cli_commands/__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. From 069138de58261117d2d709d905a5ea59f3c2b422 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Rapha=C3=ABl=20Vandon?= Date: Tue, 22 Aug 2023 10:30:37 -0700 Subject: [PATCH 08/11] static method to get cli commands + import optims --- airflow/auth/managers/base_auth_manager.py | 5 +++-- airflow/auth/managers/fab/fab_auth_manager.py | 15 ++++++++++---- airflow/cli/cli_parser.py | 4 ++-- airflow/www/extensions/init_auth_manager.py | 20 ++++++++++++++----- 4 files changed, 31 insertions(+), 13 deletions(-) diff --git a/airflow/auth/managers/base_auth_manager.py b/airflow/auth/managers/base_auth_manager.py index 53805a5fb1e86..a512804b4ca94 100644 --- a/airflow/auth/managers/base_auth_manager.py +++ b/airflow/auth/managers/base_auth_manager.py @@ -20,11 +20,11 @@ from abc import abstractmethod from typing import TYPE_CHECKING -from airflow.auth.managers.models.base_user import BaseUser from airflow.exceptions import AirflowException from airflow.utils.log.logging_mixin import LoggingMixin if TYPE_CHECKING: + from airflow.auth.managers.models.base_user import BaseUser from airflow.cli.cli_config import CLICommand from airflow.www.security import AirflowSecurityManager @@ -39,7 +39,8 @@ class BaseAuthManager(LoggingMixin): def __init__(self): self._security_manager: AirflowSecurityManager | None = None - def get_cli_commands(self) -> list[CLICommand]: + @staticmethod + def get_cli_commands() -> list[CLICommand]: """Vends CLI commands to be included in Airflow CLI. Override this method to expose commands via Airflow CLI to manage this auth manager. diff --git a/airflow/auth/managers/fab/fab_auth_manager.py b/airflow/auth/managers/fab/fab_auth_manager.py index c1b1958f78b05..718ac774586f0 100644 --- a/airflow/auth/managers/fab/fab_auth_manager.py +++ b/airflow/auth/managers/fab/fab_auth_manager.py @@ -17,8 +17,9 @@ # under the License. from __future__ import annotations +from typing import TYPE_CHECKING + from flask import url_for -from flask_login import current_user from airflow import AirflowException from airflow.auth.managers.base_auth_manager import BaseAuthManager @@ -27,13 +28,14 @@ SYNC_PERM_COMMAND, USERS_COMMANDS, ) -from airflow.auth.managers.fab.models import User -from airflow.auth.managers.fab.security_manager.override import FabAirflowSecurityManagerOverride from airflow.cli.cli_config import ( CLICommand, GroupCommand, ) +if TYPE_CHECKING: + from airflow.auth.managers.fab.models import User + class FabAuthManager(BaseAuthManager): """ @@ -42,7 +44,8 @@ class FabAuthManager(BaseAuthManager): This auth manager is responsible for providing a backward compatible user management experience to users. """ - def get_cli_commands(self) -> list[CLICommand]: + @staticmethod + def get_cli_commands() -> list[CLICommand]: """Vends CLI commands to be included in Airflow CLI.""" return [ GroupCommand( @@ -72,6 +75,8 @@ def get_user_name(self) -> str: def get_user(self) -> User: """Return the user associated to the user in session.""" + from flask_login import current_user + return current_user def get_user_id(self) -> str: @@ -84,6 +89,8 @@ def is_logged_in(self) -> bool: def get_security_manager_override_class(self) -> type: """Return the security manager override.""" + from airflow.auth.managers.fab.security_manager.override import FabAirflowSecurityManagerOverride + return FabAirflowSecurityManagerOverride def get_url_login(self, **kwargs) -> str: diff --git a/airflow/cli/cli_parser.py b/airflow/cli/cli_parser.py index 6ffbd9dbd674f..132560014c67e 100644 --- a/airflow/cli/cli_parser.py +++ b/airflow/cli/cli_parser.py @@ -44,7 +44,7 @@ from airflow.exceptions import AirflowException from airflow.executors.executor_loader import ExecutorLoader from airflow.utils.helpers import partition -from airflow.www.extensions.init_auth_manager import get_auth_manager +from airflow.www.extensions.init_auth_manager import get_auth_manager_cls airflow_commands = core_commands @@ -64,7 +64,7 @@ # other commands. try: - auth_mgr = get_auth_manager() + auth_mgr = get_auth_manager_cls() airflow_commands.extend(auth_mgr.get_cli_commands()) except Exception: log.exception("cannot load CLI commands from auth manager") diff --git a/airflow/www/extensions/init_auth_manager.py b/airflow/www/extensions/init_auth_manager.py index a53fdf304befc..24ae020862dc9 100644 --- a/airflow/www/extensions/init_auth_manager.py +++ b/airflow/www/extensions/init_auth_manager.py @@ -26,12 +26,10 @@ from airflow.auth.managers.base_auth_manager import BaseAuthManager -@cache -def get_auth_manager() -> BaseAuthManager: - """ - Initialize auth manager. +def get_auth_manager_cls() -> type[BaseAuthManager]: + """Returns just the auth manager class without initializing it. - Import the user manager class, instantiate it and return it. + Useful to save execution time if only static methods need to be called. """ auth_manager_cls = conf.getimport(section="core", key="auth_manager") @@ -41,4 +39,16 @@ def get_auth_manager() -> BaseAuthManager: "Please specify one using section/key [core/auth_manager]." ) + return auth_manager_cls + + +@cache +def get_auth_manager() -> BaseAuthManager: + """ + Initialize auth manager. + + Import the user manager class, instantiate it and return it. + """ + auth_manager_cls = get_auth_manager_cls() + return auth_manager_cls() From 9822bb36604c236a72db6bc1e9ffca1fc36468e6 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Rapha=C3=ABl=20Vandon?= Date: Tue, 22 Aug 2023 15:26:05 -0700 Subject: [PATCH 09/11] =?UTF-8?q?add=20missing=20init=20file=20?= =?UTF-8?q?=F0=9F=A4=A6?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- .../auth/managers/fab/cli_commands/__init__.py | 17 +++++++++++++++++ 1 file changed, 17 insertions(+) create mode 100644 airflow/auth/managers/fab/cli_commands/__init__.py diff --git a/airflow/auth/managers/fab/cli_commands/__init__.py b/airflow/auth/managers/fab/cli_commands/__init__.py new file mode 100644 index 0000000000000..217e5db960782 --- /dev/null +++ b/airflow/auth/managers/fab/cli_commands/__init__.py @@ -0,0 +1,17 @@ +# +# 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 18dd23a43375f3716aa81ead0df40874d6fe5283 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Rapha=C3=ABl=20Vandon?= Date: Wed, 23 Aug 2023 12:41:26 -0700 Subject: [PATCH 10/11] move import in method to save exec time --- airflow/auth/managers/fab/fab_auth_manager.py | 8 ++++++-- 1 file changed, 6 insertions(+), 2 deletions(-) diff --git a/airflow/auth/managers/fab/fab_auth_manager.py b/airflow/auth/managers/fab/fab_auth_manager.py index 718ac774586f0..64fa7784746f6 100644 --- a/airflow/auth/managers/fab/fab_auth_manager.py +++ b/airflow/auth/managers/fab/fab_auth_manager.py @@ -19,8 +19,6 @@ from typing import TYPE_CHECKING -from flask import url_for - from airflow import AirflowException from airflow.auth.managers.base_auth_manager import BaseAuthManager from airflow.auth.managers.fab.cli_commands.definition import ( @@ -95,6 +93,8 @@ def get_security_manager_override_class(self) -> type: def get_url_login(self, **kwargs) -> str: """Return the login page url.""" + from flask import url_for + if not self.security_manager.auth_view: raise AirflowException("`auth_view` not defined in the security manager.") if "next_url" in kwargs and kwargs["next_url"]: @@ -104,12 +104,16 @@ def get_url_login(self, **kwargs) -> str: def get_url_logout(self): """Return the logout page url.""" + from flask import url_for + if not self.security_manager.auth_view: raise AirflowException("`auth_view` not defined in the security manager.") return url_for(f"{self.security_manager.auth_view.endpoint}.logout") def get_url_user_profile(self) -> str | None: """Return the url to a page displaying info about the current user.""" + from flask import url_for + if not self.security_manager.user_view: return None return url_for(f"{self.security_manager.user_view.endpoint}.userinfo") From 33f8e22dd3ce7758e418994dec5b3318e05cdf8d Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Rapha=C3=ABl=20Vandon?= Date: Wed, 23 Aug 2023 13:48:17 -0700 Subject: [PATCH 11/11] fix the mocking --- airflow/auth/managers/fab/fab_auth_manager.py | 20 +++++++++---------- .../managers/fab/test_fab_auth_manager.py | 8 ++++---- 2 files changed, 14 insertions(+), 14 deletions(-) diff --git a/airflow/auth/managers/fab/fab_auth_manager.py b/airflow/auth/managers/fab/fab_auth_manager.py index 64fa7784746f6..def0590b1f5c9 100644 --- a/airflow/auth/managers/fab/fab_auth_manager.py +++ b/airflow/auth/managers/fab/fab_auth_manager.py @@ -91,29 +91,29 @@ def get_security_manager_override_class(self) -> type: return FabAirflowSecurityManagerOverride - def get_url_login(self, **kwargs) -> str: - """Return the login page url.""" + def url_for(self, *args, **kwargs): + """Wrapper to allow mocking without having to import at the top of the file.""" from flask import url_for + return url_for(*args, **kwargs) + + def get_url_login(self, **kwargs) -> str: + """Return the login page url.""" if not self.security_manager.auth_view: raise AirflowException("`auth_view` not defined in the security manager.") if "next_url" in kwargs and kwargs["next_url"]: - return url_for(f"{self.security_manager.auth_view.endpoint}.login", next=kwargs["next_url"]) + return self.url_for(f"{self.security_manager.auth_view.endpoint}.login", next=kwargs["next_url"]) else: - return url_for(f"{self.security_manager.auth_view.endpoint}.login") + return self.url_for(f"{self.security_manager.auth_view.endpoint}.login") def get_url_logout(self): """Return the logout page url.""" - from flask import url_for - if not self.security_manager.auth_view: raise AirflowException("`auth_view` not defined in the security manager.") - return url_for(f"{self.security_manager.auth_view.endpoint}.logout") + return self.url_for(f"{self.security_manager.auth_view.endpoint}.logout") def get_url_user_profile(self) -> str | None: """Return the url to a page displaying info about the current user.""" - from flask import url_for - if not self.security_manager.user_view: return None - return url_for(f"{self.security_manager.user_view.endpoint}.userinfo") + return self.url_for(f"{self.security_manager.user_view.endpoint}.userinfo") diff --git a/tests/auth/managers/fab/test_fab_auth_manager.py b/tests/auth/managers/fab/test_fab_auth_manager.py index cb4ef8ebc2113..66c279510e55e 100644 --- a/tests/auth/managers/fab/test_fab_auth_manager.py +++ b/tests/auth/managers/fab/test_fab_auth_manager.py @@ -85,14 +85,14 @@ def test_get_url_login_when_auth_view_not_defined(self, auth_manager): with pytest.raises(AirflowException, match="`auth_view` not defined in the security manager."): auth_manager.get_url_login() - @mock.patch("airflow.auth.managers.fab.fab_auth_manager.url_for") + @mock.patch.object(FabAuthManager, "url_for") def test_get_url_login(self, mock_url_for, auth_manager): auth_manager.security_manager.auth_view = Mock() auth_manager.security_manager.auth_view.endpoint = "test_endpoint" auth_manager.get_url_login() mock_url_for.assert_called_once_with("test_endpoint.login") - @mock.patch("airflow.auth.managers.fab.fab_auth_manager.url_for") + @mock.patch.object(FabAuthManager, "url_for") def test_get_url_login_with_next(self, mock_url_for, auth_manager): auth_manager.security_manager.auth_view = Mock() auth_manager.security_manager.auth_view.endpoint = "test_endpoint" @@ -103,7 +103,7 @@ def test_get_url_logout_when_auth_view_not_defined(self, auth_manager): with pytest.raises(AirflowException, match="`auth_view` not defined in the security manager."): auth_manager.get_url_logout() - @mock.patch("airflow.auth.managers.fab.fab_auth_manager.url_for") + @mock.patch.object(FabAuthManager, "url_for") def test_get_url_logout(self, mock_url_for, auth_manager): auth_manager.security_manager.auth_view = Mock() auth_manager.security_manager.auth_view.endpoint = "test_endpoint" @@ -113,7 +113,7 @@ def test_get_url_logout(self, mock_url_for, auth_manager): def test_get_url_user_profile_when_auth_view_not_defined(self, auth_manager): assert auth_manager.get_url_user_profile() is None - @mock.patch("airflow.auth.managers.fab.fab_auth_manager.url_for") + @mock.patch.object(FabAuthManager, "url_for") def test_get_url_user_profile(self, mock_url_for, auth_manager): expected_url = "test_url" mock_url_for.return_value = expected_url