pgtoolkit/tools/runquery.py
changeset 104 d8ff52a0390f
parent 103 24e94a3da209
child 105 10551741f61f
equal deleted inserted replaced
103:24e94a3da209 104:d8ff52a0390f
     1 from pgtoolkit.toolbase import SimpleTool
       
     2 
       
     3 import logging.handlers
       
     4 import time
       
     5 from datetime import datetime, timedelta
       
     6 from psycopg2 import ProgrammingError
       
     7 
       
     8 
       
     9 class RunQueryTool(ToolBase):
       
    10 
       
    11     """
       
    12     Execute configured queries in target database.
       
    13     """
       
    14 
       
    15     def __init__(self):
       
    16         SimpleTool.__init__(self, name='runquery')
       
    17         self.target_isolation_level = 'autocommit'
       
    18 
       
    19     def specify_args(self):
       
    20         SimpleTool.specify_args(self)
       
    21         self.parser.add_argument('-q', dest='queries', metavar='QUERY', nargs='*', help='Queries to run.')
       
    22         self.parser.add_argument('-f', dest='file', metavar='FILE', help='Read query from file.')
       
    23         self.parser.add_argument('--one-query-per-line', action='store_true', help='When reading queries from file, consider each line as separate query.')
       
    24         self.parser.add_argument('-p', '--parameter', dest='parameters', metavar='PARAM=VALUE', nargs='*',
       
    25             help="If query should be used as format template, these parameters will be substituted.")
       
    26         self.parser.add_argument('--output-file', dest='output_file', metavar='OUTPUT_FILE', help='Write query result in specified file.')
       
    27         self.parser.add_argument('--format', dest='format', metavar='FORMAT', help='Format string for each line in output file (using Python\'s format()).')
       
    28 
       
    29         self.config.add_option('queries', type=list, default=[])
       
    30         self.config.add_option('log_path', type=str)
       
    31 
       
    32     def load_args(self, args=None, config_file=None):
       
    33         SimpleTool.load_args(self, args, config_file)
       
    34         self.queries = self.args.queries or self.config.queries
       
    35         # read query from file
       
    36         if self.args.file:
       
    37             with open(self.args.file, 'r', encoding='utf8') as f:
       
    38                 data = f.read()
       
    39             if self.args.one_query_per_line:
       
    40                 file_queries = [ln for ln in data.splitlines() if not ln.lstrip().startswith('--')]
       
    41                 self.queries = file_queries + self.queries
       
    42             else:
       
    43                 self.queries.insert(0, data)
       
    44         # prepare parameters
       
    45         self._prepare_parameters(self.args.parameters)
       
    46 
       
    47     def init_logging(self):
       
    48         ToolBase.init_logging(self)
       
    49         if self.config.log_path:
       
    50             self.init_file_logs(self.config.log_path)
       
    51 
       
    52     def init_file_logs(self, path):
       
    53         format = logging.Formatter('%(asctime)s %(levelname)-5s %(message)s', '%y-%m-%d %H:%M:%S')
       
    54         handler = logging.handlers.TimedRotatingFileHandler(path+'/main.log', when='midnight', backupCount=5)
       
    55         handler.setFormatter(format)
       
    56         handler.setLevel(logging.DEBUG)
       
    57         logging.getLogger('main').addHandler(handler)
       
    58 
       
    59         format = logging.Formatter('%(asctime)s %(message)s', '%y-%m-%d %H:%M:%S')
       
    60         handler = logging.handlers.TimedRotatingFileHandler(path+'/pgnotices.log', when='midnight', backupCount=5)
       
    61         handler.setFormatter(format)
       
    62         handler.setLevel(logging.DEBUG)
       
    63         logging.getLogger('pgmanager_notices').addHandler(handler)
       
    64 
       
    65     def main(self):
       
    66         """Execute the queries."""
       
    67         for q in self.queries:
       
    68             if self.parameters:
       
    69                 q = q.format(**self.parameters)
       
    70             self.log.info('%s', q if len(q) < 100 else q[:100]+'...')
       
    71             with self.pgm.cursor('target') as curs:
       
    72                 curs.execute(q)
       
    73                 self.log.info('Rows affected: %d', curs.rowcount)
       
    74                 try:
       
    75                     rows = curs.fetchall_dict()
       
    76                     self._write_output_file(rows)
       
    77                 except ProgrammingError:
       
    78                     pass
       
    79         self.log.info('Done')
       
    80 
       
    81     def _write_output_file(self, rows):
       
    82         if not self.args.output_file:
       
    83             return
       
    84         with open(self.args.output_file, 'w', encoding='utf8') as f:
       
    85             for row in rows:
       
    86                 print(self.args.format.format(row), file=f)
       
    87 
       
    88     def _prepare_parameters(self, parameters):
       
    89         self.parameters = {}
       
    90         for parameter in parameters or ():
       
    91             name, value = parameter.split('=', 1)
       
    92             self.parameters[name] = value
       
    93 
       
    94 
       
    95 cls = RunQueryTool
       
    96