[AIRFLOW-3752] Add/remove user from role via CLI (#4572)
* [AIRFLOW-3752] Add/remove user from role via the CLI Update the `users` subcommand to enable 2 new actions: - `--add-role`: Make the user a member of the given role - `--remove-role`: Remove the user's membership in the given role For installations that use an external identity provider (e.g., Google OAuth) the username is typically a long ID string. For the sake of convenience, we allow the CLI operator to reference the target user via either their `username` or their `email` (but not both). * Update argparse spec Accidentally left off this update to the argparse spec in the last commit. * Add unit tests * Fix lint failures
This commit is contained in:
Andrew Stahlman
Tao Feng
Родитель
89dd877d93
Коммит
9f8ca32da4
10
UPDATING.md
10
UPDATING.md
|
|
@ -153,6 +153,16 @@ To delete a user:
|
|||
airflow users --delete --username jondoe
|
||||
```
|
||||
|
||||
To add a user to a role:
|
||||
```bash
|
||||
airflow users --add-role --username jondoe --role Public
|
||||
```
|
||||
|
||||
To remove a user from a role:
|
||||
```bash
|
||||
airflow users --remove-role --username jondoe --role Public
|
||||
```
|
||||
|
||||
### User model changes
|
||||
This patch changes the `User.superuser` field from a hardcoded boolean to a `Boolean()` database column. `User.superuser` will default to `False`, which means that this privilege will have to be granted manually to any users that may require it.
|
||||
|
||||
|
|
|
|||
|
|
@ -1339,7 +1339,7 @@ def users(args):
|
|||
|
||||
return
|
||||
|
||||
if args.create:
|
||||
elif args.create:
|
||||
fields = {
|
||||
'role': args.role,
|
||||
'username': args.username,
|
||||
|
|
@ -1377,7 +1377,7 @@ def users(args):
|
|||
else:
|
||||
raise SystemExit('Failed to create user.')
|
||||
|
||||
if args.delete:
|
||||
elif args.delete:
|
||||
if not args.username:
|
||||
raise SystemExit('Required arguments are missing: username')
|
||||
|
||||
|
|
@ -1394,6 +1394,54 @@ def users(args):
|
|||
else:
|
||||
raise SystemExit('Failed to delete user.')
|
||||
|
||||
elif args.add_role or args.remove_role:
|
||||
if args.add_role and args.remove_role:
|
||||
raise SystemExit('Conflicting args: --add-role and --remove-role'
|
||||
' are mutually exclusive')
|
||||
|
||||
if not args.username and not args.email:
|
||||
raise SystemExit('Missing args: must supply one of --username or --email')
|
||||
|
||||
if args.username and args.email:
|
||||
raise SystemExit('Conflicting args: must supply either --username'
|
||||
' or --email, but not both')
|
||||
if not args.role:
|
||||
raise SystemExit('Required args are missing: role')
|
||||
|
||||
appbuilder = cached_appbuilder()
|
||||
user = (appbuilder.sm.find_user(username=args.username) or
|
||||
appbuilder.sm.find_user(email=args.email))
|
||||
if not user:
|
||||
raise SystemExit('User "{}" does not exist'.format(
|
||||
args.username or args.email))
|
||||
|
||||
role = appbuilder.sm.find_role(args.role)
|
||||
if not role:
|
||||
raise SystemExit('"{}" is not a valid role.'.format(args.role))
|
||||
|
||||
if args.remove_role:
|
||||
if role in user.roles:
|
||||
user.roles = [r for r in user.roles if r != role]
|
||||
appbuilder.sm.update_user(user)
|
||||
print('User "{}" removed from role "{}".'.format(
|
||||
user,
|
||||
args.role))
|
||||
else:
|
||||
raise SystemExit('User "{}" is not a member of role "{}".'.format(
|
||||
user,
|
||||
args.role))
|
||||
elif args.add_role:
|
||||
if role in user.roles:
|
||||
raise SystemExit('User "{}" is already a member of role "{}".'.format(
|
||||
user,
|
||||
args.role))
|
||||
else:
|
||||
user.roles.append(role)
|
||||
appbuilder.sm.update_user(user)
|
||||
print('User "{}" added to role "{}".'.format(
|
||||
user,
|
||||
args.role))
|
||||
|
||||
|
||||
@cli_utils.action_logging
|
||||
def list_dag_runs(args, dag=None):
|
||||
|
|
@ -1903,6 +1951,14 @@ class CLIFactory(object):
|
|||
('-d', '--delete'),
|
||||
help='Delete a user',
|
||||
action='store_true'),
|
||||
'add_role': Arg(
|
||||
('--add-role',),
|
||||
help='Add user to a role',
|
||||
action='store_true'),
|
||||
'remove_role': Arg(
|
||||
('--remove-role',),
|
||||
help='Remove user from a role',
|
||||
action='store_true'),
|
||||
'autoscale': Arg(
|
||||
('-a', '--autoscale'),
|
||||
help="Minimum and Maximum number of worker to autoscale"),
|
||||
|
|
@ -2065,8 +2121,9 @@ class CLIFactory(object):
|
|||
'conn_id', 'conn_uri', 'conn_extra') + tuple(alternative_conn_specs),
|
||||
}, {
|
||||
'func': users,
|
||||
'help': "List/Create/Delete users",
|
||||
'help': "List/Create/Delete/Update users",
|
||||
'args': ('list_users', 'create_user', 'delete_user',
|
||||
'add_role', 'remove_role',
|
||||
'username', 'email', 'firstname', 'lastname', 'role',
|
||||
'password', 'use_random_password'),
|
||||
},
|
||||
|
|
|
|||
|
|
@ -30,4 +30,4 @@ and click ``List Roles`` in the new UI.
|
|||
|
||||
|
||||
The image shows a role which could only write to example_python_operator is created.
|
||||
And we could assign the given role to a new user using ``airflow users --role`` cli command.
|
||||
And we could assign the given role to a new user using ``airflow users --add-role`` cli command.
|
||||
|
|
|
|||
|
|
@ -1061,6 +1061,8 @@ class CoreTest(unittest.TestCase):
|
|||
|
||||
class CliTests(unittest.TestCase):
|
||||
|
||||
TEST_USER_EMAIL = 'test-user@example.com'
|
||||
|
||||
@classmethod
|
||||
def setUpClass(cls):
|
||||
super(CliTests, cls).setUpClass()
|
||||
|
|
@ -1080,6 +1082,9 @@ class CliTests(unittest.TestCase):
|
|||
|
||||
def tearDown(self):
|
||||
self._cleanup(session=self.session)
|
||||
test_user = self.appbuilder.sm.find_user(email=CliTests.TEST_USER_EMAIL)
|
||||
if test_user:
|
||||
self.appbuilder.sm.del_register_user(test_user)
|
||||
super(CliTests, self).tearDown()
|
||||
|
||||
@staticmethod
|
||||
|
|
@ -1148,6 +1153,64 @@ class CliTests(unittest.TestCase):
|
|||
for i in range(0, 3):
|
||||
self.assertIn('user{}'.format(i), stdout)
|
||||
|
||||
def _does_user_belong_to_role(self, email, rolename):
|
||||
user = self.appbuilder.sm.find_user(email=email)
|
||||
role = self.appbuilder.sm.find_role(rolename)
|
||||
if user and role:
|
||||
return role in user.roles
|
||||
|
||||
return False
|
||||
|
||||
def test_cli_add_user_role(self):
|
||||
args = self.parser.parse_args([
|
||||
'users', '-c', '--username', 'test4', '--lastname', 'doe',
|
||||
'--firstname', 'jon',
|
||||
'--email', self.TEST_USER_EMAIL, '--role', 'Viewer', '--use_random_password'
|
||||
])
|
||||
cli.users(args)
|
||||
|
||||
self.assertFalse(
|
||||
self._does_user_belong_to_role(email=self.TEST_USER_EMAIL,
|
||||
rolename='Op'),
|
||||
"User should not yet be a member of role 'Op'"
|
||||
)
|
||||
|
||||
args = self.parser.parse_args([
|
||||
'users', '--add-role', '--username', 'test4', '--role', 'Op'
|
||||
])
|
||||
cli.users(args)
|
||||
|
||||
self.assertTrue(
|
||||
self._does_user_belong_to_role(email=self.TEST_USER_EMAIL,
|
||||
rolename='Op'),
|
||||
"User should have been added to role 'Op'"
|
||||
)
|
||||
|
||||
def test_cli_remove_user_role(self):
|
||||
args = self.parser.parse_args([
|
||||
'users', '-c', '--username', 'test4', '--lastname', 'doe',
|
||||
'--firstname', 'jon',
|
||||
'--email', self.TEST_USER_EMAIL, '--role', 'Viewer', '--use_random_password'
|
||||
])
|
||||
cli.users(args)
|
||||
|
||||
self.assertTrue(
|
||||
self._does_user_belong_to_role(email=self.TEST_USER_EMAIL,
|
||||
rolename='Viewer'),
|
||||
"User should have been created with role 'Viewer'"
|
||||
)
|
||||
|
||||
args = self.parser.parse_args([
|
||||
'users', '--remove-role', '--username', 'test4', '--role', 'Viewer'
|
||||
])
|
||||
cli.users(args)
|
||||
|
||||
self.assertFalse(
|
||||
self._does_user_belong_to_role(email=self.TEST_USER_EMAIL,
|
||||
rolename='Viewer'),
|
||||
"User should have been removed from role 'Viewer'"
|
||||
)
|
||||
|
||||
def test_cli_sync_perm(self):
|
||||
# test whether sync_perm cli will throw exceptions or not
|
||||
args = self.parser.parse_args([
|
||||
|
|
|
|||
Загрузка…
Ссылка в новой задаче