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 ]) 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 148 class 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 213 class 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 235 def _prompt(prompt): 236 """Prompt for input, return result.""" 237 return raw_input(prompt) 238 239 240 class 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
