diff -r bd0beda49bcb -r 98c7809af415 pgtoolkit/toolbase.py --- 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) +