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 12 import argparse 13 import json 14 import os 15 import shlex 16 import shutil 17 import subprocess 18 import sys 19 20 INFRA_BOTS_DIR = os.path.abspath(os.path.realpath(os.path.join( 21 os.path.dirname(os.path.abspath(__file__)), os.pardir))) 22 sys.path.insert(0, INFRA_BOTS_DIR) 23 import utils 24 import zip_utils 25 26 27 ASSETS_DIR = os.path.join(INFRA_BOTS_DIR, 'assets') 28 SKIA_DIR = os.path.abspath(os.path.join(INFRA_BOTS_DIR, os.pardir, os.pardir)) 29 30 CIPD_PACKAGE_NAME_TMPL = 'skia/bots/%s' 31 DEFAULT_CIPD_SERVICE_URL = 'https://chrome-infra-packages.appspot.com' 32 33 DEFAULT_GS_BUCKET = 'skia-buildbots' 34 GS_SUBDIR_TMPL = 'gs://%s/assets/%s' 35 GS_PATH_TMPL = '%s/%s.zip' 36 37 TAG_PROJECT_SKIA = 'project:skia' 38 TAG_VERSION_PREFIX = 'version:' 39 TAG_VERSION_TMPL = '%s%%s' % TAG_VERSION_PREFIX 40 41 WHICH = 'where' if sys.platform.startswith('win') else 'which' 42 43 VERSION_FILENAME = 'VERSION' 44 ZIP_BLACKLIST = ['.git', '.svn', '*.pyc', '.DS_STORE'] 45 46 47 class 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', '1', 126 '-verification-timeout', '30m0s', 127 ]) 128 129 def download(self, name, version, target_dir): 130 """Download a CIPD package.""" 131 pkg_name = CIPD_PACKAGE_NAME_TMPL % name 132 version_tag = TAG_VERSION_TMPL % version 133 target_dir = os.path.abspath(target_dir) 134 with utils.tmp_dir(): 135 infile = os.path.join(os.getcwd(), 'input') 136 with open(infile, 'w') as f: 137 f.write('%s %s' % (pkg_name, version_tag)) 138 self._run([ 139 'ensure', 140 '--root', target_dir, 141 '--list', infile, 142 ]) 143 144 def delete_contents(self, name): 145 """Delete data for the given asset.""" 146 self._run(['pkg-delete', CIPD_PACKAGE_NAME_TMPL % name]) 147 148 149 class GSStore(object): 150 """Wrapper object for interacting with Google Storage.""" 151 def __init__(self, gsutil=None, bucket=DEFAULT_GS_BUCKET): 152 if gsutil: 153 gsutil = os.path.abspath(gsutil) 154 else: 155 gsutil = subprocess.check_output([WHICH, 'gsutil']).rstrip() 156 self._gsutil = [gsutil] 157 if gsutil.endswith('.py'): 158 self._gsutil = ['python', gsutil] 159 self._gs_bucket = bucket 160 161 def copy(self, src, dst): 162 """Copy src to dst.""" 163 subprocess.check_call(self._gsutil + ['cp', src, dst]) 164 165 def list(self, path): 166 """List objects in the given path.""" 167 try: 168 return subprocess.check_output(self._gsutil + ['ls', path]).splitlines() 169 except subprocess.CalledProcessError: 170 # If the prefix does not exist, we'll get an error, which is okay. 171 return [] 172 173 def get_available_versions(self, name): 174 """Return the existing version numbers for the asset.""" 175 files = self.list(GS_SUBDIR_TMPL % (self._gs_bucket, name)) 176 bnames = [os.path.basename(f) for f in files] 177 suffix = '.zip' 178 versions = [int(f[:-len(suffix)]) for f in bnames if f.endswith(suffix)] 179 versions.sort() 180 return versions 181 182 def upload(self, name, version, target_dir): 183 """Upload to GS.""" 184 target_dir = os.path.abspath(target_dir) 185 with utils.tmp_dir(): 186 zip_file = os.path.join(os.getcwd(), '%d.zip' % version) 187 zip_utils.zip(target_dir, zip_file, blacklist=ZIP_BLACKLIST) 188 gs_path = GS_PATH_TMPL % (GS_SUBDIR_TMPL % (self._gs_bucket, name), 189 str(version)) 190 self.copy(zip_file, gs_path) 191 192 def download(self, name, version, target_dir): 193 """Download from GS.""" 194 gs_path = GS_PATH_TMPL % (GS_SUBDIR_TMPL % (self._gs_bucket, name), 195 str(version)) 196 target_dir = os.path.abspath(target_dir) 197 with utils.tmp_dir(): 198 zip_file = os.path.join(os.getcwd(), '%d.zip' % version) 199 self.copy(gs_path, zip_file) 200 zip_utils.unzip(zip_file, target_dir) 201 202 def delete_contents(self, name): 203 """Delete data for the given asset.""" 204 gs_path = GS_SUBDIR_TMPL % (self._gs_bucket, name) 205 attempt_delete = True 206 try: 207 subprocess.check_call(self._gsutil + ['ls', gs_path]) 208 except subprocess.CalledProcessError: 209 attempt_delete = False 210 if attempt_delete: 211 subprocess.check_call(self._gsutil + ['rm', '-rf', gs_path]) 212 213 214 class MultiStore(object): 215 """Wrapper object which uses CIPD as the primary store and GS for backup.""" 216 def __init__(self, cipd_url=DEFAULT_CIPD_SERVICE_URL, 217 gsutil=None, gs_bucket=DEFAULT_GS_BUCKET): 218 self._cipd = CIPDStore(cipd_url=cipd_url) 219 self._gs = GSStore(gsutil=gsutil, bucket=gs_bucket) 220 221 def get_available_versions(self, name): 222 return self._cipd.get_available_versions(name) 223 224 def upload(self, name, version, target_dir): 225 self._cipd.upload(name, version, target_dir) 226 self._gs.upload(name, version, target_dir) 227 228 def download(self, name, version, target_dir): 229 self._gs.download(name, version, target_dir) 230 231 def delete_contents(self, name): 232 self._cipd.delete_contents(name) 233 self._gs.delete_contents(name) 234 235 236 def _prompt(prompt): 237 """Prompt for input, return result.""" 238 return raw_input(prompt) 239 240 241 class Asset(object): 242 def __init__(self, name, store): 243 self._store = store 244 self._name = name 245 self._dir = os.path.join(ASSETS_DIR, self._name) 246 247 @property 248 def version_file(self): 249 """Return the path to the version file for this asset.""" 250 return os.path.join(self._dir, VERSION_FILENAME) 251 252 def get_current_version(self): 253 """Obtain the current version of the asset.""" 254 if not os.path.isfile(self.version_file): 255 return -1 256 with open(self.version_file) as f: 257 return int(f.read()) 258 259 def get_available_versions(self): 260 """Return the existing version numbers for this asset.""" 261 return self._store.get_available_versions(self._name) 262 263 def get_next_version(self): 264 """Find the next available version number for the asset.""" 265 versions = self.get_available_versions() 266 if len(versions) == 0: 267 return 0 268 return versions[-1] + 1 269 270 def download_version(self, version, target_dir): 271 """Download the specified version of the asset.""" 272 self._store.download(self._name, version, target_dir) 273 274 def download_current_version(self, target_dir): 275 """Download the version of the asset specified in its version file.""" 276 v = self.get_current_version() 277 self.download_version(v, target_dir) 278 279 def upload_new_version(self, target_dir, commit=False): 280 """Upload a new version and update the version file for the asset.""" 281 version = self.get_next_version() 282 self._store.upload(self._name, version, target_dir) 283 284 def _write_version(): 285 with open(self.version_file, 'w') as f: 286 f.write(str(version)) 287 subprocess.check_call([utils.GIT, 'add', self.version_file]) 288 289 with utils.chdir(SKIA_DIR): 290 if commit: 291 with utils.git_branch(): 292 _write_version() 293 subprocess.check_call([ 294 utils.GIT, 'commit', '-m', 'Update %s version' % self._name]) 295 subprocess.check_call([utils.GIT, 'cl', 'upload', '--bypass-hooks']) 296 else: 297 _write_version() 298 299 @classmethod 300 def add(cls, name, store): 301 """Add an asset.""" 302 asset = cls(name, store) 303 if os.path.isdir(asset._dir): 304 raise Exception('Asset %s already exists!' % asset._name) 305 306 print 'Creating asset in %s' % asset._dir 307 os.mkdir(asset._dir) 308 def copy_script(script): 309 src = os.path.join(ASSETS_DIR, 'scripts', script) 310 dst = os.path.join(asset._dir, script) 311 print 'Creating %s' % dst 312 shutil.copy(src, dst) 313 subprocess.check_call([utils.GIT, 'add', dst]) 314 315 for script in ('download.py', 'upload.py', 'common.py'): 316 copy_script(script) 317 resp = _prompt('Add script to automate creation of this asset? (y/n) ') 318 if resp == 'y': 319 copy_script('create.py') 320 copy_script('create_and_upload.py') 321 print 'You will need to add implementation to the creation script.' 322 print 'Successfully created asset %s.' % asset._name 323 return asset 324 325 def remove(self, remove_in_store=False): 326 """Remove this asset.""" 327 # Ensure that the asset exists. 328 if not os.path.isdir(self._dir): 329 raise Exception('Asset %s does not exist!' % self._name) 330 331 # Cleanup the store. 332 if remove_in_store: 333 self._store.delete_contents(self._name) 334 335 # Remove the asset. 336 subprocess.check_call([utils.GIT, 'rm', '-rf', self._dir]) 337 if os.path.isdir(self._dir): 338 shutil.rmtree(self._dir) 339
