pgtoolkit/toolbase.py
changeset 56 94e091c23ebb
parent 51 bdc44f96cb0b
child 62 af637235ca81
--- a/pgtoolkit/toolbase.py	Thu Dec 13 17:15:10 2012 +0100
+++ b/pgtoolkit/toolbase.py	Mon Dec 17 16:48:12 2012 +0100
@@ -16,7 +16,7 @@
 
 
 class ToolBase:
-    def __init__(self, name, desc):
+    def __init__(self, name, desc, **kwargs):
         self.parser = argparse.ArgumentParser(description=desc)
         self.parser.add_argument('-d', dest='debug', action='store_true',
             help='Debug mode - print database queries.')
@@ -97,8 +97,8 @@
 
 
 class SimpleTool(ToolBase):
-    def __init__(self, name, desc):
-        ToolBase.__init__(self, name, desc)
+    def __init__(self, name, desc, **kwargs):
+        ToolBase.__init__(self, name, desc, **kwargs)
         self.parser.add_argument('target', metavar='target', type=str, help='Target database')
 
     def init(self):
@@ -107,19 +107,26 @@
 
 
 class SrcDstTool(ToolBase):
-    def __init__(self, name, desc):
-        ToolBase.__init__(self, name, desc)
+    def __init__(self, name, desc, **kwargs):
+        ToolBase.__init__(self, name, desc, **kwargs)
         self.parser.add_argument('src', metavar='source', type=str, help='Source database')
         self.parser.add_argument('dst', metavar='destination', type=str, help='Destination database')
+        if 'allow_reverse' in kwargs and kwargs['allow_reverse']:
+            self.parser.add_argument('-r', '--reverse', action='store_true', help='Reverse operation. Swap source and destination.')
 
     def init(self):
         ToolBase.init(self)
+        if self.is_reversed():
+            self.args.src, self.args.dst = self.args.dst, self.args.src
         self.prepare_conns_from_cmdline_args('src', 'dst')
 
+    def is_reversed(self):
+        return 'reverse' in self.args and self.args.reverse
+
 
 class SrcDstTablesTool(SrcDstTool):
-    def __init__(self, name, desc):
-        SrcDstTool.__init__(self, name, desc)
+    def __init__(self, name, desc, **kwargs):
+        SrcDstTool.__init__(self, name, desc, **kwargs)
         self.parser.add_argument('-t', '--src-table', metavar='source_table',
             dest='srctable', type=str, default='', help='Source table name.')
         self.parser.add_argument('-s', '--src-schema', metavar='source_schema',
@@ -152,6 +159,11 @@
         if not self.table2:
             self.table2 = self.table1
 
+        # swap src, dst when in reverse mode
+        if self.is_reversed():
+            self.schema1, self.schema2 = self.schema2, self.schema1
+            self.table1, self.table2 = self.table2, self.table1
+
     def tables(self):
         '''Generator. Yields schema1, table1, schema2, table2.'''
         srcconn = self.pgm.get_conn('src')