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