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_ALIASES = {
  16     'public': cloud_storage.PUBLIC_BUCKET,
  17     'partner': cloud_storage.PARTNER_BUCKET,
  18     'google-only': cloud_storage.INTERNAL_BUCKET,
  19 }
  20 BUCKETS = {bucket: easy_bucket_name for easy_bucket_name, bucket
  21            in BUCKET_ALIASES.iteritems()}
  22
  23
  24 def _GetPaths(path):
  25   root, ext = os.path.splitext(path)
  26   if ext == '.sha1':
  27     file_path = root
  28     hash_path = path
  29   else:
  30     file_path = path
  31     hash_path = path + '.sha1'
  32   return file_path, hash_path
  33
  34
  35 def _FindFilesInCloudStorage(files):
  36   """Returns a dict of all files and which buckets they're in."""
  37   # Preprocessing: get the contents of all buckets.
  38   bucket_contents = {}
  39   for bucket in BUCKETS:
  40     try:
  41       bucket_contents[bucket] = cloud_storage.List(bucket)
  42     except (cloud_storage.PermissionError, cloud_storage.CredentialsError):
  43       pass
  44
  45   # Check if each file is in the bucket contents.
  46   file_buckets = {}
  47   for path in files:
  48     file_path, hash_path = _GetPaths(path)
  49
  50     if file_path in file_buckets:
  51       # Ignore duplicates, if both data and sha1 file were in the file list.
  52       continue
  53     if not os.path.exists(hash_path):
  54       # Probably got some non-Cloud Storage files in the file list. Ignore.
  55       continue
  56
  57     file_hash = cloud_storage.ReadHash(hash_path)
  58     file_buckets[file_path] = []
  59     for bucket in BUCKETS:
  60       if bucket in bucket_contents and file_hash in bucket_contents[bucket]:
  61         file_buckets[file_path].append(bucket)
  62
  63   return file_buckets
  64
  65
  66 class Ls(command_line.Command):
  67   """List which bucket each file is in."""
  68
  69   @classmethod
  70   def AddCommandLineArgs(cls, parser):
  71     parser.add_argument('-r', '--recursive', action='store_true')
  72     parser.add_argument('paths', nargs='+')
  73
  74   @classmethod
  75   def ProcessCommandLineArgs(cls, parser, args):
  76     for path in args.paths:
  77       if not os.path.exists(path):
  78         parser.error('Path not found: %s' % path)
  79
  80   def Run(self, args):
  81     def GetFilesInPaths(paths, recursive):
  82       """If path is a dir, yields all files in path, otherwise just yields path.
  83
  84       If recursive is true, walks subdirectories recursively."""
  85       for path in paths:
  86         if not os.path.isdir(path):
  87           yield path
  88           continue
  89
  90         if recursive:
  91           for root, _, filenames in os.walk(path):
  92             for filename in filenames:
  93               yield os.path.join(root, filename)
  94         else:
  95           for filename in os.listdir(path):
  96             yield os.path.join(path, filename)
  97
  98     files = _FindFilesInCloudStorage(GetFilesInPaths(args.paths, args.recursive))
  99
 100     if not files:
 101       print 'No files in Cloud Storage.'
 102       return
 103
 104     for file_path, buckets in sorted(files.iteritems()):
 105       if buckets:
 106         buckets = [BUCKETS[bucket] for bucket in buckets]
 107         print '%-11s  %s' % (','.join(buckets), file_path)
 108       else:
 109         print '%-11s  %s' % ('not found', file_path)
 110
 111
 112 class Mv(command_line.Command):
 113   """Move files to the given bucket."""
 114
 115   @classmethod
 116   def AddCommandLineArgs(cls, parser):
 117     parser.add_argument('files', nargs='+')
 118     parser.add_argument('bucket', choices=BUCKET_ALIASES)
 119
 120   @classmethod
 121   def ProcessCommandLineArgs(cls, parser, args):
 122     args.bucket = BUCKET_ALIASES[args.bucket]
 123
 124   def Run(self, args):
 125     files = _FindFilesInCloudStorage(args.files)
 126
 127     for file_path, buckets in sorted(files.iteritems()):
 128       if not buckets:
 129         raise IOError('%s not found in Cloud Storage.' % file_path)
 130
 131     for file_path, buckets in sorted(files.iteritems()):
 132       if args.bucket in buckets:
 133         buckets.remove(args.bucket)
 134       if not buckets:
 135         logging.info('Skipping %s, no action needed.' % file_path)
 136         continue
 137
 138       # Move to the target bucket.
 139       file_hash = cloud_storage.ReadHash(file_path + '.sha1')
 140       cloud_storage.Move(buckets.pop(), args.bucket, file_hash)
 141
 142       # Delete all additional copies.
 143       for bucket in buckets:
 144         cloud_storage.Delete(bucket, file_hash)
 145
 146
 147 class Rm(command_line.Command):
 148   """Remove files from Cloud Storage."""
 149
 150   @classmethod
 151   def AddCommandLineArgs(cls, parser):
 152     parser.add_argument('files', nargs='+')
 153
 154   def Run(self, args):
 155     files = _FindFilesInCloudStorage(args.files)
 156     for file_path, buckets in sorted(files.iteritems()):
 157       file_hash = cloud_storage.ReadHash(file_path + '.sha1')
 158       for bucket in buckets:
 159         cloud_storage.Delete(bucket, file_hash)
 160
 161
 162 class Upload(command_line.Command):
 163   """Upload files to Cloud Storage."""
 164
 165   @classmethod
 166   def AddCommandLineArgs(cls, parser):
 167     parser.add_argument('files', nargs='+')
 168     parser.add_argument('bucket', choices=BUCKET_ALIASES)
 169
 170   @classmethod
 171   def ProcessCommandLineArgs(cls, parser, args):
 172     args.bucket = BUCKET_ALIASES[args.bucket]
 173
 174     for path in args.files:
 175       if not os.path.exists(path):
 176         parser.error('File not found: %s' % path)
 177
 178   def Run(self, args):
 179     for file_path in args.files:
 180       file_hash = cloud_storage.CalculateHash(file_path)
 181
 182       # Create or update the hash file.
 183       hash_path = file_path + '.sha1'
 184       with open(hash_path, 'wb') as f:
 185         f.write(file_hash)
 186         f.flush()
 187
 188       # Add the data to Cloud Storage.
 189       cloud_storage.Insert(args.bucket, file_hash, file_path)
 190
 191       # Add the hash file to the branch, for convenience. :)
 192       subprocess.call(['git', 'add', hash_path])
 193
 194
 195 COMMANDS = (Ls, Mv, Rm, Upload)
 196
 197
 198 def main():
 199   logging.getLogger().setLevel(logging.INFO)
 200
 201   parser = argparse.ArgumentParser()
 202   subparsers = parser.add_subparsers()
 203
 204   for command in COMMANDS:
 205     subparser = subparsers.add_parser(
 206         command.Name(), help=command.Description())
 207     subparser.set_defaults(command=command)
 208     command.AddCommandLineArgs(subparser)
 209
 210   args = parser.parse_args()
 211   args.command.ProcessCommandLineArgs(parser, args)
 212   args.command().Run(args)
 213
 214
 215 if __name__ == '__main__':
 216   main()