cumanage.py

"""
Run cumanage operations in an AWS hosted stack.
"""

import argparse
import shellish
import boto3
from . import base, s3


class CUManageCommand(s3.S3BucketMixin, base.HostingCommand):
    """ Run cumanage operations in an AWS hosted stack. """

    name = 'cumanage'

    def setup_args(self, parser):
        self.add_stack_argument()
        self.add_argument('proxyargs', nargs=argparse.REMAINDER)
        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 '
            'environment variables.')
        self.add_argument(
            '--check-exit-code',
            action='store_true',
            help='Block and wait for the exit status of the cumanage task')

    def run(self, args):
        stack = self.get_stack(args.stack)
        if stack is None:
            raise SystemExit('Stack not found: %s' % args.stack)
        passphrase = stack.check_passphrase(args.passphrase)
        wait_until_complete = args.check_exit_code

        exit(run_cumanage_task(
            stack,
            args.proxyargs,
            wait_until_complete,
            passphrase
        ))


def task_exit_code(client, cluster, task_arns):
    '''Check tasks exit codes on `cluster`'''
    response = client.describe_tasks(cluster=cluster, tasks=task_arns)
    tasks = response['tasks']
    containers = [container for containers in (x['containers'] for x in tasks)
                  for container in containers]
    exit_codes = [x.get('exitCode', -1) for x in containers]
    for code in exit_codes:
        if code != 0:
            return code
    return 0


def run_cumanage_task(stack,
                      proxyargs,
                      wait_until_complete=False,
                      passphrase=None):
    '''Submit and run cumanage'''
    if passphrase is None:
        passphrase = ''
    client = boto3.client('ecs', region_name=stack.region)
    taskdef = 'odyssey-%s-cumanage-tool' % stack.name
    shellish.vtmlprint("<b>Submitting:</b> <blue>%s</blue>" %
                       ' '.join(proxyargs))
    overrides = {
        "containerOverrides": [{
            "name": "cumanage",
            "command": proxyargs
        }]
    }
    response = client.run_task(cluster=stack.cluster,
                               taskDefinition=taskdef,
                               overrides=overrides)
    stack.save('Ran cumanage: %s' % ' '.join(proxyargs), passphrase)
    if wait_until_complete:
        task_arns = [x['taskArn'] for x in response['tasks']]
        waiter = client.get_waiter('tasks_stopped')
        if waiter.wait(cluster=stack.cluster, tasks=task_arns):
            raise SystemExit('Task waiter timed out!')
        return task_exit_code(client, stack.cluster, task_arns)
    return 0