"""
Suspend an auto scaling group's scaling processes that can interfere with CodeDeploy deploys.
It assumes a single ASG per deployment group.

To use this:
* create a lambda function with this code, then hook up it up to an SNS topic that receives all deployment events (but not host events).
* attach that topic as a trigger in your deployment groups.

Unlike AWS's in-appspec approach, this supports arbitrary deploy concurrency.

License: MIT, copyright Venmo.
"""

import json

import boto3

codedeploy = boto3.client('codedeploy')
autoscaling = boto3.client('autoscaling')

def lambda_handler(event, context):
    print "Received event:", event
    message = json.loads(event['Records'][0]['Sns']['Message'])
    
    deployment = codedeploy.get_deployment(deploymentId=message['deploymentId'])
    application = deployment['deploymentInfo']['applicationName']
    group = deployment['deploymentInfo']['deploymentGroupName']
    
    group = codedeploy.get_deployment_group(
        applicationName=application,
        deploymentGroupName=group,
    )
    
    asg = group['deploymentGroupInfo']['autoScalingGroups'][0]['name']
    
    suspend_or_resume = autoscaling.suspend_processes
    if message['status'] != 'CREATED':
        # One of stopped, failed, succeeded.
        suspend_or_resume = autoscaling.resume_processes
    
    # See https://github.com/awslabs/aws-codedeploy-samples/commit/fbd80c744bce49fc3efb7b686699f813b7964c85#diff-a6b6b2ae2199a57d1849b5b68115616aR114.
    res = suspend_or_resume(
        AutoScalingGroupName=asg,
        ScalingProcesses=['AZRebalance', 'AlarmNotification', 'ScheduledActions', 'ReplaceUnhealthy'],
    )
    
    # The returned value doesn't otherwise go to logs.
    print suspend_or_resume.__name__, asg, res
    
    return suspend_or_resume.__name__, asg, res