tablecopy.py
author Radek Brich <brich.radek@ifortuna.cz>
Mon, 14 Apr 2014 22:28:12 +0200
changeset 96 acf63df539b7
parent 93 b72591087495
child 104 d8ff52a0390f
permissions -rwxr-xr-x
Update tablecopy tool: Add --disable-triggers option.
Ignore whitespace changes - Everywhere: Within whitespace: At end of lines:
93
b72591087495 Change python3.2 to python3 in scripts.
Radek Brich <brich.radek@ifortuna.cz>
parents: 56
diff changeset
     1
#!/usr/bin/env python3
34
98c7809af415 Add PgDataCopy. Add TableCopyTool.Add SrcDstTablesTool class to toolbase, use in tablecopy, tablediff.
Radek Brich <radek.brich@devl.cz>
parents:
diff changeset
     2
#
35
e7f79c4a27ce TableCopyTool: order tables with respect to references.
Radek Brich <radek.brich@devl.cz>
parents: 34
diff changeset
     3
# Copy data between tables with same table schema.
e7f79c4a27ce TableCopyTool: order tables with respect to references.
Radek Brich <radek.brich@devl.cz>
parents: 34
diff changeset
     4
#
56
94e091c23ebb Add reverse parameter for diff tools.
Radek Brich <radek.brich@devl.cz>
parents: 35
diff changeset
     5
# Copies full table, target table must be empty.
35
e7f79c4a27ce TableCopyTool: order tables with respect to references.
Radek Brich <radek.brich@devl.cz>
parents: 34
diff changeset
     6
# Can copy multiple tables in one run.
e7f79c4a27ce TableCopyTool: order tables with respect to references.
Radek Brich <radek.brich@devl.cz>
parents: 34
diff changeset
     7
# Sorts the tables according to references.
e7f79c4a27ce TableCopyTool: order tables with respect to references.
Radek Brich <radek.brich@devl.cz>
parents: 34
diff changeset
     8
#
e7f79c4a27ce TableCopyTool: order tables with respect to references.
Radek Brich <radek.brich@devl.cz>
parents: 34
diff changeset
     9
# This may be used instead of dump/restore.
34
98c7809af415 Add PgDataCopy. Add TableCopyTool.Add SrcDstTablesTool class to toolbase, use in tablecopy, tablediff.
Radek Brich <radek.brich@devl.cz>
parents:
diff changeset
    10
#
98c7809af415 Add PgDataCopy. Add TableCopyTool.Add SrcDstTablesTool class to toolbase, use in tablecopy, tablediff.
Radek Brich <radek.brich@devl.cz>
parents:
diff changeset
    11
98c7809af415 Add PgDataCopy. Add TableCopyTool.Add SrcDstTablesTool class to toolbase, use in tablecopy, tablediff.
Radek Brich <radek.brich@devl.cz>
parents:
diff changeset
    12
import io
98c7809af415 Add PgDataCopy. Add TableCopyTool.Add SrcDstTablesTool class to toolbase, use in tablecopy, tablediff.
Radek Brich <radek.brich@devl.cz>
parents:
diff changeset
    13
35
e7f79c4a27ce TableCopyTool: order tables with respect to references.
Radek Brich <radek.brich@devl.cz>
parents: 34
diff changeset
    14
from pgtoolkit import toolbase, pgmanager, pgbrowser, pgdatacopy
34
98c7809af415 Add PgDataCopy. Add TableCopyTool.Add SrcDstTablesTool class to toolbase, use in tablecopy, tablediff.
Radek Brich <radek.brich@devl.cz>
parents:
diff changeset
    15
from pgtoolkit.progresswrapper import ProgressWrapper
98c7809af415 Add PgDataCopy. Add TableCopyTool.Add SrcDstTablesTool class to toolbase, use in tablecopy, tablediff.
Radek Brich <radek.brich@devl.cz>
parents:
diff changeset
    16
98c7809af415 Add PgDataCopy. Add TableCopyTool.Add SrcDstTablesTool class to toolbase, use in tablecopy, tablediff.
Radek Brich <radek.brich@devl.cz>
parents:
diff changeset
    17
98c7809af415 Add PgDataCopy. Add TableCopyTool.Add SrcDstTablesTool class to toolbase, use in tablecopy, tablediff.
Radek Brich <radek.brich@devl.cz>
parents:
diff changeset
    18
class TableCopyTool(toolbase.SrcDstTablesTool):
98c7809af415 Add PgDataCopy. Add TableCopyTool.Add SrcDstTablesTool class to toolbase, use in tablecopy, tablediff.
Radek Brich <radek.brich@devl.cz>
parents:
diff changeset
    19
    def __init__(self):
98c7809af415 Add PgDataCopy. Add TableCopyTool.Add SrcDstTablesTool class to toolbase, use in tablecopy, tablediff.
Radek Brich <radek.brich@devl.cz>
parents:
diff changeset
    20
        toolbase.SrcDstTablesTool.__init__(self, name='tablecopy', desc='Table copy tool.')
56
94e091c23ebb Add reverse parameter for diff tools.
Radek Brich <radek.brich@devl.cz>
parents: 35
diff changeset
    21
34
98c7809af415 Add PgDataCopy. Add TableCopyTool.Add SrcDstTablesTool class to toolbase, use in tablecopy, tablediff.
Radek Brich <radek.brich@devl.cz>
parents:
diff changeset
    22
        self.parser.add_argument('-n', '--no-action', dest='noaction', action='store_true',
98c7809af415 Add PgDataCopy. Add TableCopyTool.Add SrcDstTablesTool class to toolbase, use in tablecopy, tablediff.
Radek Brich <radek.brich@devl.cz>
parents:
diff changeset
    23
            help="Do nothing, just print tables to be copied. Useful in combination with --regex.")
35
e7f79c4a27ce TableCopyTool: order tables with respect to references.
Radek Brich <radek.brich@devl.cz>
parents: 34
diff changeset
    24
        self.parser.add_argument('--no-sort', dest='nosort', action='store_true',
e7f79c4a27ce TableCopyTool: order tables with respect to references.
Radek Brich <radek.brich@devl.cz>
parents: 34
diff changeset
    25
            help="Do not sort. By default, tables are sorted by foreign key references.")
96
acf63df539b7 Update tablecopy tool: Add --disable-triggers option.
Radek Brich <brich.radek@ifortuna.cz>
parents: 93
diff changeset
    26
        self.parser.add_argument('--disable-triggers', dest='notriggers', action='store_true',
acf63df539b7 Update tablecopy tool: Add --disable-triggers option.
Radek Brich <brich.radek@ifortuna.cz>
parents: 93
diff changeset
    27
            help="Disable all triggers for the session.")
56
94e091c23ebb Add reverse parameter for diff tools.
Radek Brich <radek.brich@devl.cz>
parents: 35
diff changeset
    28
34
98c7809af415 Add PgDataCopy. Add TableCopyTool.Add SrcDstTablesTool class to toolbase, use in tablecopy, tablediff.
Radek Brich <radek.brich@devl.cz>
parents:
diff changeset
    29
        self.init()
98c7809af415 Add PgDataCopy. Add TableCopyTool.Add SrcDstTablesTool class to toolbase, use in tablecopy, tablediff.
Radek Brich <radek.brich@devl.cz>
parents:
diff changeset
    30
98c7809af415 Add PgDataCopy. Add TableCopyTool.Add SrcDstTablesTool class to toolbase, use in tablecopy, tablediff.
Radek Brich <radek.brich@devl.cz>
parents:
diff changeset
    31
    def main(self):
35
e7f79c4a27ce TableCopyTool: order tables with respect to references.
Radek Brich <radek.brich@devl.cz>
parents: 34
diff changeset
    32
        self.srcconn = self.pgm.get_conn('src')
e7f79c4a27ce TableCopyTool: order tables with respect to references.
Radek Brich <radek.brich@devl.cz>
parents: 34
diff changeset
    33
        self.dstconn = self.pgm.get_conn('dst')
56
94e091c23ebb Add reverse parameter for diff tools.
Radek Brich <radek.brich@devl.cz>
parents: 35
diff changeset
    34
96
acf63df539b7 Update tablecopy tool: Add --disable-triggers option.
Radek Brich <brich.radek@ifortuna.cz>
parents: 93
diff changeset
    35
        if self.args.notriggers:
acf63df539b7 Update tablecopy tool: Add --disable-triggers option.
Radek Brich <brich.radek@ifortuna.cz>
parents: 93
diff changeset
    36
            curs = self.dstconn.cursor()
acf63df539b7 Update tablecopy tool: Add --disable-triggers option.
Radek Brich <brich.radek@ifortuna.cz>
parents: 93
diff changeset
    37
            curs.execute('SET session_replication_role = replica;')
acf63df539b7 Update tablecopy tool: Add --disable-triggers option.
Radek Brich <brich.radek@ifortuna.cz>
parents: 93
diff changeset
    38
            curs.close()
acf63df539b7 Update tablecopy tool: Add --disable-triggers option.
Radek Brich <brich.radek@ifortuna.cz>
parents: 93
diff changeset
    39
            self.dstconn.commit()
acf63df539b7 Update tablecopy tool: Add --disable-triggers option.
Radek Brich <brich.radek@ifortuna.cz>
parents: 93
diff changeset
    40
35
e7f79c4a27ce TableCopyTool: order tables with respect to references.
Radek Brich <radek.brich@devl.cz>
parents: 34
diff changeset
    41
        dc = pgdatacopy.PgDataCopy(self.srcconn, self.dstconn)
56
94e091c23ebb Add reverse parameter for diff tools.
Radek Brich <radek.brich@devl.cz>
parents: 35
diff changeset
    42
35
e7f79c4a27ce TableCopyTool: order tables with respect to references.
Radek Brich <radek.brich@devl.cz>
parents: 34
diff changeset
    43
        if self.args.nosort:
e7f79c4a27ce TableCopyTool: order tables with respect to references.
Radek Brich <radek.brich@devl.cz>
parents: 34
diff changeset
    44
            for table in self.tables():
e7f79c4a27ce TableCopyTool: order tables with respect to references.
Radek Brich <radek.brich@devl.cz>
parents: 34
diff changeset
    45
                self.copy_table(dc, *table)
e7f79c4a27ce TableCopyTool: order tables with respect to references.
Radek Brich <radek.brich@devl.cz>
parents: 34
diff changeset
    46
        else:
e7f79c4a27ce TableCopyTool: order tables with respect to references.
Radek Brich <radek.brich@devl.cz>
parents: 34
diff changeset
    47
            # sort tables with respect to references
e7f79c4a27ce TableCopyTool: order tables with respect to references.
Radek Brich <radek.brich@devl.cz>
parents: 34
diff changeset
    48
            details = dict()
e7f79c4a27ce TableCopyTool: order tables with respect to references.
Radek Brich <radek.brich@devl.cz>
parents: 34
diff changeset
    49
            pending = set()
e7f79c4a27ce TableCopyTool: order tables with respect to references.
Radek Brich <radek.brich@devl.cz>
parents: 34
diff changeset
    50
            references = dict()
56
94e091c23ebb Add reverse parameter for diff tools.
Radek Brich <radek.brich@devl.cz>
parents: 35
diff changeset
    51
35
e7f79c4a27ce TableCopyTool: order tables with respect to references.
Radek Brich <radek.brich@devl.cz>
parents: 34
diff changeset
    52
            # build list of all table to be copied (pending) and references map
e7f79c4a27ce TableCopyTool: order tables with respect to references.
Radek Brich <radek.brich@devl.cz>
parents: 34
diff changeset
    53
            for table in self.tables():
e7f79c4a27ce TableCopyTool: order tables with respect to references.
Radek Brich <radek.brich@devl.cz>
parents: 34
diff changeset
    54
                srcschema, srctable, dstschema, dsttable = table
e7f79c4a27ce TableCopyTool: order tables with respect to references.
Radek Brich <radek.brich@devl.cz>
parents: 34
diff changeset
    55
                name = dstschema + '.' + dsttable
e7f79c4a27ce TableCopyTool: order tables with respect to references.
Radek Brich <radek.brich@devl.cz>
parents: 34
diff changeset
    56
                details[name] = table
e7f79c4a27ce TableCopyTool: order tables with respect to references.
Radek Brich <radek.brich@devl.cz>
parents: 34
diff changeset
    57
                pending.add(name)
e7f79c4a27ce TableCopyTool: order tables with respect to references.
Radek Brich <radek.brich@devl.cz>
parents: 34
diff changeset
    58
                references[name] = self.get_references(dstschema, dsttable)
56
94e091c23ebb Add reverse parameter for diff tools.
Radek Brich <radek.brich@devl.cz>
parents: 35
diff changeset
    59
35
e7f79c4a27ce TableCopyTool: order tables with respect to references.
Radek Brich <radek.brich@devl.cz>
parents: 34
diff changeset
    60
            # copy files with fulfilled references, repeat until all done
e7f79c4a27ce TableCopyTool: order tables with respect to references.
Radek Brich <radek.brich@devl.cz>
parents: 34
diff changeset
    61
            while pending:
e7f79c4a27ce TableCopyTool: order tables with respect to references.
Radek Brich <radek.brich@devl.cz>
parents: 34
diff changeset
    62
                for name in list(pending):
e7f79c4a27ce TableCopyTool: order tables with respect to references.
Radek Brich <radek.brich@devl.cz>
parents: 34
diff changeset
    63
                    ok = True
e7f79c4a27ce TableCopyTool: order tables with respect to references.
Radek Brich <radek.brich@devl.cz>
parents: 34
diff changeset
    64
                    for ref in references[name]:
e7f79c4a27ce TableCopyTool: order tables with respect to references.
Radek Brich <radek.brich@devl.cz>
parents: 34
diff changeset
    65
                        if ref in pending:
e7f79c4a27ce TableCopyTool: order tables with respect to references.
Radek Brich <radek.brich@devl.cz>
parents: 34
diff changeset
    66
                            ok = False
e7f79c4a27ce TableCopyTool: order tables with respect to references.
Radek Brich <radek.brich@devl.cz>
parents: 34
diff changeset
    67
                    if ok:
e7f79c4a27ce TableCopyTool: order tables with respect to references.
Radek Brich <radek.brich@devl.cz>
parents: 34
diff changeset
    68
                        table = details[name]
e7f79c4a27ce TableCopyTool: order tables with respect to references.
Radek Brich <radek.brich@devl.cz>
parents: 34
diff changeset
    69
                        self.copy_table(dc, *table)
e7f79c4a27ce TableCopyTool: order tables with respect to references.
Radek Brich <radek.brich@devl.cz>
parents: 34
diff changeset
    70
                        pending.remove(name)
e7f79c4a27ce TableCopyTool: order tables with respect to references.
Radek Brich <radek.brich@devl.cz>
parents: 34
diff changeset
    71
e7f79c4a27ce TableCopyTool: order tables with respect to references.
Radek Brich <radek.brich@devl.cz>
parents: 34
diff changeset
    72
    def get_references(self, schema, table):
e7f79c4a27ce TableCopyTool: order tables with respect to references.
Radek Brich <radek.brich@devl.cz>
parents: 34
diff changeset
    73
        browser = pgbrowser.PgBrowser(self.dstconn)
e7f79c4a27ce TableCopyTool: order tables with respect to references.
Radek Brich <radek.brich@devl.cz>
parents: 34
diff changeset
    74
        cons = browser.list_constraints(table, schema)
e7f79c4a27ce TableCopyTool: order tables with respect to references.
Radek Brich <radek.brich@devl.cz>
parents: 34
diff changeset
    75
        return [con['fschema'] + '.' + con['fname'] for con in cons if con['type'] == 'f']
e7f79c4a27ce TableCopyTool: order tables with respect to references.
Radek Brich <radek.brich@devl.cz>
parents: 34
diff changeset
    76
e7f79c4a27ce TableCopyTool: order tables with respect to references.
Radek Brich <radek.brich@devl.cz>
parents: 34
diff changeset
    77
    def copy_table(self, dc, srcschema, srctable, dstschema, dsttable):
e7f79c4a27ce TableCopyTool: order tables with respect to references.
Radek Brich <radek.brich@devl.cz>
parents: 34
diff changeset
    78
        print('Copying [%s] %s.%s --> [%s] %s.%s' % (
e7f79c4a27ce TableCopyTool: order tables with respect to references.
Radek Brich <radek.brich@devl.cz>
parents: 34
diff changeset
    79
            self.args.src, srcschema, srctable,
e7f79c4a27ce TableCopyTool: order tables with respect to references.
Radek Brich <radek.brich@devl.cz>
parents: 34
diff changeset
    80
            self.args.dst, dstschema, dsttable))
56
94e091c23ebb Add reverse parameter for diff tools.
Radek Brich <radek.brich@devl.cz>
parents: 35
diff changeset
    81
35
e7f79c4a27ce TableCopyTool: order tables with respect to references.
Radek Brich <radek.brich@devl.cz>
parents: 34
diff changeset
    82
        if self.args.noaction:
e7f79c4a27ce TableCopyTool: order tables with respect to references.
Radek Brich <radek.brich@devl.cz>
parents: 34
diff changeset
    83
            return
56
94e091c23ebb Add reverse parameter for diff tools.
Radek Brich <radek.brich@devl.cz>
parents: 35
diff changeset
    84
35
e7f79c4a27ce TableCopyTool: order tables with respect to references.
Radek Brich <radek.brich@devl.cz>
parents: 34
diff changeset
    85
        dc.set_source(srctable, srcschema)
e7f79c4a27ce TableCopyTool: order tables with respect to references.
Radek Brich <radek.brich@devl.cz>
parents: 34
diff changeset
    86
        dc.set_destination(dsttable, dstschema)
56
94e091c23ebb Add reverse parameter for diff tools.
Radek Brich <radek.brich@devl.cz>
parents: 35
diff changeset
    87
35
e7f79c4a27ce TableCopyTool: order tables with respect to references.
Radek Brich <radek.brich@devl.cz>
parents: 34
diff changeset
    88
        try:
e7f79c4a27ce TableCopyTool: order tables with respect to references.
Radek Brich <radek.brich@devl.cz>
parents: 34
diff changeset
    89
            dc.check()
e7f79c4a27ce TableCopyTool: order tables with respect to references.
Radek Brich <radek.brich@devl.cz>
parents: 34
diff changeset
    90
        except pgdatacopy.TargetNotEmptyError as e:
e7f79c4a27ce TableCopyTool: order tables with respect to references.
Radek Brich <radek.brich@devl.cz>
parents: 34
diff changeset
    91
            print(' - error:', str(e))
e7f79c4a27ce TableCopyTool: order tables with respect to references.
Radek Brich <radek.brich@devl.cz>
parents: 34
diff changeset
    92
            return
56
94e091c23ebb Add reverse parameter for diff tools.
Radek Brich <radek.brich@devl.cz>
parents: 35
diff changeset
    93
35
e7f79c4a27ce TableCopyTool: order tables with respect to references.
Radek Brich <radek.brich@devl.cz>
parents: 34
diff changeset
    94
        print(' - read                           ')
e7f79c4a27ce TableCopyTool: order tables with respect to references.
Radek Brich <radek.brich@devl.cz>
parents: 34
diff changeset
    95
        buf = io.BytesIO()
e7f79c4a27ce TableCopyTool: order tables with respect to references.
Radek Brich <radek.brich@devl.cz>
parents: 34
diff changeset
    96
        wrapped = ProgressWrapper(buf)
e7f79c4a27ce TableCopyTool: order tables with respect to references.
Radek Brich <radek.brich@devl.cz>
parents: 34
diff changeset
    97
        dc.read(wrapped)
e7f79c4a27ce TableCopyTool: order tables with respect to references.
Radek Brich <radek.brich@devl.cz>
parents: 34
diff changeset
    98
        data = buf.getvalue()
e7f79c4a27ce TableCopyTool: order tables with respect to references.
Radek Brich <radek.brich@devl.cz>
parents: 34
diff changeset
    99
        buf.close()
56
94e091c23ebb Add reverse parameter for diff tools.
Radek Brich <radek.brich@devl.cz>
parents: 35
diff changeset
   100
35
e7f79c4a27ce TableCopyTool: order tables with respect to references.
Radek Brich <radek.brich@devl.cz>
parents: 34
diff changeset
   101
        print(' - write                          ')
e7f79c4a27ce TableCopyTool: order tables with respect to references.
Radek Brich <radek.brich@devl.cz>
parents: 34
diff changeset
   102
        buf = io.BytesIO(data)
e7f79c4a27ce TableCopyTool: order tables with respect to references.
Radek Brich <radek.brich@devl.cz>
parents: 34
diff changeset
   103
        wrapped = ProgressWrapper(buf, len(data))
e7f79c4a27ce TableCopyTool: order tables with respect to references.
Radek Brich <radek.brich@devl.cz>
parents: 34
diff changeset
   104
        dc.write(wrapped)
e7f79c4a27ce TableCopyTool: order tables with respect to references.
Radek Brich <radek.brich@devl.cz>
parents: 34
diff changeset
   105
        buf.close()
56
94e091c23ebb Add reverse parameter for diff tools.
Radek Brich <radek.brich@devl.cz>
parents: 35
diff changeset
   106
35
e7f79c4a27ce TableCopyTool: order tables with respect to references.
Radek Brich <radek.brich@devl.cz>
parents: 34
diff changeset
   107
        print(' - analyze                        ')
e7f79c4a27ce TableCopyTool: order tables with respect to references.
Radek Brich <radek.brich@devl.cz>
parents: 34
diff changeset
   108
        dc.analyze()
34
98c7809af415 Add PgDataCopy. Add TableCopyTool.Add SrcDstTablesTool class to toolbase, use in tablecopy, tablediff.
Radek Brich <radek.brich@devl.cz>
parents:
diff changeset
   109
98c7809af415 Add PgDataCopy. Add TableCopyTool.Add SrcDstTablesTool class to toolbase, use in tablecopy, tablediff.
Radek Brich <radek.brich@devl.cz>
parents:
diff changeset
   110
98c7809af415 Add PgDataCopy. Add TableCopyTool.Add SrcDstTablesTool class to toolbase, use in tablecopy, tablediff.
Radek Brich <radek.brich@devl.cz>
parents:
diff changeset
   111
tool = TableCopyTool()
98c7809af415 Add PgDataCopy. Add TableCopyTool.Add SrcDstTablesTool class to toolbase, use in tablecopy, tablediff.
Radek Brich <radek.brich@devl.cz>
parents:
diff changeset
   112
tool.main()
98c7809af415 Add PgDataCopy. Add TableCopyTool.Add SrcDstTablesTool class to toolbase, use in tablecopy, tablediff.
Radek Brich <radek.brich@devl.cz>
parents:
diff changeset
   113