diff -r adc1615d8fc5 -r 94e091c23ebb pgtoolkit/toolbase.py --- 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')