1 from pgtoolkit.toolbase import ToolBase

     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         ToolBase.__init__(self, name='runquery')

    17         self.target_isolation_level = 'autocommit'

    18 

    19     def specify_args(self):

    20         ToolBase.specify_args(self)

    21         self.parser.add_argument('target', nargs='?', metavar='target', type=str, help='Target database')

    22         self.parser.add_argument('-c', dest='config', type=str, help='Additional config file (besides pgtoolkit.conf).')

    23         self.parser.add_argument('-q', dest='queries', metavar='QUERY', nargs='*', help='Queries to run.')

    24         self.parser.add_argument('-f', dest='file', metavar='FILE', help='Read query from file.')

    25         self.parser.add_argument('--one-query-per-line', action='store_true', help='When reading queries from file, consider each line as separate query.')

    26         self.parser.add_argument('-p', '--parameter', dest='parameters', metavar='PARAM=VALUE', nargs='*',

    27             help="If query should be used as format template, these parameters will be substituted.")

    28         self.parser.add_argument('--output-file', dest='output_file', metavar='OUTPUT_FILE', help='Write query result in specified file.')

    29         self.parser.add_argument('--format', dest='format', metavar='FORMAT', help='Format string for each line in output file (using Python\'s format()).')

    30 

    31         self.config.add_option('target', type=str, default=None)

    32         self.config.add_option('queries', type=list, default=[])

    33         self.config.add_option('log_path', type=str)

    34 

    35     def setup(self, args=None):

    36         ToolBase.setup(self, args)

    37         self.prepare_conns(target=self.args.target or self.config.target)

    38 

    39     def load_args(self, args=None, config_file=None):

    40         ToolBase.load_args(self, args, config_file)

    41         if self.args.config:

    42             self.config.load(self.args.config)

    43         self.queries = self.args.queries or self.config.queries

    44         # read query from file

    45         if self.args.file:

    46             with open(self.args.file, 'r', encoding='utf8') as f:

    47                 data = f.read()

    48             if self.args.one_query_per_line:

    49                 file_queries = [ln for ln in data.splitlines() if not ln.lstrip().startswith('--')]

    50                 self.queries = file_queries + self.queries

    51             else:

    52                 self.queries.insert(0, data)

    53         # prepare parameters

    54         self._prepare_parameters(self.args.parameters)

    55 

    56     def init_logging(self):

    57         ToolBase.init_logging(self)

    58         if self.config.log_path:

    59             self.init_file_logs(self.config.log_path)

    60 

    61     def init_file_logs(self, path):

    62         format = logging.Formatter('%(asctime)s %(levelname)-5s %(message)s', '%y-%m-%d %H:%M:%S')

    63         handler = logging.handlers.TimedRotatingFileHandler(path+'/main.log', when='midnight', backupCount=5)

    64         handler.setFormatter(format)

    65         handler.setLevel(logging.DEBUG)

    66         logging.getLogger('main').addHandler(handler)

    67 

    68         format = logging.Formatter('%(asctime)s %(message)s', '%y-%m-%d %H:%M:%S')

    69         handler = logging.handlers.TimedRotatingFileHandler(path+'/pgnotices.log', when='midnight', backupCount=5)

    70         handler.setFormatter(format)

    71         handler.setLevel(logging.DEBUG)

    72         logging.getLogger('pgmanager_notices').addHandler(handler)

    73 

    74     def main(self):

    75         """Execute the queries."""

    76         print(self.queries)

    77         for q in self.queries:

    78             if self.parameters:

    79                 q = q.format(**self.parameters)

    80             self.log.info('%s', q if len(q) < 100 else q[:100]+'...')

    81             with self.pgm.cursor('target') as curs:

    82                 curs.execute(q)

    83                 self.log.info('Rows affected: %d', curs.rowcount)

    84                 try:

    85                     rows = curs.fetchall_dict()

    86                     self._write_output_file(rows)

    87                 except ProgrammingError:

    88                     pass

    89         self.log.info('Done')

    90 

    91     def _write_output_file(self, rows):

    92         if not self.args.output_file:

    93             return

    94         with open(self.args.output_file, 'w', encoding='utf8') as f:

    95             for row in rows:

    96                 print(self.args.format.format(row), file=f)

    97 

    98     def _prepare_parameters(self, parameters):

    99         self.parameters = {}

   100         for parameter in parameters or ():

   101             name, value = parameter.split('=', 1)

   102             self.parameters[name] = value

   103 

   104 

   105 cls = RunQueryTool

   106