author | Radek Brich <brich.radek@ifortuna.cz> |
Tue, 06 May 2014 18:37:43 +0200 | |
changeset 101 | 2a2d0d5df03b |
parent 98 | batchcopy.py@024299702087 |
permissions | -rw-r--r-- |
101
2a2d0d5df03b
Refactor ToolBase to allow tool composition. Add TableSync tool (composited). Move more tools under pgtool.
Radek Brich <brich.radek@ifortuna.cz>
parents:
98
diff
changeset
|
1 |
from pgtoolkit.toolbase import SrcDstTablesTool |
98
024299702087
Update batchcopy: When target record exists, allow to ignore / update the error (--dst-exists parameter).
Radek Brich <brich.radek@ifortuna.cz>
parents:
97
diff
changeset
|
2 |
from pgtoolkit.pgmanager import IntegrityError |
97 | 3 |
|
4 |
||
101
2a2d0d5df03b
Refactor ToolBase to allow tool composition. Add TableSync tool (composited). Move more tools under pgtool.
Radek Brich <brich.radek@ifortuna.cz>
parents:
98
diff
changeset
|
5 |
class BatchCopyTool(SrcDstTablesTool): |
2a2d0d5df03b
Refactor ToolBase to allow tool composition. Add TableSync tool (composited). Move more tools under pgtool.
Radek Brich <brich.radek@ifortuna.cz>
parents:
98
diff
changeset
|
6 |
|
2a2d0d5df03b
Refactor ToolBase to allow tool composition. Add TableSync tool (composited). Move more tools under pgtool.
Radek Brich <brich.radek@ifortuna.cz>
parents:
98
diff
changeset
|
7 |
""" |
2a2d0d5df03b
Refactor ToolBase to allow tool composition. Add TableSync tool (composited). Move more tools under pgtool.
Radek Brich <brich.radek@ifortuna.cz>
parents:
98
diff
changeset
|
8 |
Copy data from one table to another, filtering by specified condition. |
2a2d0d5df03b
Refactor ToolBase to allow tool composition. Add TableSync tool (composited). Move more tools under pgtool.
Radek Brich <brich.radek@ifortuna.cz>
parents:
98
diff
changeset
|
9 |
|
2a2d0d5df03b
Refactor ToolBase to allow tool composition. Add TableSync tool (composited). Move more tools under pgtool.
Radek Brich <brich.radek@ifortuna.cz>
parents:
98
diff
changeset
|
10 |
""" |
2a2d0d5df03b
Refactor ToolBase to allow tool composition. Add TableSync tool (composited). Move more tools under pgtool.
Radek Brich <brich.radek@ifortuna.cz>
parents:
98
diff
changeset
|
11 |
|
97 | 12 |
def __init__(self): |
101
2a2d0d5df03b
Refactor ToolBase to allow tool composition. Add TableSync tool (composited). Move more tools under pgtool.
Radek Brich <brich.radek@ifortuna.cz>
parents:
98
diff
changeset
|
13 |
SrcDstTablesTool.__init__(self, name='batchcopy', desc='') |
97 | 14 |
|
101
2a2d0d5df03b
Refactor ToolBase to allow tool composition. Add TableSync tool (composited). Move more tools under pgtool.
Radek Brich <brich.radek@ifortuna.cz>
parents:
98
diff
changeset
|
15 |
def specify_args(self): |
2a2d0d5df03b
Refactor ToolBase to allow tool composition. Add TableSync tool (composited). Move more tools under pgtool.
Radek Brich <brich.radek@ifortuna.cz>
parents:
98
diff
changeset
|
16 |
SrcDstTablesTool.specify_args(self) |
97 | 17 |
self.parser.add_argument('--table-name', type=str, help='Table to be copied.') |
18 |
self.parser.add_argument('--src-filter', type=str, help='WHERE condition for source query.') |
|
19 |
self.parser.add_argument('--file-with-ids', type=str, help='Read source IDs from file (each ID on new line). Use these in --src-filter as {ids}') |
|
98
024299702087
Update batchcopy: When target record exists, allow to ignore / update the error (--dst-exists parameter).
Radek Brich <brich.radek@ifortuna.cz>
parents:
97
diff
changeset
|
20 |
self.parser.add_argument('--dst-exists', choices=['rollback', 'ignore', 'update'], default='rollback', help='What to do when destination record already exists.') |
97 | 21 |
|
22 |
def main(self): |
|
23 |
# read list of IDs from file |
|
24 |
ids = '<no IDs read>' |
|
25 |
if self.args.file_with_ids: |
|
26 |
with open(self.args.file_with_ids, 'r') as f: |
|
27 |
ids = ','.join(ln.rstrip() for ln in f.readlines()) |
|
28 |
||
29 |
# read source data |
|
30 |
with self.pgm.cursor('src') as src_curs: |
|
31 |
condition = self.args.src_filter.format(ids=ids) or 'true' |
|
32 |
src_curs.execute('SELECT * FROM {} WHERE {}'.format(self.args.table_name, condition)) |
|
33 |
#TODO: ORDER BY id OFFSET 0 LIMIT 100 |
|
34 |
data = src_curs.fetchall_dict() |
|
35 |
src_curs.connection.commit() |
|
36 |
||
37 |
with self.pgm.cursor('dst') as dst_curs: |
|
98
024299702087
Update batchcopy: When target record exists, allow to ignore / update the error (--dst-exists parameter).
Radek Brich <brich.radek@ifortuna.cz>
parents:
97
diff
changeset
|
38 |
copied = 0 |
97 | 39 |
for row in data: |
40 |
keys = ', '.join(row.keys()) |
|
41 |
values_mask = ', '.join(['%s'] * len(row)) |
|
42 |
query = 'INSERT INTO {} ({}) VALUES ({})'.format(self.args.table_name, keys, values_mask) |
|
98
024299702087
Update batchcopy: When target record exists, allow to ignore / update the error (--dst-exists parameter).
Radek Brich <brich.radek@ifortuna.cz>
parents:
97
diff
changeset
|
43 |
try: |
024299702087
Update batchcopy: When target record exists, allow to ignore / update the error (--dst-exists parameter).
Radek Brich <brich.radek@ifortuna.cz>
parents:
97
diff
changeset
|
44 |
dst_curs.execute('SAVEPOINT the_query;') |
024299702087
Update batchcopy: When target record exists, allow to ignore / update the error (--dst-exists parameter).
Radek Brich <brich.radek@ifortuna.cz>
parents:
97
diff
changeset
|
45 |
dst_curs.execute(query, list(row.values())) |
024299702087
Update batchcopy: When target record exists, allow to ignore / update the error (--dst-exists parameter).
Radek Brich <brich.radek@ifortuna.cz>
parents:
97
diff
changeset
|
46 |
dst_curs.execute('RELEASE SAVEPOINT the_query;') |
024299702087
Update batchcopy: When target record exists, allow to ignore / update the error (--dst-exists parameter).
Radek Brich <brich.radek@ifortuna.cz>
parents:
97
diff
changeset
|
47 |
copied += 1 |
024299702087
Update batchcopy: When target record exists, allow to ignore / update the error (--dst-exists parameter).
Radek Brich <brich.radek@ifortuna.cz>
parents:
97
diff
changeset
|
48 |
except IntegrityError: |
024299702087
Update batchcopy: When target record exists, allow to ignore / update the error (--dst-exists parameter).
Radek Brich <brich.radek@ifortuna.cz>
parents:
97
diff
changeset
|
49 |
if self.args.dst_exists == 'rollback': |
024299702087
Update batchcopy: When target record exists, allow to ignore / update the error (--dst-exists parameter).
Radek Brich <brich.radek@ifortuna.cz>
parents:
97
diff
changeset
|
50 |
dst_curs.connection.rollback() |
024299702087
Update batchcopy: When target record exists, allow to ignore / update the error (--dst-exists parameter).
Radek Brich <brich.radek@ifortuna.cz>
parents:
97
diff
changeset
|
51 |
break |
024299702087
Update batchcopy: When target record exists, allow to ignore / update the error (--dst-exists parameter).
Radek Brich <brich.radek@ifortuna.cz>
parents:
97
diff
changeset
|
52 |
elif self.args.dst_exists == 'ignore': |
024299702087
Update batchcopy: When target record exists, allow to ignore / update the error (--dst-exists parameter).
Radek Brich <brich.radek@ifortuna.cz>
parents:
97
diff
changeset
|
53 |
dst_curs.execute('ROLLBACK TO SAVEPOINT the_query;') |
024299702087
Update batchcopy: When target record exists, allow to ignore / update the error (--dst-exists parameter).
Radek Brich <brich.radek@ifortuna.cz>
parents:
97
diff
changeset
|
54 |
elif self.args.dst_exists == 'update': |
024299702087
Update batchcopy: When target record exists, allow to ignore / update the error (--dst-exists parameter).
Radek Brich <brich.radek@ifortuna.cz>
parents:
97
diff
changeset
|
55 |
raise NotImplementedError() |
97 | 56 |
dst_curs.connection.commit() |
57 |
||
98
024299702087
Update batchcopy: When target record exists, allow to ignore / update the error (--dst-exists parameter).
Radek Brich <brich.radek@ifortuna.cz>
parents:
97
diff
changeset
|
58 |
self.log.info('Copied %s rows.', copied) |
97 | 59 |
|
60 |
||
101
2a2d0d5df03b
Refactor ToolBase to allow tool composition. Add TableSync tool (composited). Move more tools under pgtool.
Radek Brich <brich.radek@ifortuna.cz>
parents:
98
diff
changeset
|
61 |
cls = BatchCopyTool |