pgtoolkit/tools/loopquery.py
changeset 102 fda45bdfd68d
parent 101 2a2d0d5df03b
equal deleted inserted replaced
101:2a2d0d5df03b 102:fda45bdfd68d
     1 from pgtoolkit.toolbase import ToolBase
     1 from pgtoolkit.toolbase import SimpleTool
     2 
     2 
     3 import logging.handlers
     3 import logging.handlers
     4 import time
     4 import time
     5 from datetime import datetime, timedelta
     5 from datetime import datetime, timedelta
     6 
     6 
     7 
     7 
     8 class LoopQueryTool(ToolBase):
     8 class LoopQueryTool(SimpleTool):
     9 
     9 
    10     """
    10     """
    11     Execute queries in loop, with configurable interval.
    11     Execute queries in loop, with configurable interval.
    12 
    12 
    13     """
    13     """
    14 
    14 
    15     def __init__(self):
    15     def __init__(self):
    16         ToolBase.__init__(self, name='loopquery')
    16         SimpleTool.__init__(self, name='loopquery')
    17         self.target_isolation_level = 'autocommit'
    17         self.target_isolation_level = 'autocommit'
    18 
    18 
    19     def specify_args(self):
    19     def specify_args(self):
    20         ToolBase.specify_args(self)
    20         SimpleTool.specify_args(self)
    21         self.parser.add_argument('target', nargs='?', metavar='target', type=str, help='Target database')
       
    22         self.parser.add_argument('-c', dest='config', type=str, help='Additional config file (besides pgtoolkit.conf).')
       
    23         self.parser.add_argument('-q', dest='queries', metavar='QUERY', nargs='*', help='Queries to run.')
    21         self.parser.add_argument('-q', dest='queries', metavar='QUERY', nargs='*', help='Queries to run.')
    24         self.parser.add_argument('--mins', dest='delay_mins', type=int, help='Delay between queries in minutes.')
    22         self.parser.add_argument('--mins', dest='delay_mins', type=int, help='Delay between queries in minutes.')
    25         self.parser.add_argument('--secs', dest='delay_secs', type=int, help='Delay between queries in seconds.')
    23         self.parser.add_argument('--secs', dest='delay_secs', type=int, help='Delay between queries in seconds.')
    26 
    24 
    27         self.config.add_option('target', type=str, default=None)
       
    28         self.config.add_option('queries', type=list, default=[])
    25         self.config.add_option('queries', type=list, default=[])
    29         self.config.add_option('delay_mins', type=int, default=0)
    26         self.config.add_option('delay_mins', type=int, default=0)
    30         self.config.add_option('delay_secs', type=int, default=0)
    27         self.config.add_option('delay_secs', type=int, default=0)
    31         self.config.add_option('log_path', type=str)
    28         self.config.add_option('log_path', type=str)
    32 
    29 
    33     def load_args(self, args=None, config_file=None):
    30     def load_args(self, args=None, config_file=None):
    34         ToolBase.load_args(self, args, config_file)
    31         SimpleTool.load_args(self, args, config_file)
    35         if self.args.config:
       
    36             self.config.load(self.args.config)
       
    37         self.queries = self.args.queries or self.config.queries
    32         self.queries = self.args.queries or self.config.queries
    38         self.delay_mins = self.args.delay_mins or self.config.delay_mins
    33         self.delay_mins = self.args.delay_mins or self.config.delay_mins
    39         self.delay_secs = self.args.delay_secs or self.config.delay_secs
    34         self.delay_secs = self.args.delay_secs or self.config.delay_secs
    40 
    35 
    41     def init_logging(self):
    36     def init_logging(self):
    42         ToolBase.init_logging(self)
    37         SimpleTool.init_logging(self)
    43         if self.config.log_path:
    38         if self.config.log_path:
    44             self.init_file_logs(self.config.log_path)
    39             self.init_file_logs(self.config.log_path)
    45 
       
    46     def setup(self, args=None):
       
    47         ToolBase.setup(self, args)
       
    48         self.prepare_conns(target=self.args.target or self.config.target)
       
    49 
    40 
    50     def init_file_logs(self, path):
    41     def init_file_logs(self, path):
    51         format = logging.Formatter('%(asctime)s %(levelname)-5s %(message)s', '%y-%m-%d %H:%M:%S')
    42         format = logging.Formatter('%(asctime)s %(levelname)-5s %(message)s', '%y-%m-%d %H:%M:%S')
    52         handler = logging.handlers.TimedRotatingFileHandler(path+'/main.log', when='midnight', backupCount=5)
    43         handler = logging.handlers.TimedRotatingFileHandler(path+'/main.log', when='midnight', backupCount=5)
    53         handler.setFormatter(format)
    44         handler.setFormatter(format)