S3-website-cache-control-setter

headers.py

# -*- coding: utf-8 -*-

from flask import Flask
import logging

import boto3
import os

from copy import deepcopy

app = Flask(__name__)

s3 = boto3.client('s3')

# Because of boto3, these need to be strings, not ints...
TXT_TTL = os.environ.get('TXT_TTL', '300')
HTML_TTL = os.environ.get('HTML_TTL', '300')
ASSET_TTL = os.environ.get('ASSET_TTL', '300')

# In case we're called as an HTTP request, handle gracefully
@app.route('/')
def index():
    return 'Nothing to see here. You can go about your business. Move along.', 200


def reset_cache_control(event, context):
    """
    Process a file upload.
    """

    # Get the uploaded file's information
    for record in event['Records']:
        bucket = record['s3']['bucket']['name'] # Will be `my-bucket`
        key = record['s3']['object']['key'] # Will be the file path of whatever file was uploaded.

        if key.endswith('.html'):
            set_cache_control(bucket, key, ttl=HTML_TTL)
        elif key.endswith('.txt'):
            set_cache_control(bucket, key, ttl=TXT_TTL)
        else:
            set_cache_control(bucket, key, ttl=ASSET_TTL)

def set_cache_control(bucket, key, ttl=1800):
    cache_control = 'max-age={}, public'.format(ttl)
    obj = s3.head_object(Bucket=bucket, Key=key)

    if 'CacheControl' in obj and obj['CacheControl'] == cache_control:
        logging.debug('Cache-Control header already has the appropriate value.')
        return

    m = deepcopy(obj['Metadata'])
    s3.copy_object(
        Bucket=bucket,
        Key=key,
        CopySource='{}/{}'.format(bucket, key),
        CacheControl=cache_control,
        ContentType=obj['ContentType'],
        Metadata=m,
        MetadataDirective='REPLACE'
    )