Scripts to download files from google storage based on sha1 sums

upload_to_google_storage.py

+#!/usr/bin/env python
+# Copyright (c) 2012 The Chromium Authors. All rights reserved.
+# Use of this source code is governed by a BSD-style license that can be
+# found in the LICENSE file.
+
+"""Uploads files to Google Storage content addressed."""
+
+import optparse
+import os
+import Queue
+import re
+import sys
+import threading
+import time
+
+import gstools
+
+GSUTIL_DEFAULT_PATH = os.path.join(
+    os.path.dirname(os.path.abspath(__file__)),
+    'third_party', 'gsutil', 'gsutil')
+
+USAGE_STRING = """%prog [options] target [target2 ...].
+Target is the file intended to be uploaded to Google Storage.
+If target is "-", then a list of files will be taken from standard input
+
+This script will generate a file (original filename).sha1 containing the
+sha1 sum of the uploaded file.
+It is recommended that the .sha1 file is checked into the repository,
+the original file removed from the repository, and a hook added to the
+DEPS file to call download_from_google_storage.py.
+
+Example usages
+--------------
+
+Scan the current directory and upload all files larger than 1MB:
+find . -name .svn -prune -o -size +1000k -type f -print0 | %prog -0 -
+"""
+
+
+def _upload_worker(thread_num, q, base_url, gsutil, options, md5_lock):
+  while True:
+    filename, sha1_sum = q.get()
+    if filename is None:

[M-A Ruel, 2013/02/27 21:52:07]

if not filename:
  break

optional style nit: Personally, I wouldn't keep the comment.

[Ryan Tseng, 2013/02/27 23:34:41]

Done.

+      break  # A (None, None) item is inserted for each thread to mark EOF.
+    file_url = '%s/%s' % (base_url, sha1_sum)
+    if gsutil.check_call('ls', file_url)[0] == 0 and not options.force:
+      # File exists, check MD5 hash.
+      _, out, _ = gsutil.check_call('ls', '-L', file_url)
+      etag_match = re.search('ETag:\s+([a-z0-9]{32})', out)
+      if etag_match:
+        remote_md5 = etag_match.group(1)
+        # Calculate the MD5 checksum to match it to Google Storage's ETag.
+        if options.use_md5:
+          local_md5 = gstools.GetMD5Cached(filename, md5_lock)
+        else:
+          local_md5 = gstools.GetMD5(filename, md5_lock)
+        if local_md5 == remote_md5:
+          print ('File %s already exists at %s and MD5 matches, exiting' %
+                 (filename, file_url))
+          continue
+    print 'Uploading %s to %s' % (filename, file_url)
+    code = gsutil.call('cp', '-q', filename, file_url)

[M-A Ruel, 2013/02/27 21:52:07]

Does this library throws exceptions when receiving an HTTP 500, TCP RST, etc?

[Ryan Tseng, 2013/02/27 23:34:41]

It doesn't throw an exception, it just returns a non-zero code.
It can be made to throw exceptions for non-whitelisted errors though, is that
preferable?

[M-A Ruel, 2013/02/28 14:53:56]

No, my point was more that occasionally, exceptions like IOError can leak
through a library even if they try to handle "most" of them. So it's fine to
leave it as-is and look for issues in practice.

+    if code != 0:
+      print >> sys.stderr, gsutil.stderr
+      continue
+
+
+def get_targets(options, args, parser):
+  if not args:
+    parser.error('Missing target.')
+  elif len(args) == 1 and args[0] == '-':

[M-A Ruel, 2013/02/27 21:52:07]

s/elif/if/
and add an empty line above. It's clearer.

[Ryan Tseng, 2013/02/27 23:34:41]

Done.

+    # Take stdin as a newline or null seperated list of files.
+    if options.use_null_terminator:
+      input_filenames = sys.stdin.read().split('\0')

[M-A Ruel, 2013/02/27 21:52:07]

return sys.stdin.read().split('\0')

[Ryan Tseng, 2013/02/27 23:34:41]

Done.

+    else:
+      input_filenames = sys.stdin.read().splitlines()

[M-A Ruel, 2013/02/27 21:52:07]

return 

[Ryan Tseng, 2013/02/27 23:34:41]

Done.

+  else:

[M-A Ruel, 2013/02/27 21:52:07]

return args
remove the "else:" line, not needed.

[Ryan Tseng, 2013/02/27 23:34:41]

Done.

+    input_filenames = args
+
+  return input_filenames
+
+
+def upload_to_google_storage(input_filenames, base_url, gsutil, options):
+  # We only want one MD5 calculation happening at a time to avoid HD thrashing.
+  md5_lock = threading.Lock()
+
+  # Start up all the worker threads.
+  all_threads = []
+  upload_queue = Queue.Queue()
+  upload_timer = time.time()
+  for thread_num in range(options.num_threads):
+    t = threading.Thread(
+        target=_upload_worker,
+        args=[thread_num, upload_queue, base_url, 
+              gsutil.clone(), options, md5_lock])
+    t.daemon = True
+    t.start()
+    all_threads.append(t)
+
+  # We want to hash everything in a single thread since its faster.
+  # The bottleneck is in disk IO, not CPU.
+  hash_timer = time.time()  # For timing statistics.
+  for filename in input_filenames:
+    if not os.path.exists(filename):
+      print 'Error: %s not found, skipping.' % filename
+      continue
+    if os.path.exists('%s.sha1' % filename) and options.skip_hashing:
+      print 'Found hash for %s, skipping.' % filename
+      upload_queue.put((filename, open('%s.sha1' % filename).read()))
+      continue
+    print 'Calculating hash for %s...' % filename,
+    sha1_sum = gstools.GetSHA1(filename)
+    with open(filename + '.sha1', 'wb') as f:
+      f.write(sha1_sum)
+    print 'done'
+    upload_queue.put((filename, sha1_sum))
+  hash_time = time.time() - hash_timer
+
+  # Wait for everything to finish.
+  for _ in all_threads:
+    upload_queue.put((None, None))  # To mark the end of the work queue.
+  for t in all_threads:
+    t.join()
+
+  print 'Success.'
+  print 'Hashing %s files took %1f seconds' % (len(input_filenames), hash_time)
+  print 'Uploading took %1f seconds' % (time.time() - upload_timer)
+  return 0
+
+
+def main(args):
+  parser = optparse.OptionParser(USAGE_STRING)
+  parser.add_option('-b', '--bucket',
+                    help='Google Storage bucket to upload to.')
+  parser.add_option('-e', '--boto', help='Specify a custom boto file.')
+  parser.add_option('-f', '--force', action='store_true',
+                    help='Force upload even if remote file exists.')
+  parser.add_option('-g', '--gsutil_path', default=GSUTIL_DEFAULT_PATH,
+                    help='Path to the gsutil script.')
+  parser.add_option('-m', '--use_md5', action='store_true',
+                    help='Generate MD5 files when scanning, and don\'t check '
+                    'the MD5 checksum if a .md5 file is found.')
+  parser.add_option('-t', '--num_threads', default=1, type='int',
+                    help='Number of uploader threads to run.')
+  parser.add_option('-s', '--skip_hashing', action='store_true',
+                    help='Skip hashing if .sha1 file exists.')
+  parser.add_option('-0', '--use_null_terminator', action='store_true',
+                    help='Use \\0 instead of \\n when parsing '
+                    'the file list from stdin.  This is useful if the input '
+                    'is coming from "find ... -print0".')
+  (options, args) = parser.parse_args()
+
+  # Enumerate our inputs.
+  input_filenames = get_targets(options, args, parser)
+
+  # Make sure we can find a working instance of gsutil.
+  if os.path.exists(GSUTIL_DEFAULT_PATH):
+    gsutil = gstools.Gsutil(GSUTIL_DEFAULT_PATH)
+  else:
+    print >> sys.stderr, ('gsutil not found in %s, bad depot_tools checkout?' %
+                          GSUTIL_DEFAULT_PATH)
+    return 1
+
+  # Check we have a valid bucket with valid permissions.
+  base_url, code = gstools.CheckBucketPermissions(options.bucket, gsutil)
+  if code:
+    return code
+
+  return upload_to_google_storage(input_filenames, base_url, gsutil, options)
+
+
+if __name__ == '__main__':
+  sys.exit(main(sys.argv))