loopquery.py
author Radek Brich <radek.brich@devl.cz>
Thu, 31 Jan 2013 11:02:04 +0100
changeset 62 af637235ca81
parent 51 bdc44f96cb0b
child 65 ea79c89e9517
permissions -rwxr-xr-x
Update loopquery: allow any number of queries, support reading parameters from config file.

#!/usr/bin/env python3.2
"""
loopquery

Execute some queries in loop. Execute interval can be set.

"""

from pgtoolkit import toolbase, config

import logging.handlers
import time
from datetime import datetime, timedelta


class LoopQueryTool(toolbase.ToolBase):
    def __init__(self):
        toolbase.ToolBase.__init__(self, name='loopquery', desc='Run query in loop.')
        self.parser.add_argument('target', nargs='?', metavar='target', type=str, help='Target database')
        self.parser.add_argument('-c', dest='config', type=str, help='Additional config file (besides pgtoolkit.conf).')
        self.parser.add_argument('-q', dest='queries', nargs='*', help='Queries to run.')
        self.parser.add_argument('--mins', dest='delay_mins', type=int, help='Delay between queries in minutes.')
        self.parser.add_argument('--secs', dest='delay_secs', type=int, help='Delay between queries in seconds.')

        self.config.add_argument('target', type=str, default=None)
        self.config.add_argument('queries', type=list, default=[])
        self.config.add_argument('delay_mins', type=int, default=0)
        self.config.add_argument('delay_secs', type=int, default=0)
        self.config.add_argument('log_path', type=str)

        self.target_isolation_level = 'autocommit'

        self.init()

    def init(self):
        toolbase.ToolBase.init(self)
        if self.args.config:
            self.config.load(self.args.config)
        self.queries = self.args.queries or self.config.queries
        self.delay_mins = self.config.delay_mins
        self.delay_secs = self.args.delay_secs or self.config.delay_secs
        if self.config.log_path:
            self.init_file_logs(self.config.log_path)
        self.prepare_conns(target = self.args.target or self.config.target)

    def init_file_logs(self, path):
        format = logging.Formatter('%(asctime)s %(levelname)-5s %(message)s', '%y-%m-%d %H:%M:%S')
        handler = logging.handlers.TimedRotatingFileHandler(path+'/main.log', when='midnight', backupCount=5)
        handler.setFormatter(format)
        handler.setLevel(logging.DEBUG)
        logging.getLogger('main').addHandler(handler)

        format = logging.Formatter('%(asctime)s %(message)s', '%y-%m-%d %H:%M:%S')
        handler = logging.handlers.TimedRotatingFileHandler(path+'/pgnotices.log', when='midnight', backupCount=5)
        handler.setFormatter(format)
        handler.setLevel(logging.DEBUG)
        logging.getLogger('pgmanager_notices').addHandler(handler)

    def main(self):
        self.prepare()
        while True:
            self.wait()
            self.action()

    def prepare(self):
        """Check current time, set next action time."""
        dt = datetime.today()
        dt = dt.replace(second = 0, microsecond = 0)
        self.next_action_time = dt + timedelta(minutes = self.delay_mins,
                                               seconds = self.delay_secs)

    def wait(self):
        """Wait for action time, compute next action time."""
        now = datetime.today()
        self.log.debug('Next run %s', self.next_action_time)
        if self.next_action_time > now:
            td = self.next_action_time - now
            self.log.debug('Sleep %ds', td.seconds + td.microseconds/1e6)
            time.sleep(td.seconds + td.microseconds/1e6)
        self.next_action_time += timedelta(minutes = self.delay_mins,
                                           seconds = self.delay_secs)

    def action(self):
        """Execute the queries."""
        for q in self.queries:
            self.log.info('%s', q)
            with self.pgm.cursor('target') as curs:
                curs.execute(q)
        self.log.info('Done')


tool = LoopQueryTool()
tool.main()