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-- |
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 | 2 |
|
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 | 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 | 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 | 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.') |
|
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 | 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: |
|
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 | 61 |
with open(self.args.output, 'w', encoding='utf8') as f: |
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 | 69 |
f.write('\n') |
70 |
||
71 |
||
105
10551741f61f
Move batchquery, listtables tools under pgtool.
Radek Brich <radek.brich@devl.cz>
parents:
104
diff
changeset
|
72 |
cls = BatchQueryTool |
29 | 73 |