[AIRFLOW-3051] Change CLI to make users ops similar to connections

The ability to manipulate users from the  command line is a bit clunky.  Currently 'airflow create_user' and 'airflow delete_user' and 'airflow list_users'.  It seems that these ought to be made more like connections, so that it becomes 'airflow users list ...', 'airflow users delete ...' and 'airflow users create ...'

Author: kaxil

airflow/bin/cli.py

@@ -1347,74 +1347,75 @@ def kerberos(args):  # noqa
 
 
 @cli_utils.action_logging
-def create_user(args):
-    fields = {
-        'role': args.role,
-        'username': args.username,
-        'email': args.email,
-        'firstname': args.firstname,
-        'lastname': args.lastname,
-    }
-    empty_fields = [k for k, v in fields.items() if not v]
-    if empty_fields:
-        raise SystemExit('Required arguments are missing: {}.'.format(
-            ', '.join(empty_fields)))
-
-    appbuilder = cached_appbuilder()
-    role = appbuilder.sm.find_role(args.role)
-    if not role:
-        raise SystemExit('{} is not a valid role.'.format(args.role))
-
-    if args.use_random_password:
-        password = ''.join(random.choice(string.printable) for _ in range(16))
-    elif args.password:
-        password = args.password
-    else:
-        password = getpass.getpass('Password:')
-        password_confirmation = getpass.getpass('Repeat for confirmation:')
-        if password != password_confirmation:
-            raise SystemExit('Passwords did not match!')
+def users(args):
+    if args.list:
+
+        appbuilder = cached_appbuilder()
+        users = appbuilder.sm.get_all_users()
+        fields = ['id', 'username', 'email', 'first_name', 'last_name', 'roles']
+        users = [[user.__getattribute__(field) for field in fields] for user in users]
+        msg = tabulate(users, [field.capitalize().replace('_', ' ') for field in fields],
+                       tablefmt="fancy_grid")
+        if sys.version_info[0] < 3:
+            msg = msg.encode('utf-8')
+        print(msg)
 
-    if appbuilder.sm.find_user(args.username):
-        print('{} already exist in the db'.format(args.username))
         return
-    user = appbuilder.sm.add_user(args.username, args.firstname, args.lastname,
-                                  args.email, role, password)
-    if user:
-        print('{} user {} created.'.format(args.role, args.username))
-    else:
-        raise SystemExit('Failed to create user.')
 
+    if args.create:
+        fields = {
+            'role': args.role,
+            'username': args.username,
+            'email': args.email,
+            'firstname': args.firstname,
+            'lastname': args.lastname,
+        }
+        empty_fields = [k for k, v in fields.items() if not v]
+        if empty_fields:
+            raise SystemExit('Required arguments are missing: {}.'.format(
+                ', '.join(empty_fields)))
 
-@cli_utils.action_logging
-def delete_user(args):
-    if not args.username:
-        raise SystemExit('Required arguments are missing: username')
+        appbuilder = cached_appbuilder()
+        role = appbuilder.sm.find_role(args.role)
+        if not role:
+            raise SystemExit('{} is not a valid role.'.format(args.role))
+
+        if args.use_random_password:
+            password = ''.join(random.choice(string.printable) for _ in range(16))
+        elif args.password:
+            password = args.password
+        else:
+            password = getpass.getpass('Password:')
+            password_confirmation = getpass.getpass('Repeat for confirmation:')
+            if password != password_confirmation:
+                raise SystemExit('Passwords did not match!')
 
-    appbuilder = cached_appbuilder()
+        if appbuilder.sm.find_user(args.username):
+            print('{} already exist in the db'.format(args.username))
+            return
+        user = appbuilder.sm.add_user(args.username, args.firstname, args.lastname,
+                                      args.email, role, password)
+        if user:
+            print('{} user {} created.'.format(args.role, args.username))
+        else:
+            raise SystemExit('Failed to create user.')
 
-    try:
-        u = next(u for u in appbuilder.sm.get_all_users() if u.username == args.username)
-    except StopIteration:
-        raise SystemExit('{} is not a valid user.'.format(args.username))
+    if args.delete:
+        if not args.username:
+            raise SystemExit('Required arguments are missing: username')
 
-    if appbuilder.sm.del_register_user(u):
-        print('User {} deleted.'.format(args.username))
-    else:
-        raise SystemExit('Failed to delete user.')
+        appbuilder = cached_appbuilder()
 
+        try:
+            u = next(u for u in appbuilder.sm.get_all_users()
+                     if u.username == args.username)
+        except StopIteration:
+            raise SystemExit('{} is not a valid user.'.format(args.username))
 
-@cli_utils.action_logging
-def list_users(args):
-    appbuilder = cached_appbuilder()
-    users = appbuilder.sm.get_all_users()
-    fields = ['id', 'username', 'email', 'first_name', 'last_name', 'roles']
-    users = [[user.__getattribute__(field) for field in fields] for user in users]
-    msg = tabulate(users, [field.capitalize().replace('_', ' ') for field in fields],
-                   tablefmt="fancy_grid")
-    if sys.version_info[0] < 3:
-        msg = msg.encode('utf-8')
-    print(msg)
+        if appbuilder.sm.del_register_user(u):
+            print('User {} deleted.'.format(args.username))
+        else:
+            raise SystemExit('Failed to delete user.')
 
 
 @cli_utils.action_logging
@@ -1527,10 +1528,6 @@ class CLIFactory(object):
             "Do not prompt to confirm reset. Use with care!",
             "store_true",
             default=False),
-        'username': Arg(
-            ('-u', '--username',),
-            help='Username of the user',
-            type=str),
 
         # list_dag_runs
         'no_backfill': Arg(
@@ -1886,33 +1883,52 @@ class CLIFactory(object):
             ('--conn_extra',),
             help='Connection `Extra` field, optional when adding a connection',
             type=str),
-        # create_user
-        'role': Arg(
-            ('-r', '--role',),
-            help='Role of the user. Existing roles include Admin, '
-                 'User, Op, Viewer, and Public',
+        # users
+        'username': Arg(
+            ('--username',),
+            help='Username of the user, required to create/delete a user',
             type=str),
         'firstname': Arg(
-            ('-f', '--firstname',),
-            help='First name of the user',
+            ('--firstname',),
+            help='First name of the user, required to create a user',
             type=str),
         'lastname': Arg(
-            ('-l', '--lastname',),
-            help='Last name of the user',
+            ('--lastname',),
+            help='Last name of the user, required to create a user',
+            type=str),
+        'role': Arg(
+            ('--role',),
+            help='Role of the user. Existing roles include Admin, '
+                 'User, Op, Viewer, and Public. Required to create a user',
             type=str),
         'email': Arg(
-            ('-e', '--email',),
-            help='Email of the user',
+            ('--email',),
+            help='Email of the user, required to create a user',
             type=str),
         'password': Arg(
-            ('-p', '--password',),
-            help='Password of the user',
+            ('--password',),
+            help='Password of the user, required to create a user '
+                 'without --use_random_password',
             type=str),
         'use_random_password': Arg(
             ('--use_random_password',),
-            help='Do not prompt for password.  Use random string instead',
+            help='Do not prompt for password. Use random string instead.'
+                 ' Required to create a user without --password ',
             default=False,
             action='store_true'),
+        'list_users': Arg(
+            ('-l', '--list'),
+            help='List all users',
+            action='store_true'),
+        'create_user': Arg(
+            ('-c', '--create'),
+            help='Create a user',
+            action='store_true'),
+        'delete_user': Arg(
+            ('-d', '--delete'),
+            help='Delete a user',
+            action='store_true'),
+
     }
     subparsers = (
         {
@@ -2070,18 +2086,11 @@ class CLIFactory(object):
             'args': ('list_connections', 'add_connection', 'delete_connection',
                      'conn_id', 'conn_uri', 'conn_extra') + tuple(alternative_conn_specs),
         }, {
-            'func': create_user,
-            'help': "Create an account for the Web UI (FAB-based)",
-            'args': ('role', 'username', 'email', 'firstname', 'lastname',
+            'func': users,
+            'help': "List/Create/Delete users",
+            'args': ('list_users', 'create_user', 'delete_user',
+                     'username', 'email', 'firstname', 'lastname', 'role',
                      'password', 'use_random_password'),
-        }, {
-            'func': delete_user,
-            'help': "Delete an account for the Web UI",
-            'args': ('username',),
-        }, {
-            'func': list_users,
-            'help': "List accounts for the Web UI",
-            'args': tuple(),
         },
         {
             'func': sync_perm,

tests/core.py

@@ -1089,40 +1089,44 @@ def test_cli_list_dag_runs(self):
 
     def test_cli_create_user_random_password(self):
         args = self.parser.parse_args([
-            'create_user', '-u', 'test1', '-l', 'doe', '-f', 'jon',
-            '-e', 'jdoe@foo.com', '-r', 'Viewer', '--use_random_password'
+            'users', '-c', '--username', 'test1', '--lastname', 'doe',
+            '--firstname', 'jon',
+            '--email', 'jdoe@foo.com', '--role', 'Viewer', '--use_random_password'
         ])
-        cli.create_user(args)
+        cli.users(args)
 
     def test_cli_create_user_supplied_password(self):
         args = self.parser.parse_args([
-            'create_user', '-u', 'test2', '-l', 'doe', '-f', 'jon',
-            '-e', 'jdoe@apache.org', '-r', 'Viewer', '-p', 'test'
+            'users', '-c', '--username', 'test2', '--lastname', 'doe',
+            '--firstname', 'jon',
+            '--email', 'jdoe@apache.org', '--role', 'Viewer', '--password', 'test'
         ])
-        cli.create_user(args)
+        cli.users(args)
 
     def test_cli_delete_user(self):
         args = self.parser.parse_args([
-            'create_user', '-u', 'test3', '-l', 'doe', '-f', 'jon',
-            '-e', 'jdoe@example.com', '-r', 'Viewer', '--use_random_password'
+            'users', '-c', '--username', 'test3', '--lastname', 'doe',
+            '--firstname', 'jon',
+            '--email', 'jdoe@example.com', '--role', 'Viewer', '--use_random_password'
         ])
-        cli.create_user(args)
+        cli.users(args)
         args = self.parser.parse_args([
-            'delete_user', '-u', 'test3',
+            'users', '-d', '--username', 'test3',
         ])
-        cli.delete_user(args)
+        cli.users(args)
 
     def test_cli_list_users(self):
         for i in range(0, 3):
             args = self.parser.parse_args([
-                'create_user', '-u', 'user{}'.format(i), '-l', 'doe', '-f', 'jon',
-                '-e', 'jdoe+{}@gmail.com'.format(i), '-r', 'Viewer',
+                'users', '-c', '--username', 'user{}'.format(i), '--lastname',
+                'doe', '--firstname', 'jon',
+                '--email', 'jdoe+{}@gmail.com'.format(i), '--role', 'Viewer',
                 '--use_random_password'
             ])
-            cli.create_user(args)
+            cli.users(args)
         with mock.patch('sys.stdout',
                         new_callable=six.StringIO) as mock_stdout:
-            cli.list_users(self.parser.parse_args(['list_users']))
+            cli.users(self.parser.parse_args(['users', '-l']))
             stdout = mock_stdout.getvalue()
         for i in range(0, 3):
             self.assertIn('user{}'.format(i), stdout)