pgtoolkit/tools/loopquery.py
changeset 104 d8ff52a0390f
parent 103 24e94a3da209
child 105 10551741f61f
equal deleted inserted replaced
103:24e94a3da209 104:d8ff52a0390f
     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