pgtoolkit/tools/tablediff.py
changeset 101 2a2d0d5df03b
parent 100 d6088dba8fea
equal deleted inserted replaced
100:d6088dba8fea 101:2a2d0d5df03b
     1 from pgtoolkit import toolbase, pgmanager, pgdatadiff
     1 from pgtoolkit import toolbase, pgmanager, pgdatadiff
     2 from pgtoolkit.toolbase import SrcDstTablesTool
     2 from pgtoolkit.toolbase import SrcDstTablesTool
     3 from pycolib.ansicolor import highlight, BOLD, YELLOW
     3 from pycolib.ansicolor import highlight, BOLD, YELLOW
     4 
     4 
       
     5 import sys
     5 
     6 
     6 class TableDiffTool(toolbase.SrcDstTablesTool):
     7 
       
     8 class TableDiffTool(SrcDstTablesTool):
     7 
     9 
     8     """
    10     """
     9     Print differencies between data in tables.
    11     Print differences between data in tables.
    10 
    12 
    11     Requirements:
    13     Requirements:
    12     * Source table must have defined PRIMARY KEY.
    14     * Source table must have defined PRIMARY KEY.
    13     * Destination table must contain all columns from source table.
    15     * Destination table must contain all columns from source table.
    14       Order is not important.
    16       Order is not important.
    15 
    17 
    16     """
    18     """
    17 
    19 
    18     def __init__(self):
    20     def __init__(self):
    19         SrcDstTablesTool.__init__(self, name='tablediff', desc=self.__doc__, allow_reverse=True)
    21         SrcDstTablesTool.__init__(self, name='tablediff', desc=self.__doc__, allow_reverse=True)
       
    22 
       
    23     def specify_args(self):
       
    24         SrcDstTablesTool.specify_args(self)
    20         self.parser.add_argument('--sql', action='store_true', help='Output is SQL script.')
    25         self.parser.add_argument('--sql', action='store_true', help='Output is SQL script.')
    21         self.parser.add_argument('--rowcount', action='store_true', help='Compare number of rows.')
    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.')
    22 
    28 
    23     def main(self):
    29     def main(self):
    24         srcconn = self.pgm.get_conn('src')
    30         srcconn = self.pgm.get_conn('src')
    25         dstconn = self.pgm.get_conn('dst')
    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
    26 
    37 
    27         dd = pgdatadiff.PgDataDiff(srcconn, dstconn)
    38         dd = pgdatadiff.PgDataDiff(srcconn, dstconn)
    28 
    39 
    29         for srcschema, srctable, dstschema, dsttable in self.tables():
    40         for srcschema, srctable, dstschema, dsttable in self.tables():
    30             print('-- Diff from [%s] %s.%s to [%s] %s.%s' % (
    41             print('-- Diff from [%s] %s.%s to [%s] %s.%s' % (
    31                 self.args.src, srcschema, srctable,
    42                 self.args.src, srcschema, srctable,
    32                 self.args.dst, dstschema, dsttable))
    43                 self.args.dst, dstschema, dsttable),
       
    44                 file=output_file)
    33 
    45 
    34             if self.args.rowcount:
    46             if self.args.rowcount:
    35                 with self.pgm.cursor('src') as curs:
    47                 with self.pgm.cursor('src') as curs:
    36                     curs.execute('''SELECT count(*) FROM "%s"."%s"''' % (srcschema, srctable))
    48                     curs.execute('''SELECT count(*) FROM "%s"."%s"''' % (srcschema, srctable))
    37                     srccount = curs.fetchone()[0]
    49                     srccount = curs.fetchone()[0]
    38                 with self.pgm.cursor('dst') as curs:
    50                 with self.pgm.cursor('dst') as curs:
    39                     curs.execute('''SELECT count(*) FROM "%s"."%s"''' % (dstschema, dsttable))
    51                     curs.execute('''SELECT count(*) FROM "%s"."%s"''' % (dstschema, dsttable))
    40                     dstcount = curs.fetchone()[0]
    52                     dstcount = curs.fetchone()[0]
    41                 if srccount != dstcount:
    53                 if srccount != dstcount:
    42                     print(highlight(1, BOLD | YELLOW), "Row count differs: src=%s dst=%s" % (srccount, dstcount), highlight(0), sep='')
    54                     print(highlight(1, BOLD | YELLOW),
       
    55                         "Row count differs: src=%s dst=%s" % (srccount, dstcount),
       
    56                         highlight(0), sep='', file=output_file)
    43                 continue
    57                 continue
    44 
    58 
    45             dd.settable1(srctable, srcschema)
    59             dd.settable1(srctable, srcschema)
    46             dd.settable2(dsttable, dstschema)
    60             dd.settable2(dsttable, dstschema)
    47 
    61 
    48             if self.args.sql:
    62             if self.args.sql:
    49                 dd.print_patch()
    63                 dd.print_patch(file=output_file)
    50             else:
    64             else:
    51                 dd.print_diff()
    65                 dd.print_diff(file=output_file)
    52 
    66 
    53 
    67 
    54 cls = TableDiffTool
    68 cls = TableDiffTool
    55 
    69