TableCopyTool: order tables with respect to references.
authorRadek Brich <radek.brich@devl.cz>
Fri, 11 May 2012 14:16:36 +0200
changeset 35 e7f79c4a27ce
parent 34 98c7809af415
child 36 e67101c22e83
TableCopyTool: order tables with respect to references.
pgtoolkit/pgbrowser.py
pgtoolkit/pgdatacopy.py
tablecopy.py
tablediff.py
--- a/pgtoolkit/pgbrowser.py	Thu May 10 08:42:21 2012 +0200
+++ b/pgtoolkit/pgbrowser.py	Fri May 11 14:16:36 2012 +0200
@@ -228,12 +228,16 @@
     def list_constraints(self, table, schema='public'):
         return self._query('''
             SELECT
-                conname as "name",
-                r.contype as "type",
+                r.conname AS "name",
+                r.contype AS "type",
+                cf.relname AS "fname",
+                nf.nspname AS "fschema",
                 pg_catalog.pg_get_constraintdef(r.oid, true) as "definition"
             FROM pg_catalog.pg_constraint r
             JOIN pg_catalog.pg_class c ON r.conrelid = c.oid
             JOIN pg_catalog.pg_namespace n ON n.oid = c.relnamespace
+            LEFT JOIN pg_catalog.pg_class cf ON r.confrelid = cf.oid
+            LEFT JOIN pg_catalog.pg_namespace nf ON nf.oid = cf.relnamespace
             WHERE n.nspname = %s AND c.relname = %s
             ORDER BY 1
             ''', [schema, table])
--- a/pgtoolkit/pgdatacopy.py	Thu May 10 08:42:21 2012 +0200
+++ b/pgtoolkit/pgdatacopy.py	Fri May 11 14:16:36 2012 +0200
@@ -28,7 +28,9 @@
 
 class TargetNotEmptyError(Exception):
     
-    pass
+    def __init__(self, msg, table):
+        Exception.__init__(self, msg)
+        self.table = table
 
 
 class PgDataCopy:
@@ -74,7 +76,7 @@
         curs.execute(q)
         curs.connection.commit()
         if curs.rowcount > 0:
-            raise TargetNotEmptyError('Target table contains data: %s' % self.fulltable2)
+            raise TargetNotEmptyError('Target table contains data.', self.fulltable2)
         self.cols = [desc[0] for desc in curs.description]
     
     def read(self, tmpfile):
--- a/tablecopy.py	Thu May 10 08:42:21 2012 +0200
+++ b/tablecopy.py	Fri May 11 14:16:36 2012 +0200
@@ -1,11 +1,16 @@
 #!/usr/bin/env python3.2
 #
-# Copy data from one table to another table with same schema.
+# Copy data between tables with same table schema.
+#
+# Can copy multiple tables in one run.
+# Sorts the tables according to references.
+#
+# This may be used instead of dump/restore.
 #
 
 import io
 
-from pgtoolkit import toolbase, pgmanager, pgdatacopy
+from pgtoolkit import toolbase, pgmanager, pgbrowser, pgdatacopy
 from pgtoolkit.progresswrapper import ProgressWrapper
 
 
@@ -15,40 +20,83 @@
         
         self.parser.add_argument('-n', '--no-action', dest='noaction', action='store_true',
             help="Do nothing, just print tables to be copied. Useful in combination with --regex.")
+        self.parser.add_argument('--no-sort', dest='nosort', action='store_true',
+            help="Do not sort. By default, tables are sorted by foreign key references.")
         
         self.init()
 
     def main(self):
-        srcconn = self.pgm.get_conn('src')
-        dstconn = self.pgm.get_conn('dst')
+        self.srcconn = self.pgm.get_conn('src')
+        self.dstconn = self.pgm.get_conn('dst')
         
-        dc = pgdatacopy.PgDataCopy(srcconn, dstconn)
+        dc = pgdatacopy.PgDataCopy(self.srcconn, self.dstconn)
         
-        for srcschema, srctable, dstschema, dsttable in self.tables():
-            print('Copying [%s] %s.%s --> [%s] %s.%s' % (
-                self.args.src, srcschema, srctable,
-                self.args.dst, dstschema, dsttable))
+        if self.args.nosort:
+            for table in self.tables():
+                self.copy_table(dc, *table)
+        else:
+            # sort tables with respect to references
+            details = dict()
+            pending = set()
+            references = dict()
             
-            if self.args.noaction:
-                continue
+            # build list of all table to be copied (pending) and references map
+            for table in self.tables():
+                srcschema, srctable, dstschema, dsttable = table
+                name = dstschema + '.' + dsttable
+                details[name] = table
+                pending.add(name)
+                references[name] = self.get_references(dstschema, dsttable)
             
-            dc.set_source(srctable, srcschema)
-            dc.set_destination(dsttable, dstschema)
+            # copy files with fulfilled references, repeat until all done
+            while pending:
+                for name in list(pending):
+                    ok = True
+                    for ref in references[name]:
+                        if ref in pending:
+                            ok = False
+                    if ok:
+                        table = details[name]
+                        self.copy_table(dc, *table)
+                        pending.remove(name)
+
+    def get_references(self, schema, table):
+        browser = pgbrowser.PgBrowser(self.dstconn)
+        cons = browser.list_constraints(table, schema)
+        return [con['fschema'] + '.' + con['fname'] for con in cons if con['type'] == 'f']
+
+    def copy_table(self, dc, srcschema, srctable, dstschema, dsttable):
+        print('Copying [%s] %s.%s --> [%s] %s.%s' % (
+            self.args.src, srcschema, srctable,
+            self.args.dst, dstschema, dsttable))
         
-            dc.check()
+        if self.args.noaction:
+            return
         
-            buf = io.BytesIO()
-            wrapped = ProgressWrapper(buf)
-            dc.read(wrapped)
-            data = buf.getvalue()
-            buf.close()
-            
-            buf = io.BytesIO(data)
-            wrapped = ProgressWrapper(buf, len(data))
-            dc.write(wrapped)
-            buf.close()
+        dc.set_source(srctable, srcschema)
+        dc.set_destination(dsttable, dstschema)
+        
+        try:
+            dc.check()
+        except pgdatacopy.TargetNotEmptyError as e:
+            print(' - error:', str(e))
+            return
         
-            dc.analyze()
+        print(' - read                           ')
+        buf = io.BytesIO()
+        wrapped = ProgressWrapper(buf)
+        dc.read(wrapped)
+        data = buf.getvalue()
+        buf.close()
+        
+        print(' - write                          ')
+        buf = io.BytesIO(data)
+        wrapped = ProgressWrapper(buf, len(data))
+        dc.write(wrapped)
+        buf.close()
+        
+        print(' - analyze                        ')
+        dc.analyze()
 
 
 tool = TableCopyTool()
--- a/tablediff.py	Thu May 10 08:42:21 2012 +0200
+++ b/tablediff.py	Fri May 11 14:16:36 2012 +0200
@@ -28,7 +28,7 @@
         dd = pgdatadiff.PgDataDiff(srcconn, dstconn)
         
         for srcschema, srctable, dstschema, dsttable in self.tables():
-            print('Diff from [%s] %s.%s to [%s] %s.%s' % (
+            print('-- Diff from [%s] %s.%s to [%s] %s.%s' % (
                 self.args.src, srcschema, srctable,
                 self.args.dst, dstschema, dsttable))