s3_cache.py 3.5 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108
  1. #!/usr/bin/env python2.7
  2. from __future__ import absolute_import, unicode_literals, print_function, division
  3. from sys import argv
  4. from os import environ, stat, remove as _delete_file
  5. from os.path import isfile, dirname, basename, abspath
  6. from hashlib import sha256
  7. from subprocess import check_call as run
  8. from boto.s3.connection import S3Connection
  9. from boto.s3.key import Key
  10. from boto.exception import S3ResponseError
  11. NEED_TO_UPLOAD_MARKER = '.need-to-upload'
  12. BYTES_PER_MB = 1024 * 1024
  13. try:
  14. BUCKET_NAME = environ['TWBS_S3_BUCKET']
  15. except KeyError:
  16. raise SystemExit("TWBS_S3_BUCKET environment variable not set!")
  17. def _sha256_of_file(filename):
  18. hasher = sha256()
  19. with open(filename, 'rb') as input_file:
  20. hasher.update(input_file.read())
  21. file_hash = hasher.hexdigest()
  22. print('sha256({}) = {}'.format(filename, file_hash))
  23. return file_hash
  24. def _delete_file_quietly(filename):
  25. try:
  26. _delete_file(filename)
  27. except (OSError, IOError):
  28. pass
  29. def _tarball_size(directory):
  30. kib = stat(_tarball_filename_for(directory)).st_size // BYTES_PER_MB
  31. return "{} MiB".format(kib)
  32. def _tarball_filename_for(directory):
  33. return abspath('./{}.tar.gz'.format(basename(directory)))
  34. def _create_tarball(directory):
  35. print("Creating tarball of {}...".format(directory))
  36. run(['tar', '-czf', _tarball_filename_for(directory), '-C', dirname(directory), basename(directory)])
  37. def _extract_tarball(directory):
  38. print("Extracting tarball of {}...".format(directory))
  39. run(['tar', '-xzf', _tarball_filename_for(directory), '-C', dirname(directory)])
  40. def download(directory):
  41. _delete_file_quietly(NEED_TO_UPLOAD_MARKER)
  42. try:
  43. print("Downloading {} tarball from S3...".format(friendly_name))
  44. key.get_contents_to_filename(_tarball_filename_for(directory))
  45. except S3ResponseError as err:
  46. open(NEED_TO_UPLOAD_MARKER, 'a').close()
  47. print(err)
  48. raise SystemExit("Cached {} download failed!".format(friendly_name))
  49. print("Downloaded {}.".format(_tarball_size(directory)))
  50. _extract_tarball(directory)
  51. print("{} successfully installed from cache.".format(friendly_name))
  52. def upload(directory):
  53. _create_tarball(directory)
  54. print("Uploading {} tarball to S3... ({})".format(friendly_name, _tarball_size(directory)))
  55. key.set_contents_from_filename(_tarball_filename_for(directory))
  56. print("{} cache successfully updated.".format(friendly_name))
  57. _delete_file_quietly(NEED_TO_UPLOAD_MARKER)
  58. if __name__ == '__main__':
  59. # Uses environment variables:
  60. # AWS_ACCESS_KEY_ID -- AWS Access Key ID
  61. # AWS_SECRET_ACCESS_KEY -- AWS Secret Access Key
  62. argv.pop(0)
  63. if len(argv) != 4:
  64. raise SystemExit("USAGE: s3_cache.py <download | upload> <friendly name> <dependencies file> <directory>")
  65. mode, friendly_name, dependencies_file, directory = argv
  66. conn = S3Connection()
  67. bucket = conn.lookup(BUCKET_NAME, validate=False)
  68. if bucket is None:
  69. raise SystemExit("Could not access bucket!")
  70. dependencies_file_hash = _sha256_of_file(dependencies_file)
  71. key = Key(bucket, dependencies_file_hash)
  72. key.storage_class = 'REDUCED_REDUNDANCY'
  73. if mode == 'download':
  74. download(directory)
  75. elif mode == 'upload':
  76. if isfile(NEED_TO_UPLOAD_MARKER): # FIXME
  77. upload(directory)
  78. else:
  79. print("No need to upload anything.")
  80. else:
  81. raise SystemExit("Unrecognized mode {!r}".format(mode))