#!/usr/bin/env python3.2

#

# Copy data from one table to another table with same schema.

#

import io

from pgtoolkit import toolbase, pgmanager, pgdatacopy

from pgtoolkit.progresswrapper import ProgressWrapper

class TableCopyTool(toolbase.SrcDstTablesTool):

    def __init__(self):

        toolbase.SrcDstTablesTool.__init__(self, name='tablecopy', desc='Table copy tool.')

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

    def main(self):

        srcconn = self.pgm.get_conn('src')

        dstconn = self.pgm.get_conn('dst')

        dc = pgdatacopy.PgDataCopy(srcconn, 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.noaction:

                continue

            dc.set_source(srctable, srcschema)

            dc.set_destination(dsttable, dstschema)

            dc.check()

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

tool = TableCopyTool()

tool.main()