OLD | NEW |
1 #!/usr/bin/env python | 1 #!/usr/bin/env python |
2 # Copyright (c) 2012 The Chromium Authors. All rights reserved. | 2 # Copyright (c) 2012 The Chromium Authors. All rights reserved. |
3 # Use of this source code is governed by a BSD-style license that can be | 3 # Use of this source code is governed by a BSD-style license that can be |
4 # found in the LICENSE file. | 4 # found in the LICENSE file. |
5 | 5 |
6 """Uploads files to Google Storage content addressed.""" | 6 """Uploads files to Google Storage content addressed.""" |
7 | 7 |
8 import hashlib | 8 import hashlib |
9 import optparse | 9 import optparse |
10 import os | 10 import os |
11 import Queue | 11 import Queue |
12 import re | 12 import re |
13 import stat | 13 import stat |
14 import sys | 14 import sys |
| 15 import tarfile |
15 import threading | 16 import threading |
16 import time | 17 import time |
17 | 18 |
18 from download_from_google_storage import check_bucket_permissions | 19 from download_from_google_storage import check_bucket_permissions |
19 from download_from_google_storage import get_sha1 | 20 from download_from_google_storage import get_sha1 |
20 from download_from_google_storage import Gsutil | 21 from download_from_google_storage import Gsutil |
21 from download_from_google_storage import printer_worker | 22 from download_from_google_storage import printer_worker |
22 from download_from_google_storage import GSUTIL_DEFAULT_PATH | 23 from download_from_google_storage import GSUTIL_DEFAULT_PATH |
23 | 24 |
24 USAGE_STRING = """%prog [options] target [target2 ...]. | 25 USAGE_STRING = """%prog [options] target [target2 ...]. |
(...skipping 98 matching lines...) Expand 10 before | Expand all | Expand 10 after Loading... |
123 if use_null_terminator: | 124 if use_null_terminator: |
124 return sys.stdin.read().split('\0') | 125 return sys.stdin.read().split('\0') |
125 else: | 126 else: |
126 return sys.stdin.read().splitlines() | 127 return sys.stdin.read().splitlines() |
127 else: | 128 else: |
128 return args | 129 return args |
129 | 130 |
130 | 131 |
131 def upload_to_google_storage( | 132 def upload_to_google_storage( |
132 input_filenames, base_url, gsutil, force, | 133 input_filenames, base_url, gsutil, force, |
133 use_md5, num_threads, skip_hashing): | 134 use_md5, num_threads, skip_hashing, archive): |
134 # We only want one MD5 calculation happening at a time to avoid HD thrashing. | 135 # We only want one MD5 calculation happening at a time to avoid HD thrashing. |
135 md5_lock = threading.Lock() | 136 md5_lock = threading.Lock() |
136 | 137 |
137 # Start up all the worker threads plus the printer thread. | 138 # Start up all the worker threads plus the printer thread. |
138 all_threads = [] | 139 all_threads = [] |
139 ret_codes = Queue.Queue() | 140 ret_codes = Queue.Queue() |
140 ret_codes.put((0, None)) | 141 ret_codes.put((0, None)) |
141 upload_queue = Queue.Queue() | 142 upload_queue = Queue.Queue() |
142 upload_timer = time.time() | 143 upload_timer = time.time() |
143 stdout_queue = Queue.Queue() | 144 stdout_queue = Queue.Queue() |
(...skipping 52 matching lines...) Expand 10 before | Expand all | Expand 10 after Loading... |
196 for ret_code, message in ret_codes.queue: | 197 for ret_code, message in ret_codes.queue: |
197 max_ret_code = max(ret_code, max_ret_code) | 198 max_ret_code = max(ret_code, max_ret_code) |
198 if message: | 199 if message: |
199 print >> sys.stderr, message | 200 print >> sys.stderr, message |
200 | 201 |
201 if not max_ret_code: | 202 if not max_ret_code: |
202 print 'Success!' | 203 print 'Success!' |
203 | 204 |
204 return max_ret_code | 205 return max_ret_code |
205 | 206 |
| 207 def create_archives(dirs): |
| 208 archive_names = [] |
| 209 for name in dirs: |
| 210 tarname = '%s.tar.gz' % name |
| 211 with tarfile.open(tarname, 'w:gz') as tar: |
| 212 tar.add(name) |
| 213 archive_names.append(tarname) |
| 214 return archive_names |
206 | 215 |
| 216 def validate_archive_dirs(dirs): |
| 217 def just_below_cwd(name): |
| 218 return name in next(os.walk('.'))[1] |
| 219 return reduce(lambda x, y: x + 1 if not just_below_cwd(y) else x, dirs, 0) |
| 220 |
207 def main(args): | 221 def main(args): |
208 parser = optparse.OptionParser(USAGE_STRING) | 222 parser = optparse.OptionParser(USAGE_STRING) |
209 parser.add_option('-b', '--bucket', | 223 parser.add_option('-b', '--bucket', |
210 help='Google Storage bucket to upload to.') | 224 help='Google Storage bucket to upload to.') |
211 parser.add_option('-e', '--boto', help='Specify a custom boto file.') | 225 parser.add_option('-e', '--boto', help='Specify a custom boto file.') |
| 226 parser.add_option('-z', '--archive', action='store_true', |
| 227 help='Archive directory as a tar.gz file') |
212 parser.add_option('-f', '--force', action='store_true', | 228 parser.add_option('-f', '--force', action='store_true', |
213 help='Force upload even if remote file exists.') | 229 help='Force upload even if remote file exists.') |
214 parser.add_option('-g', '--gsutil_path', default=GSUTIL_DEFAULT_PATH, | |
215 help='Path to the gsutil script.') | |
216 parser.add_option('-m', '--use_md5', action='store_true', | 230 parser.add_option('-m', '--use_md5', action='store_true', |
217 help='Generate MD5 files when scanning, and don\'t check ' | 231 help='Generate MD5 files when scanning, and don\'t check ' |
218 'the MD5 checksum if a .md5 file is found.') | 232 'the MD5 checksum if a .md5 file is found.') |
219 parser.add_option('-t', '--num_threads', default=1, type='int', | 233 parser.add_option('-t', '--num_threads', default=1, type='int', |
220 help='Number of uploader threads to run.') | 234 help='Number of uploader threads to run.') |
221 parser.add_option('-s', '--skip_hashing', action='store_true', | 235 parser.add_option('-s', '--skip_hashing', action='store_true', |
222 help='Skip hashing if .sha1 file exists.') | 236 help='Skip hashing if .sha1 file exists.') |
223 parser.add_option('-0', '--use_null_terminator', action='store_true', | 237 parser.add_option('-0', '--use_null_terminator', action='store_true', |
224 help='Use \\0 instead of \\n when parsing ' | 238 help='Use \\0 instead of \\n when parsing ' |
225 'the file list from stdin. This is useful if the input ' | 239 'the file list from stdin. This is useful if the input ' |
226 'is coming from "find ... -print0".') | 240 'is coming from "find ... -print0".') |
227 (options, args) = parser.parse_args() | 241 (options, args) = parser.parse_args() |
228 | 242 |
229 # Enumerate our inputs. | 243 # Enumerate our inputs. |
230 input_filenames = get_targets(args, parser, options.use_null_terminator) | 244 input_filenames = get_targets(args, parser, options.use_null_terminator) |
231 | 245 |
| 246 |
| 247 if options.archive: |
| 248 if validate_archive_dirs(input_filenames): |
| 249 parser.error('Only directories just below cwd are valid entries when ' |
| 250 'using the --archive argument. Entries was %s' % |
| 251 input_filenames) |
| 252 return 1 |
| 253 input_filenames = create_archives(input_filenames) |
| 254 |
232 # Make sure we can find a working instance of gsutil. | 255 # Make sure we can find a working instance of gsutil. |
233 if os.path.exists(GSUTIL_DEFAULT_PATH): | 256 if os.path.exists(GSUTIL_DEFAULT_PATH): |
234 gsutil = Gsutil(GSUTIL_DEFAULT_PATH, boto_path=options.boto) | 257 gsutil = Gsutil(GSUTIL_DEFAULT_PATH, boto_path=options.boto) |
235 else: | 258 else: |
236 gsutil = None | 259 gsutil = None |
237 for path in os.environ["PATH"].split(os.pathsep): | 260 for path in os.environ["PATH"].split(os.pathsep): |
238 if os.path.exists(path) and 'gsutil' in os.listdir(path): | 261 if os.path.exists(path) and 'gsutil' in os.listdir(path): |
239 gsutil = Gsutil(os.path.join(path, 'gsutil'), boto_path=options.boto) | 262 gsutil = Gsutil(os.path.join(path, 'gsutil'), boto_path=options.boto) |
240 if not gsutil: | 263 if not gsutil: |
241 parser.error('gsutil not found in %s, bad depot_tools checkout?' % | 264 parser.error('gsutil not found in %s, bad depot_tools checkout?' % |
242 GSUTIL_DEFAULT_PATH) | 265 GSUTIL_DEFAULT_PATH) |
243 | 266 |
244 base_url = 'gs://%s' % options.bucket | 267 base_url = 'gs://%s' % options.bucket |
245 | 268 |
246 # Check we have a valid bucket with valid permissions. | 269 # Check we have a valid bucket with valid permissions. |
247 code = check_bucket_permissions(base_url, gsutil) | 270 code = check_bucket_permissions(base_url, gsutil) |
248 if code: | 271 if code: |
249 return code | 272 return code |
250 | 273 |
251 return upload_to_google_storage( | 274 return upload_to_google_storage( |
252 input_filenames, base_url, gsutil, options.force, options.use_md5, | 275 input_filenames, base_url, gsutil, options.force, options.use_md5, |
253 options.num_threads, options.skip_hashing) | 276 options.num_threads, options.skip_hashing, options.archive) |
254 | 277 |
255 | 278 |
256 if __name__ == '__main__': | 279 if __name__ == '__main__': |
257 sys.exit(main(sys.argv)) | 280 sys.exit(main(sys.argv)) |
OLD | NEW |