pydbkit/tools/runquery.py
author Radek Brich <radek.brich@devl.cz>
Wed, 09 Jul 2014 18:03:54 +0200
changeset 104 d8ff52a0390f
parent 102 pgtoolkit/tools/runquery.py@fda45bdfd68d
child 105 10551741f61f
permissions -rw-r--r--
Rename to pydbkit.
Ignore whitespace changes - Everywhere: Within whitespace: At end of lines:
104
d8ff52a0390f Rename to pydbkit.
Radek Brich <radek.brich@devl.cz>
parents: 102
diff changeset
     1
from pydbkit.toolbase import SimpleTool
51
bdc44f96cb0b Add loopquery tool. Log notices to console from all tools (toolbase).
Radek Brich <radek.brich@devl.cz>
parents:
diff changeset
     2
62
af637235ca81 Update loopquery: allow any number of queries, support reading parameters from config file.
Radek Brich <radek.brich@devl.cz>
parents: 51
diff changeset
     3
import logging.handlers
51
bdc44f96cb0b Add loopquery tool. Log notices to console from all tools (toolbase).
Radek Brich <radek.brich@devl.cz>
parents:
diff changeset
     4
import time
62
af637235ca81 Update loopquery: allow any number of queries, support reading parameters from config file.
Radek Brich <radek.brich@devl.cz>
parents: 51
diff changeset
     5
from datetime import datetime, timedelta
92
Radek Brich <radek.brich@devl.cz>
parents: 91
diff changeset
     6
from psycopg2 import ProgrammingError
Radek Brich <radek.brich@devl.cz>
parents: 91
diff changeset
     7
51
bdc44f96cb0b Add loopquery tool. Log notices to console from all tools (toolbase).
Radek Brich <radek.brich@devl.cz>
parents:
diff changeset
     8
101
2a2d0d5df03b Refactor ToolBase to allow tool composition. Add TableSync tool (composited). Move more tools under pgtool.
Radek Brich <brich.radek@ifortuna.cz>
parents: 99
diff changeset
     9
class RunQueryTool(ToolBase):
2a2d0d5df03b Refactor ToolBase to allow tool composition. Add TableSync tool (composited). Move more tools under pgtool.
Radek Brich <brich.radek@ifortuna.cz>
parents: 99
diff changeset
    10
2a2d0d5df03b Refactor ToolBase to allow tool composition. Add TableSync tool (composited). Move more tools under pgtool.
Radek Brich <brich.radek@ifortuna.cz>
parents: 99
diff changeset
    11
    """
2a2d0d5df03b Refactor ToolBase to allow tool composition. Add TableSync tool (composited). Move more tools under pgtool.
Radek Brich <brich.radek@ifortuna.cz>
parents: 99
diff changeset
    12
    Execute configured queries in target database.
2a2d0d5df03b Refactor ToolBase to allow tool composition. Add TableSync tool (composited). Move more tools under pgtool.
Radek Brich <brich.radek@ifortuna.cz>
parents: 99
diff changeset
    13
    """
91
fba99a54ec49 Add to runquery tool: Parameters for query, output file and format for output file. Fix loopquery tool.
Radek Brich <radek.brich@devl.cz>
parents: 90
diff changeset
    14
51
bdc44f96cb0b Add loopquery tool. Log notices to console from all tools (toolbase).
Radek Brich <radek.brich@devl.cz>
parents:
diff changeset
    15
    def __init__(self):
102
fda45bdfd68d Update ToolBase: Load cascade of config files, add -c parameter.
Radek Brich <brich.radek@ifortuna.cz>
parents: 101
diff changeset
    16
        SimpleTool.__init__(self, name='runquery')
101
2a2d0d5df03b Refactor ToolBase to allow tool composition. Add TableSync tool (composited). Move more tools under pgtool.
Radek Brich <brich.radek@ifortuna.cz>
parents: 99
diff changeset
    17
        self.target_isolation_level = 'autocommit'
2a2d0d5df03b Refactor ToolBase to allow tool composition. Add TableSync tool (composited). Move more tools under pgtool.
Radek Brich <brich.radek@ifortuna.cz>
parents: 99
diff changeset
    18
2a2d0d5df03b Refactor ToolBase to allow tool composition. Add TableSync tool (composited). Move more tools under pgtool.
Radek Brich <brich.radek@ifortuna.cz>
parents: 99
diff changeset
    19
    def specify_args(self):
102
fda45bdfd68d Update ToolBase: Load cascade of config files, add -c parameter.
Radek Brich <brich.radek@ifortuna.cz>
parents: 101
diff changeset
    20
        SimpleTool.specify_args(self)
65
ea79c89e9517 Update loopquery tool.
Radek Brich <radek.brich@devl.cz>
parents: 62
diff changeset
    21
        self.parser.add_argument('-q', dest='queries', metavar='QUERY', nargs='*', help='Queries to run.')
90
43e5d2cf4ddb Update runquery: Add parameter to read queries from file.
Radek Brich <radek.brich@devl.cz>
parents: 86
diff changeset
    22
        self.parser.add_argument('-f', dest='file', metavar='FILE', help='Read query from file.')
99
245646538743 Update runquery tool: Add --one-query-per-line parameter.
Radek Brich <brich.radek@ifortuna.cz>
parents: 93
diff changeset
    23
        self.parser.add_argument('--one-query-per-line', action='store_true', help='When reading queries from file, consider each line as separate query.')
91
fba99a54ec49 Add to runquery tool: Parameters for query, output file and format for output file. Fix loopquery tool.
Radek Brich <radek.brich@devl.cz>
parents: 90
diff changeset
    24
        self.parser.add_argument('-p', '--parameter', dest='parameters', metavar='PARAM=VALUE', nargs='*',
fba99a54ec49 Add to runquery tool: Parameters for query, output file and format for output file. Fix loopquery tool.
Radek Brich <radek.brich@devl.cz>
parents: 90
diff changeset
    25
            help="If query should be used as format template, these parameters will be substituted.")
fba99a54ec49 Add to runquery tool: Parameters for query, output file and format for output file. Fix loopquery tool.
Radek Brich <radek.brich@devl.cz>
parents: 90
diff changeset
    26
        self.parser.add_argument('--output-file', dest='output_file', metavar='OUTPUT_FILE', help='Write query result in specified file.')
fba99a54ec49 Add to runquery tool: Parameters for query, output file and format for output file. Fix loopquery tool.
Radek Brich <radek.brich@devl.cz>
parents: 90
diff changeset
    27
        self.parser.add_argument('--format', dest='format', metavar='FORMAT', help='Format string for each line in output file (using Python\'s format()).')
62
af637235ca81 Update loopquery: allow any number of queries, support reading parameters from config file.
Radek Brich <radek.brich@devl.cz>
parents: 51
diff changeset
    28
86
b61b54aa9f96 Fix: ConfigParser add_argument -> add_option.
Radek Brich <brich.radek@ifortuna.cz>
parents: 78
diff changeset
    29
        self.config.add_option('queries', type=list, default=[])
b61b54aa9f96 Fix: ConfigParser add_argument -> add_option.
Radek Brich <brich.radek@ifortuna.cz>
parents: 78
diff changeset
    30
        self.config.add_option('log_path', type=str)
62
af637235ca81 Update loopquery: allow any number of queries, support reading parameters from config file.
Radek Brich <radek.brich@devl.cz>
parents: 51
diff changeset
    31
101
2a2d0d5df03b Refactor ToolBase to allow tool composition. Add TableSync tool (composited). Move more tools under pgtool.
Radek Brich <brich.radek@ifortuna.cz>
parents: 99
diff changeset
    32
    def load_args(self, args=None, config_file=None):
102
fda45bdfd68d Update ToolBase: Load cascade of config files, add -c parameter.
Radek Brich <brich.radek@ifortuna.cz>
parents: 101
diff changeset
    33
        SimpleTool.load_args(self, args, config_file)
62
af637235ca81 Update loopquery: allow any number of queries, support reading parameters from config file.
Radek Brich <radek.brich@devl.cz>
parents: 51
diff changeset
    34
        self.queries = self.args.queries or self.config.queries
90
43e5d2cf4ddb Update runquery: Add parameter to read queries from file.
Radek Brich <radek.brich@devl.cz>
parents: 86
diff changeset
    35
        # read query from file
43e5d2cf4ddb Update runquery: Add parameter to read queries from file.
Radek Brich <radek.brich@devl.cz>
parents: 86
diff changeset
    36
        if self.args.file:
43e5d2cf4ddb Update runquery: Add parameter to read queries from file.
Radek Brich <radek.brich@devl.cz>
parents: 86
diff changeset
    37
            with open(self.args.file, 'r', encoding='utf8') as f:
43e5d2cf4ddb Update runquery: Add parameter to read queries from file.
Radek Brich <radek.brich@devl.cz>
parents: 86
diff changeset
    38
                data = f.read()
99
245646538743 Update runquery tool: Add --one-query-per-line parameter.
Radek Brich <brich.radek@ifortuna.cz>
parents: 93
diff changeset
    39
            if self.args.one_query_per_line:
245646538743 Update runquery tool: Add --one-query-per-line parameter.
Radek Brich <brich.radek@ifortuna.cz>
parents: 93
diff changeset
    40
                file_queries = [ln for ln in data.splitlines() if not ln.lstrip().startswith('--')]
245646538743 Update runquery tool: Add --one-query-per-line parameter.
Radek Brich <brich.radek@ifortuna.cz>
parents: 93
diff changeset
    41
                self.queries = file_queries + self.queries
245646538743 Update runquery tool: Add --one-query-per-line parameter.
Radek Brich <brich.radek@ifortuna.cz>
parents: 93
diff changeset
    42
            else:
90
43e5d2cf4ddb Update runquery: Add parameter to read queries from file.
Radek Brich <radek.brich@devl.cz>
parents: 86
diff changeset
    43
                self.queries.insert(0, data)
91
fba99a54ec49 Add to runquery tool: Parameters for query, output file and format for output file. Fix loopquery tool.
Radek Brich <radek.brich@devl.cz>
parents: 90
diff changeset
    44
        # prepare parameters
fba99a54ec49 Add to runquery tool: Parameters for query, output file and format for output file. Fix loopquery tool.
Radek Brich <radek.brich@devl.cz>
parents: 90
diff changeset
    45
        self._prepare_parameters(self.args.parameters)
101
2a2d0d5df03b Refactor ToolBase to allow tool composition. Add TableSync tool (composited). Move more tools under pgtool.
Radek Brich <brich.radek@ifortuna.cz>
parents: 99
diff changeset
    46
2a2d0d5df03b Refactor ToolBase to allow tool composition. Add TableSync tool (composited). Move more tools under pgtool.
Radek Brich <brich.radek@ifortuna.cz>
parents: 99
diff changeset
    47
    def init_logging(self):
2a2d0d5df03b Refactor ToolBase to allow tool composition. Add TableSync tool (composited). Move more tools under pgtool.
Radek Brich <brich.radek@ifortuna.cz>
parents: 99
diff changeset
    48
        ToolBase.init_logging(self)
62
af637235ca81 Update loopquery: allow any number of queries, support reading parameters from config file.
Radek Brich <radek.brich@devl.cz>
parents: 51
diff changeset
    49
        if self.config.log_path:
af637235ca81 Update loopquery: allow any number of queries, support reading parameters from config file.
Radek Brich <radek.brich@devl.cz>
parents: 51
diff changeset
    50
            self.init_file_logs(self.config.log_path)
af637235ca81 Update loopquery: allow any number of queries, support reading parameters from config file.
Radek Brich <radek.brich@devl.cz>
parents: 51
diff changeset
    51
af637235ca81 Update loopquery: allow any number of queries, support reading parameters from config file.
Radek Brich <radek.brich@devl.cz>
parents: 51
diff changeset
    52
    def init_file_logs(self, path):
af637235ca81 Update loopquery: allow any number of queries, support reading parameters from config file.
Radek Brich <radek.brich@devl.cz>
parents: 51
diff changeset
    53
        format = logging.Formatter('%(asctime)s %(levelname)-5s %(message)s', '%y-%m-%d %H:%M:%S')
af637235ca81 Update loopquery: allow any number of queries, support reading parameters from config file.
Radek Brich <radek.brich@devl.cz>
parents: 51
diff changeset
    54
        handler = logging.handlers.TimedRotatingFileHandler(path+'/main.log', when='midnight', backupCount=5)
af637235ca81 Update loopquery: allow any number of queries, support reading parameters from config file.
Radek Brich <radek.brich@devl.cz>
parents: 51
diff changeset
    55
        handler.setFormatter(format)
af637235ca81 Update loopquery: allow any number of queries, support reading parameters from config file.
Radek Brich <radek.brich@devl.cz>
parents: 51
diff changeset
    56
        handler.setLevel(logging.DEBUG)
af637235ca81 Update loopquery: allow any number of queries, support reading parameters from config file.
Radek Brich <radek.brich@devl.cz>
parents: 51
diff changeset
    57
        logging.getLogger('main').addHandler(handler)
af637235ca81 Update loopquery: allow any number of queries, support reading parameters from config file.
Radek Brich <radek.brich@devl.cz>
parents: 51
diff changeset
    58
af637235ca81 Update loopquery: allow any number of queries, support reading parameters from config file.
Radek Brich <radek.brich@devl.cz>
parents: 51
diff changeset
    59
        format = logging.Formatter('%(asctime)s %(message)s', '%y-%m-%d %H:%M:%S')
af637235ca81 Update loopquery: allow any number of queries, support reading parameters from config file.
Radek Brich <radek.brich@devl.cz>
parents: 51
diff changeset
    60
        handler = logging.handlers.TimedRotatingFileHandler(path+'/pgnotices.log', when='midnight', backupCount=5)
af637235ca81 Update loopquery: allow any number of queries, support reading parameters from config file.
Radek Brich <radek.brich@devl.cz>
parents: 51
diff changeset
    61
        handler.setFormatter(format)
af637235ca81 Update loopquery: allow any number of queries, support reading parameters from config file.
Radek Brich <radek.brich@devl.cz>
parents: 51
diff changeset
    62
        handler.setLevel(logging.DEBUG)
af637235ca81 Update loopquery: allow any number of queries, support reading parameters from config file.
Radek Brich <radek.brich@devl.cz>
parents: 51
diff changeset
    63
        logging.getLogger('pgmanager_notices').addHandler(handler)
af637235ca81 Update loopquery: allow any number of queries, support reading parameters from config file.
Radek Brich <radek.brich@devl.cz>
parents: 51
diff changeset
    64
51
bdc44f96cb0b Add loopquery tool. Log notices to console from all tools (toolbase).
Radek Brich <radek.brich@devl.cz>
parents:
diff changeset
    65
    def main(self):
62
af637235ca81 Update loopquery: allow any number of queries, support reading parameters from config file.
Radek Brich <radek.brich@devl.cz>
parents: 51
diff changeset
    66
        """Execute the queries."""
af637235ca81 Update loopquery: allow any number of queries, support reading parameters from config file.
Radek Brich <radek.brich@devl.cz>
parents: 51
diff changeset
    67
        for q in self.queries:
92
Radek Brich <radek.brich@devl.cz>
parents: 91
diff changeset
    68
            if self.parameters:
Radek Brich <radek.brich@devl.cz>
parents: 91
diff changeset
    69
                q = q.format(**self.parameters)
90
43e5d2cf4ddb Update runquery: Add parameter to read queries from file.
Radek Brich <radek.brich@devl.cz>
parents: 86
diff changeset
    70
            self.log.info('%s', q if len(q) < 100 else q[:100]+'...')
51
bdc44f96cb0b Add loopquery tool. Log notices to console from all tools (toolbase).
Radek Brich <radek.brich@devl.cz>
parents:
diff changeset
    71
            with self.pgm.cursor('target') as curs:
91
fba99a54ec49 Add to runquery tool: Parameters for query, output file and format for output file. Fix loopquery tool.
Radek Brich <radek.brich@devl.cz>
parents: 90
diff changeset
    72
                curs.execute(q)
67
78478a49c0bd Add runquery tool. Update setup.py.
Radek Brich <radek.brich@devl.cz>
parents: 66
diff changeset
    73
                self.log.info('Rows affected: %d', curs.rowcount)
92
Radek Brich <radek.brich@devl.cz>
parents: 91
diff changeset
    74
                try:
Radek Brich <radek.brich@devl.cz>
parents: 91
diff changeset
    75
                    rows = curs.fetchall_dict()
Radek Brich <radek.brich@devl.cz>
parents: 91
diff changeset
    76
                    self._write_output_file(rows)
Radek Brich <radek.brich@devl.cz>
parents: 91
diff changeset
    77
                except ProgrammingError:
Radek Brich <radek.brich@devl.cz>
parents: 91
diff changeset
    78
                    pass
62
af637235ca81 Update loopquery: allow any number of queries, support reading parameters from config file.
Radek Brich <radek.brich@devl.cz>
parents: 51
diff changeset
    79
        self.log.info('Done')
51
bdc44f96cb0b Add loopquery tool. Log notices to console from all tools (toolbase).
Radek Brich <radek.brich@devl.cz>
parents:
diff changeset
    80
91
fba99a54ec49 Add to runquery tool: Parameters for query, output file and format for output file. Fix loopquery tool.
Radek Brich <radek.brich@devl.cz>
parents: 90
diff changeset
    81
    def _write_output_file(self, rows):
fba99a54ec49 Add to runquery tool: Parameters for query, output file and format for output file. Fix loopquery tool.
Radek Brich <radek.brich@devl.cz>
parents: 90
diff changeset
    82
        if not self.args.output_file:
fba99a54ec49 Add to runquery tool: Parameters for query, output file and format for output file. Fix loopquery tool.
Radek Brich <radek.brich@devl.cz>
parents: 90
diff changeset
    83
            return
fba99a54ec49 Add to runquery tool: Parameters for query, output file and format for output file. Fix loopquery tool.
Radek Brich <radek.brich@devl.cz>
parents: 90
diff changeset
    84
        with open(self.args.output_file, 'w', encoding='utf8') as f:
fba99a54ec49 Add to runquery tool: Parameters for query, output file and format for output file. Fix loopquery tool.
Radek Brich <radek.brich@devl.cz>
parents: 90
diff changeset
    85
            for row in rows:
fba99a54ec49 Add to runquery tool: Parameters for query, output file and format for output file. Fix loopquery tool.
Radek Brich <radek.brich@devl.cz>
parents: 90
diff changeset
    86
                print(self.args.format.format(row), file=f)
fba99a54ec49 Add to runquery tool: Parameters for query, output file and format for output file. Fix loopquery tool.
Radek Brich <radek.brich@devl.cz>
parents: 90
diff changeset
    87
fba99a54ec49 Add to runquery tool: Parameters for query, output file and format for output file. Fix loopquery tool.
Radek Brich <radek.brich@devl.cz>
parents: 90
diff changeset
    88
    def _prepare_parameters(self, parameters):
fba99a54ec49 Add to runquery tool: Parameters for query, output file and format for output file. Fix loopquery tool.
Radek Brich <radek.brich@devl.cz>
parents: 90
diff changeset
    89
        self.parameters = {}
fba99a54ec49 Add to runquery tool: Parameters for query, output file and format for output file. Fix loopquery tool.
Radek Brich <radek.brich@devl.cz>
parents: 90
diff changeset
    90
        for parameter in parameters or ():
fba99a54ec49 Add to runquery tool: Parameters for query, output file and format for output file. Fix loopquery tool.
Radek Brich <radek.brich@devl.cz>
parents: 90
diff changeset
    91
            name, value = parameter.split('=', 1)
fba99a54ec49 Add to runquery tool: Parameters for query, output file and format for output file. Fix loopquery tool.
Radek Brich <radek.brich@devl.cz>
parents: 90
diff changeset
    92
            self.parameters[name] = value
fba99a54ec49 Add to runquery tool: Parameters for query, output file and format for output file. Fix loopquery tool.
Radek Brich <radek.brich@devl.cz>
parents: 90
diff changeset
    93
51
bdc44f96cb0b Add loopquery tool. Log notices to console from all tools (toolbase).
Radek Brich <radek.brich@devl.cz>
parents:
diff changeset
    94
101
2a2d0d5df03b Refactor ToolBase to allow tool composition. Add TableSync tool (composited). Move more tools under pgtool.
Radek Brich <brich.radek@ifortuna.cz>
parents: 99
diff changeset
    95
cls = RunQueryTool
51
bdc44f96cb0b Add loopquery tool. Log notices to console from all tools (toolbase).
Radek Brich <radek.brich@devl.cz>
parents:
diff changeset
    96