annotated_run.py: Add LogDog bootstrapping.

scripts/slave/cipd.py

Index: scripts/slave/cipd.py
diff --git a/scripts/slave/cipd.py b/scripts/slave/cipd.py
new file mode 100755
index 0000000000000000000000000000000000000000..3849c323ea7149fea99331a22d4f6f9b8b03470a
--- /dev/null
+++ b/scripts/slave/cipd.py
@@ -0,0 +1,159 @@
+#!/usr/bin/env python
+# Copyright 2015 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.
+
+"""cipd.py bootstraps a CIPD client and installs named CIPD packages to a
+directory.
+"""
+
+import argparse
+import collections
+import json
+import logging
+import os
+import platform
+import subprocess
+import sys
+import tempfile
+
+# Install Infra build environment.
+BUILD_ROOT = os.path.dirname(os.path.dirname(os.path.dirname(
+                             os.path.abspath(__file__))))
+sys.path.insert(0, os.path.join(BUILD_ROOT, 'scripts'))
+
+import common.env
+
+# Instance-wide logger.
+LOGGER = logging.getLogger('cipd')
+
+# Used to contain a CIPD package specification.
+CipdPackage = collections.namedtuple('CipdPackage', ('name', 'version'))
+
+
+def get_normalized_platform(plat, machine):
+  if plat.startswith('linux'):
+    plat = 'linux'
+  elif plat.startswith(('win', 'cygwin')):
+    plat = 'win'
+  elif plat.startswith(('darwin', 'mac')):
+    plat = 'mac'
+  else:  # pragma: no cover
+    raise ValueError("Don't understand platform [%s]" % (plat,))
+
+  if machine == 'x86_64':
+    bits = 64
+  elif machine == 'i386':
+    bits = 32
+  else:  # pragma: no cover
+    raise ValueError("Don't understand machine [%s]" % (machine,))
+
+  return plat, bits
+
+def bootstrap(path):
+  bootstrap_path = os.path.join(common.env.Build, 'scripts', 'slave',
+                                'recipe_modules', 'cipd', 'resources',
+                                'bootstrap.py')
+
+  plat, bits = get_normalized_platform(sys.platform, platform.machine())
+  plat = '%s-%s' % (
+      plat.replace('win', 'windows'),
+      {
+          32: '386',
+          64: 'amd64',
+      }[bits]
+  )
+  json_output = os.path.join(path, 'cipd_bootstrap.json')
+
+  cmd = [
+    sys.executable,
+    bootstrap_path,
+    '--platform', plat,
+    '--dest-directory', path,
+    '--json-output', json_output,
+  ]
+  LOGGER.debug('Installing CIPD client: %s', cmd)
+  subprocess.check_call(cmd)
+
+  with open(json_output, 'r') as fd:
+    return json.load(fd)['executable']
+
+
+def cipd_ensure(client, path, packages, service_account_json=None,
+                json_output=None):
+  manifest_path = os.path.join(path, 'cipd_manifest.txt')
+  with open(manifest_path, 'w') as fd:
+    fd.write('# This file is autogenerated by %s\n\n' % (
+             os.path.abspath(__file__),))
+    for pkg in packages:
+      LOGGER.debug('Adding package [%s] version [%s] to manifest.',
+                   pkg.name, pkg.version)
+      fd.write('%s  %s\n' % (pkg.name, pkg.version))
+
+  cmd = [
+      client,
+      'ensure',
+      '-list', manifest_path,
+      '-root', path,
+  ]
+  if service_account_json:
+    cmd += ['-service-account-json', service_account_json]
+  if json_output:
+    cmd += ['-json-output', json_output]
+  LOGGER.debug('Ensuring CIPD packages: %s', cmd)
+  subprocess.check_call(cmd)
+
+
+def parse_cipd_package(v):
+  idx = v.index('@')
+  if idx > 0:
+    package, version = v[:idx], v[idx+1:]
+    if package and version:
+      return CipdPackage(package, version)
+  raise argparse.ArgumentTypeError('Package must be expressed using the format '
+                                   '"PACKAGE@VERSION".')
+
+
+def main(argv):
+  parser = argparse.ArgumentParser()
+  parser.add_argument('-v', '--verbose', action='count',
+                      help='Increase logging. Can be specified multiple times.')
+  parser.add_argument('-P', '--package', metavar='PACKAGE@VERSION',
+                      action='append', default=[], required=True,
+                      type=parse_cipd_package,
+                      help='Add a <package:version> to the CIPD manifest.')
+  parser.add_argument('-d', '--dest-directory', required=True,
+                      help='The CIPD package destination directory.')
+  parser.add_argument('-o', '--json-output',
+                      help='Output package results to a JSON file.')
+  parser.add_argument('--service-account-json',
+                      help='If specified, use this service account JSON.')
+  opts = parser.parse_args(argv[1:])
+
+  # Verbosity.
+  if opts.verbose == 0:
+    level = logging.WARNING
+  elif opts.verbose == 1:
+    level = logging.INFO
+  else:
+    level = logging.DEBUG
+  logging.getLogger().setLevel(level)
+
+  if not os.path.isdir(opts.dest_directory):
+    LOGGER.info('Creating destination directory: %s', opts.dest_directory)
+    os.makedirs(opts.dest_directory)
+
+  LOGGER.debug('Bootstrapping CIPD client...')
+  client = bootstrap(opts.dest_directory)
+  LOGGER.info('Using CIPD client at [%s].', client)
+
+  LOGGER.info('CIPD ensuring %d package(s)...', len(opts.package))
+  cipd_ensure(client, opts.dest_directory, opts.package,
+              service_account_json=opts.service_account_json,
+              json_output=opts.json_output)
+  return 0
+
+
+if __name__ == '__main__':
+  logging.basicConfig(level=logging.WARNING)
+  sys.exit(main(sys.argv))