pgtoolkit/tools/tablediff.py
changeset 101 2a2d0d5df03b
parent 100 d6088dba8fea
--- a/pgtoolkit/tools/tablediff.py	Tue May 06 18:37:41 2014 +0200
+++ b/pgtoolkit/tools/tablediff.py	Tue May 06 18:37:43 2014 +0200
@@ -2,11 +2,13 @@
 from pgtoolkit.toolbase import SrcDstTablesTool
 from pycolib.ansicolor import highlight, BOLD, YELLOW
 
+import sys
 
-class TableDiffTool(toolbase.SrcDstTablesTool):
+
+class TableDiffTool(SrcDstTablesTool):
 
     """
-    Print differencies between data in tables.
+    Print differences between data in tables.
 
     Requirements:
     * Source table must have defined PRIMARY KEY.
@@ -17,19 +19,29 @@
 
     def __init__(self):
         SrcDstTablesTool.__init__(self, name='tablediff', desc=self.__doc__, allow_reverse=True)
+
+    def specify_args(self):
+        SrcDstTablesTool.specify_args(self)
         self.parser.add_argument('--sql', action='store_true', help='Output is SQL script.')
         self.parser.add_argument('--rowcount', action='store_true', help='Compare number of rows.')
+        self.parser.add_argument('-o', '--output-file', help='Output file for sql queries.')
 
     def main(self):
         srcconn = self.pgm.get_conn('src')
         dstconn = self.pgm.get_conn('dst')
 
+        if self.args.output_file:
+            output_file = open(self.args.output_file, 'w')
+        else:
+            output_file = sys.stdout
+
         dd = pgdatadiff.PgDataDiff(srcconn, dstconn)
 
         for srcschema, srctable, dstschema, dsttable in self.tables():
             print('-- Diff from [%s] %s.%s to [%s] %s.%s' % (
                 self.args.src, srcschema, srctable,
-                self.args.dst, dstschema, dsttable))
+                self.args.dst, dstschema, dsttable),
+                file=output_file)
 
             if self.args.rowcount:
                 with self.pgm.cursor('src') as curs:
@@ -39,16 +51,18 @@
                     curs.execute('''SELECT count(*) FROM "%s"."%s"''' % (dstschema, dsttable))
                     dstcount = curs.fetchone()[0]
                 if srccount != dstcount:
-                    print(highlight(1, BOLD | YELLOW), "Row count differs: src=%s dst=%s" % (srccount, dstcount), highlight(0), sep='')
+                    print(highlight(1, BOLD | YELLOW),
+                        "Row count differs: src=%s dst=%s" % (srccount, dstcount),
+                        highlight(0), sep='', file=output_file)
                 continue
 
             dd.settable1(srctable, srcschema)
             dd.settable2(dsttable, dstschema)
 
             if self.args.sql:
-                dd.print_patch()
+                dd.print_patch(file=output_file)
             else:
-                dd.print_diff()
+                dd.print_diff(file=output_file)
 
 
 cls = TableDiffTool