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
6from recipe_engine import recipe_api
7
8UPLOAD_ATTEMPTS = 5
9
10class GSUtilApi(recipe_api.RecipeApi):
11  def __call__(self, step_name, *args):
12    """Run gsutil with the given args."""
13    if 'Win' in self.m.vars.builder_cfg.get('os', ''):
14      return self.m.run(self.m.python, step_name, script=str(self.m.vars.workdir.join('cipd_bin_packages').join('gsutil')), args=args)
15    return self.m.step(step_name, cmd=['gsutil'] + list(args))
16
17  def cp(self, name, src, dst, extra_args=None, multithread=False):
18    """Attempt to upload or download files to/from Google Cloud Storage (GCS).
19
20    Args:
21      name: string. Will be used to fill out the step name.
22      src: string. Absolute path for a local file or gcs file (e.g. gs://...)
23      dst: string. Same as src.
24      extra_args: optional list of args to be passed to gsutil. e.g. [-Z] asks
25        all files be compressed with gzip after upload and before download.
26      multi_thread: if the -m argument should be used to copy multiple items
27        at once (e.g. gsutil -m cp foo* gs://bar/dir)
28
29    If the operation fails, it will be retried multiple times.
30    """
31    cmd = ['cp']
32    if multithread:
33      cmd = ['-m'] + cmd
34    if extra_args:
35      cmd.extend(extra_args)
36    cmd.extend([src, dst])
37
38    name = 'upload %s' % name
39    for i in range(UPLOAD_ATTEMPTS):
40      step_name = name
41      if i > 0:
42        step_name += ' (attempt %d)' % (i+1)
43      try:
44        self(step_name, *cmd)
45        break
46      except self.m.step.StepFailure:
47        if i == UPLOAD_ATTEMPTS - 1:
48          raise
49