diff -r d6088dba8fea -r 2a2d0d5df03b pgtoolkit/tools/runquery.py --- /dev/null Thu Jan 01 00:00:00 1970 +0000 +++ b/pgtoolkit/tools/runquery.py Tue May 06 18:37:43 2014 +0200 @@ -0,0 +1,106 @@ +from pgtoolkit.toolbase import ToolBase + +import logging.handlers +import time +from datetime import datetime, timedelta +from psycopg2 import ProgrammingError + + +class RunQueryTool(ToolBase): + + """ + Execute configured queries in target database. + """ + + def __init__(self): + ToolBase.__init__(self, name='runquery') + self.target_isolation_level = 'autocommit' + + def specify_args(self): + ToolBase.specify_args(self) + 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.parser.add_argument('--one-query-per-line', action='store_true', help='When reading queries from file, consider each line as separate query.') + self.parser.add_argument('-p', '--parameter', dest='parameters', metavar='PARAM=VALUE', nargs='*', + help="If query should be used as format template, these parameters will be substituted.") + self.parser.add_argument('--output-file', dest='output_file', metavar='OUTPUT_FILE', help='Write query result in specified file.') + self.parser.add_argument('--format', dest='format', metavar='FORMAT', help='Format string for each line in output file (using Python\'s format()).') + + 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) + + def setup(self, args=None): + ToolBase.setup(self, args) + self.prepare_conns(target=self.args.target or self.config.target) + + def load_args(self, args=None, config_file=None): + ToolBase.load_args(self, args, config_file) + 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() + if self.args.one_query_per_line: + file_queries = [ln for ln in data.splitlines() if not ln.lstrip().startswith('--')] + self.queries = file_queries + self.queries + else: + self.queries.insert(0, data) + # prepare parameters + self._prepare_parameters(self.args.parameters) + + def init_logging(self): + ToolBase.init_logging(self) + if self.config.log_path: + self.init_file_logs(self.config.log_path) + + 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.""" + print(self.queries) + for q in self.queries: + if self.parameters: + q = q.format(**self.parameters) + self.log.info('%s', q if len(q) < 100 else q[:100]+'...') + with self.pgm.cursor('target') as curs: + curs.execute(q) + self.log.info('Rows affected: %d', curs.rowcount) + try: + rows = curs.fetchall_dict() + self._write_output_file(rows) + except ProgrammingError: + pass + self.log.info('Done') + + def _write_output_file(self, rows): + if not self.args.output_file: + return + with open(self.args.output_file, 'w', encoding='utf8') as f: + for row in rows: + print(self.args.format.format(row), file=f) + + def _prepare_parameters(self, parameters): + self.parameters = {} + for parameter in parameters or (): + name, value = parameter.split('=', 1) + self.parameters[name] = value + + +cls = RunQueryTool +