1 #!/usr/bin/env python3.2 |
1 #!/usr/bin/env python3.2 |
2 # |
2 # |
3 # Copy data from one table to another table with same schema. |
3 # Copy data between tables with same table schema. |
|
4 # |
|
5 # Can copy multiple tables in one run. |
|
6 # Sorts the tables according to references. |
|
7 # |
|
8 # This may be used instead of dump/restore. |
4 # |
9 # |
5 |
10 |
6 import io |
11 import io |
7 |
12 |
8 from pgtoolkit import toolbase, pgmanager, pgdatacopy |
13 from pgtoolkit import toolbase, pgmanager, pgbrowser, pgdatacopy |
9 from pgtoolkit.progresswrapper import ProgressWrapper |
14 from pgtoolkit.progresswrapper import ProgressWrapper |
10 |
15 |
11 |
16 |
12 class TableCopyTool(toolbase.SrcDstTablesTool): |
17 class TableCopyTool(toolbase.SrcDstTablesTool): |
13 def __init__(self): |
18 def __init__(self): |
14 toolbase.SrcDstTablesTool.__init__(self, name='tablecopy', desc='Table copy tool.') |
19 toolbase.SrcDstTablesTool.__init__(self, name='tablecopy', desc='Table copy tool.') |
15 |
20 |
16 self.parser.add_argument('-n', '--no-action', dest='noaction', action='store_true', |
21 self.parser.add_argument('-n', '--no-action', dest='noaction', action='store_true', |
17 help="Do nothing, just print tables to be copied. Useful in combination with --regex.") |
22 help="Do nothing, just print tables to be copied. Useful in combination with --regex.") |
|
23 self.parser.add_argument('--no-sort', dest='nosort', action='store_true', |
|
24 help="Do not sort. By default, tables are sorted by foreign key references.") |
18 |
25 |
19 self.init() |
26 self.init() |
20 |
27 |
21 def main(self): |
28 def main(self): |
22 srcconn = self.pgm.get_conn('src') |
29 self.srcconn = self.pgm.get_conn('src') |
23 dstconn = self.pgm.get_conn('dst') |
30 self.dstconn = self.pgm.get_conn('dst') |
24 |
31 |
25 dc = pgdatacopy.PgDataCopy(srcconn, dstconn) |
32 dc = pgdatacopy.PgDataCopy(self.srcconn, self.dstconn) |
26 |
33 |
27 for srcschema, srctable, dstschema, dsttable in self.tables(): |
34 if self.args.nosort: |
28 print('Copying [%s] %s.%s --> [%s] %s.%s' % ( |
35 for table in self.tables(): |
29 self.args.src, srcschema, srctable, |
36 self.copy_table(dc, *table) |
30 self.args.dst, dstschema, dsttable)) |
37 else: |
|
38 # sort tables with respect to references |
|
39 details = dict() |
|
40 pending = set() |
|
41 references = dict() |
31 |
42 |
32 if self.args.noaction: |
43 # build list of all table to be copied (pending) and references map |
33 continue |
44 for table in self.tables(): |
|
45 srcschema, srctable, dstschema, dsttable = table |
|
46 name = dstschema + '.' + dsttable |
|
47 details[name] = table |
|
48 pending.add(name) |
|
49 references[name] = self.get_references(dstschema, dsttable) |
34 |
50 |
35 dc.set_source(srctable, srcschema) |
51 # copy files with fulfilled references, repeat until all done |
36 dc.set_destination(dsttable, dstschema) |
52 while pending: |
|
53 for name in list(pending): |
|
54 ok = True |
|
55 for ref in references[name]: |
|
56 if ref in pending: |
|
57 ok = False |
|
58 if ok: |
|
59 table = details[name] |
|
60 self.copy_table(dc, *table) |
|
61 pending.remove(name) |
|
62 |
|
63 def get_references(self, schema, table): |
|
64 browser = pgbrowser.PgBrowser(self.dstconn) |
|
65 cons = browser.list_constraints(table, schema) |
|
66 return [con['fschema'] + '.' + con['fname'] for con in cons if con['type'] == 'f'] |
|
67 |
|
68 def copy_table(self, dc, srcschema, srctable, dstschema, dsttable): |
|
69 print('Copying [%s] %s.%s --> [%s] %s.%s' % ( |
|
70 self.args.src, srcschema, srctable, |
|
71 self.args.dst, dstschema, dsttable)) |
37 |
72 |
|
73 if self.args.noaction: |
|
74 return |
|
75 |
|
76 dc.set_source(srctable, srcschema) |
|
77 dc.set_destination(dsttable, dstschema) |
|
78 |
|
79 try: |
38 dc.check() |
80 dc.check() |
|
81 except pgdatacopy.TargetNotEmptyError as e: |
|
82 print(' - error:', str(e)) |
|
83 return |
39 |
84 |
40 buf = io.BytesIO() |
85 print(' - read ') |
41 wrapped = ProgressWrapper(buf) |
86 buf = io.BytesIO() |
42 dc.read(wrapped) |
87 wrapped = ProgressWrapper(buf) |
43 data = buf.getvalue() |
88 dc.read(wrapped) |
44 buf.close() |
89 data = buf.getvalue() |
45 |
90 buf.close() |
46 buf = io.BytesIO(data) |
|
47 wrapped = ProgressWrapper(buf, len(data)) |
|
48 dc.write(wrapped) |
|
49 buf.close() |
|
50 |
91 |
51 dc.analyze() |
92 print(' - write ') |
|
93 buf = io.BytesIO(data) |
|
94 wrapped = ProgressWrapper(buf, len(data)) |
|
95 dc.write(wrapped) |
|
96 buf.close() |
|
97 |
|
98 print(' - analyze ') |
|
99 dc.analyze() |
52 |
100 |
53 |
101 |
54 tool = TableCopyTool() |
102 tool = TableCopyTool() |
55 tool.main() |
103 tool.main() |
56 |
104 |