This is an automated email from the ASF dual-hosted git repository.
onikolas pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/airflow.git
The following commit(s) were added to refs/heads/main by this push:
new b1a3b42880 Make auth managers provide their own airflow CLI commands
(#33481)
b1a3b42880 is described below
commit b1a3b4288022c67db22cbc7d24b0c4b2b122453b
Author: Raphaël Vandon <[email protected]>
AuthorDate: Wed Aug 23 14:57:53 2023 -0700
Make auth managers provide their own airflow CLI commands (#33481)
* add a way for auth managers to define their own CLI commands
* extract command definition to a different file for readability
* static method to get cli commands + import optims
* move import in method to save exec time
---
airflow/auth/managers/base_auth_manager.py | 11 +-
.../managers/fab/cli_commands/__init__.py} | 29 +--
.../auth/managers/fab/cli_commands/definition.py | 220 +++++++++++++++++++++
.../managers/fab/cli_commands}/role_command.py | 12 +-
.../fab/cli_commands}/sync_perm_command.py | 2 +-
.../managers/fab/cli_commands}/user_command.py | 15 +-
.../managers/fab/cli_commands/utils.py} | 0
airflow/auth/managers/fab/fab_auth_manager.py | 52 ++++-
airflow/cli/cli_config.py | 206 +------------------
airflow/cli/cli_parser.py | 8 +
airflow/cli/commands/standalone_command.py | 2 +-
.../providers/celery/executors/celery_executor.py | 2 +-
airflow/www/extensions/init_auth_manager.py | 20 +-
.../auth/managers/fab/cli_commands/__init__.py | 28 ---
.../fab/cli_commands}/test_role_command.py | 10 +-
.../fab/cli_commands}/test_sync_perm_command.py | 6 +-
.../fab/cli_commands}/test_user_command.py | 8 +-
tests/auth/managers/fab/test_fab_auth_manager.py | 8 +-
18 files changed, 325 insertions(+), 314 deletions(-)
diff --git a/airflow/auth/managers/base_auth_manager.py
b/airflow/auth/managers/base_auth_manager.py
index bcc5e13892..a512804b4c 100644
--- a/airflow/auth/managers/base_auth_manager.py
+++ b/airflow/auth/managers/base_auth_manager.py
@@ -20,11 +20,12 @@ from __future__ import annotations
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/www/extensions/init_auth_manager.py
b/airflow/auth/managers/fab/cli_commands/__init__.py
similarity index 50%
copy from airflow/www/extensions/init_auth_manager.py
copy to airflow/auth/managers/fab/cli_commands/__init__.py
index a53fdf304b..217e5db960 100644
--- a/airflow/www/extensions/init_auth_manager.py
+++ b/airflow/auth/managers/fab/cli_commands/__init__.py
@@ -1,3 +1,4 @@
+#
# 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
@@ -14,31 +15,3 @@
# KIND, either express or implied. See the License for the
# specific language governing permissions and limitations
# under the License.
-from __future__ import annotations
-
-from typing import TYPE_CHECKING
-
-from airflow.compat.functools import cache
-from airflow.configuration import conf
-from airflow.exceptions import AirflowConfigException
-
-if TYPE_CHECKING:
- from airflow.auth.managers.base_auth_manager import BaseAuthManager
-
-
-@cache
-def get_auth_manager() -> BaseAuthManager:
- """
- Initialize auth manager.
-
- Import the user manager class, instantiate it and return it.
- """
- auth_manager_cls = conf.getimport(section="core", key="auth_manager")
-
- if not auth_manager_cls:
- raise AirflowConfigException(
- "No auth manager defined in the config. "
- "Please specify one using section/key [core/auth_manager]."
- )
-
- return auth_manager_cls()
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 0000000000..478f6d8d30
--- /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": "[email protected]",
+ "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 [email protected]"
+ ),
+ ),
+ 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 a582b33195..34ea8fb9d3 100644
--- a/airflow/cli/commands/role_command.py
+++ b/airflow/auth/managers/fab/cli_commands/role_command.py
@@ -23,6 +23,7 @@ import itertools
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 @@ from airflow.www.security import EXISTING_ROLES
@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 4d4e280637..14b6e58bbb 100644
--- a/airflow/cli/commands/sync_perm_command.py
+++ b/airflow/auth/managers/fab/cli_commands/sync_perm_command.py
@@ -26,7 +26,7 @@ from airflow.utils.providers_configuration_loader import
providers_configuration
@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 bc982719c9..84e6318e40 100644
--- a/airflow/cli/commands/user_command.py
+++ b/airflow/auth/managers/fab/cli_commands/user_command.py
@@ -29,6 +29,7 @@ import re2
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 848f5ff188..def0590b1f 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 @@ class FabAuthManager(BaseAuthManager):
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 @@ class FabAuthManager(BaseAuthManager):
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 de1d43be6f..c06bd32d96 100644
--- a/airflow/cli/cli_config.py
+++ b/airflow/cli/cli_config.py
@@ -229,6 +229,12 @@ ARG_REVISION_RANGE = Arg(
),
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 @@ ARG_FULL = Arg(
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": "[email protected]",
- "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 @@ ARG_ALLOW_MULTIPLE = Arg(
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 @@ PROVIDERS_COMMANDS = (
)
-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 [email protected]"
- ),
- ),
- 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 @@ core_commands: list[CLICommand] = [
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 ee1e60f1b2..07c7695f2d 100644
--- a/airflow/cli/cli_parser.py
+++ b/airflow/cli/cli_parser.py
@@ -46,6 +46,7 @@ from airflow.cli.utils import CliConflictError
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 @@ except Exception:
# 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 ceae1c6dca..0beacb71d1 100644
--- a/airflow/cli/commands/standalone_command.py
+++ b/airflow/cli/commands/standalone_command.py
@@ -182,7 +182,7 @@ class StandaloneCommand:
# 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 8bdff2a25e..cc1b6e8122 100644
--- a/airflow/providers/celery/executors/celery_executor.py
+++ b/airflow/providers/celery/executors/celery_executor.py
@@ -37,7 +37,6 @@ from celery import states as celery_states
try:
from airflow.cli.cli_config import (
- ARG_AUTOSCALE,
ARG_DAEMON,
ARG_LOG_FILE,
ARG_PID,
@@ -143,6 +142,7 @@ ARG_FLOWER_BASIC_AUTH = Arg(
)
# 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 a53fdf304b..24ae020862 100644
--- a/airflow/www/extensions/init_auth_manager.py
+++ b/airflow/www/extensions/init_auth_manager.py
@@ -26,12 +26,10 @@ if TYPE_CHECKING:
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/airflow/www/extensions/init_auth_manager.py
b/tests/auth/managers/fab/cli_commands/__init__.py
similarity index 50%
copy from airflow/www/extensions/init_auth_manager.py
copy to tests/auth/managers/fab/cli_commands/__init__.py
index a53fdf304b..13a83393a9 100644
--- a/airflow/www/extensions/init_auth_manager.py
+++ b/tests/auth/managers/fab/cli_commands/__init__.py
@@ -14,31 +14,3 @@
# KIND, either express or implied. See the License for the
# specific language governing permissions and limitations
# under the License.
-from __future__ import annotations
-
-from typing import TYPE_CHECKING
-
-from airflow.compat.functools import cache
-from airflow.configuration import conf
-from airflow.exceptions import AirflowConfigException
-
-if TYPE_CHECKING:
- from airflow.auth.managers.base_auth_manager import BaseAuthManager
-
-
-@cache
-def get_auth_manager() -> BaseAuthManager:
- """
- Initialize auth manager.
-
- Import the user manager class, instantiate it and return it.
- """
- auth_manager_cls = conf.getimport(section="core", key="auth_manager")
-
- if not auth_manager_cls:
- raise AirflowConfigException(
- "No auth manager defined in the config. "
- "Please specify one using section/key [core/auth_manager]."
- )
-
- return auth_manager_cls()
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 544d8e9560..5259ca8482 100644
--- a/tests/cli/commands/test_role_command.py
+++ b/tests/auth/managers/fab/cli_commands/test_role_command.py
@@ -23,10 +23,11 @@ from contextlib import redirect_stdout
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 = "[email protected]"
TEST_USER2_EMAIL = "[email protected]"
@@ -34,9 +35,8 @@ TEST_USER2_EMAIL = "[email protected]"
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 d55d4e3724..d59fb34a8c 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 __future__ import annotations
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 @@ class TestCliSyncPerm:
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 32f2e5db88..4daf0e6e36 100644
--- a/tests/cli/commands/test_user_command.py
+++ b/tests/auth/managers/fab/cli_commands/test_user_command.py
@@ -25,7 +25,8 @@ from contextlib import redirect_stdout
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 = "[email protected]"
@@ -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 cb4ef8ebc2..66c279510e 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 @@ class TestFabAuthManager:
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 @@ class TestFabAuthManager:
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 @@ class TestFabAuthManager:
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
