diff --git a/airflow/auth/managers/base_auth_manager.py b/airflow/auth/managers/base_auth_manager.py index bcc5e13892e2b..a512804b4ca94 100644 --- a/airflow/auth/managers/base_auth_manager.py +++ b/airflow/auth/managers/base_auth_manager.py @@ -20,11 +20,12 @@ 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 @@ -38,6 +39,14 @@ class BaseAuthManager(LoggingMixin): def __init__(self): self._security_manager: AirflowSecurityManager | None = None + @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. + """ + return [] + @abstractmethod def get_user_name(self) -> str: """Return the username associated to the user in session.""" 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. 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/cli/commands/role_command.py b/airflow/auth/managers/fab/cli_commands/role_command.py similarity index 94% rename from airflow/cli/commands/role_command.py rename to airflow/auth/managers/fab/cli_commands/role_command.py index a582b33195320..34ea8fb9d30e0 100644 --- a/airflow/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/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/cli/commands/user_command.py b/airflow/auth/managers/fab/cli_commands/user_command.py similarity index 95% rename from airflow/cli/commands/user_command.py rename to airflow/auth/managers/fab/cli_commands/user_command.py index bc982719c94df..84e6318e40537 100644 --- a/airflow/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 diff --git a/airflow/auth/managers/fab/fab_auth_manager.py b/airflow/auth/managers/fab/fab_auth_manager.py index 848f5ff188d6c..def0590b1f5c9 100644 --- a/airflow/auth/managers/fab/fab_auth_manager.py +++ b/airflow/auth/managers/fab/fab_auth_manager.py @@ -17,13 +17,22 @@ # under the License. from __future__ import annotations -from flask import url_for -from flask_login import current_user +from typing import TYPE_CHECKING from airflow import AirflowException 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.auth.managers.fab.cli_commands.definition import ( + ROLES_COMMANDS, + SYNC_PERM_COMMAND, + USERS_COMMANDS, +) +from airflow.cli.cli_config import ( + CLICommand, + GroupCommand, +) + +if TYPE_CHECKING: + from airflow.auth.managers.fab.models import User class FabAuthManager(BaseAuthManager): @@ -33,6 +42,23 @@ class FabAuthManager(BaseAuthManager): This auth manager is responsible for providing a backward compatible user management experience to users. """ + @staticmethod + def get_cli_commands() -> list[CLICommand]: + """Vends CLI commands to be included in Airflow CLI.""" + return [ + GroupCommand( + name="users", + help="Manage users", + subcommands=USERS_COMMANDS, + ), + GroupCommand( + name="roles", + help="Manage roles", + subcommands=ROLES_COMMANDS, + ), + SYNC_PERM_COMMAND, # not in a command group + ] + def get_user_name(self) -> str: """ Return the username associated to the user in session. @@ -47,6 +73,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: @@ -59,25 +87,33 @@ 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 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.""" 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.""" 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/airflow/cli/cli_config.py b/airflow/cli/cli_config.py index de1d43be6ffe5..c06bd32d96cbb 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",), @@ -1024,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",), @@ -1839,115 +1770,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,22 +1993,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", - 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/cli_parser.py b/airflow/cli/cli_parser.py index ee1e60f1b210a..07c7695f2dddb 100644 --- a/airflow/cli/cli_parser.py +++ b/airflow/cli/cli_parser.py @@ -46,6 +46,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_cls airflow_commands = core_commands.copy() # make a copy to prevent bad interactions in tests @@ -64,6 +65,13 @@ # Do not re-raise the exception since we want the CLI to still function for # other commands. +try: + 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") + # do not re-raise for the same reason as above + ALL_COMMANDS_DICT: dict[str, CLICommand] = {sp.name: sp for sp in airflow_commands} 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", 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() 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. 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 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