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