--- 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)
-