Update loopquery: allow any number of queries, support reading parameters from config file.
#!/usr/bin/env python3.2
from pgtoolkit import toolbase
from pgtoolkit.highlight import highlight
class BatchQueryTool(toolbase.SimpleTool):
def __init__(self):
toolbase.SimpleTool.__init__(self, name='batchquery', desc='Run a query using columns from CSV file as arguments.')
self.parser.add_argument('--query', dest='query', type=str, help='Query to run. Use %%s for arguments, or %%(name)s for named arguments (see --header).')
self.parser.add_argument('--file', dest='file', type=str, help='CSV file with data to use as arguments.')
self.parser.add_argument('--init', dest='init', type=str, help='Query which initialize database session (eg. temporary function).')
self.parser.add_argument('--output', dest='output', type=str, help='File name for results.')
self.parser.add_argument('--outputfunc', dest='outputfunc', type=str, help='Python function which will format results (format_row(args, rows)).')
self.parser.add_argument('--header', dest='header', action='store_true', help='First line of CSV is header with names for columns. These name can be used in query.')
self.init()
def _split_line(self, line):
return [x.strip() for x in line.split(',')]
def main(self):
results = []
# load query from file
with open(self.args.query, 'r', encoding='utf8') as f:
query = f.read()
# connect DB
with self.pgm.cursor('target') as curs:
# run init query
if self.args.init:
with open(self.args.init, 'r', encoding='utf8') as f:
curs.execute(f.read(), [])
# read CSV file
with open(self.args.file, 'r', encoding='utf8') as f:
# read header
names = None
if self.args.header:
line = f.readline()
names = self._split_line(line)
# read and process lines
for line in f:
args = self._split_line(line)
if names:
args = dict(zip(names, args))
curs.execute(query, args)
rows = curs.fetchall()
results.append((args, rows))
curs.connection.commit()
# write results to output file
if self.args.output:
format_row = None
if self.args.outputfunc:
with open(self.args.outputfunc, 'r', encoding='utf8') as f:
d = dict()
exec(f.read(), d)
format_row = d['format_row']
with open(self.args.output, 'w', encoding='utf8') as f:
for args, rows in results:
if format_row:
f.write(format_row(args, rows))
else:
f.write(repr(args))
f.write(' -> ')
f.write(repr(rows))
f.write('\n')
tool = BatchQueryTool()
tool.main()