forked from vasilvv/mrtools
-
Notifications
You must be signed in to change notification settings - Fork 0
/
mrquery.py
executable file
·107 lines (87 loc) · 3.38 KB
/
mrquery.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
#!/usr/bin/python
#
## mrquery
##
## This program provides an interface to send arbitrary queries to Moira.
#
import common
import re, argparse, json
import pymoira
class QueryInfo:
def __init__(self, name):
info, = client.query('_help', (name,))
# FIXME: the format Moira server uses is horrible.
# There should be a saner way to do this
description = ' '.join(info)
match = re.match( r"\s*(\S+), (\S+) \((.+)\)( => (.+))?", description )
if not match:
raise pymoira.UserError( "Unable to parse the Moira query description" )
self.name, self.shortname, inputs, output_whole, outputs = match.groups()
self.inputs = inputs.split(' ')
self.outputs = outputs.split(' ') if outputs else []
def setup_arguments(argparser):
"""Sets up the arguments."""
argparser.add_argument('query', help = 'The Moira query name to execute')
argparser.add_argument('arg', nargs = '*', help = 'The arguments to the query')
argparser.add_argument('-j', '--json', action = 'store_true', help = 'Output the results of the query in JSON')
def do_query():
"""Runs the query and returns the result with query information."""
info = QueryInfo(args.query)
if len(info.inputs) != len(args.arg):
raise pymoira.UserError( "Query argument count mismatch (%i expected, %i supplied)" % (len(info.inputs), len(args.arg)) )
output = client.query(info.name, args.arg)
result = []
for row in output:
new_row = {}
if len(row) != len(info.outputs):
raise pymoira.UserError("Moira server returned unexpected amount of columns in a row")
for name, value in zip(info.outputs, row):
new_row[name] = value
result.append(new_row)
return result, info
def show_help():
try:
query = args.arg[0]
except IndexError:
common.error( "Query name was not specified" )
return
info, = client.query('_help', (query,))
print ' '.join(info).strip()
def show_queries_list():
queries = client.query('_list_queries', ())
for query, in queries:
print query
def show_user_list():
users = client.query('_list_users', ())
for user in users:
print "%s from %s %s, since %s %s" % user
def handle_query():
"""Outputs the results of the query or an error message."""
if args.json:
try:
result, info = do_query()
print json.dumps({ 'status' : 'ok', 'result' : result })
except pymoira.BaseError as e:
print json.dumps({ 'status' : 'error', 'message' : str(e) })
else:
try:
# Those are not real queries, and QueryInfo() would fail for them
if args.query == '_help':
show_help()
return
if args.query == '_list_queries':
show_queries_list()
return
if args.query == '_list_users':
show_user_list()
return
result, info = do_query()
for row in result:
fields = [ (field_name, row[field_name]) for field_name in info.outputs ]
common.show_fields(*fields)
print ""
except pymoira.BaseError as err:
common.error(err)
if __name__ == '__main__':
client, args = common.init('mrquery', 'Send raw queries to Moira', setup_arguments)
handle_query()