def users_create(args):
"""Creates new user in the DB"""
appbuilder = cached_appbuilder()
role = appbuilder.sm.find_role(args.role)
if not role:
valid_roles = appbuilder.sm.get_all_roles()
raise SystemExit('{} is not a valid role. Valid roles are: {}'.format(
args.role, valid_roles))
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:'******'Repeat for confirmation:')
if password != password_confirmation:
raise SystemExit('Passwords did not match!')
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.')
def _import_users(users_list): # pylint: disable=redefined-outer-name
appbuilder = cached_appbuilder()
users_created = []
users_updated = []
for user in users_list:
roles = []
for rolename in user['roles']:
role = appbuilder.sm.find_role(rolename)
if not role:
valid_roles = appbuilder.sm.get_all_roles()
print("Error: '{}' is not a valid role. Valid roles are: {}".
format(rolename, valid_roles))
sys.exit(1)
else:
roles.append(role)
required_fields = [
'username', 'firstname', 'lastname', 'email', 'roles'
]
for field in required_fields:
if not user.get(field):
print("Error: '{}' is a required field, but was not "
"specified".format(field))
sys.exit(1)
existing_user = appbuilder.sm.find_user(email=user['email'])
if existing_user:
print("Found existing user with email '{}'".format(user['email']))
existing_user.roles = roles
existing_user.first_name = user['firstname']
existing_user.last_name = user['lastname']
if existing_user.username != user['username']:
print("Error: Changing the username is not allowed - "
"please delete and recreate the user with "
"email '{}'".format(user['email']))
sys.exit(1)
appbuilder.sm.update_user(existing_user)
users_updated.append(user['email'])
else:
print("Creating new user with email '{}'".format(user['email']))
appbuilder.sm.add_user(
username=user['username'],
first_name=user['firstname'],
last_name=user['lastname'],
email=user['email'],
role=roles[0], # add_user() requires exactly 1 role
)
if len(roles) > 1:
new_user = appbuilder.sm.find_user(email=user['email'])
new_user.roles = roles
appbuilder.sm.update_user(new_user)
users_created.append(user['email'])
return users_created, users_updated
def roles_list(args):
"""Lists all existing roles"""
appbuilder = cached_appbuilder()
roles = appbuilder.sm.get_all_roles()
print("Existing roles:\n")
role_names = sorted([[r.name] for r in roles])
msg = tabulate(role_names, headers=['Role'], tablefmt=args.output)
print(msg)
def users_list(args):
"""Lists users at the command line"""
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=args.output)
print(msg)
def test_dag_link(self):
from airflow.www.app import cached_appbuilder
with cached_appbuilder(testing=True).app.test_request_context():
html = str(utils.dag_link({
'dag_id': '<a&1>',
'execution_date': datetime.now()
}))
self.assertIn('%3Ca%261%3E', html)
self.assertNotIn('<a&1>', html)
def test_dag_link(self):
from airflow.www.app import cached_appbuilder
with cached_appbuilder(testing=True).app.test_request_context():
html = str(utils.dag_link({
'dag_id': '<a&1>',
'execution_date': datetime.now()
}))
self.assertIn('%3Ca%261%3E', html)
self.assertNotIn('<a&1>', html)
def sync_perm(args):
"""Updates permissions for existing roles and DAGs"""
appbuilder = cached_appbuilder()
print('Updating permission, view-menu for all existing roles')
appbuilder.sm.sync_roles()
print('Updating permission on all DAG views')
dags = DagBag().dags.values()
for dag in dags:
appbuilder.sm.sync_perm_for_dag(
dag.dag_id,
dag.access_control)
def users_delete(args):
"""Deletes user from DB"""
appbuilder = cached_appbuilder()
try:
user = 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 appbuilder.sm.del_register_user(user):
print('User {} deleted.'.format(args.username))
else:
raise SystemExit('Failed to delete user.')
def users_manage_role(args, remove=False):
"""Deletes or appends user roles"""
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')
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:
valid_roles = appbuilder.sm.get_all_roles()
raise SystemExit('{} is not a valid role. Valid roles are: {}'.format(args.role, valid_roles))
if remove:
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))
else:
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))
def users_export(args):
"""Exports all users to the json file"""
appbuilder = cached_appbuilder()
users = appbuilder.sm.get_all_users()
fields = ['id', 'username', 'email', 'first_name', 'last_name', 'roles']
# In the User model the first and last name fields have underscores,
# but the corresponding parameters in the CLI don't
def remove_underscores(s):
return re.sub("_", "", s)
users = [{
remove_underscores(field): user.__getattribute__(field)
if field != 'roles' else [r.name for r in user.roles]
for field in fields
} for user in users]
with open(args.export, 'w') as file:
file.write(json.dumps(users, sort_keys=True, indent=4))
print("{} users successfully exported to {}".format(
len(users), file.name))
def roles_create(args):
"""Creates new empty role in DB"""
appbuilder = cached_appbuilder()
for role_name in args.role:
appbuilder.sm.add_role(role_name)
