97
|
1 |
#!/usr/bin/env python3
|
|
2 |
|
|
3 |
from pgtoolkit import toolbase
|
|
4 |
|
|
5 |
|
|
6 |
class BatchCopyTool(toolbase.SrcDstTablesTool):
|
|
7 |
def __init__(self):
|
|
8 |
toolbase.SrcDstTablesTool.__init__(self, name='batchcopy', desc='Copy data from one table to another.')
|
|
9 |
|
|
10 |
self.parser.add_argument('--table-name', type=str, help='Table to be copied.')
|
|
11 |
self.parser.add_argument('--src-filter', type=str, help='WHERE condition for source query.')
|
|
12 |
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}')
|
|
13 |
#TODO: duplicates=rollback|ignore|update
|
|
14 |
|
|
15 |
self.init()
|
|
16 |
|
|
17 |
def main(self):
|
|
18 |
# read list of IDs from file
|
|
19 |
ids = '<no IDs read>'
|
|
20 |
if self.args.file_with_ids:
|
|
21 |
with open(self.args.file_with_ids, 'r') as f:
|
|
22 |
ids = ','.join(ln.rstrip() for ln in f.readlines())
|
|
23 |
|
|
24 |
# read source data
|
|
25 |
with self.pgm.cursor('src') as src_curs:
|
|
26 |
condition = self.args.src_filter.format(ids=ids) or 'true'
|
|
27 |
src_curs.execute('SELECT * FROM {} WHERE {}'.format(self.args.table_name, condition))
|
|
28 |
#TODO: ORDER BY id OFFSET 0 LIMIT 100
|
|
29 |
data = src_curs.fetchall_dict()
|
|
30 |
src_curs.connection.commit()
|
|
31 |
|
|
32 |
with self.pgm.cursor('dst') as dst_curs:
|
|
33 |
for row in data:
|
|
34 |
keys = ', '.join(row.keys())
|
|
35 |
values_mask = ', '.join(['%s'] * len(row))
|
|
36 |
query = 'INSERT INTO {} ({}) VALUES ({})'.format(self.args.table_name, keys, values_mask)
|
|
37 |
dst_curs.execute(query, list(row.values()))
|
|
38 |
dst_curs.connection.commit()
|
|
39 |
|
|
40 |
self.log.info('Copied {} rows.' % len(rows))
|
|
41 |
|
|
42 |
|
|
43 |
tool = BatchCopyTool()
|
|
44 |
tool.main()
|