Simplifying CodePipeline lambda jobs

[JordonPhillips]

Problem

One of the possible action types in CodePipeline allows you to invoke external code to do some work. Pointing this action type at a Lambda function requires that your function implement some boilerplate code to handle the interaction. It would be great if Chalice could abstract as much of that away as possible.

I propose a decorator for a function which does the following:

Handles Input Artifacts

CodePipeline will pass in input artifacts as part of its event. The lambda function must then create an s3 client with given credentials, pull them down, and then unzip them.

Suggestion

For each input artifact, provide an easy method to get a corresponding file-like object or ZipFile.

Handles Input Parameters

CodePipeline allows for passing in a configuration string to a given invoke action. While this can be anything, it often ends up being a JSON string. The lambda function must grab and parse this.

Suggestion

Automatically decode the parameters and pass them to a function. Bonus points for passing them as arguments to the function.

Handles Output Artifacts

Like input artifacts, output artifacts must be handled manually in the same way each time.

Suggestion

For each output artifact, provide a file-like object that users can write to which handles uploading to S3.

Handles Errors

CodePipeline won't recognize the function execution failing as a failure state for the job, so users have to handle them manually.

Suggestion

Chalice should catch errors, call put_job_failure_result, and then reraise the errors. Otherwise provide error response type objects similar to what it does for HTTP errors.

Handles Timeouts

A timeout will not be registered by CodePipeline as a failure, so a lambda function must be incredibly sure that it can run in the allotted time. If it times out, CodePipe will wait for a while before giving up on the job.

Suggestion

Provide a configurable manual timeout that will call put_job_failure_result before the lambda function actually times out.

Handles Returns

A lambda function must explicitly say it's done with a job for CodePipeline to recognize it as done. Otherwise CodePipeline will continue to wait and eventually register a failure.

Suggestion

Chalice should call put_job_success_result using the function's return value as the message, and also provide return objects similar to what it does for HTTP return types.

Passes continuation tokens

When you call put_job_success_result you can pass a continuation token. If you do, CodePipeline will continue to wait until you call that function without a continuation token. It also allows you to provide things like messages and progress for each round.

Suggestion

Chalice should pass in the continuation token and ensure that it is part of any success return objects it provides.

Examples

Artifact Merge Example

To show an example of some of the savings this could provide, I'll show a CodePipeline function that merges two input artifact into one output artifact both how one would implement it now and how it could look in the future.

lambda_function example

This is how one might write it currently.

import io
import json
import zipfile

import boto3
import botocore

def _get_user_parameters(job_data):
    configuration = job_data['actionConfiguration']['configuration']
    user_parameters = configuration['UserParameters']
    decoded_parameters = json.loads(user_parameters)
    return decoded_parameters

def create_s3_client(job_data):
    key_id = job_data['artifactCredentials']['accessKeyId']
    key_secret = job_data['artifactCredentials']['secretAccessKey']
    session_token = job_data['artifactCredentials']['sessionToken']

    session = boto3.Session(
        aws_access_key_id=key_id,
        aws_secret_access_key=key_secret,
        aws_session_token=session_token
    )
    client = session.client(
        's3',
        config=botocore.client.Config(signature_version='s3v4')
    )
    return client

def _fetch_artifact(s3, artifact):
    bucket = artifact['location']['s3Location']['bucketName']
    key = artifact['location']['s3Location']['objectKey']

    args = {
        'Key': key,
        'Bucket': bucket
    }
    response = s3.get_object(**args)

    body = response['Body']
    body_bytes = io.BytesIO(body.read())
    name = artifact['name']
    return name, body_bytes

def _add_to_zipfile(src_zipfile, dst_zipfile, prefix=None):
    for name in src_zipfile.namelist():
        data = src_zipfile.read(name)
        arc_name = name
        if prefix is not None:
            arc_name = '%s/%s' % (prefix, name)
        dst_zipfile.writestr(arc_name, data)

def _get_output_artifact_location(job_data):
    location = job_data['outputArtifacts'][0]['location']['s3Location']
    bucket = location['bucketName']
    key = location['objectKey']
    return bucket, key

def _write_to_bucket(s3, bucket, key, data):
    data.seek(0)
    s3.put_object(
        Bucket=bucket,
        Key=key,
        Body=data
    )

def lambda_handler(event, context):
    code_pipeline = boto3.client('codepipeline')
    job_id = event['CodePipeline.job']['id']
    try:
        job_data = event['CodePipeline.job']['data']
        params = _get_user_parameters(job_data)
        s3 = create_s3_client(job_data)
        print(job_data)

        artifacts = job_data['inputArtifacts']
        final_zipfile_data = io.BytesIO()
        with zipfile.ZipFile(final_zipfile_data, 'w',
                             compression=zipfile.ZIP_DEFLATED) as dst_zfile:
            for artifact in artifacts:
                artifact_name, artifact_data = _fetch_artifact(s3, artifact)
                with zipfile.ZipFile(artifact_data, 'r') as src_zfile:
                    if artifact_name in params['wrap']:
                        prefix = artifact_name
                    else:
                        prefix = None
                    _add_to_zipfile(src_zfile, dst_zfile, prefix=prefix)

        bucket, key = _get_output_artifact_location(job_data)
        _write_to_bucket(s3, bucket, key, final_zipfile_data)
        code_pipeline.put_job_success_result(
            jobId=job_id,
        )
    except Exception as e:
        code_pipeline.put_job_failure_result(
            jobId=job_id,
            failureDetails={
                'type': 'JobFailed',
                'message': str(e)
            }
        )
        raise e
    return "Complete."

pipeline_function example

This is how one might write it with Chalice support.

import io
import zipfile

from chalice import Chalice

app = Chalice(app_name='CodePipelineExample')

# Timeout 5 seconds before the lambda invocation itself times out.
@app.pipeline_function(name='MergeArtifacts', timeout=5)
def pipeline_handler(input_artifacts, output_artifacts, configuration):
    final_zipfile_data = io.BytesIO()
    with zipfile.ZipFile(final_zipfile_data, 'w',
                         compression=zipfile.ZIP_DEFLATED) as dst_zfile:
        for artifact_name, artifact in input_artifacts:
            with zipfile.ZipFile(artifact_data, 'r') as src_zfile:
                if artifact_name in configuration['wrap']:
                    prefix = artifact_name
                else:
                    prefix = None
                _add_to_zipfile(src_zfile, dst_zfile, prefix=prefix)

    output_artifacts[0].write(final_zipfile_data.getvalue())

def _add_to_zipfile(src_zipfile, dst_zipfile, prefix=None):
    for name in src_zipfile.namelist():
        data = src_zipfile.read(name)
        arc_name = name
        if prefix is not None:
            arc_name = '%s/%s' % (prefix, name)
        dst_zipfile.writestr(arc_name, data)

The function went from 110 lines to 33 lines, and added timeout safety.

[stealthycoin]

How about something like this

import io
import zipfile

from chalice import Chalice

app = Chalice(app_name='merge')

@app.pipeline_function()
def second_stage(input_artifacts, user_parameters):
    merged_data = input_artifacts.get('Merged')
    with zipfile.ZipFile(merged_data, 'r') as zfile:
        for name in zfile.namelist():
            print(name)

def _add_zipfile_to_zipfile(src_zipfile, dst_zipfile, prefix=None):
    for name in src_zipfile.namelist():
        data = src_zipfile.read(name)
        arc_name = name
        if prefix is not None:
            arc_name = '%s/%s' % (prefix, name)
        dst_zipfile.writestr(arc_name, data)

def _add_raw_to_zipfile(name, src_data, dst_zipfile, prefix=None):
    arc_name = name
    if prefix is not None:
        arc_name = '%s/%s' % (prefix, name)
    dst_zipfile.writestr(arc_name, src_data.read())

@app.pipeline_function()
def merge(input_artifacts, user_parameters):
    merged_content = io.BytesIO()
    with zipfile.ZipFile(merged_content, 'w',
                         compression=zipfile.ZIP_DEFLATED) as dst_zfile:
        for artifact_name, artifact_data in input_artifacts.items():
            if artifact_name in user_parameters.get('wrap', []):
                prefix = artifact_name
            else:
                prefix = None
            try:
                with zipfile.ZipFile(artifact_data, 'r') as src_zfile:
                    _add_zipfile_to_zipfile(src_zfile, dst_zfile,
                                            prefix=prefix)
            except zipfile.BadZipFile:
                _add_raw_to_zipfile(artifact_name, artifact_data, dst_zfile,
                                    prefix=prefix)

    output_name = user_parameters.get('output_name')
    if output_name is None:
        output_name = 'Merged'
    return {
        output_name: merged_content
    }

Using the branch https://github.com/stealthycoin/chalice/tree/code-pipeline-integration

Would need a lot of parameterization of the input and output artifacts to capture all the cases like custom SSE keys etc...