Bootstrap/test-infra/s3_cache.py

#!/usr/bin/env python2.7
from __future__ import absolute_import, unicode_literals, print_function, division

from sys import argv
from os import environ, stat, chdir, remove as _delete_file
from os.path import isfile, dirname, basename, abspath, realpath, expandvars
from hashlib import sha256
from subprocess import check_call as run
from json import load
from contextlib import contextmanager
from datetime import datetime

from boto.s3.connection import S3Connection
from boto.s3.key import Key
from boto.exception import S3ResponseError


CONFIG_FILE = './S3Cachefile.json'
NEED_TO_UPLOAD_MARKER = '.need-to-upload'
BYTES_PER_MB = 1024 * 1024


@contextmanager
def timer():
    start = datetime.utcnow()
    yield
    end = datetime.utcnow()
    elapsed = end - start
    print("\tDone. Took", int(elapsed.total_seconds()), "second(s).")


def _sha256_of_file(filename):
    hasher = sha256()
    with open(filename, 'rb') as input_file:
        hasher.update(input_file.read())
    file_hash = hasher.hexdigest()
    print('sha256({}) = {}'.format(filename, file_hash))
    return file_hash


def _delete_file_quietly(filename):
    try:
        _delete_file(filename)
    except (OSError, IOError):
        pass


def _tarball_size(directory):
    kib = stat(_tarball_filename_for(directory)).st_size // BYTES_PER_MB
    return "{} MiB".format(kib)


def _tarball_filename_for(directory):
    return abspath('./{}.tar.gz'.format(basename(directory)))


def _create_tarball(directory):
    print("Creating tarball of {}...".format(directory))
    with timer():
        run(['tar', '-czf', _tarball_filename_for(directory), '-C', dirname(directory), basename(directory)])


def _extract_tarball(directory):
    print("Extracting tarball of {}...".format(directory))
    with timer():
        run(['tar', '-xzf', _tarball_filename_for(directory), '-C', dirname(directory)])


def download(directory):
    _delete_file_quietly(NEED_TO_UPLOAD_MARKER)
    try:
        print("Downloading {} tarball from S3...".format(cache_name))
        with timer():
            key.get_contents_to_filename(_tarball_filename_for(directory))
    except S3ResponseError as err:
        open(NEED_TO_UPLOAD_MARKER, 'a').close()
        print(err)
        raise SystemExit("Cached {} download failed!".format(cache_name))
    print("Downloaded {}.".format(_tarball_size(directory)))
    _extract_tarball(directory)
    print("{} successfully installed from cache.".format(cache_name))


def upload(directory):
    _create_tarball(directory)
    print("Uploading {} tarball to S3... ({})".format(cache_name, _tarball_size(directory)))
    with timer():
        key.set_contents_from_filename(_tarball_filename_for(directory))
    print("{} cache successfully updated.".format(cache_name))
    _delete_file_quietly(NEED_TO_UPLOAD_MARKER)


if __name__ == '__main__':
    # Uses environment variables:
    #   AWS_ACCESS_KEY_ID -- AWS Access Key ID
    #   AWS_SECRET_ACCESS_KEY -- AWS Secret Access Key
    argv.pop(0)
    if len(argv) != 2:
        raise SystemExit("USAGE: s3_cache.py <download | upload> <cache name>")
    mode, cache_name = argv
    script_dir = dirname(realpath(__file__))
    chdir(script_dir)
    try:
        with open(CONFIG_FILE, 'rt') as config_file:
            config = load(config_file)
    except (IOError, OSError, ValueError) as config_err:
        print(config_err)
        raise SystemExit("Error when trying to load config from JSON file!")

    try:
        cache_info = config[cache_name]
        key_file = expandvars(cache_info["key"])
        fallback_cmd = cache_info["generate"]
        directory = expandvars(cache_info["cache"])
    except (TypeError, KeyError) as load_err:
        print(load_err)
        raise SystemExit("Config for cache named {!r} is missing or malformed!".format(cache_name))

    try:
        try:
            BUCKET_NAME = environ['TWBS_S3_BUCKET']
        except KeyError:
            raise SystemExit("TWBS_S3_BUCKET environment variable not set!")

        conn = S3Connection()
        bucket = conn.lookup(BUCKET_NAME)
        if bucket is None:
            raise SystemExit("Could not access bucket!")

        key_file_hash = _sha256_of_file(key_file)

        key = Key(bucket, key_file_hash)
        key.storage_class = 'REDUCED_REDUNDANCY'

        if mode == 'download':
            download(directory)
        elif mode == 'upload':
            if isfile(NEED_TO_UPLOAD_MARKER):  # FIXME
                upload(directory)
            else:
                print("No need to upload anything.")
        else:
            raise SystemExit("Unrecognized mode {!r}".format(mode))
    except BaseException as exc:
        if mode != 'download':
            raise
        print("Error!:", exc)
        print("Unable to download from cache.")
        print("Running fallback command to generate cache directory {!r}: {}".format(directory, fallback_cmd))
        with timer():
            run(fallback_cmd, shell=True)