1 from pgtoolkit.toolbase import SrcDstTool

     2 from pgtoolkit.tools.tablediff import TableDiffTool

     3 from pgtoolkit.tools.runquery import RunQueryTool

     4 

     5 

     6 class TableSyncTool(SrcDstTool):

     7 

     8     """

     9     Synchronize tables between two databases (tablediff + runquery).

    10 

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

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

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

    14 

    15     """

    16 

    17     def __init__(self):

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

    19         self.tablediff = TableDiffTool()

    20         self.tablediff.specify_args()

    21         self.runquery = RunQueryTool()

    22         self.runquery.specify_args()

    23 

    24     def specify_args(self):

    25         SrcDstTool.specify_args(self)

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

    27             help="Tables to be synchronized.")

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

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

    30 

    31     def init_logging(self):

    32         SrcDstTool.init_logging(self)

    33         self.runquery.log = self.log

    34 

    35     def setup(self, args=None):

    36         SrcDstTool.setup(self, args)

    37         self.target_isolation_level = 'autocommit'

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

    39 

    40     def main(self):

    41         for table in self.args.tables:

    42             self.sync(table)

    43 

    44     def sync(self, table):

    45         if '.' in table:

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

    47         else:

    48             schema = self.args.schema

    49         # Call tablediff

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

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

    52         self.tablediff.main()

    53         # Call runquery

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

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

    56         self.runquery.main()

    57 

    58 

    59 cls = TableSyncTool

    60