pgtoolkit/tools/runquery.py
changeset 102 fda45bdfd68d
parent 101 2a2d0d5df03b
equal deleted inserted replaced
101:2a2d0d5df03b 102:fda45bdfd68d
     1 from pgtoolkit.toolbase import ToolBase
     1 from pgtoolkit.toolbase import SimpleTool
     2 
     2 
     3 import logging.handlers
     3 import logging.handlers
     4 import time
     4 import time
     5 from datetime import datetime, timedelta
     5 from datetime import datetime, timedelta
     6 from psycopg2 import ProgrammingError
     6 from psycopg2 import ProgrammingError
    11     """
    11     """
    12     Execute configured queries in target database.
    12     Execute configured queries in target database.
    13     """
    13     """
    14 
    14 
    15     def __init__(self):
    15     def __init__(self):
    16         ToolBase.__init__(self, name='runquery')
    16         SimpleTool.__init__(self, name='runquery')
    17         self.target_isolation_level = 'autocommit'
    17         self.target_isolation_level = 'autocommit'
    18 
    18 
    19     def specify_args(self):
    19     def specify_args(self):
    20         ToolBase.specify_args(self)
    20         SimpleTool.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.')
    21         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.')
    22         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.')
    23         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='*',
    24         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.")
    25             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.')
    26         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()).')
    27         self.parser.add_argument('--format', dest='format', metavar='FORMAT', help='Format string for each line in output file (using Python\'s format()).')
    30 
    28 
    31         self.config.add_option('target', type=str, default=None)
       
    32         self.config.add_option('queries', type=list, default=[])
    29         self.config.add_option('queries', type=list, default=[])
    33         self.config.add_option('log_path', type=str)
    30         self.config.add_option('log_path', type=str)
    34 
    31 
    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):
    32     def load_args(self, args=None, config_file=None):
    40         ToolBase.load_args(self, args, config_file)
    33         SimpleTool.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
    34         self.queries = self.args.queries or self.config.queries
    44         # read query from file
    35         # read query from file
    45         if self.args.file:
    36         if self.args.file:
    46             with open(self.args.file, 'r', encoding='utf8') as f:
    37             with open(self.args.file, 'r', encoding='utf8') as f:
    47                 data = f.read()
    38                 data = f.read()
    71         handler.setLevel(logging.DEBUG)
    62         handler.setLevel(logging.DEBUG)
    72         logging.getLogger('pgmanager_notices').addHandler(handler)
    63         logging.getLogger('pgmanager_notices').addHandler(handler)
    73 
    64 
    74     def main(self):
    65     def main(self):
    75         """Execute the queries."""
    66         """Execute the queries."""
    76         print(self.queries)
       
    77         for q in self.queries:
    67         for q in self.queries:
    78             if self.parameters:
    68             if self.parameters:
    79                 q = q.format(**self.parameters)
    69                 q = q.format(**self.parameters)
    80             self.log.info('%s', q if len(q) < 100 else q[:100]+'...')
    70             self.log.info('%s', q if len(q) < 100 else q[:100]+'...')
    81             with self.pgm.cursor('target') as curs:
    71             with self.pgm.cursor('target') as curs: