12 import time |
12 import time |
13 from datetime import datetime, timedelta |
13 from datetime import datetime, timedelta |
14 |
14 |
15 |
15 |
16 class RunQueryTool(toolbase.ToolBase): |
16 class RunQueryTool(toolbase.ToolBase): |
|
17 |
17 def __init__(self): |
18 def __init__(self): |
18 toolbase.ToolBase.__init__(self, name='runquery', desc='Run configured queries.') |
19 toolbase.ToolBase.__init__(self, name='runquery', desc='Run configured queries.') |
19 self.parser.add_argument('target', nargs='?', metavar='target', type=str, help='Target database') |
20 self.parser.add_argument('target', nargs='?', metavar='target', type=str, help='Target database') |
20 self.parser.add_argument('-c', dest='config', type=str, help='Additional config file (besides pgtoolkit.conf).') |
21 self.parser.add_argument('-c', dest='config', type=str, help='Additional config file (besides pgtoolkit.conf).') |
21 self.parser.add_argument('-q', dest='queries', metavar='QUERY', nargs='*', help='Queries to run.') |
22 self.parser.add_argument('-q', dest='queries', metavar='QUERY', nargs='*', help='Queries to run.') |
22 self.parser.add_argument('-f', dest='file', metavar='FILE', help='Read query from file.') |
23 self.parser.add_argument('-f', dest='file', metavar='FILE', help='Read query from file.') |
|
24 self.parser.add_argument('-p', '--parameter', dest='parameters', metavar='PARAM=VALUE', nargs='*', |
|
25 help="If query should be used as format template, these parameters will be substituted.") |
|
26 self.parser.add_argument('--output-file', dest='output_file', metavar='OUTPUT_FILE', help='Write query result in specified file.') |
|
27 self.parser.add_argument('--format', dest='format', metavar='FORMAT', help='Format string for each line in output file (using Python\'s format()).') |
23 |
28 |
24 self.config.add_option('target', type=str, default=None) |
29 self.config.add_option('target', type=str, default=None) |
25 self.config.add_option('queries', type=list, default=[]) |
30 self.config.add_option('queries', type=list, default=[]) |
26 self.config.add_option('log_path', type=str) |
31 self.config.add_option('log_path', type=str) |
27 |
32 |
37 # read query from file |
42 # read query from file |
38 if self.args.file: |
43 if self.args.file: |
39 with open(self.args.file, 'r', encoding='utf8') as f: |
44 with open(self.args.file, 'r', encoding='utf8') as f: |
40 data = f.read() |
45 data = f.read() |
41 self.queries.insert(0, data) |
46 self.queries.insert(0, data) |
|
47 # prepare parameters |
|
48 self._prepare_parameters(self.args.parameters) |
42 if self.config.log_path: |
49 if self.config.log_path: |
43 self.init_file_logs(self.config.log_path) |
50 self.init_file_logs(self.config.log_path) |
44 self.prepare_conns(target = self.args.target or self.config.target) |
51 self.prepare_conns(target = self.args.target or self.config.target) |
45 |
52 |
46 def init_file_logs(self, path): |
53 def init_file_logs(self, path): |
57 logging.getLogger('pgmanager_notices').addHandler(handler) |
64 logging.getLogger('pgmanager_notices').addHandler(handler) |
58 |
65 |
59 def main(self): |
66 def main(self): |
60 """Execute the queries.""" |
67 """Execute the queries.""" |
61 for q in self.queries: |
68 for q in self.queries: |
|
69 q = q.format(**self.parameters) |
62 self.log.info('%s', q if len(q) < 100 else q[:100]+'...') |
70 self.log.info('%s', q if len(q) < 100 else q[:100]+'...') |
63 with self.pgm.cursor('target') as curs: |
71 with self.pgm.cursor('target') as curs: |
64 rows = curs.execute(q) |
72 curs.execute(q) |
65 self.log.info('Rows affected: %d', curs.rowcount) |
73 self.log.info('Rows affected: %d', curs.rowcount) |
66 self.log.info('Result: %r', rows) |
74 rows = curs.fetchall_dict() |
|
75 self._write_output_file(rows) |
67 self.log.info('Done') |
76 self.log.info('Done') |
|
77 |
|
78 def _write_output_file(self, rows): |
|
79 if not self.args.output_file: |
|
80 return |
|
81 with open(self.args.output_file, 'w', encoding='utf8') as f: |
|
82 for row in rows: |
|
83 print(self.args.format.format(row), file=f) |
|
84 |
|
85 def _prepare_parameters(self, parameters): |
|
86 self.parameters = {} |
|
87 for parameter in parameters or (): |
|
88 name, value = parameter.split('=', 1) |
|
89 self.parameters[name] = value |
68 |
90 |
69 |
91 |
70 tool = RunQueryTool() |
92 tool = RunQueryTool() |
71 tool.main() |
93 tool.main() |
72 |
94 |