1 from pydbkit.toolbase import SimpleTool

     2 

     3 

     4 class BatchQueryTool(SimpleTool):

     5 

     6     """

     7     Run a query using columns from CSV file as arguments.

     8     """

     9 

    10     def __init__(self):

    11         SimpleTool.__init__(self, name='batchquery')

    12 

    13     def specify_args(self):

    14         SimpleTool.specify_args(self)

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

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

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

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

    19         self.parser.add_argument('--outputfunc', dest='outputfunc', type=str, help='Python function which will format results (format_row(args, rows)).')

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

    21 

    22     def _split_line(self, line):

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

    24 

    25     def main(self):

    26         results = []

    27         # load query from file

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

    29             query = f.read()

    30         # connect DB

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

    32             # run init query

    33             if self.args.init:

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

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

    36             # read CSV file

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

    38                 # read header

    39                 names = None

    40                 if self.args.header:

    41                     line = f.readline()

    42                     names = self._split_line(line)

    43                 # read and process lines

    44                 for line in f:

    45                     args = self._split_line(line)

    46                     if names:

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

    48                     curs.execute(query, args)

    49                     rows = curs.fetchall()

    50                     results.append((args, rows))

    51             curs.connection.commit()

    52         # write results to output file

    53         if self.args.output:

    54             format_row = None

    55             if self.args.outputfunc:

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

    57                     d = dict()

    58                     exec(f.read(), d)

    59                     format_row = d['format_row']

    60 

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

    62                 for args, rows in results:

    63                     if format_row:

    64                         f.write(format_row(args, rows))

    65                     else:

    66                         f.write(repr(args))

    67                         f.write(' -> ')

    68                         f.write(repr(rows))

    69                     f.write('\n')

    70 

    71 

    72 cls = BatchQueryTool

    73