• Home
  • Line#
  • Scopes#
  • Navigate#
  • Raw
  • Download
1#!/usr/bin/env python
2#
3# Copyright 2016 Google Inc.
4#
5# Use of this source code is governed by a BSD-style license that can be
6# found in the LICENSE file.
7
8
9"""Utilities for managing assets."""
10
11
12import argparse
13import json
14import os
15import shlex
16import shutil
17import subprocess
18import sys
19
20INFRA_BOTS_DIR = os.path.abspath(os.path.realpath(os.path.join(
21    os.path.dirname(os.path.abspath(__file__)), os.pardir)))
22sys.path.insert(0, INFRA_BOTS_DIR)
23import utils
24import zip_utils
25
26
27ASSETS_DIR = os.path.join(INFRA_BOTS_DIR, 'assets')
28SKIA_DIR = os.path.abspath(os.path.join(INFRA_BOTS_DIR, os.pardir, os.pardir))
29
30CIPD_PACKAGE_NAME_TMPL = 'skia/bots/%s'
31DEFAULT_CIPD_SERVICE_URL = 'https://chrome-infra-packages.appspot.com'
32
33DEFAULT_GS_BUCKET = 'skia-assets'
34GS_SUBDIR_TMPL = 'gs://%s/assets/%s'
35GS_PATH_TMPL = '%s/%s.zip'
36
37TAG_PROJECT_SKIA = 'project:skia'
38TAG_VERSION_PREFIX = 'version:'
39TAG_VERSION_TMPL = '%s%%s' % TAG_VERSION_PREFIX
40
41VERSION_FILENAME = 'VERSION'
42ZIP_BLACKLIST = ['.git', '.svn', '*.pyc', '.DS_STORE']
43
44
45class CIPDStore(object):
46  """Wrapper object for CIPD."""
47  def __init__(self, cipd_url=DEFAULT_CIPD_SERVICE_URL):
48    self._cipd = 'cipd'
49    if sys.platform == 'win32':
50      self._cipd = 'cipd.bat'
51    self._cipd_url = cipd_url
52    self._check_setup()
53
54  def _check_setup(self):
55    """Verify that we have the CIPD binary and that we're authenticated."""
56    try:
57      self._run(['auth-info'], specify_service_url=False)
58    except OSError:
59      raise Exception('CIPD binary not found on your path (typically in '
60                      'depot_tools). You may need to update depot_tools.')
61    except subprocess.CalledProcessError:
62      raise Exception('CIPD not authenticated. You may need to run:\n\n'
63                      '$ %s auth-login' % self._cipd)
64
65  def _run(self, cmd, specify_service_url=True):
66    """Run the given command."""
67    cipd_args = []
68    if specify_service_url:
69      cipd_args.extend(['--service-url', self._cipd_url])
70    if os.getenv('USE_CIPD_GCE_AUTH'):
71      # Enable automatic GCE authentication. For context see
72      # https://bugs.chromium.org/p/skia/issues/detail?id=6385#c3
73      cipd_args.extend(['-service-account-json', ':gce'])
74    return subprocess.check_output(
75        [self._cipd] + cmd + cipd_args,
76        stderr=subprocess.STDOUT)
77
78  def _json_output(self, cmd):
79    """Run the given command, return the JSON output."""
80    with utils.tmp_dir():
81      json_output = os.path.join(os.getcwd(), 'output.json')
82      self._run(cmd + ['--json-output', json_output])
83      with open(json_output) as f:
84        parsed = json.load(f)
85    return parsed.get('result', [])
86
87  def _search(self, pkg_name):
88    try:
89      res = self._json_output(['search', pkg_name, '--tag', TAG_PROJECT_SKIA])
90    except subprocess.CalledProcessError as e:
91      if 'no such package' in e.output:
92        return []
93      raise
94    return [r['instance_id'] for r in res or []]
95
96  def _describe(self, pkg_name, instance_id):
97    """Obtain details about the given package and instance ID."""
98    return self._json_output(['describe', pkg_name, '--version', instance_id])
99
100  def get_available_versions(self, name):
101    """List available versions of the asset."""
102    pkg_name = CIPD_PACKAGE_NAME_TMPL % name
103    versions = []
104    for instance_id in self._search(pkg_name):
105      details = self._describe(pkg_name, instance_id)
106      for tag in details.get('tags'):
107        tag_name = tag.get('tag', '')
108        if tag_name.startswith(TAG_VERSION_PREFIX):
109          trimmed = tag_name[len(TAG_VERSION_PREFIX):]
110          try:
111            versions.append(int(trimmed))
112          except ValueError:
113            raise ValueError('Found package instance with invalid version '
114                             'tag: %s' % tag_name)
115    versions.sort()
116    return versions
117
118  def upload(self, name, version, target_dir, extra_tags=None):
119    """Create a CIPD package."""
120    cmd = [
121        'create',
122        '--name', CIPD_PACKAGE_NAME_TMPL % name,
123        '--in', target_dir,
124        '--tag', TAG_PROJECT_SKIA,
125        '--tag', TAG_VERSION_TMPL % version,
126        '--compression-level', '1',
127        '-verification-timeout', '30m0s',
128    ]
129    if extra_tags:
130      for tag in extra_tags:
131        cmd.extend(['--tag', tag])
132    self._run(cmd)
133
134  def download(self, name, version, target_dir):
135    """Download a CIPD package."""
136    pkg_name = CIPD_PACKAGE_NAME_TMPL % name
137    version_tag = TAG_VERSION_TMPL % version
138    target_dir = os.path.abspath(target_dir)
139    with utils.tmp_dir():
140      infile = os.path.join(os.getcwd(), 'input')
141      with open(infile, 'w') as f:
142        f.write('%s %s' % (pkg_name, version_tag))
143      self._run([
144          'ensure',
145          '--root', target_dir,
146          '--list', infile,
147      ])
148
149  def delete_contents(self, name):
150    """Delete data for the given asset."""
151    self._run(['pkg-delete', CIPD_PACKAGE_NAME_TMPL % name])
152
153
154class GSStore(object):
155  """Wrapper object for interacting with Google Storage."""
156  def __init__(self, gsutil=None, bucket=DEFAULT_GS_BUCKET):
157    if gsutil:
158      gsutil = os.path.abspath(gsutil)
159    else:
160      gsutils = subprocess.check_output([
161          utils.WHICH, 'gsutil']).rstrip().splitlines()
162      for g in gsutils:
163        ok = True
164        try:
165          subprocess.check_call([g, 'version'])
166        except OSError:
167          ok = False
168        if ok:
169          gsutil = g
170          break
171    self._gsutil = [gsutil]
172    if gsutil.endswith('.py'):
173      self._gsutil = ['python', gsutil]
174    self._gs_bucket = bucket
175
176  def copy(self, src, dst):
177    """Copy src to dst."""
178    subprocess.check_call(self._gsutil + ['cp', src, dst])
179
180  def list(self, path):
181    """List objects in the given path."""
182    try:
183      return subprocess.check_output(self._gsutil + ['ls', path]).splitlines()
184    except subprocess.CalledProcessError:
185      # If the prefix does not exist, we'll get an error, which is okay.
186      return []
187
188  def get_available_versions(self, name):
189    """Return the existing version numbers for the asset."""
190    files = self.list(GS_SUBDIR_TMPL % (self._gs_bucket, name))
191    bnames = [os.path.basename(f) for f in files]
192    suffix = '.zip'
193    versions = [int(f[:-len(suffix)]) for f in bnames if f.endswith(suffix)]
194    versions.sort()
195    return versions
196
197  # pylint: disable=unused-argument
198  def upload(self, name, version, target_dir, extra_tags=None):
199    """Upload to GS."""
200    target_dir = os.path.abspath(target_dir)
201    with utils.tmp_dir():
202      zip_file = os.path.join(os.getcwd(), '%d.zip' % version)
203      zip_utils.zip(target_dir, zip_file, blacklist=ZIP_BLACKLIST)
204      gs_path = GS_PATH_TMPL % (GS_SUBDIR_TMPL % (self._gs_bucket, name),
205                                str(version))
206      self.copy(zip_file, gs_path)
207
208  def download(self, name, version, target_dir):
209    """Download from GS."""
210    gs_path = GS_PATH_TMPL % (GS_SUBDIR_TMPL % (self._gs_bucket, name),
211                              str(version))
212    target_dir = os.path.abspath(target_dir)
213    with utils.tmp_dir():
214      zip_file = os.path.join(os.getcwd(), '%d.zip' % version)
215      self.copy(gs_path, zip_file)
216      zip_utils.unzip(zip_file, target_dir)
217
218  def delete_contents(self, name):
219    """Delete data for the given asset."""
220    gs_path = GS_SUBDIR_TMPL % (self._gs_bucket, name)
221    attempt_delete = True
222    try:
223      subprocess.check_call(self._gsutil + ['ls', gs_path])
224    except subprocess.CalledProcessError:
225      attempt_delete = False
226    if attempt_delete:
227      subprocess.check_call(self._gsutil + ['rm', '-rf', gs_path])
228
229
230class MultiStore(object):
231  """Wrapper object which uses CIPD as the primary store and GS for backup."""
232  def __init__(self, cipd_url=DEFAULT_CIPD_SERVICE_URL,
233               gsutil=None, gs_bucket=DEFAULT_GS_BUCKET):
234    self._cipd = CIPDStore(cipd_url=cipd_url)
235    self._gs = GSStore(gsutil=gsutil, bucket=gs_bucket)
236
237  def get_available_versions(self, name):
238    return self._cipd.get_available_versions(name)
239
240  def upload(self, name, version, target_dir, extra_tags=None):
241    self._cipd.upload(name, version, target_dir, extra_tags=extra_tags)
242    self._gs.upload(name, version, target_dir, extra_tags=extra_tags)
243
244  def download(self, name, version, target_dir):
245    self._gs.download(name, version, target_dir)
246
247  def delete_contents(self, name):
248    self._cipd.delete_contents(name)
249    self._gs.delete_contents(name)
250
251
252def _prompt(prompt):
253  """Prompt for input, return result."""
254  return raw_input(prompt)
255
256
257class Asset(object):
258  def __init__(self, name, store):
259    self._store = store
260    self._name = name
261    self._dir = os.path.join(ASSETS_DIR, self._name)
262
263  @property
264  def version_file(self):
265    """Return the path to the version file for this asset."""
266    return os.path.join(self._dir, VERSION_FILENAME)
267
268  def get_current_version(self):
269    """Obtain the current version of the asset."""
270    if not os.path.isfile(self.version_file):
271      return -1
272    with open(self.version_file) as f:
273      return int(f.read())
274
275  def get_available_versions(self):
276    """Return the existing version numbers for this asset."""
277    return self._store.get_available_versions(self._name)
278
279  def get_next_version(self):
280    """Find the next available version number for the asset."""
281    versions = self.get_available_versions()
282    if len(versions) == 0:
283      return 0
284    return versions[-1] + 1
285
286  def download_version(self, version, target_dir):
287    """Download the specified version of the asset."""
288    self._store.download(self._name, version, target_dir)
289
290  def download_current_version(self, target_dir):
291    """Download the version of the asset specified in its version file."""
292    v = self.get_current_version()
293    self.download_version(v, target_dir)
294
295  def upload_new_version(self, target_dir, commit=False, extra_tags=None):
296    """Upload a new version and update the version file for the asset."""
297    version = self.get_next_version()
298    self._store.upload(self._name, version, target_dir, extra_tags=extra_tags)
299
300    def _write_version():
301      with open(self.version_file, 'w') as f:
302        f.write(str(version))
303      subprocess.check_call([utils.GIT, 'add', self.version_file])
304
305    with utils.chdir(SKIA_DIR):
306      if commit:
307        with utils.git_branch():
308          _write_version()
309          subprocess.check_call([
310              utils.GIT, 'commit', '-m', 'Update %s version' % self._name])
311          subprocess.check_call([utils.GIT, 'cl', 'upload', '--bypass-hooks'])
312      else:
313        _write_version()
314
315  @classmethod
316  def add(cls, name, store):
317    """Add an asset."""
318    asset = cls(name, store)
319    if os.path.isdir(asset._dir):
320      raise Exception('Asset %s already exists!' % asset._name)
321
322    print 'Creating asset in %s' % asset._dir
323    os.mkdir(asset._dir)
324    def copy_script(script):
325      src = os.path.join(ASSETS_DIR, 'scripts', script)
326      dst = os.path.join(asset._dir, script)
327      print 'Creating %s' % dst
328      shutil.copy(src, dst)
329      subprocess.check_call([utils.GIT, 'add', dst])
330
331    for script in ('download.py', 'upload.py', 'common.py'):
332      copy_script(script)
333    resp = _prompt('Add script to automate creation of this asset? (y/n) ')
334    if resp == 'y':
335      copy_script('create.py')
336      copy_script('create_and_upload.py')
337      print 'You will need to add implementation to the creation script.'
338    print 'Successfully created asset %s.' % asset._name
339    return asset
340
341  def remove(self, remove_in_store=False):
342    """Remove this asset."""
343    # Ensure that the asset exists.
344    if not os.path.isdir(self._dir):
345      raise Exception('Asset %s does not exist!' % self._name)
346
347    # Cleanup the store.
348    if remove_in_store:
349      self._store.delete_contents(self._name)
350
351    # Remove the asset.
352    subprocess.check_call([utils.GIT, 'rm', '-rf', self._dir])
353    if os.path.isdir(self._dir):
354      shutil.rmtree(self._dir)
355