from pgtoolkit.toolbase import SrcDstTool

from pgtoolkit.tools.tablediff import TableDiffTool

from pgtoolkit.tools.runquery import RunQueryTool

class TableSyncTool(SrcDstTool):

    """

    Synchronize tables between two databases (tablediff + runquery).

    This will essentially call following commands on each table from list:

    * pgtool tablediff <source> <target> -r -s <schema> -t <table> --sql -o /tmp/diff.sql

    * pgtool runquery <target> -f /tmp/diff.sql

    """

    def __init__(self):

        SrcDstTool.__init__(self, name='tablesync', force_reverse=True)

        self.tablediff = TableDiffTool()

        self.tablediff.specify_args()

        self.runquery = RunQueryTool()

        self.runquery.specify_args()

    def specify_args(self):

        SrcDstTool.specify_args(self)

        self.parser.add_argument('-t', dest='tables', metavar='table', nargs='*',

            help="Tables to be synchronized.")

        self.parser.add_argument('-s', '--schema', metavar='default_schema',

            dest='schema', type=str, default='public', help='Default schema name.')

    def init_logging(self):

        SrcDstTool.init_logging(self)

        self.runquery.log = self.log

    def setup(self, args=None):

        SrcDstTool.setup(self, args)

        self.target_isolation_level = 'autocommit'

        self.prepare_conns(target=self.args.src)

    def main(self):

        for table in self.args.tables:

            self.sync(table)

    def sync(self, table):

        if '.' in table:

            schema, table = table.split('.', 1)

        else:

            schema = self.args.schema

        # Call tablediff

        self.tablediff.load_args([self.args.src, self.args.dst,

            '-r', '-s', schema, '-t', table, '--sql', '-o', '/tmp/diff.sql'])

        self.tablediff.main()

        # Call runquery

        self.runquery.load_args([self.args.src, '--one-query-per-line',

            '-f', '/tmp/diff.sql'])

        self.runquery.main()

cls = TableSyncTool