1 from pgtoolkit.toolbase import SimpleTool |
|
2 |
|
3 import logging.handlers |
|
4 import time |
|
5 from datetime import datetime, timedelta |
|
6 |
|
7 |
|
8 class LoopQueryTool(SimpleTool): |
|
9 |
|
10 """ |
|
11 Execute queries in loop, with configurable interval. |
|
12 |
|
13 """ |
|
14 |
|
15 def __init__(self): |
|
16 SimpleTool.__init__(self, name='loopquery') |
|
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('--mins', dest='delay_mins', type=int, help='Delay between queries in minutes.') |
|
23 self.parser.add_argument('--secs', dest='delay_secs', type=int, help='Delay between queries in seconds.') |
|
24 |
|
25 self.config.add_option('queries', type=list, default=[]) |
|
26 self.config.add_option('delay_mins', type=int, default=0) |
|
27 self.config.add_option('delay_secs', type=int, default=0) |
|
28 self.config.add_option('log_path', type=str) |
|
29 |
|
30 def load_args(self, args=None, config_file=None): |
|
31 SimpleTool.load_args(self, args, config_file) |
|
32 self.queries = self.args.queries or self.config.queries |
|
33 self.delay_mins = self.args.delay_mins or self.config.delay_mins |
|
34 self.delay_secs = self.args.delay_secs or self.config.delay_secs |
|
35 |
|
36 def init_logging(self): |
|
37 SimpleTool.init_logging(self) |
|
38 if self.config.log_path: |
|
39 self.init_file_logs(self.config.log_path) |
|
40 |
|
41 def init_file_logs(self, path): |
|
42 format = logging.Formatter('%(asctime)s %(levelname)-5s %(message)s', '%y-%m-%d %H:%M:%S') |
|
43 handler = logging.handlers.TimedRotatingFileHandler(path+'/main.log', when='midnight', backupCount=5) |
|
44 handler.setFormatter(format) |
|
45 handler.setLevel(logging.DEBUG) |
|
46 logging.getLogger('main').addHandler(handler) |
|
47 |
|
48 format = logging.Formatter('%(asctime)s %(message)s', '%y-%m-%d %H:%M:%S') |
|
49 handler = logging.handlers.TimedRotatingFileHandler(path+'/pgnotices.log', when='midnight', backupCount=5) |
|
50 handler.setFormatter(format) |
|
51 handler.setLevel(logging.DEBUG) |
|
52 logging.getLogger('pgmanager_notices').addHandler(handler) |
|
53 |
|
54 def main(self): |
|
55 self.reset() |
|
56 while True: |
|
57 self.wait() |
|
58 self.action() |
|
59 |
|
60 def reset(self): |
|
61 """Check current time, set next action time.""" |
|
62 dt = datetime.today() |
|
63 dt = dt.replace(microsecond = 0) |
|
64 self.next_action_time = dt + timedelta(minutes = self.delay_mins, |
|
65 seconds = self.delay_secs) |
|
66 |
|
67 def wait(self): |
|
68 """Wait for action time, compute next action time.""" |
|
69 now = datetime.today() |
|
70 self.log.debug('Next run %s', self.next_action_time) |
|
71 if self.next_action_time > now: |
|
72 td = self.next_action_time - now |
|
73 self.log.debug('Sleep %ds', td.seconds + td.microseconds/1e6) |
|
74 time.sleep(td.seconds + td.microseconds/1e6) |
|
75 self.next_action_time += timedelta(minutes = self.delay_mins, |
|
76 seconds = self.delay_secs) |
|
77 # in case that action took too long and next planned time would |
|
78 # be in past -> reset planner |
|
79 if self.next_action_time < now: |
|
80 self.reset() |
|
81 |
|
82 def action(self): |
|
83 """Execute the queries.""" |
|
84 for q in self.queries: |
|
85 self.log.info('%s', q) |
|
86 with self.pgm.cursor('target') as curs: |
|
87 curs.execute(q) |
|
88 self.log.info('Done') |
|
89 |
|
90 |
|
91 cls = LoopQueryTool |
|
92 |
|