def move_user_to_unaffiliated(user_id, **kwargs):
"""
Moves a user to the Unaffiliated organization and removes access
to all data except for the reference data set.
"""
import sfa_api.utils.storage_interface as storage
storage._call_procedure('move_user_to_unaffiliated',
str(user_id),
with_current_user=False)
click.echo(f'User {user_id} moved to unaffiliated organization.')
def add_job_role(user_id, role_name, **kwargs):
"""
Add a job role(s) (ROLE_NAME) to a job execution user (USER_ID)
"""
import sfa_api.utils.storage_interface as storage
for role in role_name:
storage._call_procedure('grant_job_role',
str(user_id),
role,
with_current_user=False)
click.echo(f'Added role {role} to user {user_id}')
def delete_job(job_id, **kwargs):
"""
Delete JOB_ID from the database
"""
import sfa_api.utils.storage_interface as storage
try:
storage._call_procedure('delete_job',
str(job_id),
with_current_user=False)
except pymysql.err.InternalError as e:
if e.args[0] == 1305:
fail(e.args[1])
else: # pragma: no cover
raise
else:
click.echo(f'Job {job_id} deleted successfully.')
def create_job_user(organization_name, encryption_key, **kwargs):
"""
Creates a new user in Auth0 to run background jobs for the organization.
Make sure AUTH0_CLIENT_ID and AUTH0_CLIENT_SECRET are properly set.
"""
from sfa_api.utils import auth0_info
import sfa_api.utils.storage_interface as storage
org_id = None
for org in storage._call_procedure('list_all_organizations',
with_current_user=False):
if org['name'] == organization_name:
org_id = org['id']
break
if org_id is None:
fail(f'Organization {organization_name} not found')
username = ('job-execution@' +
organization_name.lower().replace(" ", "-") +
'.solarforecastarbiter.org')
passwd = auth0_info.random_password()
user_id, auth0_id = storage.create_job_user(username, passwd, org_id,
encryption_key)
click.echo(f'Created user {username} with Auth0 ID {auth0_id}')
def delete_user(user_id, **kwargs):
"""
Remove a user from the framework.
"""
import sfa_api.utils.storage_interface as storage
try:
storage._call_procedure('delete_user',
str(user_id),
with_current_user=False)
except pymysql.err.InternalError as e:
if e.args[0] == 1305:
fail(e.args[1])
else: # pragma: no cover
raise
else:
click.echo(f'User {user_id} deleted successfully.')
def exchange_token(user_id):
"""
Get the refresh token from MySQL for the user_id, decrypt it, and
exchange it for an access token. This requires the same
TOKEN_ENCRYPTION_KEY that was used to encrypt the token along with
the same AUTH0_CLIENT_ID and AUTH0_CLIENT_SECRET to do anything
useful with the refresh token.
Parameters
----------
user_id : str
Retrieve an access token for this user
Returns
-------
HiddenToken
The access token that can be accessed at the .token property
Raises
------
KeyError
If no token is found for user_id
"""
try:
enc_token = storage._call_procedure(
'fetch_token',
(user_id, ),
with_current_user=False,
)[0]['token'].encode()
except IndexError:
raise KeyError(f'No token for {user_id} found')
f = Fernet(current_app.config['TOKEN_ENCRYPTION_KEY'])
refresh_token = f.decrypt(enc_token).decode()
access_token = exchange_refresh_token(refresh_token)
return HiddenToken(access_token)
def create_organization(organization_name, **kwargs):
"""Creates a new organization.
"""
import sfa_api.utils.storage_interface as storage
try:
storage._call_procedure('create_organization',
organization_name,
with_current_user=False)
except pymysql.err.DataError:
fail("Organization name must be 32 characters or fewer.")
except pymysql.err.IntegrityError as e:
if e.args[0] == 1062:
fail(f'Organization {organization_name} already exists.')
else: # pragma: no cover
raise
else:
click.echo(f'Created organization {organization_name}.')
def set_org_accepted_tou(organization_id, **kwargs):
"""
Sets an organizaiton's accepted terms of use field to true.
"""
import sfa_api.utils.storage_interface as storage
try:
storage._call_procedure('set_org_accepted_tou',
str(organization_id),
with_current_user=False)
except pymysql.err.InternalError as e:
if e.args[0] == 1305:
fail(e.args[1])
else: # pragma: no cover
raise
else:
click.echo(f'Organization {organization_id} has been marked '
'as accepting the terms of use.')
def list_jobs(**kwargs):
"""
List information for all jobs
"""
import pprint
import sfa_api.utils.storage_interface as storage
jobs = storage._call_procedure('list_jobs', with_current_user=False)
# a table would be too wide, so pretty print instead
click.echo(pprint.pformat(jobs))
def promote_to_admin(user_id, organization_id, **kwargs):
"""
Grants a user admin permissions in the organizations.
"""
import sfa_api.utils.storage_interface as storage
try:
storage._call_procedure('promote_user_to_org_admin',
str(user_id),
str(organization_id),
with_current_user=False)
except pymysql.err.IntegrityError as e:
if e.args[0] == 1062:
fail('User already granted admin permissions.')
else: # pragma: no cover
raise
except StorageAuthError as e:
click.echo(e.args[0])
else:
click.echo(f'Promoted user {user_id} to administrate '
f'organization {organization_id}')
def add_user_to_org(user_id, organization_id, **kwargs):
"""
Adds a user to an organization. The user must currently be
unaffiliated.
"""
import sfa_api.utils.storage_interface as storage
try:
storage._call_procedure('add_user_to_org',
str(user_id),
str(organization_id),
with_current_user=False)
except pymysql.err.IntegrityError as e:
if e.args[0] == 1452:
fail('Organization does not exist')
else: # pragma: no cover
raise
except StorageAuthError as e:
fail(e.args[0])
else:
click.echo(f'Added user {user_id} to organization {organization_id}')
def list_organizations(**kwargs):
"""
Prints a table of organization names and ids.
"""
import sfa_api.utils.storage_interface as storage
organizations = storage._call_procedure('list_all_organizations',
with_current_user=False)
table_format = '{:<34}|{:<38}|{:<12}'
headers = table_format.format('Name', 'Organization ID', 'Accepted TOU')
click.echo(headers)
click.echo('-' * len(headers))
for org in organizations:
click.echo(
table_format.format(org['name'], org["id"],
str(bool(org['accepted_tou']))))
def schedule_jobs(scheduler):
"""
Sync jobs between MySQL and RQ scheduler, adding new jobs
from MySQL, updating jobs if they have changed, and remove
RQ jobs that have been removed from MySQL
Parameters
----------
scheduler : rq_scheduler.Scheduler
The scheduler instance to compare MySQL jobs with
"""
logger.debug('Syncing MySQL and RQ jobs...')
sql_jobs = storage._call_procedure('list_jobs', with_current_user=False)
rq_jobs = scheduler.get_jobs()
sql_dict = {k['id']: k for k in sql_jobs}
rq_dict = {j.id: j for j in rq_jobs}
for to_cancel in set(rq_dict.keys()) - set(sql_dict.keys()):
logger.info('Removing extra RQ jobs %s',
rq_dict[to_cancel].meta.get('job_name', to_cancel))
scheduler.cancel(to_cancel)
# make sure job removed from redis
rq_dict[to_cancel].delete()
for job_id, sql_job in sql_dict.items():
if job_id in rq_dict:
if (sql_job['modified_at'] !=
rq_dict[job_id].meta['last_modified_in_sql']):
logger.info('Removing job %s', sql_job['name'])
scheduler.cancel(job_id)
else:
continue # pragma: no cover
try:
convert_sql_to_rq_job(sql_job, scheduler)
except (ValueError, json.JSONDecodeError, KeyError) as e:
logger.error('Failed to schedule job %s with error %s', sql_job, e)
def list_users(**kwargs):
"""
Prints a table of user information including auth0 id, user id,
organization and organization id. AUTH0_CLIENT_ID and AUTH0_CLIENT_SECRET
must be properly set to retrieve emails.
"""
import sfa_api.utils.storage_interface as storage
from sfa_api.utils.auth0_info import list_user_emails
import logging
logging.getLogger('sfa_api.utils.auth0_info').setLevel('CRITICAL')
users = storage._call_procedure('list_all_users', with_current_user=False)
emails = list_user_emails([u['auth0_id'] for u in users])
table_format = '{:<34}|{:<38}|{:<44}|{:<34}|{:<38}'
headers = table_format.format('auth0_id', 'User ID', 'User Email',
'Organization Name', 'Organization ID')
click.echo(headers)
click.echo('-' * len(headers))
for user in users:
click.echo(
table_format.format(user['auth0_id'], user['id'],
emails[user['auth0_id']],
user['organization_name'],
user['organization_id']))
def create_job(job_type, name, user_id, cron_string, timeout=None, **kwargs):
"""
Create a job in the database
Parameters
----------
job_type : str
Type of background job. This determines what kwargs are expected
name : str
Name for the job
user_id : str
ID of the user to execute this job
cron_string : str
Crontab string to schedule job
timeout : str
Maximum runtime before the job is killed. An integer default units is
seconds, otherwise, can specify the unit e.g. 1h, 2m, 10s
**kwargs
Keyword arguments that will be passed along when the job is executed
Returns
-------
str
ID of the MySQL job
Raises
------
ValueError
If the job type is not supported
"""
logger.info('Creating %s job', job_type)
if job_type == 'daily_observation_validation':
keys = ('start_td', 'end_td')
elif job_type == 'reference_nwp':
keys = ('issue_time_buffer', )
elif job_type == 'periodic_report':
# report must already exist
keys = ('report_id', )
elif job_type in ('reference_persistence',
'reference_probabilistic_persistence'):
keys = ()
else:
raise ValueError(f'Job type {job_type} is not supported')
params = {}
if 'base_url' in kwargs:
params['base_url'] = kwargs['base_url']
for k in keys:
params[k] = kwargs[k]
schedule = {'type': 'cron', 'cron_string': cron_string}
if timeout is not None:
schedule['timeout'] = timeout
id_ = storage.generate_uuid()
storage._call_procedure('store_job',
id_,
str(user_id),
name,
job_type,
json.dumps(params),
json.dumps(schedule),
0,
with_current_user=False)
return id_