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.
+
+"""Script to upload files to Google Storage."""

[M-A Ruel, 2013/02/22 01:15:56]

"""Uploads files to Google Storage."""

[Ryan Tseng, 2013/02/22 02:38:00]

Done.

+
+import optparse
+import os
+import Queue
+import re
+import sys
+import threading
+import time
+
+from common import Gsutil
+from common import GetSHA1
+from common import GetMD5
+
+# TODO(hinoka): This is currently incorrect. Should find a better default.
+GSUTIL_DEFAULT_PATH = os.path.join(os.path.dirname(os.path.normpath(__file__)),

[M-A Ruel, 2013/02/22 01:15:56]

s/normpath/abspath/

But you can't commit this as-is, unless you move this file into a subdirectory.

[Ryan Tseng, 2013/02/22 02:38:00]

Done.

+    '..', '..', '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 -
+"""
+
+

[M-A Ruel, 2013/02/22 01:15:56]

Remove one line

[Ryan Tseng, 2013/02/22 02:38:00]

Done.

+
+def _upload_worker(thread_num, q, base_url, gsutil, options, md5_lock):
+  while True:
+    try:
+      filename, sha1_sum = q.get_nowait()
+      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.groups()[0]

[M-A Ruel, 2013/02/22 01:15:56]

remote_md5 = etag_match.group(1)

[Ryan Tseng, 2013/02/22 02:38:00]

Done.

+          # Calculate the MD5 checksum to match it to Google Storage's ETag.
+          local_md5 = GetMD5(filename, md5_lock, options.use_md5)
+          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)
+      if code != 0:
+        print >> sys.stderr, gsutil.stderr
+        continue
+    except Queue.Empty:
+      return
+

[M-A Ruel, 2013/02/22 01:15:56]

2 lines

[Ryan Tseng, 2013/02/22 02:38:00]

Done.

+def main(args):
+  parser = optparse.OptionParser(USAGE_STRING)
+  parser.add_option('-b', '--bucket', default='chrome-artifacts',

[M-A Ruel, 2013/02/22 01:15:56]

I prefer no default, at least not if this file is going to live in depot_tools

[Ryan Tseng, 2013/02/22 02:38:00]

Done.

+                    help='Google Storage bucket to upload to.')
+  parser.add_option('-e', '--boto', default=None,

[M-A Ruel, 2013/02/22 01:15:56]

No need for default=None

[Ryan Tseng, 2013/02/22 02:38:00]

Done.

+                    help='Specify a custom boto file.')
+  parser.add_option('-f', '--force', action='store_true', default=False,

[M-A Ruel, 2013/02/22 01:15:56]

No need for default=False, same below.

[Ryan Tseng, 2013/02/22 02:38:00]

Done.

+                    help='Force upload even if remote file exists.')
+  parser.add_option('-g', '--gsutil_path', default=GSUTIL_DEFAULT_PATH,

[M-A Ruel, 2013/02/22 01:15:56]

Why this argument at all if gsutil is included in depot_tools?

[Ryan Tseng, 2013/02/22 02:38:00]

Removed.

+                    help='Path to the gsutil script.')
+  parser.add_option('-m', '--use_md5', action='store_true', default=False,
+                    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', default=False,

[M-A Ruel, 2013/02/22 01:15:56]

Why not the default? Same for --use_md5.

[Ryan Tseng, 2013/02/22 02:38:00]

I'm avoiding the situation where you modify the file, but forget to remove the
.sha1 file.  This way the default always ensure that the .sha1 file is correct. 
The flag should only be turned on if the user is confident that the .sha1 files
accurate reflect the latest changes.

+                    help='Skip hashing if .sha1 file exists.')
+  parser.add_option('-0', '--use_null_terminator', action='store_true',
+                    default=False, 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()
+
+  if len(args) < 1:
+    parser.error('Missing target.')
+  elif len(args) == 1 and args[0] == '-':
+    # Take stdin as a newline or null seperated list of files.
+    if options.use_null_terminator:
+      input_filenames = [line for line in sys.stdin.read().split('\0')]

[M-A Ruel, 2013/02/22 01:15:56]

input_filenames = sys.stdin.read().split('\0')

[Ryan Tseng, 2013/02/22 02:38:00]

Done.

+    else:
+      input_filenames = [line.strip() for line in sys.stdin.readlines()]

[M-A Ruel, 2013/02/22 01:15:56]

Technically, you would want to have it work with a file ending it a whitespace,
you'd want
input_files = sys.stdin.read().splitlines()

[Ryan Tseng, 2013/02/22 02:38:00]

Done.

+  else:
+    input_filenames = args
+  base_url = 'gs://%s' % options.bucket
+
+  # Make sure we can find a working instance of gsutil.
+  if os.path.exists(options.gsutil_path):
+    gsutil = Gsutil(options.gsutil_path)
+  else:
+    for path in os.environ["PATH"].split(os.pathsep):
+      if os.path.exists(path) and 'gsutil' in os.listdir(path):
+        gsutil = Gsutil(os.path.join(path, 'gsutil'))
+
+  # Check if we have permissions to the Google Storage bucket.
+  code, _, ls_err = gsutil.check_call('ls', base_url)
+  if code == 403:
+    code, _, _ = gsutil.call('config')
+    if code != 0:
+      print >> sys.stderr, 'Error while authenticating to %s.' % base_url
+      return 403
+  elif code == 404:
+    print >> sys.stderr, '%s not found.' % base_url
+    return 404
+  elif code != 0:
+    print >> sys.stderr, ls_err
+    return code
+
+  # We want to hash everything in a single thread since its faster.

[M-A Ruel, 2013/02/22 01:15:56]

I don't understand why it'd be faster since it's CPU bound.

[Ryan Tseng, 2013/02/22 02:38:00]

We are most definitely IO bound at harddrive read speed when calculating MD5
sums.

[M-A Ruel, 2013/02/25 15:15:06]

Err right, sorry.

+  # The bottleneck is in disk IO, not CPU.
+  upload_queue = Queue.Queue()
+  hash_timer = time.time()
+  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 = GetSHA1(filename)
+    with open(filename + '.sha1', 'w') as f:

[M-A Ruel, 2013/02/22 01:15:56]

'wb'

[Ryan Tseng, 2013/02/22 02:38:00]

Done.

+      f.write(sha1_sum)
+    print 'done'
+    upload_queue.put((filename, sha1_sum))
+  hash_time = time.time() - hash_timer
+
+  # Start up all the worker threads.
+  all_threads = []
+
+  # We only want one MD5 calculation happening at a time.

[M-A Ruel, 2013/02/22 01:15:56]

Why?

[Ryan Tseng, 2013/02/22 02:38:00]

Harddrive IO bound.  The harddrive read head jitters if we try to read two files
at a time.  I found that reading two files at a time doesn't actually increase
read speed.  In fact running 10 threads is significantly slower than 1 thread at
this stage.

+  md5_lock = threading.Lock()
+  upload_timer = time.time()
+
+  for thread_num in range(options.num_threads):
+    t = threading.Thread(target=_upload_worker, args=[thread_num,

[M-A Ruel, 2013/02/22 01:15:56]

Don't split arguments like that

[Ryan Tseng, 2013/02/22 02:38:00]

??? What did I do?

[M-A Ruel, 2013/02/25 15:15:06]

t = threading.Thread(
    target=_upload_worker,
    args=[
      thread_num, upload_queue, base_url,
      gsutil.clone(), options, md5_lock])

or as you prefer, but don't split an argument on two list without first putting
it first on the line.

[Ryan Tseng, 2013/02/27 02:06:55]

Done.

+        upload_queue, base_url, gsutil.clone(), options, md5_lock])
+    t.daemon = True
+    t.start()
+    all_threads.append(t)
+
+  # Wait for everything to finish.
+  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
+

[M-A Ruel, 2013/02/22 01:15:56]

two lines

[Ryan Tseng, 2013/02/22 02:38:00]

Done.

+if __name__ == '__main__':
+  sys.exit(main(sys.argv))