1 from pgtoolkit.toolbase import SimpleTool |
|
2 |
|
3 import logging.handlers |
|
4 import time |
|
5 from datetime import datetime, timedelta |
|
6 from psycopg2 import ProgrammingError |
|
7 |
|
8 |
|
9 class RunQueryTool(ToolBase): |
|
10 |
|
11 """ |
|
12 Execute configured queries in target database. |
|
13 """ |
|
14 |
|
15 def __init__(self): |
|
16 SimpleTool.__init__(self, name='runquery') |
|
17 self.target_isolation_level = 'autocommit' |
|
18 |
|
19 def specify_args(self): |
|
20 SimpleTool.specify_args(self) |
|
21 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('--one-query-per-line', action='store_true', help='When reading queries from file, consider each line as separate query.') |
|
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()).') |
|
28 |
|
29 self.config.add_option('queries', type=list, default=[]) |
|
30 self.config.add_option('log_path', type=str) |
|
31 |
|
32 def load_args(self, args=None, config_file=None): |
|
33 SimpleTool.load_args(self, args, config_file) |
|
34 self.queries = self.args.queries or self.config.queries |
|
35 # read query from file |
|
36 if self.args.file: |
|
37 with open(self.args.file, 'r', encoding='utf8') as f: |
|
38 data = f.read() |
|
39 if self.args.one_query_per_line: |
|
40 file_queries = [ln for ln in data.splitlines() if not ln.lstrip().startswith('--')] |
|
41 self.queries = file_queries + self.queries |
|
42 else: |
|
43 self.queries.insert(0, data) |
|
44 # prepare parameters |
|
45 self._prepare_parameters(self.args.parameters) |
|
46 |
|
47 def init_logging(self): |
|
48 ToolBase.init_logging(self) |
|
49 if self.config.log_path: |
|
50 self.init_file_logs(self.config.log_path) |
|
51 |
|
52 def init_file_logs(self, path): |
|
53 format = logging.Formatter('%(asctime)s %(levelname)-5s %(message)s', '%y-%m-%d %H:%M:%S') |
|
54 handler = logging.handlers.TimedRotatingFileHandler(path+'/main.log', when='midnight', backupCount=5) |
|
55 handler.setFormatter(format) |
|
56 handler.setLevel(logging.DEBUG) |
|
57 logging.getLogger('main').addHandler(handler) |
|
58 |
|
59 format = logging.Formatter('%(asctime)s %(message)s', '%y-%m-%d %H:%M:%S') |
|
60 handler = logging.handlers.TimedRotatingFileHandler(path+'/pgnotices.log', when='midnight', backupCount=5) |
|
61 handler.setFormatter(format) |
|
62 handler.setLevel(logging.DEBUG) |
|
63 logging.getLogger('pgmanager_notices').addHandler(handler) |
|
64 |
|
65 def main(self): |
|
66 """Execute the queries.""" |
|
67 for q in self.queries: |
|
68 if self.parameters: |
|
69 q = q.format(**self.parameters) |
|
70 self.log.info('%s', q if len(q) < 100 else q[:100]+'...') |
|
71 with self.pgm.cursor('target') as curs: |
|
72 curs.execute(q) |
|
73 self.log.info('Rows affected: %d', curs.rowcount) |
|
74 try: |
|
75 rows = curs.fetchall_dict() |
|
76 self._write_output_file(rows) |
|
77 except ProgrammingError: |
|
78 pass |
|
79 self.log.info('Done') |
|
80 |
|
81 def _write_output_file(self, rows): |
|
82 if not self.args.output_file: |
|
83 return |
|
84 with open(self.args.output_file, 'w', encoding='utf8') as f: |
|
85 for row in rows: |
|
86 print(self.args.format.format(row), file=f) |
|
87 |
|
88 def _prepare_parameters(self, parameters): |
|
89 self.parameters = {} |
|
90 for parameter in parameters or (): |
|
91 name, value = parameter.split('=', 1) |
|
92 self.parameters[name] = value |
|
93 |
|
94 |
|
95 cls = RunQueryTool |
|
96 |
|