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