runquery.py
author Radek Brich <radek.brich@devl.cz>
Fri, 04 Oct 2013 16:38:24 +0200
changeset 90 43e5d2cf4ddb
parent 86 b61b54aa9f96
child 91 fba99a54ec49
permissions -rwxr-xr-x
Update runquery: Add parameter to read queries from file.

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

Execute configured queries once.

"""

from pgtoolkit import toolbase

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


class RunQueryTool(toolbase.ToolBase):
    def __init__(self):
        toolbase.ToolBase.__init__(self, name='runquery', desc='Run configured queries.')
        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', metavar='QUERY', nargs='*', help='Queries to run.')
        self.parser.add_argument('-f', dest='file', metavar='FILE', help='Read query from file.')

        self.config.add_option('target', type=str, default=None)
        self.config.add_option('queries', type=list, default=[])
        self.config.add_option('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
        # read query from file
        if self.args.file:
            with open(self.args.file, 'r', encoding='utf8') as f:
                data = f.read()
                self.queries.insert(0, data)
        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):
        """Execute the queries."""
        for q in self.queries:
            self.log.info('%s', q if len(q) < 100 else q[:100]+'...')
            with self.pgm.cursor('target') as curs:
                rows = curs.execute(q)
                self.log.info('Rows affected: %d', curs.rowcount)
                self.log.info('Result: %r', rows)
        self.log.info('Done')


tool = RunQueryTool()
tool.main()