recipe_modules/gsutil/api.py

# Copyright 2016 The Chromium Authors. All rights reserved.
# Use of this source code is governed by a BSD-style license that can be
# found in the LICENSE file.

from recipe_engine import recipe_api

import re


class GSUtilApi(recipe_api.RecipeApi):
  """GSUtilApi provides support for GSUtil."""

  def __init__(self, *args, **kwargs):
    super(GSUtilApi, self).__init__(*args, **kwargs)
    self._gsutil_tool = None

  def __call__(self, *args, **kwargs):
    """Return a step to run arbitrary gsutil command."""
    assert self._gsutil_tool
    name = kwargs.pop('name', 'gsutil ' + args[0])
    infra_step = kwargs.pop('infra_step', True)
    cmd_prefix = []
    # Note that metadata arguments have to be passed before the command.
    metadata = kwargs.pop('metadata', [])
    if metadata:
      for k, v in sorted(metadata.iteritems(), key=lambda (k, _): k):
        field = self._get_metadata_field(k)
        param = (field) if v is None else ('%s:%s' % (field, v))
        cmd_prefix.extend(['-h', param])
    if kwargs.pop('parallel_upload', False):
      cmd_prefix.extend([
        '-o',
        'GSUtil:parallel_composite_upload_threshold=50M'
      ])
    if kwargs.pop('multithreaded', False):
      cmd_prefix.extend(['-m'])
    cmd_prefix.extend([
      '-o',
      'GSUtil:software_update_check_period=0',
    ])
    return self.m.python(name, self._gsutil_tool, cmd_prefix + list(args),
                         venv=self.resource('gsutil.vpython'),
                         infra_step=infra_step, **kwargs)

  @recipe_api.non_step
  def normalize(self, url):
    """Normalize a GS URL using the gs:// URL prefix."""
    gs_prefix = 'gs://'
    # Defines the regex that matches a normalized URL.
    for prefix in (
      gs_prefix,
      'https://storage.cloud.google.com/',
      'https://storage.googleapis.com/',
      ):
      if url.startswith(prefix):
        return gs_prefix + url[len(prefix):]
    raise AssertionError("%s cannot be normalized" % url)

  @recipe_api.non_step
  def join(self, *parts):
    """Constructs a GS path from composite parts."""
    return '/'.join(p.strip('/') for p in parts)

  @classmethod
  def _http_url(cls, bucket, dest, unauthenticated_url=False):
    if unauthenticated_url:
      base = 'https://storage.googleapis.com/%s/%s'
    else:
      base = 'https://storage.cloud.google.com/%s/%s'
    return base % (bucket, dest)

  @staticmethod
  def _get_metadata_field(name, provider_prefix=None):
    """Returns: (str) the metadata field to use with Google Storage

    The Google Storage specification for metadata can be found at:
    https://developers.google.com/storage/docs/gsutil/addlhelp/WorkingWithObjectMetadata
    """
    # Already contains custom provider prefix
    if name.lower().startswith('x-'):
      return name

    # See if it's innately supported by Google Storage
    if name in (
      'Cache-Control',
      'Content-Disposition',
      'Content-Encoding',
      'Content-Language',
      'Content-MD5',
      'Content-Type',
      ):
      return name

    # Add provider prefix
    if not provider_prefix:
      provider_prefix = 'x-goog-meta'
    return '%s-%s' % (provider_prefix, name)

  def ensure_gsutil(self, version=None):
    with self.m.step.nest('ensure_gsutil'):
      with self.m.context(infra_steps=True):
        gsutil_dir = self.m.path['start_dir'].join('cipd', 'gsutil')

        self.m.cipd.ensure(
            gsutil_dir, {'infra/tools/gsutil': version or 'latest'})
        self._gsutil_tool = gsutil_dir.join('gsutil')

        return self._gsutil_tool

  def upload(self, bucket, src, dst, link_name='gsutil.upload',
             unauthenticated_url=False, **kwargs):
    step = self('cp', src, 'gs://%s/%s' % (bucket, dst), **kwargs)
    if link_name:
      step.presentation.links[link_name] = self._http_url(
          bucket, dst, unauthenticated_url=unauthenticated_url)
    return step

  def copy(self, src_bucket, src, dst_bucket, dst, link_name='gsutil.copy',
           unauthenticated_url=False, **kwargs):
    step = self('cp',
        'gs://%s/%s' % (src_bucket, src),
        'gs://%s/%s' % (dst_bucket, dst),
        **kwargs)
    if link_name:
      step.presentation.links[link_name] = self._http_url(
          dst_bucket, dst, unauthenticated_url=unauthenticated_url)
    return step