use S3-based node_modules caching

1 files changed, 87 insertions, 0 deletions

diff --git a/test-infra/node_modules_cache.py b/test-infra/node_modules_cache.py
new file mode 100755
index 000000000..9d9272fcc
--- /dev/null
+++ b/test-infra/node_modules_cache.py
@@ -0,0 +1,87 @@
+#!/usr/bin/env python2.7
+from __future__ import absolute_import, unicode_literals, print_function, division
+
+from sys import argv
+from os import environ, stat, remove as _delete_file
+from os.path import isfile
+from hashlib import sha256
+from subprocess import check_call as run
+
+from boto.s3.connection import S3Connection
+from boto.s3.key import Key
+from boto.exception import S3ResponseError
+
+
+NODE_MODULES_TARBALL = 'node_modules.tar.gz'
+NEED_TO_UPLOAD_MARKER = '.need-to-upload'
+BYTES_PER_MB = 1024 * 1024
+try:
+    BUCKET_NAME = environ['TWBS_S3_BUCKET']
+except KeyError:
+    raise SystemExit("TWBS_S3_BUCKET environment variable not set!")
+
+
+def _sha256_of_file(filename):
+    hasher = sha256()
+    with open(filename, 'rb') as input_file:
+        hasher.update(input_file.read())
+    return hasher.hexdigest()
+
+
+def _delete_file_quietly(filename):
+    try:
+        _delete_file(filename)
+    except (OSError, IOError):
+        pass
+
+
+def _tarball_size():
+    kib = stat(NODE_MODULES_TARBALL).st_size // BYTES_PER_MB
+    return "{} MiB".format(kib)
+
+
+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) != 1:
+        raise SystemExit("USAGE: node_modules_cache.py <download | upload>")
+    mode = argv.pop()
+
+    conn = S3Connection()
+    bucket = conn.lookup(BUCKET_NAME)
+    if bucket is None:
+        raise SystemExit("Could not access bucket!")
+
+    package_json_hash = _sha256_of_file('package.json')
+    print('sha256(package.json) = ' + package_json_hash)
+
+    key = Key(bucket, package_json_hash)
+    key.storage_class = 'REDUCED_REDUNDANCY'
+
+    if mode == 'download':
+        _delete_file_quietly(NEED_TO_UPLOAD_MARKER)
+        try:
+            print("Downloading tarball from S3...")
+            key.get_contents_to_filename(NODE_MODULES_TARBALL)
+        except S3ResponseError as err:
+            open(NEED_TO_UPLOAD_MARKER, 'a').close()
+            print(err)
+            raise SystemExit("Cached node_modules download failed!")
+        print("Downloaded {}.".format(_tarball_size()))
+        print("Extracting tarball...")
+        run(['tar', 'xzf', NODE_MODULES_TARBALL])
+        print("node_modules successfully installed from cache.")
+    elif mode == 'upload':
+        if isfile(NEED_TO_UPLOAD_MARKER):
+            print("Creating tarball...")
+            run(['tar', 'czf', NODE_MODULES_TARBALL, 'node_modules'])
+            print("Uploading tarball to S3... ({})".format(_tarball_size()))
+            key.set_contents_from_filename(NODE_MODULES_TARBALL)
+            print("node_modules cache successfully updated.")
+            _delete_file_quietly(NEED_TO_UPLOAD_MARKER)
+        else:
+            print("No need to upload anything.")
+    else:
+        raise SystemExit("Unrecognized mode {!r}".format(mode))