Rietveld Code Review Tool
Help | Bug tracker | Discussion group | Source code

Side by Side Diff: run.py

Issue 6572117575335936: Issue 2200 - PART II/II - Introduce --local and --remote parameters in {run,kick}.py (Closed)
Patch Set: Created April 8, 2015, 12:31 p.m.
Left:
Right:
Use n/p to move between diff chunks; N/P to move between comments.
Jump to:
View unified diff | Download patch
« no previous file with comments | « kick.py ('k') | no next file » | no next file with comments »
Toggle Intra-line Diffs ('i') | Expand Comments ('e') | Collapse Comments ('c') | Show Comments Hide Comments ('s')
OLDNEW
1 #!/usr/bin/env python 1 #!/usr/bin/env python
2 # coding: utf-8 2 # coding: utf-8
3 3
4 import argparse 4 import argparse
5 import sys 5 import sys
6 import os 6 import os
7 import re 7 import re
8 import subprocess 8 import subprocess
9 import yaml 9 import yaml
10 10
11 def createArgumentParser(**kwargs): 11 def createArgumentParser(**kwargs):
12 parser = argparse.ArgumentParser(**kwargs) 12 parser = argparse.ArgumentParser(**kwargs)
13 parser.add_argument( 13 parser.add_argument(
14 '-u', '--user', metavar='user', dest='user', type=str, default='vagrant', 14 '-u', '--user', metavar='user', dest='user', type=str, default='vagrant',
15 help='user name for use with SSH, must exist on all hosts' 15 help='user name for use with SSH, must exist on all hosts'
16 ) 16 )
17 17
18 parser.add_argument(
19 '-l', '--local', action='store_false', dest='remote', default=None,
20 help='use the local version of hosts.yaml'
21 )
22
23 parser.add_argument(
24 '-r', '--remote', metavar='<master>', dest='remote', type=str,
25 help='use a remote (puppet-master) version of hosts.yaml'
26 )
27
18 return parser 28 return parser
19 29
20 def parseOptions(args): 30 def parseOptions(args):
21 description = 'Run a command on the given hosts or groups of hosts' 31 description = 'Run a command on the given hosts or groups of hosts'
22 parser = createArgumentParser(description=description) 32 parser = createArgumentParser(description=description)
23 parser.add_argument( 33 parser.add_argument(
24 '-i', '--ignore-errors', action='store_true', dest='ignore_errors', 34 '-i', '--ignore-errors', action='store_true', dest='ignore_errors',
25 help='continue execution on next host in case of an error' 35 help='continue execution on next host in case of an error'
26 ) 36 )
27 37
28 hosts = set() 38 hosts = set()
29 parser.add_argument( 39 parser.add_argument(
30 '-t', '--target', metavar='host|group', 40 '-t', '--target', metavar='host|group',
31 help='target host or group, can be specified multiple times', 41 help='target host or group, can be specified multiple times',
32 type=lambda value: hosts.update([value]) 42 type=lambda value: hosts.update([value])
33 ) 43 )
34 44
35 parser.add_argument( 45 parser.add_argument(
36 'args', metavar='command', type=str, nargs='+', 46 'args', metavar='command', type=str, nargs='+',
37 help='the command to run on the specified hosts' 47 help='the command to run on the specified hosts'
38 ) 48 )
39 49
40 options = parser.parse_args(args) 50 options = parser.parse_args(args)
41 options.hosts = hosts 51 options.hosts = hosts
42 return options 52 return options
43 53
44 def getValidHosts(): 54 def getValidHosts(options):
45 dirname = os.path.dirname(sys.argv[0]) 55 path_canonical = ('modules', 'private', 'hiera', 'hosts.yaml')
46 path_name = os.path.join(dirname, 'modules', 'private', 'hiera', 'hosts.yaml') 56
47 with open(path_name, 'rb') as handle: 57 if options.remote:
48 config = yaml.load(handle) 58 login = ['-l', options.user] if options.user else []
59 path_name = '/etc/puppet/infrastructure/%s' % ("/".join(path_canonical),)
60 command = ['ssh'] + login + [options.remote, '--', 'sudo', 'cat', path_name]
61 child = subprocess.Popen(command, stderr=sys.stderr, stdout=subprocess.PIPE)
62 config = yaml.load(child.stdout)
63 child.stdout.close()
64 child.wait()
65 elif options.remote is False:
66 dirname = os.path.dirname(sys.argv[0])
67 path_name = os.path.join(dirname, *path_canonical)
68 with open(path_name, 'rb') as handle:
69 config = yaml.load(handle)
70 else:
71 print >>sys.stderr, 'Please either specify a --remote host or use --local'
72 sys.exit(1)
73
49 servers = config.get('servers', {}) 74 servers = config.get('servers', {})
50 return servers 75 return servers
51 76
52 def resolveHostList(hosts): 77 def resolveHostList(options):
53 78
54 result = set() 79 result = set()
55 80
56 try: 81 try:
57 valid_hosts = getValidHosts() 82 valid_hosts = getValidHosts(options)
58 except Warning as error: 83 except Warning as error:
59 print >>sys.stderr, 'Warning: failed to determine valid hosts:', error 84 print >>sys.stderr, 'Warning: failed to determine valid hosts:', error
60 result.update(hosts) 85 result.update(options.hosts)
61 else: 86 else:
62 for name in hosts: 87 for name in options.hosts:
63 chunk = [ 88 chunk = [
64 value.get('dns', key) for (key, value) in valid_hosts.items() 89 value.get('dns', key) for (key, value) in valid_hosts.items()
65 90
66 if name == key 91 if name == key
67 or name == '*' 92 or name == '*'
68 or name == value.get('dns', None) 93 or name == value.get('dns', None)
69 or name in value.get('groups', ()) 94 or name in value.get('groups', ())
70 ] 95 ]
71 96
72 if len(chunk) == 0: 97 if len(chunk) == 0:
73 print >>sys.stderr, 'Warning: failed to recognize host or group', name 98 print >>sys.stderr, 'Warning: failed to recognize host or group', name
74 else: 99 else:
75 result.update(chunk) 100 result.update(chunk)
76 101
77 return result 102 return result
78 103
79 def runCommand(user, host, command, ignore_errors=False): 104 def runCommand(user, host, command, ignore_errors=False):
80 if not isinstance(command, list): 105 if not isinstance(command, list):
81 command = [command] 106 command = [command]
82 command = ['ssh'] + (['-l', user] if user else []) + [host] + command 107 command = ['ssh'] + (['-l', user] if user else []) + [host] + command
83 if ignore_errors: 108 if ignore_errors:
84 subprocess.call(command) 109 subprocess.call(command)
85 else: 110 else:
86 subprocess.check_call(command) 111 subprocess.check_call(command)
87 112
88 if __name__ == '__main__': 113 if __name__ == '__main__':
89 options = parseOptions(sys.argv[1:]) 114 options = parseOptions(sys.argv[1:])
90 selectedHosts = resolveHostList(options.hosts) 115 selectedHosts = resolveHostList(options)
91 if len(selectedHosts) == 0: 116 if len(selectedHosts) == 0:
92 print >>sys.stderr, 'No valid hosts or groups specified, nothing to do' 117 print >>sys.stderr, 'No valid hosts or groups specified, nothing to do'
93 sys.exit(0) 118 sys.exit(0)
94 for host in selectedHosts: 119 for host in selectedHosts:
95 print >>sys.stderr, 'Running on %s...' % host 120 print >>sys.stderr, 'Running on %s...' % host
96 runCommand(options.user, host, options.args, options.ignore_errors) 121 runCommand(options.user, host, options.args, options.ignore_errors)
OLDNEW
« no previous file with comments | « kick.py ('k') | no next file » | no next file with comments »

Powered by Google App Engine
This is Rietveld