OLD | NEW |
(Empty) | |
| 1 # Copyright 2017 The Chromium Authors. All rights reserved. |
| 2 # Use of this source code is governed by a BSD-style license that can be |
| 3 # found in the LICENSE file. |
| 4 |
| 5 import hashlib |
| 6 import os |
| 7 |
| 8 from pylib.base import output_manager |
| 9 from pylib.utils import logdog_helper |
| 10 from pylib.utils import google_storage_helper |
| 11 |
| 12 |
| 13 class RemoteOutputManager(output_manager.OutputManager): |
| 14 |
| 15 def __init__(self, bucket): |
| 16 """Uploads output files to Google Storage or LogDog. |
| 17 |
| 18 Files will either be uploaded directly to Google Storage or LogDog |
| 19 depending on the datatype. |
| 20 |
| 21 Args |
| 22 bucket: Bucket to use when saving to Google Storage. |
| 23 """ |
| 24 super(RemoteOutputManager, self).__init__() |
| 25 self._bucket = bucket |
| 26 |
| 27 #override |
| 28 def _CreateArchiveJob(self, in_filepath, out_filename, out_subdir, datatype): |
| 29 if datatype == output_manager.Datatype.TEXT: |
| 30 return LogdogArchiveJob(in_filepath, out_filename, out_subdir, datatype) |
| 31 else: |
| 32 if self._bucket is None: |
| 33 return None |
| 34 return GoogleStorageArchiveJob( |
| 35 in_filepath, out_filename, out_subdir, datatype, self._bucket) |
| 36 |
| 37 |
| 38 class LogdogArchiveJob(output_manager.Job): |
| 39 |
| 40 def __init__(self, in_filepath, out_filename, out_subdir, datatype): |
| 41 super(LogdogArchiveJob, self).__init__( |
| 42 in_filepath, out_filename, out_subdir, datatype) |
| 43 self._stream_name = '%s_%s' % (out_subdir, out_filename) |
| 44 |
| 45 def Link(self): |
| 46 return logdog_helper.get_viewer_url(self._stream_name) |
| 47 |
| 48 def Archive(self): |
| 49 with open(self._in_filepath, 'r') as f: |
| 50 logdog_helper.text(self._stream_name, f.read()) |
| 51 |
| 52 |
| 53 class GoogleStorageArchiveJob(output_manager.Job): |
| 54 |
| 55 def __init__(self, in_filepath, out_filename, out_subdir, datatype, bucket): |
| 56 super(GoogleStorageArchiveJob, self).__init__( |
| 57 in_filepath, out_filename, out_subdir, datatype) |
| 58 self._bucket = bucket |
| 59 |
| 60 self._content_addressed = (self._datatype in ( |
| 61 output_manager.Datatype.HTML, |
| 62 output_manager.Datatype.IMAGE)) |
| 63 |
| 64 if self._content_addressed: |
| 65 sha1 = hashlib.sha1() |
| 66 with open(in_filepath, 'rb') as f: |
| 67 sha1.update(f.read()) |
| 68 self._upload_path = sha1.hexdigest() |
| 69 else: |
| 70 self._upload_path = os.path.join(out_subdir, out_filename) |
| 71 |
| 72 def Link(self): |
| 73 return google_storage_helper.get_url_link( |
| 74 self._upload_path, self._bucket) |
| 75 |
| 76 def Archive(self): |
| 77 if self._content_addressed: |
| 78 if google_storage_helper.exists(self._upload_path, self._bucket): |
| 79 return |
| 80 |
| 81 content_type = None |
| 82 if self._datatype == output_manager.Datatype.HTML: |
| 83 content_type = 'text/html' |
| 84 google_storage_helper.upload( |
| 85 self._upload_path, self._in_filepath, self._bucket, content_type) |
OLD | NEW |