loopquery.py
author Radek Brich <radek.brich@devl.cz>
Thu, 31 Jan 2013 13:41:48 +0100
changeset 64 687e18e5ca93
parent 62 af637235ca81
child 65 ea79c89e9517
permissions -rwxr-xr-x
Drop colordiff, implement it in PgDiff directly.

#!/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()