pydbkit/tools/batchquery.py
author Radek Brich <radek.brich@devl.cz>
Wed, 09 Jul 2014 18:04:11 +0200
changeset 105 10551741f61f
parent 104 batchquery.py@d8ff52a0390f
permissions -rw-r--r--
Move batchquery, listtables tools under pgtool.
Ignore whitespace changes - Everywhere: Within whitespace: At end of lines:
105
10551741f61f Move batchquery, listtables tools under pgtool.
Radek Brich <radek.brich@devl.cz>
parents: 104
diff changeset
     1
from pydbkit.toolbase import SimpleTool
29
8636719a30f6 Add batchquery tool.
Radek Brich <radek.brich@devl.cz>
parents:
diff changeset
     2
8636719a30f6 Add batchquery tool.
Radek Brich <radek.brich@devl.cz>
parents:
diff changeset
     3
105
10551741f61f Move batchquery, listtables tools under pgtool.
Radek Brich <radek.brich@devl.cz>
parents: 104
diff changeset
     4
class BatchQueryTool(SimpleTool):
10551741f61f Move batchquery, listtables tools under pgtool.
Radek Brich <radek.brich@devl.cz>
parents: 104
diff changeset
     5
10551741f61f Move batchquery, listtables tools under pgtool.
Radek Brich <radek.brich@devl.cz>
parents: 104
diff changeset
     6
    """
10551741f61f Move batchquery, listtables tools under pgtool.
Radek Brich <radek.brich@devl.cz>
parents: 104
diff changeset
     7
    Run a query using columns from CSV file as arguments.
10551741f61f Move batchquery, listtables tools under pgtool.
Radek Brich <radek.brich@devl.cz>
parents: 104
diff changeset
     8
    """
10551741f61f Move batchquery, listtables tools under pgtool.
Radek Brich <radek.brich@devl.cz>
parents: 104
diff changeset
     9
29
8636719a30f6 Add batchquery tool.
Radek Brich <radek.brich@devl.cz>
parents:
diff changeset
    10
    def __init__(self):
105
10551741f61f Move batchquery, listtables tools under pgtool.
Radek Brich <radek.brich@devl.cz>
parents: 104
diff changeset
    11
        SimpleTool.__init__(self, name='batchquery')
10551741f61f Move batchquery, listtables tools under pgtool.
Radek Brich <radek.brich@devl.cz>
parents: 104
diff changeset
    12
10551741f61f Move batchquery, listtables tools under pgtool.
Radek Brich <radek.brich@devl.cz>
parents: 104
diff changeset
    13
    def specify_args(self):
10551741f61f Move batchquery, listtables tools under pgtool.
Radek Brich <radek.brich@devl.cz>
parents: 104
diff changeset
    14
        SimpleTool.specify_args(self)
39
Radek Brich <radek.brich@devl.cz>
parents: 30
diff changeset
    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).')
29
8636719a30f6 Add batchquery tool.
Radek Brich <radek.brich@devl.cz>
parents:
diff changeset
    16
        self.parser.add_argument('--file', dest='file', type=str, help='CSV file with data to use as arguments.')
8636719a30f6 Add batchquery tool.
Radek Brich <radek.brich@devl.cz>
parents:
diff changeset
    17
        self.parser.add_argument('--init', dest='init', type=str, help='Query which initialize database session (eg. temporary function).')
8636719a30f6 Add batchquery tool.
Radek Brich <radek.brich@devl.cz>
parents:
diff changeset
    18
        self.parser.add_argument('--output', dest='output', type=str, help='File name for results.')
30
a8b7cd92f39f Fix analyzeall tool. Add user output formating to batchquery tool. Add isolation_level setting support to ToolBase.
Radek Brich <radek.brich@devl.cz>
parents: 29
diff changeset
    19
        self.parser.add_argument('--outputfunc', dest='outputfunc', type=str, help='Python function which will format results (format_row(args, rows)).')
29
8636719a30f6 Add batchquery tool.
Radek Brich <radek.brich@devl.cz>
parents:
diff changeset
    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.')
8636719a30f6 Add batchquery tool.
Radek Brich <radek.brich@devl.cz>
parents:
diff changeset
    21
8636719a30f6 Add batchquery tool.
Radek Brich <radek.brich@devl.cz>
parents:
diff changeset
    22
    def _split_line(self, line):
8636719a30f6 Add batchquery tool.
Radek Brich <radek.brich@devl.cz>
parents:
diff changeset
    23
        return [x.strip() for x in line.split(',')]
8636719a30f6 Add batchquery tool.
Radek Brich <radek.brich@devl.cz>
parents:
diff changeset
    24
8636719a30f6 Add batchquery tool.
Radek Brich <radek.brich@devl.cz>
parents:
diff changeset
    25
    def main(self):
8636719a30f6 Add batchquery tool.
Radek Brich <radek.brich@devl.cz>
parents:
diff changeset
    26
        results = []
8636719a30f6 Add batchquery tool.
Radek Brich <radek.brich@devl.cz>
parents:
diff changeset
    27
        # load query from file
8636719a30f6 Add batchquery tool.
Radek Brich <radek.brich@devl.cz>
parents:
diff changeset
    28
        with open(self.args.query, 'r', encoding='utf8') as f:
8636719a30f6 Add batchquery tool.
Radek Brich <radek.brich@devl.cz>
parents:
diff changeset
    29
            query = f.read()
8636719a30f6 Add batchquery tool.
Radek Brich <radek.brich@devl.cz>
parents:
diff changeset
    30
        # connect DB
8636719a30f6 Add batchquery tool.
Radek Brich <radek.brich@devl.cz>
parents:
diff changeset
    31
        with self.pgm.cursor('target') as curs:
8636719a30f6 Add batchquery tool.
Radek Brich <radek.brich@devl.cz>
parents:
diff changeset
    32
            # run init query
8636719a30f6 Add batchquery tool.
Radek Brich <radek.brich@devl.cz>
parents:
diff changeset
    33
            if self.args.init:
8636719a30f6 Add batchquery tool.
Radek Brich <radek.brich@devl.cz>
parents:
diff changeset
    34
                with open(self.args.init, 'r', encoding='utf8') as f:
8636719a30f6 Add batchquery tool.
Radek Brich <radek.brich@devl.cz>
parents:
diff changeset
    35
                    curs.execute(f.read(), [])
8636719a30f6 Add batchquery tool.
Radek Brich <radek.brich@devl.cz>
parents:
diff changeset
    36
            # read CSV file
8636719a30f6 Add batchquery tool.
Radek Brich <radek.brich@devl.cz>
parents:
diff changeset
    37
            with open(self.args.file, 'r', encoding='utf8') as f:
8636719a30f6 Add batchquery tool.
Radek Brich <radek.brich@devl.cz>
parents:
diff changeset
    38
                # read header
8636719a30f6 Add batchquery tool.
Radek Brich <radek.brich@devl.cz>
parents:
diff changeset
    39
                names = None
8636719a30f6 Add batchquery tool.
Radek Brich <radek.brich@devl.cz>
parents:
diff changeset
    40
                if self.args.header:
8636719a30f6 Add batchquery tool.
Radek Brich <radek.brich@devl.cz>
parents:
diff changeset
    41
                    line = f.readline()
8636719a30f6 Add batchquery tool.
Radek Brich <radek.brich@devl.cz>
parents:
diff changeset
    42
                    names = self._split_line(line)
8636719a30f6 Add batchquery tool.
Radek Brich <radek.brich@devl.cz>
parents:
diff changeset
    43
                # read and process lines
8636719a30f6 Add batchquery tool.
Radek Brich <radek.brich@devl.cz>
parents:
diff changeset
    44
                for line in f:
8636719a30f6 Add batchquery tool.
Radek Brich <radek.brich@devl.cz>
parents:
diff changeset
    45
                    args = self._split_line(line)
8636719a30f6 Add batchquery tool.
Radek Brich <radek.brich@devl.cz>
parents:
diff changeset
    46
                    if names:
8636719a30f6 Add batchquery tool.
Radek Brich <radek.brich@devl.cz>
parents:
diff changeset
    47
                        args = dict(zip(names, args))
8636719a30f6 Add batchquery tool.
Radek Brich <radek.brich@devl.cz>
parents:
diff changeset
    48
                    curs.execute(query, args)
8636719a30f6 Add batchquery tool.
Radek Brich <radek.brich@devl.cz>
parents:
diff changeset
    49
                    rows = curs.fetchall()
8636719a30f6 Add batchquery tool.
Radek Brich <radek.brich@devl.cz>
parents:
diff changeset
    50
                    results.append((args, rows))
8636719a30f6 Add batchquery tool.
Radek Brich <radek.brich@devl.cz>
parents:
diff changeset
    51
            curs.connection.commit()
8636719a30f6 Add batchquery tool.
Radek Brich <radek.brich@devl.cz>
parents:
diff changeset
    52
        # write results to output file
8636719a30f6 Add batchquery tool.
Radek Brich <radek.brich@devl.cz>
parents:
diff changeset
    53
        if self.args.output:
30
a8b7cd92f39f Fix analyzeall tool. Add user output formating to batchquery tool. Add isolation_level setting support to ToolBase.
Radek Brich <radek.brich@devl.cz>
parents: 29
diff changeset
    54
            format_row = None
a8b7cd92f39f Fix analyzeall tool. Add user output formating to batchquery tool. Add isolation_level setting support to ToolBase.
Radek Brich <radek.brich@devl.cz>
parents: 29
diff changeset
    55
            if self.args.outputfunc:
a8b7cd92f39f Fix analyzeall tool. Add user output formating to batchquery tool. Add isolation_level setting support to ToolBase.
Radek Brich <radek.brich@devl.cz>
parents: 29
diff changeset
    56
                with open(self.args.outputfunc, 'r', encoding='utf8') as f:
a8b7cd92f39f Fix analyzeall tool. Add user output formating to batchquery tool. Add isolation_level setting support to ToolBase.
Radek Brich <radek.brich@devl.cz>
parents: 29
diff changeset
    57
                    d = dict()
a8b7cd92f39f Fix analyzeall tool. Add user output formating to batchquery tool. Add isolation_level setting support to ToolBase.
Radek Brich <radek.brich@devl.cz>
parents: 29
diff changeset
    58
                    exec(f.read(), d)
a8b7cd92f39f Fix analyzeall tool. Add user output formating to batchquery tool. Add isolation_level setting support to ToolBase.
Radek Brich <radek.brich@devl.cz>
parents: 29
diff changeset
    59
                    format_row = d['format_row']
83
515fadd3d286 Add dependency on pycolib. Move common modules to pycolib. Add example table schema for meta DB.
Radek Brich <radek.brich@devl.cz>
parents: 39
diff changeset
    60
29
8636719a30f6 Add batchquery tool.
Radek Brich <radek.brich@devl.cz>
parents:
diff changeset
    61
            with open(self.args.output, 'w', encoding='utf8') as f:
8636719a30f6 Add batchquery tool.
Radek Brich <radek.brich@devl.cz>
parents:
diff changeset
    62
                for args, rows in results:
30
a8b7cd92f39f Fix analyzeall tool. Add user output formating to batchquery tool. Add isolation_level setting support to ToolBase.
Radek Brich <radek.brich@devl.cz>
parents: 29
diff changeset
    63
                    if format_row:
a8b7cd92f39f Fix analyzeall tool. Add user output formating to batchquery tool. Add isolation_level setting support to ToolBase.
Radek Brich <radek.brich@devl.cz>
parents: 29
diff changeset
    64
                        f.write(format_row(args, rows))
a8b7cd92f39f Fix analyzeall tool. Add user output formating to batchquery tool. Add isolation_level setting support to ToolBase.
Radek Brich <radek.brich@devl.cz>
parents: 29
diff changeset
    65
                    else:
a8b7cd92f39f Fix analyzeall tool. Add user output formating to batchquery tool. Add isolation_level setting support to ToolBase.
Radek Brich <radek.brich@devl.cz>
parents: 29
diff changeset
    66
                        f.write(repr(args))
a8b7cd92f39f Fix analyzeall tool. Add user output formating to batchquery tool. Add isolation_level setting support to ToolBase.
Radek Brich <radek.brich@devl.cz>
parents: 29
diff changeset
    67
                        f.write(' -> ')
a8b7cd92f39f Fix analyzeall tool. Add user output formating to batchquery tool. Add isolation_level setting support to ToolBase.
Radek Brich <radek.brich@devl.cz>
parents: 29
diff changeset
    68
                        f.write(repr(rows))
29
8636719a30f6 Add batchquery tool.
Radek Brich <radek.brich@devl.cz>
parents:
diff changeset
    69
                    f.write('\n')
8636719a30f6 Add batchquery tool.
Radek Brich <radek.brich@devl.cz>
parents:
diff changeset
    70
8636719a30f6 Add batchquery tool.
Radek Brich <radek.brich@devl.cz>
parents:
diff changeset
    71
105
10551741f61f Move batchquery, listtables tools under pgtool.
Radek Brich <radek.brich@devl.cz>
parents: 104
diff changeset
    72
cls = BatchQueryTool
29
8636719a30f6 Add batchquery tool.
Radek Brich <radek.brich@devl.cz>
parents:
diff changeset
    73