diff --git a/UPDATING.md b/UPDATING.md index 34237108e4..4ebbd9d5d4 100644 --- a/UPDATING.md +++ b/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. diff --git a/airflow/bin/cli.py b/airflow/bin/cli.py index 9db5b6eebd..999474f76b 100644 --- a/airflow/bin/cli.py +++ b/airflow/bin/cli.py @@ -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'), }, diff --git a/docs/howto/add-new-role.rst b/docs/howto/add-new-role.rst index 2374facf19..8d22103185 100644 --- a/docs/howto/add-new-role.rst +++ b/docs/howto/add-new-role.rst @@ -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. diff --git a/tests/core.py b/tests/core.py index 16f48eb5c0..6624e87df9 100644 --- a/tests/core.py +++ b/tests/core.py @@ -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([