Scripts to download files from google storage based on sha1 sums

build/download_from_google_storage.py

Index: build/download_from_google_storage.py
diff --git a/build/download_from_google_storage.py b/build/download_from_google_storage.py
new file mode 100755
index 0000000000000000000000000000000000000000..ea76339fe7240cdae0616290132269f1126a3234
--- /dev/null
+++ b/build/download_from_google_storage.py
@@ -0,0 +1,225 @@
+#!/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 download files from Google Storage."""
+
+
+import os
+import zipfile
+import tempfile
+import subprocess
+import hashlib
+import sys
+import re
+import zipfile
+import threading
+import time
+import Queue

[cmp, 2013/01/04 17:42:04]

these should be sorted alphabetically

[Ryan Tseng, 2013/01/14 21:37:12]

Done.

+from optparse import OptionParser

[cmp, 2013/01/04 17:42:04]

insert an empty line before line 20

[Ryan Tseng, 2013/01/14 21:37:12]

Done.

+GSUTIL_DEFAULT_PATH = os.path.join(os.path.dirname(os.path.normpath(__file__)),

[cmp, 2013/01/04 17:42:04]

insert an empty line before line 21

[Ryan Tseng, 2013/01/14 21:37:12]

Done.

+          '..', '..', 'third_party', 'gsutil', 'gsutil')
+
+
+class Gsutil():

[szager1, 2013/01/04 17:54:02]

I strongly recommend (but will not require) timeout functionality for this
class.  Google Storage can be flaky and transfers can hang indefinitely.  Refer
to tools/build/scripts/slave/gitzip.py for an example.

[Ryan Tseng, 2013/01/14 21:37:12]

Done.

+  def __init__(self, path):

[szager1, 2013/01/04 17:54:02]

Maybe optionally take a path to a credential file?

def __init__(self, gsutil_path, boto_path=None):
  ...

def call(self, *args):
  env = os.environ.copy()
  if self.boto_path is not None:
    env['AWS_CREDENTIAL_FILE'] = self.boto_path
  p = subprocess.Popen(..., env=env)

Refer to GSUtilSetup() in tools/build/scripts/slave/slave_utils.py

[Ryan Tseng, 2013/01/14 21:37:12]

Done.

+    if os.path.exists(path):
+      self.path = path
+    else:
+      raise IOError('GSUtil not found in %s' % path)

[szager1, 2013/01/04 17:54:02]

I'd prefer for this error to be surfaced the same way that subprocess.Popen
does, i.e., with OSError.  In fact, it's probably OK to eliminate this
error-checking entirely, and let subprocess.Popen do it.

+    stdout = None

[szager1, 2013/01/04 17:54:02]

This and the next line are no-ops

[Ryan Tseng, 2013/01/14 21:37:12]

Done.

+    stderr = None
+  def call_interactive(self, *args):

[cmp, 2013/01/04 17:42:04]

at indent of 2 spaces, there should be a single empty line before each "def"
method definition

[szager1, 2013/01/04 17:54:02]

I recommend renaming this and the next method to call() and check_call().  That
aligns with the convention of the subprocess module, and it's a bit more
intuitive.

[Ryan Tseng, 2013/01/14 21:37:12]

Done.

+    p = subprocess.Popen(('python', self.path) + args, stdout=sys.stdout,

[cmp, 2013/01/04 17:42:04]

please use sys.executable instead of 'python' since the script may have been
invoked with a particular python and we want to keep using that.  also change
line 38 to use sys.executable.

[szager1, 2013/01/04 17:54:02]

stdin, stdout, stderr args are unnecessary; default behavior is to inherit these
from the parent.

[Ryan Tseng, 2013/01/14 21:37:12]

Done.

+                         stderr=sys.stderr, stdin=sys.stdin)
+    return p.wait()
+  def call(self, *args):
+    p = subprocess.Popen(('python', self.path) + args, stdout=subprocess.PIPE,
+                         stderr=subprocess.PIPE)
+    code = p.wait()
+    out, err = p.communicate()
+    self.stdout = out
+    self.stderr = err
+
+    if code == 0:
+      return 0
+    else:

[cmp, 2013/01/04 17:42:04]

since you return from the 'code == 0' branch at line 46, you can drop the else
on line 47 and de-indent lines 48-57 by 2 spaces

[Ryan Tseng, 2013/01/14 21:37:12]

Done.

+      status_code_match = re.search('status=([0-9]+)', err)
+      if status_code_match:
+        return int(status_code_match.groups(1))
+      elif ('You are attempting to access protected data with '
+            'no configured credentials.' in err):
+        return 403
+      elif 'No such object' in err:
+        return 404
+      else:
+        return code
+
+
+def check_sha1(sha1_sum, filename):

[szager1, 2013/01/04 17:54:02]

CamelCase

[Ryan Tseng, 2013/01/14 21:37:12]

Done.

+  sha1 = hashlib.sha1()
+  sha1.update(open(filename).read())
+  return sha1_sum == sha1.hexdigest()
+
+
+def _downloader_worker_thread(thread_num, q, options, base_url, gsutil):
+  while True:
+    try:
+      input_sha1_sum, output_filename = q.get_nowait()
+      if os.path.exists(output_filename) and not options.force:
+        if check_sha1(input_sha1_sum, output_filename):
+          print 'File %s exists and SHA1 sum (%s) matches. Skipping.' % (
+              output_filename , input_sha1_sum)
+          continue
+      # Check if file exists.
+      file_url = '%s/%s' % (base_url, input_sha1_sum)
+      if gsutil.call('ls', file_url) != 0:
+        print >>sys.stderr, 'File %s for %s does not exist, skipping.' % (
+            file_url, output_filename)
+        continue
+      # Fetch the file.
+      print 'Downloading %s to %s...' % (file_url, output_filename)
+      code = gsutil.call_interactive('cp', '-q', file_url, output_filename)
+      if code != 0:
+        print >>sys.stderr, gsutil.stderr
+        return code
+      # TODO(hinoka): Delete and unzip.
+      # if options.unzip:
+      #   with zipfile.ZipFile(options.output, 'r') as source_zipfile:
+      #     source_zipfile.extractall(os.path.dirname(options.output))
+    except Queue.Empty:
+      return
+
+
+def main(args):
+  parser = OptionParser()
+  parser.add_option('-z', '--unzip', action='store_true', default=False,
+                    help='The target file is a zip file, unzip file after '
+                          'the download completes.')
+  parser.add_option('-d', '--delete', action='store_true', default=False,
+                    help='Deletes the target file after unzip.')
+  parser.add_option('-o', '--output', default=None,
+                    help='Specify the output file name.'
+                         'Defaults to the SHA1 hash.')
+  parser.add_option('-b', '--bucket', default='chrome-artifacts',
+                    help='Google Storage bucket to fetch from.')
+  parser.add_option('-f', '--force', action='store_true', default=False,
+                    help='Force download even if local file exists.')
+  parser.add_option('-r', '--recursive', action='store_true', default=False,
+                    help='Scan folders recursively for .sha1 files.')
+  parser.add_option('-t', '--num_threads', default=1, type='int',
+                    help='Number of downloader threads to run.')
+  # This file should be stored in tools/deps_scripts/ and we want the path to
+  # third_party/gsutil/gsutil
+  parser.add_option('-g', '--gsutil_path', default=GSUTIL_DEFAULT_PATH,
+          help='Path to the gsutil script.')
+
+  (options, args) = parser.parse_args()
+  if len(args) < 1:
+    print >>sys.stderr, 'Missing target.'

[szager1, 2013/01/04 17:54:02]

This is cryptic; please add a usage message.

[Ryan Tseng, 2013/01/14 21:37:12]

Done.

+    return 1
+  else:
+    input_filename = args[0]

[szager1, 2013/01/04 17:54:02]

What about handling multiple input files?  Currently, providing multiple input
files will silently ignore all but the first, which would be surprising to me.

[Ryan Tseng, 2013/01/14 21:37:12]

Added error message and TODO

+
+  # This could mean one of three things:
+  #   1. The input is a directory
+  #   2. The input is an already downloaded binary file.
+  #   3. The input is a .sha1 file

[szager1, 2013/01/04 17:54:02]

Please add an option to create a .sha1 file after download.  It's not at all
clear to me what the contents of a .sha1 file should be; there ought to be a
standard programmatic way to create it.

[Ryan Tseng, 2013/01/14 21:37:12]

Its done using the upload_to_google_storage.py script.
A sha1 file is just a file that contains a ([a-z0-9]{40}) sha1 sum.
It doesn't really make sense to create a .sha1 file after download since the
sha1 sum has to be known before anything can be downloaded, so its just as easy
to go
# echo "thesha1sumhere" > somefile.tar.gz.sha1

+  if os.path.exists(input_filename):
+    if os.path.isdir(input_filename):
+      dir_name = True
+      checked_sha1 = False
+    elif check_sha1(input_filename, input_filename):

[szager1, 2013/01/04 17:54:02]

This is the most expensive condition to check, so it should go last, if
possible.

[Ryan Tseng, 2013/01/14 21:37:12]

Done.

+      dir_name = False
+      checked_sha1 = True
+    else:
+      with open(input_filename) as f:
+        sha1_match = re.search('([A-Za-z0-9]{40})', f.read(1024))
+      if sha1_match:
+        if input_filename.endswith('.sha1'):
+          options.output = input_filename.replace('.sha1', '')

[szager1, 2013/01/04 17:54:02]

Maybe I'm paranoid, but I prefer:

options.output = input_filename[:-5]

[Ryan Tseng, 2013/01/14 21:37:12]

Done.

+        input_filename = sha1_match.groups(1)
+        dir_name = False
+        checked_sha1 = False
+      else:
+        if not re.match('[A-Za-z0-9]{40}', input_filename):
+          print >>sys.stderr, 'Input %s not recognized.' % input_filename

[szager1, 2013/01/04 17:54:02]

Please include the substance of this error message in your usage message.

[Ryan Tseng, 2013/01/14 21:37:12]

Done.

+          print >>sys.stderr, 'Input must be: '
+          print >>sys.stderr, '(1) a directory, (2) a .sha1 file or '
+          print >>sys.stderr, '(3) a sha1 sum ([A-Za-z0-9]{40})'
+          return 1
+
+  if not options.output:
+    options.output = input_filename
+  base_url = 'gs://%s' % options.bucket
+
+  if os.path.exists(options.gsutil_path):
+    gsutil = Gsutil(options.gsutil_path)

[szager1, 2013/01/04 17:54:02]

The way your code is structured, there will be a single Gsutil instance, and
multiple threads will invoke methods on that instance simultaneously.  That is a
very dangerous approach.

I strongly recommend you restructure so that there is one Gsutil instance per
task.  work_queue should contain instances of Gsutil, rather than tuples of
(input, output).

[Ryan Tseng, 2013/01/14 21:37:12]

Done by cloning the gsutil instance for each worker thread.

+  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.
+  code = gsutil.call('ls', base_url)
+  if code == 403:
+    code = gsutil.call_interactive('config')
+    if code != 0:
+      print >>sys.stderr, 'Error while authenticating to %s, exiting' % base_url
+      return 403
+  elif code == 404:
+    print >>sys.stderr, '%s not found.' % base_url
+    return 404
+  elif code != 0:
+    print >>sys.stderr, gsutil.stderr
+    return code
+
+  # Enumerate our work queue.
+  work_queue = Queue.Queue()
+  if dir_name:
+    if options.recursive:
+      for root, dirs, files in os.walk(input_filename):
+        if '.svn' in dirs:
+          dirs.remove('.svn')
+        if not options.recursive:
+          for item in dirs:
+            dirs.remove(item)
+        for filename in files:
+          full_path = os.path.join(root, filename)
+          if full_path.endswith('.sha1'):
+            with open(full_path) as f:
+              sha1_match = re.search('([A-Za-z0-9]{40})', f.read(1024))
+            if sha1_match:
+              work_queue.put((sha1_match.groups(1)[0],
+                                 full_path.replace('.sha1', '')))
+  else:
+    work_queue.put((input_filename, options.output))
+
+  # Start up all the worker threads.
+  all_threads = []
+  for thread_num in range(options.num_threads):
+    t = threading.Thread(target=_downloader_worker_thread, args=[thread_num,
+        work_queue, options, base_url, gsutil])
+    t.daemon = True
+    t.start()
+    all_threads.append(t)
+
+  def _wait_thread(threads, done):
+    for t in threads:
+      t.join()
+    print 'Now we\'re done'
+    done.set()
+
+  # Have a thread set a flag when all the tasks are done.

[szager1, 2013/01/04 17:54:02]

This is not necessary.  You can just join() all the worker threads from the main
thread, then exit.  ctrl-c will still work.

[Ryan Tseng, 2013/01/14 21:37:12]

Done.

+  done = threading.Event()
+  done_thread = threading.Thread(target=_wait_thread, args=[all_threads, done])
+  done_thread.daemon = True
+  done_thread.start()
+
+  while not done.is_set():
+    time.sleep(1)  # Do a sleep loop so we can ctrl + c out of this anytime.

[szager1, 2013/01/04 17:54:02]

My previous comment means this clause is unnecessary, but I want to emphasize
something here: no sleep() in multi-processing code, never ever ever.

[Ryan Tseng, 2013/01/14 21:37:12]

Done.

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