--- a/pgtoolkit/toolbase.py	Wed Mar 28 17:25:18 2012 +0200
+++ b/pgtoolkit/toolbase.py	Thu May 10 08:42:21 2012 +0200
@@ -1,7 +1,8 @@
 import argparse
 import logging
+import re
 
-from pgtoolkit import pgmanager, config
+from pgtoolkit import pgmanager, pgbrowser, config
 from pgtoolkit.coloredformatter import ColoredFormatter
 from pgtoolkit.highlight import highlight
 
@@ -10,6 +11,10 @@
     pass
 
 
+class BadArgsError(Exception):    
+    pass
+
+
 class ToolBase:
     def __init__(self, name, desc):
         self.parser = argparse.ArgumentParser(description=desc)
@@ -110,3 +115,84 @@
         ToolBase.init(self)
         self.prepare_conns_from_cmdline_args('src', 'dst')
 
+
+class SrcDstTablesTool(SrcDstTool):
+    def __init__(self, name, desc):
+        SrcDstTool.__init__(self, name, desc)
+        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',
+            dest='srcschema', type=str, default='', help='Source schema name (default=public).')
+        self.parser.add_argument('--dst-table', metavar='destination_table',
+            dest='dsttable', type=str, default='', help='Destination table name (default=source_table).')
+        self.parser.add_argument('--dst-schema', metavar='destination_schema',
+            dest='dstschema', type=str, default='', help='Destination schema name (default=source_schema).')
+        self.parser.add_argument('--regex', action='store_true', help="Use RE in schema or table name.")
+    
+    def init(self):
+        SrcDstTool.init(self)
+        
+        self.schema1 = self.args.srcschema
+        self.table1 = self.args.srctable
+        self.schema2 = self.args.dstschema
+        self.table2 = self.args.dsttable
+        
+        # check regex - it applies to source name, dest name must not be specified
+        # applies to only one - schema or table name
+        if self.args.regex:
+            if self.table2 or (self.schema2 and not self.table1):
+                raise BadArgsError('Cannot specify both --regex and --dst-schema, --dst-table.')
+        # schema defaults to public
+        if self.table1 and not self.schema1:
+            self.schema1 = 'public'
+        # dest defaults to source
+        if not self.schema2:
+            self.schema2 = self.schema1
+        if not self.table2:
+            self.table2 = self.table1
+
+    def tables(self):
+        '''Generator. Yields schema1, table1, schema2, table2.'''
+        srcconn = self.pgm.get_conn('src')
+        try:
+            srcbrowser = pgbrowser.PgBrowser(srcconn)
+            if self.args.regex:
+                if not self.table1:
+                    # all tables from schemas defined by regex
+                    for item in self._iter_schemas_regex(srcbrowser, self.schema1):
+                        yield item
+                else:
+                    # all tables defined by regex
+                    for item in self._iter_tables_regex(srcbrowser, self.schema1, self.schema2, self.table1):
+                        yield item
+            else:
+                if not self.table1:
+                    if not self.schema1:
+                        # all tables from all schemas
+                        for item in self._iter_schemas_regex(srcbrowser, self.schema1):
+                            yield item
+                    else:
+                        # all tables from specified schema
+                        for item in self._iter_tables_regex(srcbrowser, self.schema1, self.schema2, self.table1):
+                            yield item
+                else:
+                    # one table
+                    yield (self.schema1, self.table1, self.schema2, self.table2)
+        finally:
+            self.pgm.put_conn(srcconn, 'src')
+    
+    def _iter_schemas_regex(self, browser, regex):
+        for schema in browser.list_schemas():
+            if schema['system']:
+                continue
+            schemaname = schema['name']
+            if re.match(regex, schemaname):
+                for item in self._iter_tables_regex(browser, schemaname, schemaname, ''):
+                    yield item
+    
+    def _iter_tables_regex(self, browser, schema1, schema2, regex):
+        for table in browser.list_tables(schema1):
+            tablename = table['name']
+            if re.match(regex, tablename):
+                yield (schema1, tablename, schema2, tablename)
+