src/tools/telemetry/cloud_storage

   1 #!/usr/bin/env python
   2 # Copyright 2014 The Chromium Authors. All rights reserved.
   3 # Use of this source code is governed by a BSD-style license that can be
   4 # found in the LICENSE file.
   5
   6 import argparse
   7 import logging
   8 import os
   9 import subprocess
  10
  11 from telemetry.core import command_line
  12 from telemetry.page import cloud_storage
  13
  14
  15 BUCKET_CHOICES = {
  16     'public': cloud_storage.PUBLIC_BUCKET,
  17     #'partner': cloud_storage.PARTNER_BUCKET,
  18     'google-only': cloud_storage.INTERNAL_BUCKET,
  19 }
  20
  21
  22 def _GetPaths(path):
  23   root, ext = os.path.splitext(path)
  24   if ext == '.sha1':
  25     file_path = root
  26     hash_path = path
  27   else:
  28     file_path = path
  29     hash_path = path + '.sha1'
  30   return file_path, hash_path
  31
  32
  33 def _FindFilesInCloudStorage(files):
  34   bucket_contents = {}
  35   for easy_bucket_name, bucket in BUCKET_CHOICES.iteritems():
  36     try:
  37       bucket_contents[easy_bucket_name] = cloud_storage.List(bucket)
  38     except (cloud_storage.PermissionError, cloud_storage.CredentialsError):
  39       pass
  40
  41   file_buckets = {}
  42   for path in files:
  43     file_path, hash_path = _GetPaths(path)
  44
  45     if file_path in file_buckets:
  46       continue
  47     if not os.path.exists(hash_path):
  48       continue
  49
  50     with open(hash_path, 'rb') as f:
  51       file_hash = f.read(1024).rstrip()
  52
  53     buckets = []
  54     for bucket in BUCKET_CHOICES:
  55       if bucket not in bucket_contents:
  56         continue
  57       if file_hash in bucket_contents[bucket]:
  58         buckets.append(bucket)
  59
  60     file_buckets[file_path] = buckets
  61
  62   return file_buckets
  63
  64
  65 class Ls(command_line.ArgparseCommand):
  66   """List which bucket each file is in."""
  67
  68   def AddCommandLineOptions(self, parser):
  69     parser.add_argument('-r', '--recursive', action='store_true')
  70     parser.add_argument('paths', nargs='+')
  71
  72   def ProcessCommandLine(self, parser, args):
  73     for path in args.paths:
  74       if not os.path.exists(path):
  75         parser.error('File not found: %s' % path)
  76
  77   def Run(self, args):
  78     def GetFilesInPath(paths, recursive):
  79       """If path is a dir, yields all files in path, otherwise just yields path.
  80
  81       If recursive is true, walks subdirectories recursively."""
  82       for path in paths:
  83         if not os.path.isdir(path):
  84           yield path
  85           return
  86
  87         if recursive:
  88           for root, _, filenames in os.walk(path):
  89             for filename in filenames:
  90               yield os.path.join(root, filename)
  91         else:
  92           for filename in os.listdir(path):
  93             yield os.path.join(path, filename)
  94
  95     files = _FindFilesInCloudStorage(GetFilesInPath(args.paths, args.recursive))
  96
  97     if not files:
  98       print 'No files in Cloud Storage.'
  99       return
 100
 101     for file_path, buckets in sorted(files.iteritems()):
 102       if buckets:
 103         print '%-11s  %s' % (','.join(buckets), file_path)
 104       else:
 105         print '%-11s  %s' % ('not found', file_path)
 106
 107
 108 class Mv(command_line.ArgparseCommand):
 109   """Move files to the given bucket."""
 110
 111   def AddCommandLineOptions(self, parser):
 112     parser.add_argument('files', nargs='+')
 113     parser.add_argument('bucket', choices=BUCKET_CHOICES)
 114
 115   def ProcessCommandLine(self, parser, args):
 116     args.bucket = BUCKET_CHOICES[args.bucket]
 117
 118     for path in args.paths:
 119       _, hash_path = _GetPaths(path)
 120       if not os.path.exists(hash_path):
 121         parser.error('File not found: %s' % hash_path)
 122
 123   def Run(self, args):
 124     files = _FindFilesInCloudStorage(args.files)
 125
 126     for file_path, buckets in sorted(files.iteritems()):
 127       if not buckets:
 128         raise IOError('%s not found in Cloud Storage.' % file_path)
 129
 130     for file_path, buckets in sorted(files.iteritems()):
 131       hash_path = file_path + '.sha1'
 132       with open(hash_path, 'rb') as f:
 133         file_hash = f.read(1024).rstrip()
 134
 135       cloud_storage.Move(buckets.pop(), args.bucket, file_hash)
 136
 137       for bucket in buckets:
 138         if bucket == args.bucket:
 139           continue
 140         cloud_storage.Delete(bucket, file_hash)
 141
 142
 143 class Rm(command_line.ArgparseCommand):
 144   """Remove files from Cloud Storage."""
 145
 146   def AddCommandLineOptions(self, parser):
 147     parser.add_argument('files', nargs='+')
 148
 149   def ProcessCommandLine(self, parser, args):
 150     for path in args.paths:
 151       _, hash_path = _GetPaths(path)
 152       if not os.path.exists(hash_path):
 153         parser.error('File not found: %s' % hash_path)
 154
 155   def Run(self, args):
 156     files = _FindFilesInCloudStorage(args.files)
 157     for file_path, buckets in sorted(files.iteritems()):
 158       hash_path = file_path + '.sha1'
 159       with open(hash_path, 'rb') as f:
 160         file_hash = f.read(1024).rstrip()
 161
 162       for bucket in buckets:
 163         cloud_storage.Delete(bucket, file_hash)
 164
 165
 166 class Upload(command_line.ArgparseCommand):
 167   """Upload files to Cloud Storage."""
 168
 169   def AddCommandLineOptions(self, parser):
 170     parser.add_argument('files', nargs='+')
 171     parser.add_argument('bucket', choices=BUCKET_CHOICES)
 172
 173   def ProcessCommandLine(self, parser, args):
 174     args.bucket = BUCKET_CHOICES[args.bucket]
 175
 176     for path in args.paths:
 177       if not os.path.exists(path):
 178         parser.error('File not found: %s' % path)
 179
 180   def Run(self, args):
 181     for file_path in args.files:
 182       file_hash = cloud_storage.GetHash(file_path)
 183
 184       # Create or update the hash file.
 185       hash_path = file_path + '.sha1'
 186       with open(hash_path, 'wb') as f:
 187         f.write(file_hash)
 188         f.flush()
 189
 190       # Add the data to Cloud Storage.
 191       cloud_storage.Insert(args.bucket, file_hash, file_path)
 192
 193       # Add the hash file to the branch, for convenience. :)
 194       subprocess.call(['git', 'add', hash_path])
 195
 196
 197 COMMANDS = (Ls, Mv, Rm, Upload)
 198
 199
 200 def main():
 201   logging.getLogger().setLevel(logging.INFO)
 202
 203   parser = argparse.ArgumentParser()
 204   subparsers = parser.add_subparsers()
 205
 206   for command in COMMANDS:
 207     command = command()
 208     subparser = subparsers.add_parser(command.name, help=command.description)
 209     subparser.set_defaults(command=command)
 210     command.AddCommandLineOptions(subparser)
 211
 212   args = parser.parse_args()
 213   args.command.ProcessCommandLine(parser, args)
 214   args.command.Run(args)
 215
 216
 217 if __name__ == '__main__':
 218   main()