ecs_deploy.py

#!/usr/bin/env python
# -*- coding: utf-8 -*-
from __future__ import print_function
from __future__ import unicode_literals
from __future__ import division

import re
import signal
import sys
import time
from functools import wraps

import boto3
import click

import logging
logger = logging.getLogger(__name__)


class TimeoutException(Exception):

    def __init__(self, msg="TimeoutException"):
        self.msg = msg

    def __str__(self):
        return repr(self.msg)


def timeout(seconds, msg='TimeoutException'):
    """
    http://www.saltycrane.com/blog/2010/04/using-python-timeout-decorator-uploading-s3/
    http://code.activestate.com/recipes/307871-timing-out-function/
    https://github.com/pnpnpn/timeout-decorator
    """
    def decorate(function):
        def handler(signum, frame):
            raise TimeoutException(msg)

        @wraps(function)
        def new_function(*args, **kwargs):
            old = signal.signal(signal.SIGALRM, handler)
            signal.alarm(seconds)
            try:
                result = function(*args, **kwargs)
            finally:
                signal.signal(signal.SIGALRM, old)
            signal.alarm(0)
            return result
        return new_function
    return decorate


class DockerImage(object):

    IMAGE_REGEX = re.compile('^(?P<domain>[a-zA-Z0-9.¥-]+):?(?P<port>[0-9]+)?/(?P<repo>[a-zA-Z0-9_-]+)/?(?P<img>[a-zA-Z0-9_-]+)?:?(?P<tag>[a-zA-Z0-9¥._-]+)?$')
    ROOT_REGEX = re.compile('^(?P<img>[a-zA-Z0-9¥-]+):?(?P<tag>[a-zA-Z0-9¥.¥-]+)?$')

    def __init__(self, image):
        self.__dict__.update(image)
        self.__reassemble_image_name()

    def __reassemble_image_name(self):
        use_image = ''
        if self.domain:
            use_image += self.domain
        if self.port:
            use_image += ':' + self.port
        if self.repo:
            use_image += '/' + self.repo
        if self.img:
            if use_image:
                use_image += '/' + self.img
            else:
                use_image = self.img
        self.image_without_tag = use_image
        if self.tag:
            self.use_image = use_image + ':' + self.tag
        logger.info('Useing image name: %s', use_image)

    @staticmethod
    def validate_image(ctx, param, value):
        result = DockerImage.IMAGE_REGEX.match(value)
        if result:
            image = result.groupdict()
            if not image['domain']:
                raise click.BadParameter('Image name does not contain a domain or repo as expected.' +
                                         ' See usage for supported formats.')
            if not image['repo']:
                raise click.BadParameter('Image name is missing the actual image name.' +
                                         ' See usage for supported formats.')
            if not image['img']:
                image['img'] = image['repo']
                image['repo'] = None
        else:
            result = DockerImage.ROOT_REGEX.match(value)
            if result:
                image = result.groupdict()
                image.update(dict(domain=None, port=None, repo=None))
            else:
                raise click.BadParameter('Unable to parse image name: {}, check the format and try again.'.format(value))
        logging.debug('Image validate result: %s', image)
        return DockerImage(image)


class AwsEcsException(Exception):

    def __init__(self, msg="AwsEcsException"):
        self.msg = msg

    def __str__(self):
        return repr(self.msg)


class AwsEcs(object):

    def __init__(self, image, aws_access_key=None, aws_secret_access_key=None, region_name=None):
        self.image = image
        self.client = boto3.client('ecs',
                                   aws_access_key_id=aws_access_key,
                                   aws_secret_access_key=aws_secret_access_key,
                                   region_name=region_name)

    def get_task_definition_arn(self, cluster_name, service_name):
        response = self.client.describe_services(
            cluster=cluster_name,
            services=[
                service_name
            ]
        )
        logger.debug(response)
        if 'services' in response:
            services = response['services']
            if len(services) > 0:
                service = services[0]
                if 'taskDefinition' in service:
                    return service['taskDefinition']
        raise AwsEcsException('Task definition not found.')

    def get_task_definition(self, task_definition):
        response = self.client.describe_task_definition(
            taskDefinition=task_definition
        )
        logger.debug(response)
        if 'taskDefinition' in response:
            task_definition = response['taskDefinition']
            return task_definition
        raise AwsEcsException('Task definition not found.')

    def get_new_task_definition(self, task_definition):
        response = self.client.describe_task_definition(
            taskDefinition=task_definition
        )
        logger.debug(response)
        if 'taskDefinition' in response:
            task_definition = response['taskDefinition']
            container_definitions = task_definition['containerDefinitions']
            for container_definition in container_definitions:
                container_definition['image'] = re.sub(r'^{}:?([a-zA-Z0-9¥.¥-])+$'.format(self.image.image_without_tag),
                                                       self.image.use_image, container_definition['image'])
            return task_definition
        raise AwsEcsException('Task definition not found.')

    def update_task_definition(self, family, container_definitions, volumes):
        response = self.client.register_task_definition(
            family=family, containerDefinitions=container_definitions, volumes=volumes)
        logger.debug(response)
        return response['taskDefinition']['taskDefinitionArn']

    def update_service(self, cluster_name, service_name, task_definition, max, min):
        response = self.client.update_service(
            cluster=cluster_name,
            service=service_name,
            taskDefinition=task_definition,
            deploymentConfiguration={
                'maximumPercent': max,
                'minimumHealthyPercent': min
            }
        )
        logger.debug(response)

    def get_running_tasks(self, cluster_name, service_name):
        response = self.client.list_tasks(
            cluster=cluster_name,
            serviceName=service_name,
            desiredStatus='RUNNING'
        )
        logger.debug(response)
        if 'taskArns' in response:
            task_arns = response['taskArns']
            if task_arns and (len(task_arns) > 0):
                response = self.client.describe_tasks(
                    cluster=cluster_name,
                    tasks=task_arns
                )
                logger.debug(response)
                if 'tasks' in response:
                    return [{
                                'taskArn': t['taskArn'],
                                'taskDefinitionArn': t['taskDefinitionArn'],
                                'lastStatus': t['lastStatus'],
                                'desiredStatus': t['desiredStatus']
                            } for t in response['tasks']]
            return None
        raise AwsEcsException('Task not found.')

    def stop_task(self, cluster_name, task_arn):
        response = self.client.stop_task(
            cluster=cluster_name,
            task=task_arn
        )
        logger.debug(response)

    def deregister_task_definition(self, task_definition):
        response = self.client.deregister_task_definition(
            taskDefinition=task_definition
        )
        logger.debug(response)


@click.group()
@click.option('--aws_access_key_id',
              type=str,
              help='AWS Access Key ID.')
@click.option('--aws_secret_access_key',
              type=str,
              help='AWS Secret Access Key.')
@click.option('--region',
              type=str,
              help='AWS region name.')
@click.option('--image',
              type=str,
              required=True,
              callback=DockerImage.validate_image,
              help='Name of Docker image to run, ex: repo/image:latest.\n' +
                   'Format: [domain][:port][/repo][/][image][:tag]\n' +
                   'Examples: mariadb, mariadb:latest, silintl/mariadb, ' +
                   'silintl/mariadb:latest, private.registry.com:8000/repo/image:tag')
@click.option('--min',
              type=int,
              default=50,
              help='minumumHealthyPercent: The lower limit on the number of running tasks during a deployment.' +
                   ' (Default is 50)')
@click.option('--max',
              type=int,
              default=200,
              help='maximumPercent: The upper limit on the number of running tasks during a deployment.' +
                   ' (Default is 200)')
@click.option('--timeout',
              type=int,
              default=90,
              help='Script monitors ECS Service for new task definition to be running.' +
                   ' (Default is 90s)')
@click.option('--interval',
              type=int,
              default=5,
              help='The monitoring interval for new task definition to be running.' +
                   ' (Default is 5s)')
@click.option('--force',
              default=False,
              is_flag=True)
@click.option('--debug',
              default=False,
              is_flag=True)
@click.pass_context
def cli(ctx, aws_access_key_id, aws_secret_access_key, region, image, min, max, timeout, interval, force, debug):
    """
    Simple script for triggering blue/green deployments on Amazon EC2 Container Service
    https://github.com/silinternational/ecs-deploy
    """
    handler = logging.StreamHandler(sys.stdout)
    if debug:
        logging_level = logging.DEBUG
        logging_format = '%(asctime)s %(levelname)-9s %(name)22s %(funcName)22s:%(lineno)-4d %(message)s'
    else:
        logging_level = logging.INFO
        logging_format = '%(asctime)s %(levelname)-9s %(message)s'
    logging.root.setLevel(logging_level)
    handler.setLevel(logging_level)
    handler.setFormatter(logging.Formatter(logging_format))
    logging.root.addHandler(handler)

    ctx.obj['aws_access_key_id'] = aws_access_key_id
    ctx.obj['aws_secret_access_key'] = aws_secret_access_key
    ctx.obj['region'] = region
    ctx.obj['image'] = image
    ctx.obj['min'] = min
    ctx.obj['max'] = max
    ctx.obj['timeout'] = timeout
    ctx.obj['interval'] = interval
    ctx.obj['force'] = force


@cli.command()
@click.option('--cluster',
              type=str,
              required=True,
              help='Name of ECS cluster')
@click.argument('name')
@click.pass_context
def service(ctx, name, cluster):
    """Name of service to deploy"""
    logger.debug('Deploy service name: %s', name)
    ecs = AwsEcs(ctx.obj['image'], ctx.obj['aws_access_key_id'], ctx.obj['aws_secret_access_key'], ctx.obj['region'])

    task_definition = ecs.get_task_definition(name)
    task_definition_arn = ecs.get_task_definition_arn(cluster, name)
    logger.info('Current task definition: %s', task_definition_arn)

    new_task_definition = ecs.get_new_task_definition(name)
    new_task_definition_arn = ecs.update_task_definition(
        new_task_definition['family'],
        new_task_definition['containerDefinitions'],
        new_task_definition['volumes']
    )
    logger.info('New task definition: %s', new_task_definition_arn)

    @timeout(ctx.obj['timeout'], 'New task definition not running within {} seconds.'.format(ctx.obj['timeout']))
    def check_task_status():
        while True:
            arns = ecs.get_running_tasks(cluster, name)
            logger.info('Running task: %s', arns)
            if arns and (new_task_definition_arn in [arn['taskDefinitionArn']
                                                     for arn in arns if arn['lastStatus'] == 'RUNNING']):
                ecs.deregister_task_definition(task_definition_arn)
                if ctx.obj['force']:
                    ecs.stop_task(cluster, next(arn['taskArn']
                                                for arn in arns if arn['taskDefinitionArn'] == task_definition_arn))
                break
            time.sleep(ctx.obj['interval'])

    ecs.update_service(cluster, name, new_task_definition_arn, ctx.obj['max'], ctx.obj['min'])
    try:
        check_task_status()
    except TimeoutException as e:
        # rollback
        rollback_task_definition_arn = ecs.update_task_definition(
            task_definition['family'],
            task_definition['containerDefinitions'],
            task_definition['volumes']
        )
        logger.exception('Rollback task definition: %s', rollback_task_definition_arn)
        ecs.update_service(cluster, name, rollback_task_definition_arn, ctx.obj['max'], ctx.obj['min'])
        ecs.deregister_task_definition(task_definition_arn)
        ecs.deregister_task_definition(new_task_definition_arn)
        raise e
    logger.info('Service updated successfully, new task definition running.')


@cli.command()
@click.argument('name')
@click.pass_context
def task(ctx, name):
    """Name of task definition to deploy"""
    logger.debug('Deploy task name: %s', name)
    ecs = AwsEcs(ctx.obj['image'], ctx.obj['aws_access_key_id'], ctx.obj['aws_secret_access_key'], ctx.obj['region'])

    new_task_definition = ecs.get_new_task_definition(name)
    new_task_definition_arn = ecs.update_task_definition(
        new_task_definition['family'],
        new_task_definition['containerDefinitions'],
        new_task_definition['volumes']
    )
    logger.info('New task definition: %s', new_task_definition_arn)
    logger.info('Task definition updated successfully.')


if __name__ == '__main__':
    cli(obj={})

usage.txt

Usage: ecs_deploy.py [OPTIONS] COMMAND [ARGS]...

  Simple script for triggering blue/green deployments on Amazon EC2
  Container Service https://github.com/silinternational/ecs-deploy

Options:
  --aws_access_key_id TEXT      AWS Access Key ID.
  --aws_secret_access_key TEXT  AWS Secret Access Key.
  --region TEXT                 AWS region name.
  --image TEXT                  Name of Docker image to run, ex:
                                repo/image:latest.
                                Format:
                                [domain][:port][/repo][/][image][:tag]
                                Examples: mariadb, mariadb:latest,
                                silintl/mariadb, silintl/mariadb:latest,
                                private.registry.com:8000/repo/image:tag
                                [required]
  --min INTEGER                 minumumHealthyPercent: The lower limit on the
                                number of running tasks during a deployment.
                                (Default is 50)
  --max INTEGER                 maximumPercent: The upper limit on the number
                                of running tasks during a deployment. (Default
                                is 200)
  --timeout INTEGER             Script monitors ECS Service for new task
                                definition to be running. (Default is 90s)
  --interval INTEGER            The monitoring interval for new task
                                definition to be running. (Default is 5s)
  --debug
  --help                        Show this message and exit.

Commands:
  service  Name of service to deploy
  task     Name of task definition to deploy

Usage: ecs_deploy.py service [OPTIONS] NAME

  Name of service to deploy

Options:
  --cluster TEXT  Name of ECS cluster  [required]
  --help          Show this message and exit.

Usage: ecs_deploy.py task [OPTIONS] NAME

  Name of task definition to deploy

Options:
  --help  Show this message and exit.