sched.py

#!/usr/bin/env python
"""
HomeCU tool for interacting with schedulers for stacks
"""

import boto3
import shellish
import subprocess
from . import s3, base, cloudformation


class SchedCommand(base.HostingCommand):
    '''Manage stack scheduler'''
    name = 'sched'
    allowed_states = {'active'}

    def setup_args(self, parser):
        self.add_subcommand(ReloadCommand)


class ReloadCommand(s3.S3BucketMixin, base.HostingCommand):
    """ Reload Scheduler of an Odyssey stack"""

    name = 'reload'
    min_version = '0.8.0'
    allowed_states = {'active'}

    def setup_args(self, parser):
        self.add_stack_argument()
        self.add_argument('--verbose', '-v', action='store_true')
        self.add_argument('--passphrase', env='ODYSSEY_STACK_PASSPHRASE',
                          help='To avoid being prompted for a passphrase on '
                          'protected stacks, you can set it on the command '
                          'line or via the env.')

    def run(self, args):
        stack = self.get_stack(args.stack)
        if stack is None:
            raise SystemExit('Stack not found: %s' % args.stack)
        stack.check_passphrase(args.passphrase)
        self.check_requirements(stack)
        try:
            shellish.vtmlprint("restarting scheduler for %s..." % stack.name)
            self.restart_scheduler(args, stack)
            shellish.vtmlprint("done.")
        except Exception as ex:
            if args.verbose:
                print(ex)
            stack.update_state('partial_deploy')
            stack._save('Restart of scheduler resulted in errors')
            raise
        else:
            stack.update_state('active')
            stack._save('Restarted scheduler service')

    def restart_scheduler(self, args, stack):
        '''Restart scheduler service for stack'''
        client = boto3.client('ecs', region_name=stack.region)
        cluster = stack.cluster
        cfn_stack = cloudformation.describe_stack(stack)
        services = [x['OutputValue'] for x in cfn_stack['Outputs']
                    if 'Service' in x['OutputKey']]
        sched_service = [x for x in services if 'sched' in x.lower()][-1]
        for task in self.get_service_task_arns(client, cluster, sched_service):
            client.stop_task(
                cluster=cluster,
                task=task,
                reason='Restarting scheduler, reloading crontab'
            )

    def get_service_task_arns(self, client, cluster, service):
        '''return task arns for service'''
        response = client.list_tasks(cluster=cluster, serviceName=service)
        if response['ResponseMetadata']['HTTPStatusCode'] != 200:
            shellish.vtmlprint("<red>Error while getting tasks</red>")
            return []
        return response['taskArns']


def cmd(*args, verbose=False):
    cmd = ' '.join(map(str, args))
    shellish.vtmlprint("<b>Running:</b> <blue>%s</blue>" % cmd)
    if not verbose:
        cmd += ' >/dev/null'
    subprocess.check_call(cmd, shell=True)