run.py

#!/usr/bin/python
# Copyright 2012 Google Inc. All Rights Reserved.
#
"""Run the given shell command for each of the given hosts."""

import os
import re
import select
import subprocess
import sys
import config
from portsh import options

__author__ = 'Avery Pennarun (apenwarr@google.com)'


optspec = """
run [options...] <hosts> -- <command string...>
--
r,raw           Raw output (don't redirect stdin/stdout/stderr)
q,quiet         Quiet: don't print to stderr for nonzero exit codes
"""


def _FixNull(s):
  if s is None:
    return ''
  return str(s)


def _ArgSub(argv, host):
  for arg in argv:
    argout = []
    for part in re.split(r'(\$\w+)', arg):
      if part.startswith('$'):
        for var in config.Host.__slots__:
          if part[1:] == var:
            part = _FixNull(getattr(host, var))
            break
      argout.append(part)
    yield ''.join(argout)


class Handler(object):
  """When data is received from the given file, print it line by line."""

  def __init__(self, prefix, fileobj):
    self.prefix = prefix
    self.fileobj = fileobj
    self.buf = ''

  def __del__(self):
    if self.buf and not self.buf.endswith('\n'):
      self.buf += '\n'
    self.PrintBuf()

  def fileno(self):  #gpylint: disable-msg=C6409
    return self.fileobj.fileno()

  def PrintBuf(self):
    while '\n' in self.buf:
      line, self.buf = self.buf.split('\n', 1)
      print '%s: %s' % (self.prefix, line)

  def Run(self):
    """Read from this handler and write its results to stdout."""
    got = os.read(self.fileobj.fileno(), 65536)
    self.buf += got
    self.PrintBuf()
    return got


def main():
  o = options.Options(optspec)
  args = []
  cmd = []
  for i in range(1, len(sys.argv)):
    if sys.argv[i] == '--':
      args = sys.argv[1:i]
      cmd = sys.argv[i+1:]
      break
  opt, _, hostnames = o.parse(args)
  if not hostnames or not cmd:
    o.fatal('you must specify hosts, --, and a command string')

  try:
    hosts = [config.hosts.FindOne(name=i) for i in hostnames]
  except KeyError, e:
    o.fatal('host not configured: %s' % e)

  procs = []
  handlers = []
  for host in hosts:
    env = dict(os.environ)
    env.update(dict((key, _FixNull(getattr(host, key)))
                    for key in config.Host.__slots__))
    if len(cmd) > 1:
      argv = list(_ArgSub(cmd, host))
      shell = False
    else:
      argv = cmd[0]
      shell = True
    if opt.raw:
      p = subprocess.Popen(argv, env=env, shell=shell)
    else:
      p = subprocess.Popen(argv, env=env, shell=shell,
                           stdout=subprocess.PIPE, stderr=subprocess.PIPE)
      handlers.append(Handler(host.name, p.stdout))
      handlers.append(Handler(host.name, p.stderr))
    procs.append((host, p))

  while handlers:
    r, _, _ = select.select(handlers, [], [])
    for handler in r:
      if not handler.Run():
        handlers.remove(handler)

  final_rv = 0
  for host, p in procs:
    rv = p.wait()
    if rv != 0:
      if not opt.quiet:
        sys.stderr.write('-- %s: exited with error code %d\n'
                         % (host.name, rv))
      final_rv = 1
  return final_rv


if __name__ == '__main__':
  sys.exit(main())