1# Copyright (c) 2015 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
5import os
6
7import cloudstorage as gcs
8
9
10class EnvVarModifier(object):
11  def __init__(self, **kwargs):
12    self._vars = {}
13    self._kwargs = kwargs
14
15  def __enter__(self):
16    for k, v in self._kwargs.iteritems():
17      self._vars[k] = os.environ.get(k)
18      os.environ[k] = v
19    return self
20
21  def __exit__(self, *_):
22    for k, v in self._vars.iteritems():
23      os.environ[k] = v
24
25
26default_retry_params = gcs.RetryParams(initial_delay=0.2,
27                                       max_delay=5.0,
28                                       backoff_factor=2,
29                                       max_retry_period=15)
30gcs.set_default_retry_params(default_retry_params)
31
32
33def _remove_gcs_prefix(full_url):
34  return full_url.split('gs:/')[1]
35
36
37def WriteGCS(fullurl, data):
38  with EnvVarModifier(SERVER_SOFTWARE='') as _:
39    gcs_file = gcs.open(_remove_gcs_prefix(fullurl),
40                        'w',
41                        content_type='text/plain',
42                        options={},
43                        retry_params=default_retry_params)
44    gcs_file.write(data)
45    gcs_file.close()
46
47
48def ReadGCS(fullurl):
49  with EnvVarModifier(SERVER_SOFTWARE='') as _:
50    gcs_file = gcs.open(_remove_gcs_prefix(fullurl),
51                        'r',
52                        retry_params=default_retry_params)
53
54    contents = gcs_file.read()
55    gcs_file.close()
56
57    return contents
58
59
60def ReadGCSToFile(fullurl, output_file):
61  output_file.write(ReadGCS(fullurl))
62
63
64def StatGCS(fullurl):
65  with EnvVarModifier(SERVER_SOFTWARE='') as _:
66    try:
67      return gcs.stat(_remove_gcs_prefix(fullurl))
68    except gcs.NotFoundError:
69      return None
70