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