Skip to content

Commit e33407f

Browse files
kaxilAlice Berard
authored and
Alice Berard
committed
[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 ...'
1 parent 1c979be commit e33407f

File tree

2 files changed

+116
-103
lines changed

2 files changed

+116
-103
lines changed

airflow/bin/cli.py

+97-88
Original file line numberDiff line numberDiff line change
@@ -1347,74 +1347,75 @@ def kerberos(args): # noqa
13471347

13481348

13491349
@cli_utils.action_logging
1350-
def create_user(args):
1351-
fields = {
1352-
'role': args.role,
1353-
'username': args.username,
1354-
'email': args.email,
1355-
'firstname': args.firstname,
1356-
'lastname': args.lastname,
1357-
}
1358-
empty_fields = [k for k, v in fields.items() if not v]
1359-
if empty_fields:
1360-
raise SystemExit('Required arguments are missing: {}.'.format(
1361-
', '.join(empty_fields)))
1362-
1363-
appbuilder = cached_appbuilder()
1364-
role = appbuilder.sm.find_role(args.role)
1365-
if not role:
1366-
raise SystemExit('{} is not a valid role.'.format(args.role))
1367-
1368-
if args.use_random_password:
1369-
password = ''.join(random.choice(string.printable) for _ in range(16))
1370-
elif args.password:
1371-
password = args.password
1372-
else:
1373-
password = getpass.getpass('Password:')
1374-
password_confirmation = getpass.getpass('Repeat for confirmation:')
1375-
if password != password_confirmation:
1376-
raise SystemExit('Passwords did not match!')
1350+
def users(args):
1351+
if args.list:
1352+
1353+
appbuilder = cached_appbuilder()
1354+
users = appbuilder.sm.get_all_users()
1355+
fields = ['id', 'username', 'email', 'first_name', 'last_name', 'roles']
1356+
users = [[user.__getattribute__(field) for field in fields] for user in users]
1357+
msg = tabulate(users, [field.capitalize().replace('_', ' ') for field in fields],
1358+
tablefmt="fancy_grid")
1359+
if sys.version_info[0] < 3:
1360+
msg = msg.encode('utf-8')
1361+
print(msg)
13771362

1378-
if appbuilder.sm.find_user(args.username):
1379-
print('{} already exist in the db'.format(args.username))
13801363
return
1381-
user = appbuilder.sm.add_user(args.username, args.firstname, args.lastname,
1382-
args.email, role, password)
1383-
if user:
1384-
print('{} user {} created.'.format(args.role, args.username))
1385-
else:
1386-
raise SystemExit('Failed to create user.')
13871364

1365+
if args.create:
1366+
fields = {
1367+
'role': args.role,
1368+
'username': args.username,
1369+
'email': args.email,
1370+
'firstname': args.firstname,
1371+
'lastname': args.lastname,
1372+
}
1373+
empty_fields = [k for k, v in fields.items() if not v]
1374+
if empty_fields:
1375+
raise SystemExit('Required arguments are missing: {}.'.format(
1376+
', '.join(empty_fields)))
13881377

1389-
@cli_utils.action_logging
1390-
def delete_user(args):
1391-
if not args.username:
1392-
raise SystemExit('Required arguments are missing: username')
1378+
appbuilder = cached_appbuilder()
1379+
role = appbuilder.sm.find_role(args.role)
1380+
if not role:
1381+
raise SystemExit('{} is not a valid role.'.format(args.role))
1382+
1383+
if args.use_random_password:
1384+
password = ''.join(random.choice(string.printable) for _ in range(16))
1385+
elif args.password:
1386+
password = args.password
1387+
else:
1388+
password = getpass.getpass('Password:')
1389+
password_confirmation = getpass.getpass('Repeat for confirmation:')
1390+
if password != password_confirmation:
1391+
raise SystemExit('Passwords did not match!')
13931392

1394-
appbuilder = cached_appbuilder()
1393+
if appbuilder.sm.find_user(args.username):
1394+
print('{} already exist in the db'.format(args.username))
1395+
return
1396+
user = appbuilder.sm.add_user(args.username, args.firstname, args.lastname,
1397+
args.email, role, password)
1398+
if user:
1399+
print('{} user {} created.'.format(args.role, args.username))
1400+
else:
1401+
raise SystemExit('Failed to create user.')
13951402

1396-
try:
1397-
u = next(u for u in appbuilder.sm.get_all_users() if u.username == args.username)
1398-
except StopIteration:
1399-
raise SystemExit('{} is not a valid user.'.format(args.username))
1403+
if args.delete:
1404+
if not args.username:
1405+
raise SystemExit('Required arguments are missing: username')
14001406

1401-
if appbuilder.sm.del_register_user(u):
1402-
print('User {} deleted.'.format(args.username))
1403-
else:
1404-
raise SystemExit('Failed to delete user.')
1407+
appbuilder = cached_appbuilder()
14051408

1409+
try:
1410+
u = next(u for u in appbuilder.sm.get_all_users()
1411+
if u.username == args.username)
1412+
except StopIteration:
1413+
raise SystemExit('{} is not a valid user.'.format(args.username))
14061414

1407-
@cli_utils.action_logging
1408-
def list_users(args):
1409-
appbuilder = cached_appbuilder()
1410-
users = appbuilder.sm.get_all_users()
1411-
fields = ['id', 'username', 'email', 'first_name', 'last_name', 'roles']
1412-
users = [[user.__getattribute__(field) for field in fields] for user in users]
1413-
msg = tabulate(users, [field.capitalize().replace('_', ' ') for field in fields],
1414-
tablefmt="fancy_grid")
1415-
if sys.version_info[0] < 3:
1416-
msg = msg.encode('utf-8')
1417-
print(msg)
1415+
if appbuilder.sm.del_register_user(u):
1416+
print('User {} deleted.'.format(args.username))
1417+
else:
1418+
raise SystemExit('Failed to delete user.')
14181419

14191420

14201421
@cli_utils.action_logging
@@ -1527,10 +1528,6 @@ class CLIFactory(object):
15271528
"Do not prompt to confirm reset. Use with care!",
15281529
"store_true",
15291530
default=False),
1530-
'username': Arg(
1531-
('-u', '--username',),
1532-
help='Username of the user',
1533-
type=str),
15341531

15351532
# list_dag_runs
15361533
'no_backfill': Arg(
@@ -1886,33 +1883,52 @@ class CLIFactory(object):
18861883
('--conn_extra',),
18871884
help='Connection `Extra` field, optional when adding a connection',
18881885
type=str),
1889-
# create_user
1890-
'role': Arg(
1891-
('-r', '--role',),
1892-
help='Role of the user. Existing roles include Admin, '
1893-
'User, Op, Viewer, and Public',
1886+
# users
1887+
'username': Arg(
1888+
('--username',),
1889+
help='Username of the user, required to create/delete a user',
18941890
type=str),
18951891
'firstname': Arg(
1896-
('-f', '--firstname',),
1897-
help='First name of the user',
1892+
('--firstname',),
1893+
help='First name of the user, required to create a user',
18981894
type=str),
18991895
'lastname': Arg(
1900-
('-l', '--lastname',),
1901-
help='Last name of the user',
1896+
('--lastname',),
1897+
help='Last name of the user, required to create a user',
1898+
type=str),
1899+
'role': Arg(
1900+
('--role',),
1901+
help='Role of the user. Existing roles include Admin, '
1902+
'User, Op, Viewer, and Public. Required to create a user',
19021903
type=str),
19031904
'email': Arg(
1904-
('-e', '--email',),
1905-
help='Email of the user',
1905+
('--email',),
1906+
help='Email of the user, required to create a user',
19061907
type=str),
19071908
'password': Arg(
1908-
('-p', '--password',),
1909-
help='Password of the user',
1909+
('--password',),
1910+
help='Password of the user, required to create a user '
1911+
'without --use_random_password',
19101912
type=str),
19111913
'use_random_password': Arg(
19121914
('--use_random_password',),
1913-
help='Do not prompt for password. Use random string instead',
1915+
help='Do not prompt for password. Use random string instead.'
1916+
' Required to create a user without --password ',
19141917
default=False,
19151918
action='store_true'),
1919+
'list_users': Arg(
1920+
('-l', '--list'),
1921+
help='List all users',
1922+
action='store_true'),
1923+
'create_user': Arg(
1924+
('-c', '--create'),
1925+
help='Create a user',
1926+
action='store_true'),
1927+
'delete_user': Arg(
1928+
('-d', '--delete'),
1929+
help='Delete a user',
1930+
action='store_true'),
1931+
19161932
}
19171933
subparsers = (
19181934
{
@@ -2070,18 +2086,11 @@ class CLIFactory(object):
20702086
'args': ('list_connections', 'add_connection', 'delete_connection',
20712087
'conn_id', 'conn_uri', 'conn_extra') + tuple(alternative_conn_specs),
20722088
}, {
2073-
'func': create_user,
2074-
'help': "Create an account for the Web UI (FAB-based)",
2075-
'args': ('role', 'username', 'email', 'firstname', 'lastname',
2089+
'func': users,
2090+
'help': "List/Create/Delete users",
2091+
'args': ('list_users', 'create_user', 'delete_user',
2092+
'username', 'email', 'firstname', 'lastname', 'role',
20762093
'password', 'use_random_password'),
2077-
}, {
2078-
'func': delete_user,
2079-
'help': "Delete an account for the Web UI",
2080-
'args': ('username',),
2081-
}, {
2082-
'func': list_users,
2083-
'help': "List accounts for the Web UI",
2084-
'args': tuple(),
20852094
},
20862095
{
20872096
'func': sync_perm,

tests/core.py

+19-15
Original file line numberDiff line numberDiff line change
@@ -1089,40 +1089,44 @@ def test_cli_list_dag_runs(self):
10891089

10901090
def test_cli_create_user_random_password(self):
10911091
args = self.parser.parse_args([
1092-
'create_user', '-u', 'test1', '-l', 'doe', '-f', 'jon',
1093-
'-e', 'jdoe@foo.com', '-r', 'Viewer', '--use_random_password'
1092+
'users', '-c', '--username', 'test1', '--lastname', 'doe',
1093+
'--firstname', 'jon',
1094+
'--email', 'jdoe@foo.com', '--role', 'Viewer', '--use_random_password'
10941095
])
1095-
cli.create_user(args)
1096+
cli.users(args)
10961097

10971098
def test_cli_create_user_supplied_password(self):
10981099
args = self.parser.parse_args([
1099-
'create_user', '-u', 'test2', '-l', 'doe', '-f', 'jon',
1100-
'-e', 'jdoe@apache.org', '-r', 'Viewer', '-p', 'test'
1100+
'users', '-c', '--username', 'test2', '--lastname', 'doe',
1101+
'--firstname', 'jon',
1102+
'--email', 'jdoe@apache.org', '--role', 'Viewer', '--password', 'test'
11011103
])
1102-
cli.create_user(args)
1104+
cli.users(args)
11031105

11041106
def test_cli_delete_user(self):
11051107
args = self.parser.parse_args([
1106-
'create_user', '-u', 'test3', '-l', 'doe', '-f', 'jon',
1107-
'-e', 'jdoe@example.com', '-r', 'Viewer', '--use_random_password'
1108+
'users', '-c', '--username', 'test3', '--lastname', 'doe',
1109+
'--firstname', 'jon',
1110+
'--email', 'jdoe@example.com', '--role', 'Viewer', '--use_random_password'
11081111
])
1109-
cli.create_user(args)
1112+
cli.users(args)
11101113
args = self.parser.parse_args([
1111-
'delete_user', '-u', 'test3',
1114+
'users', '-d', '--username', 'test3',
11121115
])
1113-
cli.delete_user(args)
1116+
cli.users(args)
11141117

11151118
def test_cli_list_users(self):
11161119
for i in range(0, 3):
11171120
args = self.parser.parse_args([
1118-
'create_user', '-u', 'user{}'.format(i), '-l', 'doe', '-f', 'jon',
1119-
'-e', 'jdoe+{}@gmail.com'.format(i), '-r', 'Viewer',
1121+
'users', '-c', '--username', 'user{}'.format(i), '--lastname',
1122+
'doe', '--firstname', 'jon',
1123+
'--email', 'jdoe+{}@gmail.com'.format(i), '--role', 'Viewer',
11201124
'--use_random_password'
11211125
])
1122-
cli.create_user(args)
1126+
cli.users(args)
11231127
with mock.patch('sys.stdout',
11241128
new_callable=six.StringIO) as mock_stdout:
1125-
cli.list_users(self.parser.parse_args(['list_users']))
1129+
cli.users(self.parser.parse_args(['users', '-l']))
11261130
stdout = mock_stdout.getvalue()
11271131
for i in range(0, 3):
11281132
self.assertIn('user{}'.format(i), stdout)

0 commit comments

Comments
 (0)