--- 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()