diff -r 98c7809af415 -r e7f79c4a27ce tablecopy.py --- 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()