xref: /aosp_15_r20/external/skia/infra/bots/recipe_modules/gsutil/api.py (revision c8dee2aa9b3f27cf6c858bd81872bdeb2c07ed17)
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
14       This assumes there exists an executable called gsutil on the PATH.
15       This probably only works for Linux/Mac, but those are the only
16       hosts that we try to upload to GCS from anyway.
17    """
18    return self.m.step(step_name, cmd=['gsutil'] + list(args))
19
20  def cp(self, name, src, dst, extra_gsutil_args=None, extra_args=None,
21         multithread=False):
22    """Attempt to upload or download files to/from Google Cloud Storage (GCS).
23
24    Args:
25      name: string. Will be used to fill out the step name.
26      src: string. Absolute path for a local file or gcs file (e.g. gs://...)
27      dst: string. Same as src.
28      extra_gsutil_args: optional list of args to be passed to gsutil before the
29        cp command.
30      extra_args: optional list of args to be passed to gsutil cp. e.g. [-Z]
31        asks all files be compressed with gzip after upload and before download.
32      multi_thread: if the -m argument should be used to copy multiple items
33        at once (e.g. gsutil -m cp foo* gs://bar/dir)
34
35    If the operation fails, it will be retried multiple times.
36    """
37    cmd = []
38    if extra_gsutil_args:
39      cmd.extend(extra_gsutil_args)
40    if multithread:
41      cmd.append('-m')
42    cmd.append('cp')
43    if extra_args:
44      cmd.extend(extra_args)
45    cmd.extend([src, dst])
46
47    name = 'upload %s' % name
48    for i in range(UPLOAD_ATTEMPTS):
49      step_name = name
50      if i > 0:
51        step_name += ' (attempt %d)' % (i+1)
52      try:
53        self(step_name, *cmd)
54        break
55      except self.m.step.StepFailure:
56        if i == UPLOAD_ATTEMPTS - 1:
57          raise
58