runquery.py
changeset 91 fba99a54ec49
parent 90 43e5d2cf4ddb
child 92 8816c917979c
--- a/runquery.py	Fri Oct 04 16:38:24 2013 +0200
+++ b/runquery.py	Mon Dec 09 11:39:28 2013 +0100
@@ -14,12 +14,17 @@
 
 
 class RunQueryTool(toolbase.ToolBase):
+
     def __init__(self):
         toolbase.ToolBase.__init__(self, name='runquery', desc='Run configured queries.')
         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('-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=[])
@@ -39,6 +44,8 @@
             with open(self.args.file, 'r', encoding='utf8') as f:
                 data = f.read()
                 self.queries.insert(0, data)
+        # prepare parameters
+        self._prepare_parameters(self.args.parameters)
         if self.config.log_path:
             self.init_file_logs(self.config.log_path)
         self.prepare_conns(target = self.args.target or self.config.target)
@@ -59,13 +66,28 @@
     def main(self):
         """Execute the queries."""
         for q in self.queries:
+            q = q.format(**self.parameters)
             self.log.info('%s', q if len(q) < 100 else q[:100]+'...')
             with self.pgm.cursor('target') as curs:
-                rows = curs.execute(q)
+                curs.execute(q)
                 self.log.info('Rows affected: %d', curs.rowcount)
-                self.log.info('Result: %r', rows)
+                rows = curs.fetchall_dict()
+                self._write_output_file(rows)
         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
+
 
 tool = RunQueryTool()
 tool.main()