--- /dev/null Thu Jan 01 00:00:00 1970 +0000
+++ b/pydbkit/toolbase.py Wed Jul 09 18:03:54 2014 +0200
@@ -0,0 +1,261 @@
+from pydbkit 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 pydbkit.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/pydbkit.conf', must_exist=False)
+ # Load local config
+ self.config.load(config_file or 'pydbkit.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)
+