def increase_quota(cluster, role, cpu_str, ram_str, disk_str):
"""usage: increase_quota cluster role cpu ram[unit] disk[unit]
Increases the amount of production quota allocated to a user.
"""
cpu = float(cpu_str)
ram = parse_data(ram_str)
disk = parse_data(disk_str)
options = app.get_options()
client = AuroraClientAPI(CLUSTERS[cluster], options.verbosity == 'verbose')
resp = client.get_quota(role)
quota = resp.result.getQuotaResult.quota
log.info('Current quota for %s:\n\tCPU\t%s\n\tRAM\t%s MB\n\tDisk\t%s MB' %
(role, quota.numCpus, quota.ramMb, quota.diskMb))
new_cpu = cpu + quota.numCpus
new_ram = ram + Amount(quota.ramMb, Data.MB)
new_disk = disk + Amount(quota.diskMb, Data.MB)
log.info(
'Attempting to update quota for %s to\n\tCPU\t%s\n\tRAM\t%s MB\n\tDisk\t%s MB'
% (role, new_cpu, new_ram.as_(Data.MB), new_disk.as_(Data.MB)))
resp = client.set_quota(role, new_cpu, new_ram.as_(Data.MB),
new_disk.as_(Data.MB))
check_and_log_response(resp)
def __init__(self, cluster, role, env, jobs, ssh_user=None):
self._cluster = cluster
self._api = AuroraClientAPI(cluster=cluster)
self._role = role
self._env = env
self._jobs = jobs
self._ssh_user = ssh_user if ssh_user else self._role
def scheduler_snapshot(cluster):
"""usage: scheduler_snapshot cluster
Request that the scheduler perform a storage snapshot and block until complete.
"""
options = app.get_options()
check_and_log_response(
AuroraClientAPI(CLUSTERS['cluster'], options.verbosity).snapshot())
def scheduler_backup_now(cluster):
"""usage: scheduler_backup_now cluster
Immediately initiates a full storage backup.
"""
options = app.get_options()
check_and_log_response(
AuroraClientAPI(CLUSTERS[cluster], options.verbosity).perform_backup())
def scheduler_unload_recovery(cluster):
"""usage: scheduler_unload_recovery cluster
Unloads a staged recovery.
"""
options = app.get_options()
check_and_log_response(
AuroraClientAPI(CLUSTERS[cluster],
options.verbosity).unload_recovery())
def scheduler_stage_recovery(cluster, backup_id):
"""usage: scheduler_stage_recovery cluster backup_id
Stages a backup for recovery.
"""
options = app.get_options()
check_and_log_response(
AuroraClientAPI(CLUSTERS[cluster],
options.verbosity).stage_recovery(backup_id))
def scheduler_delete_recovery_tasks(cluster, task_ids):
"""usage: scheduler_delete_recovery_tasks cluster task_ids
Deletes a comma-separated list of task IDs from a staged recovery.
"""
ids = set(task_ids.split(','))
options = app.get_options()
check_and_log_response(
AuroraClientAPI(CLUSTERS[cluster],
options.verbosity).delete_recovery_tasks(
TaskQuery(taskIds=ids)))
def scheduler_list_backups(cluster):
"""usage: scheduler_list_backups cluster
Lists backups available for recovery.
"""
options = app.get_options()
resp = AuroraClientAPI(CLUSTERS[cluster], options.verbosity).list_backups()
check_and_log_response(resp)
backups = resp.result.listBackupsResult.backups
print('%s available backups:' % len(backups))
for backup in backups:
print(backup)
def scheduler_list_job_updates(cluster):
"""usage: scheduler_list_job_updates cluster
Lists in-flight job updates.
"""
options = app.get_options()
resp = AuroraClientAPI(CLUSTERS[cluster],
options.verbosity).get_job_updates()
check_and_log_response(resp)
print('Role\tEnv\tJob')
for update in resp.jobUpdates:
print('%s\t%s\t%s' %
(update.jobKey.role if update.jobKey else update.roleDeprecated,
update.jobKey.environment if update.jobKey else None,
update.jobKey.name if update.jobKey else update.jobDeprecated))
def set_quota(cluster, role, cpu_str, ram_mb_str, disk_mb_str):
"""usage: set_quota cluster role cpu ramMb diskMb
Alters the amount of production quota allocated to a user.
"""
try:
cpu = float(cpu_str)
ram_mb = int(ram_mb_str)
disk_mb = int(disk_mb_str)
except ValueError:
log.error('Invalid value')
options = app.get_options()
resp = AuroraClientAPI(CLUSTERS[cluster], options.verbosity).set_quota(
role, cpu, ram_mb, disk_mb)
check_and_log_response(resp)
def scheduler_print_recovery_tasks(cluster):
"""usage: scheduler_print_recovery_tasks cluster
Prints all active tasks in a staged recovery.
"""
options = app.get_options()
resp = AuroraClientAPI(CLUSTERS[cluster],
options.verbosity).query_recovery(
TaskQuery(statuses=ACTIVE_STATES))
check_and_log_response(resp)
log.info('Role\tJob\tShard\tStatus\tTask ID')
for task in resp.tasks:
assigned = task.assignedTask
conf = assigned.task
log.info('\t'.join(
(conf.owner.role, conf.jobName, str(assigned.instanceId),
ScheduleStatus._VALUES_TO_NAMES[task.status], assigned.taskId)))
def __init__(self, cluster, verbosity):
self._client = AuroraClientAPI(cluster, verbosity == 'verbose')
def query(args, options):
"""usage: query [--shards=N[,N,...]]
[--states=State[,State,...]]
cluster [role [job]]
Query Mesos about jobs and tasks.
"""
def _convert_fmt_string(fmtstr):
import re
def convert(match):
return "%%(%s)s" % match.group(1)
return re.sub(r'%(\w+)%', convert, fmtstr)
def flatten_task(t, d={}):
for key in t.__dict__.keys():
val = getattr(t, key)
try:
val.__dict__.keys()
except AttributeError:
d[key] = val
else:
flatten_task(val, d)
return d
def map_values(d):
default_value = lambda v: v
mapping = {
'status': lambda v: ScheduleStatus._VALUES_TO_NAMES[v],
}
return dict(
(k, mapping.get(k, default_value)(v)) for (k, v) in d.items())
for state in options.states.split(','):
if state not in ScheduleStatus._NAMES_TO_VALUES:
msg = "Unknown state '%s' specified. Valid states are:\n" % state
msg += ','.join(ScheduleStatus._NAMES_TO_VALUES.keys())
die(msg)
# Role, Job, Instances, States, and the listformat
if len(args) == 0:
die('Must specify at least cluster.')
cluster = args[0]
role = args[1] if len(args) > 1 else None
job = args[2] if len(args) > 2 else None
instances = set(map(
int, options.shards.split(','))) if options.shards else set()
if options.states:
states = set(
map(ScheduleStatus._NAMES_TO_VALUES.get,
options.states.split(',')))
else:
states = ACTIVE_STATES | TERMINAL_STATES
listformat = _convert_fmt_string(options.listformat)
# Figure out "expensive" queries here and bone if they do not have --force
# - Does not specify role
if role is None and not options.force:
die('--force is required for expensive queries (no role specified)')
# - Does not specify job
if job is None and not options.force:
die('--force is required for expensive queries (no job specified)')
# - Specifies status outside of ACTIVE_STATES
if not (states <= ACTIVE_STATES) and not options.force:
die('--force is required for expensive queries (states outside ACTIVE states'
)
api = AuroraClientAPI(CLUSTERS[cluster], options.verbosity)
query_info = api.query(
api.build_query(role, job, instances=instances, statuses=states))
tasks = query_info.result.scheduleStatusResult.tasks
if query_info.responseCode != ResponseCode.OK:
die('Failed to query scheduler: %s' % query_info.message)
if tasks is None:
return
try:
for task in tasks:
d = flatten_task(task)
print(listformat % map_values(d))
except KeyError:
msg = "Unknown key in format string. Valid keys are:\n"
msg += ','.join(d.keys())
die(msg)
