diff -r 24e94a3da209 -r d8ff52a0390f pgtoolkit/toolbase.py --- a/pgtoolkit/toolbase.py Mon May 26 18:18:21 2014 +0200 +++ /dev/null Thu Jan 01 00:00:00 1970 +0000 @@ -1,261 +0,0 @@ -from pgtoolkit import pgmanager, pgbrowser - -from pycolib.configparser import ConfigParser -from pycolib.coloredformatter import ColoredFormatter -from pycolib.ansicolor import highlight - -import argparse -import logging -import re -import textwrap - - -class ConnectionInfoNotFound(Exception): - pass - - -class BadArgsError(Exception): - pass - - -class ToolDescriptionFormatter(argparse.HelpFormatter): - """Help message formatter which retains any formatting in descriptions.""" - - def _fill_text(self, text, width, indent): - return textwrap.dedent(text) - - -class ToolBase: - - def __init__(self, name, desc=None, **kwargs): - self.config = ConfigParser() - self.parser = argparse.ArgumentParser(prog=name, description=desc or self.__doc__, - formatter_class=ToolDescriptionFormatter) - self.pgm = pgmanager.get_instance() - self.target_isolation_level = None - - def setup(self, args=None): - self.specify_args() - self.load_args(args) - self.init_logging() - - def specify_args(self): - self.config.add_option('databases', dict) - self.config.add_option('meta_db') - self.config.add_option('meta_query') - self.parser.add_argument('-Q', dest='show_queries', action='store_true', - help='Print database queries.') - self.parser.add_argument('-C', dest='config_file', type=str, - help='Additional config file (besides pgtoolkit.conf).') - - def load_args(self, args=None, config_file=None): - # Parse command line arguments - self.args = self.parser.parse_args(args) - # Load global config - self.config.load('/etc/pgtoolkit.conf', must_exist=False) - # Load local config - self.config.load(config_file or 'pgtoolkit.conf', must_exist=False) - # Load additional config - if self.args.config_file: - self.config.load(self.args.config_file) - - def init_logging(self): - # logging - format = ColoredFormatter(highlight(1,7,0)+'%(asctime)s %(levelname)-5s'+highlight(0)+' %(message)s', '%H:%M:%S') - handler = logging.StreamHandler() - handler.setFormatter(format) - handler.setLevel(logging.DEBUG) - self.log = logging.getLogger('main') - self.log.addHandler(handler) - self.log.setLevel(logging.DEBUG) - - log_notices = logging.getLogger('pgmanager_notices') - log_notices.addHandler(handler) - log_notices.setLevel(logging.DEBUG) - - if self.args.show_queries: - log_sql = logging.getLogger('pgmanager_sql') - log_sql.addHandler(handler) - log_sql.setLevel(logging.DEBUG) - - def prepare_conn_from_metadb(self, name, lookup_name): - """Create connection in pgmanager using meta DB. - - name -- Name for connection in pgmanager. - lookup_name -- Name of connection in meta DB. - - """ - if not self.pgm.knows_conn('meta'): - self.pgm.create_conn(name='meta', dsn=self.config.meta_db) - with self.pgm.cursor('meta') as curs: - curs.execute(self.config.meta_query, [lookup_name]) - row = curs.fetchone_dict() - curs.connection.commit() - if row: - self.pgm.create_conn(name=name, - isolation_level=self.target_isolation_level, - **row) - return True - self.pgm.close_conn('meta') - - def prepare_conn_from_config(self, name, lookup_name): - """Create connection in pgmanager using info in config.databases.""" - if self.config.databases: - if lookup_name in self.config.databases: - dsn = self.config.databases[lookup_name] - self.pgm.create_conn(name=name, - isolation_level=self.target_isolation_level, - dsn=dsn) - return True - - def prepare_conns(self, **kwargs): - """Create connections in PgManager. - - Keyword arguments meaning: - key: connection name for use in PgManager - value: connection name in config or meta DB - - """ - for name in kwargs: - lookup_name = kwargs[name] - found = self.prepare_conn_from_config(name, lookup_name) - if not found and self.config.meta_db: - found = self.prepare_conn_from_metadb(name, lookup_name) - if not found: - raise ConnectionInfoNotFound('Connection name "%s" not found in config nor in meta DB.' % lookup_name) - - -class SimpleTool(ToolBase): - - def __init__(self, name, desc=None, **kwargs): - ToolBase.__init__(self, name, desc, **kwargs) - - def specify_args(self): - ToolBase.specify_args(self) - self.config.add_option('target', type=str, default=None) - self.parser.add_argument('target', nargs='?', type=str, help='Target database') - - def load_args(self, args=None, config_file=None): - ToolBase.load_args(self, args, config_file) - self.target = self.args.target or self.config.target or 'default' - - def setup(self, args=None): - ToolBase.setup(self, args) - self.prepare_conns(target=self.target) - - -class SrcDstTool(ToolBase): - - def __init__(self, name, desc=None, *, allow_reverse=False, force_reverse=False, **kwargs): - ToolBase.__init__(self, name, desc, **kwargs) - self.allow_reverse = allow_reverse - self.force_reverse = force_reverse - - def specify_args(self): - ToolBase.specify_args(self) - 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 self.allow_reverse: - self.parser.add_argument('-r', '--reverse', action='store_true', help='Reverse operation. Swap source and destination.') - - def load_args(self, args=None, config_file=None): - ToolBase.load_args(self, args, config_file) - if self.is_reversed(): - self.args.src, self.args.dst = self.args.dst, self.args.src - - def setup(self, args=None): - ToolBase.setup(self, args) - self.prepare_conns(src=self.args.src, dst=self.args.dst) - - def is_reversed(self): - return ('reverse' in self.args and self.args.reverse) or self.force_reverse - - -class SrcDstTablesTool(SrcDstTool): - - def specify_args(self): - SrcDstTool.specify_args(self) - 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 load_args(self, args=None, config_file=None): - SrcDstTool.load_args(self, args, config_file) - self.load_table_names() - - def load_table_names(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 - - # 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') - 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) -