1 #!/usr/bin/env python3.2

     2 

     3 from pgtoolkit import toolbase

     4 from pgtoolkit.highlight import highlight

     5 

     6 

     7 class BatchQueryTool(toolbase.SimpleTool):

     8     def __init__(self):

     9         toolbase.SimpleTool.__init__(self, name='batchquery', desc='Run a query using columns from CSV file as arguments.')

    10         self.parser.add_argument('--query', dest='query', type=str, help='Query to run. Use %s for arguments, or %(name)s for named arguments (see --names).')

    11         self.parser.add_argument('--file', dest='file', type=str, help='CSV file with data to use as arguments.')

    12         self.parser.add_argument('--init', dest='init', type=str, help='Query which initialize database session (eg. temporary function).')

    13         self.parser.add_argument('--output', dest='output', type=str, help='File name for results.')

    14         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.')

    15         self.init()

    16 

    17     def _split_line(self, line):

    18         return [x.strip() for x in line.split(',')]

    19 

    20     def main(self):

    21         results = []

    22         # load query from file

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

    24             query = f.read()

    25         # connect DB

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

    27             # run init query

    28             if self.args.init:

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

    30                     curs.execute(f.read(), [])

    31             # read CSV file

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

    33                 # read header

    34                 names = None

    35                 if self.args.header:

    36                     line = f.readline()

    37                     names = self._split_line(line)

    38                 # read and process lines

    39                 for line in f:

    40                     args = self._split_line(line)

    41                     if names:

    42                         args = dict(zip(names, args))

    43                     curs.execute(query, args)

    44                     rows = curs.fetchall()

    45                     results.append((args, rows))

    46             curs.connection.commit()

    47         # write results to output file

    48         if self.args.output:

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

    50                 for args, rows in results:

    51                     f.write(repr(args))

    52                     f.write(' -> ')

    53                     f.write(repr(rows))

    54                     f.write('\n')

    55 

    56 

    57 tool = BatchQueryTool()

    58 tool.main()

    59