#!/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()