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