1 from pgtoolkit import toolbase, pgmanager, pgdatadiff

     2 from pgtoolkit.toolbase import SrcDstTablesTool

     3 from pycolib.ansicolor import highlight, BOLD, YELLOW

     4 

     5 import sys

     6 

     7 

     8 class TableDiffTool(SrcDstTablesTool):

     9 

    10     """

    11     Print differences between data in tables.

    12 

    13     Requirements:

    14     * Source table must have defined PRIMARY KEY.

    15     * Destination table must contain all columns from source table.

    16       Order is not important.

    17 

    18     """

    19 

    20     def __init__(self):

    21         SrcDstTablesTool.__init__(self, name='tablediff', desc=self.__doc__, allow_reverse=True)

    22 

    23     def specify_args(self):

    24         SrcDstTablesTool.specify_args(self)

    25         self.parser.add_argument('--sql', action='store_true', help='Output is SQL script.')

    26         self.parser.add_argument('--rowcount', action='store_true', help='Compare number of rows.')

    27         self.parser.add_argument('-o', '--output-file', help='Output file for sql queries.')

    28 

    29     def main(self):

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

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

    32 

    33         if self.args.output_file:

    34             output_file = open(self.args.output_file, 'w')

    35         else:

    36             output_file = sys.stdout

    37 

    38         dd = pgdatadiff.PgDataDiff(srcconn, dstconn)

    39 

    40         for srcschema, srctable, dstschema, dsttable in self.tables():

    41             print('-- Diff from [%s] %s.%s to [%s] %s.%s' % (

    42                 self.args.src, srcschema, srctable,

    43                 self.args.dst, dstschema, dsttable),

    44                 file=output_file)

    45 

    46             if self.args.rowcount:

    47                 with self.pgm.cursor('src') as curs:

    48                     curs.execute('''SELECT count(*) FROM "%s"."%s"''' % (srcschema, srctable))

    49                     srccount = curs.fetchone()[0]

    50                 with self.pgm.cursor('dst') as curs:

    51                     curs.execute('''SELECT count(*) FROM "%s"."%s"''' % (dstschema, dsttable))

    52                     dstcount = curs.fetchone()[0]

    53                 if srccount != dstcount:

    54                     print(highlight(1, BOLD | YELLOW),

    55                         "Row count differs: src=%s dst=%s" % (srccount, dstcount),

    56                         highlight(0), sep='', file=output_file)

    57                 continue

    58 

    59             dd.settable1(srctable, srcschema)

    60             dd.settable2(dsttable, dstschema)

    61 

    62             if self.args.sql:

    63                 dd.print_patch(file=output_file)

    64             else:

    65                 dd.print_diff(file=output_file)

    66 

    67 

    68 cls = TableDiffTool

    69