1 from pgtoolkit import pgmanager, pgbrowser

     2 

     3 from pycolib.configparser import ConfigParser

     4 from pycolib.coloredformatter import ColoredFormatter

     5 from pycolib.ansicolor import highlight

     6 

     7 import argparse

     8 import logging

     9 import re

    10 import textwrap

    11 

    12 

    13 class ConnectionInfoNotFound(Exception):

    14     pass

    15 

    16 

    17 class BadArgsError(Exception):

    18     pass

    19 

    20 

    21 class ToolDescriptionFormatter(argparse.HelpFormatter):

    22     """Help message formatter which retains any formatting in descriptions."""

    23 

    24     def _fill_text(self, text, width, indent):

    25         return textwrap.dedent(text)

    26 

    27 

    28 class ToolBase:

    29 

    30     def __init__(self, name, desc=None, **kwargs):

    31         self.config = ConfigParser()

    32         self.parser = argparse.ArgumentParser(prog=name, description=desc or self.__doc__,

    33             formatter_class=ToolDescriptionFormatter)

    34         self.pgm = pgmanager.get_instance()

    35         self.target_isolation_level = None

    36 

    37     def setup(self, args=None):

    38         self.specify_args()

    39         self.load_args(args)

    40         self.init_logging()

    41 

    42     def specify_args(self):

    43         self.config.add_option('databases', dict)

    44         self.config.add_option('meta_db')

    45         self.config.add_option('meta_query')

    46         self.parser.add_argument('-Q', dest='show_queries', action='store_true',

    47             help='Print database queries.')

    48         self.parser.add_argument('-C', dest='config_file', type=str,

    49             help='Additional config file (besides pgtoolkit.conf).')

    50 

    51     def load_args(self, args=None, config_file=None):

    52         # Parse command line arguments

    53         self.args = self.parser.parse_args(args)

    54         # Load global config

    55         self.config.load('/etc/pgtoolkit.conf', must_exist=False)

    56         # Load local config

    57         self.config.load(config_file or 'pgtoolkit.conf', must_exist=False)

    58         # Load additional config

    59         if self.args.config_file:

    60             self.config.load(self.args.config_file)

    61 

    62     def init_logging(self):

    63         # logging

    64         format = ColoredFormatter(highlight(1,7,0)+'%(asctime)s %(levelname)-5s'+highlight(0)+' %(message)s', '%H:%M:%S')

    65         handler = logging.StreamHandler()

    66         handler.setFormatter(format)

    67         handler.setLevel(logging.DEBUG)

    68         self.log = logging.getLogger('main')

    69         self.log.addHandler(handler)

    70         self.log.setLevel(logging.DEBUG)

    71 

    72         log_notices = logging.getLogger('pgmanager_notices')

    73         log_notices.addHandler(handler)

    74         log_notices.setLevel(logging.DEBUG)

    75 

    76         if self.args.show_queries:

    77             log_sql = logging.getLogger('pgmanager_sql')

    78             log_sql.addHandler(handler)

    79             log_sql.setLevel(logging.DEBUG)

    80 

    81     def prepare_conn_from_metadb(self, name, lookup_name):

    82         """Create connection in pgmanager using meta DB.

    83 

    84         name -- Name for connection in pgmanager.

    85         lookup_name -- Name of connection in meta DB.

    86 

    87         """

    88         if not self.pgm.knows_conn('meta'):

    89             self.pgm.create_conn(name='meta', dsn=self.config.meta_db)

    90         with self.pgm.cursor('meta') as curs:

    91             curs.execute(self.config.meta_query, [lookup_name])

    92             row = curs.fetchone_dict()

    93             curs.connection.commit()

    94             if row:

    95                 self.pgm.create_conn(name=name,

    96                     isolation_level=self.target_isolation_level,

    97                     **row)

    98                 return True

    99         self.pgm.close_conn('meta')

   100 

   101     def prepare_conn_from_config(self, name, lookup_name):

   102         """Create connection in pgmanager using info in config.databases."""

   103         if self.config.databases:

   104             if lookup_name in self.config.databases:

   105                 dsn = self.config.databases[lookup_name]

   106                 self.pgm.create_conn(name=name,

   107                     isolation_level=self.target_isolation_level,

   108                     dsn=dsn)

   109                 return True

   110 

   111     def prepare_conns(self, **kwargs):

   112         """Create connections in PgManager.

   113 

   114         Keyword arguments meaning:

   115             key: connection name for use in PgManager

   116             value: connection name in config or meta DB

   117 

   118         """

   119         for name in kwargs:

   120             lookup_name = kwargs[name]

   121             found = self.prepare_conn_from_config(name, lookup_name)

   122             if not found and self.config.meta_db:

   123                 found = self.prepare_conn_from_metadb(name, lookup_name)

   124             if not found:

   125                 raise ConnectionInfoNotFound('Connection name "%s" not found in config nor in meta DB.' % lookup_name)

   126 

   127 

   128 class SimpleTool(ToolBase):

   129 

   130     def __init__(self, name, desc=None, **kwargs):

   131         ToolBase.__init__(self, name, desc, **kwargs)

   132 

   133     def specify_args(self):

   134         ToolBase.specify_args(self)

   135         self.config.add_option('target', type=str, default=None)

   136         self.parser.add_argument('target', nargs='?', type=str, help='Target database')

   137 

   138     def load_args(self, args=None, config_file=None):

   139         ToolBase.load_args(self, args, config_file)

   140         self.target = self.args.target or self.config.target or 'default'

   141 

   142     def setup(self, args=None):

   143         ToolBase.setup(self, args)

   144         self.prepare_conns(target=self.target)

   145 

   146 

   147 class SrcDstTool(ToolBase):

   148 

   149     def __init__(self, name, desc=None, *, allow_reverse=False, force_reverse=False, **kwargs):

   150         ToolBase.__init__(self, name, desc, **kwargs)

   151         self.allow_reverse = allow_reverse

   152         self.force_reverse = force_reverse

   153 

   154     def specify_args(self):

   155         ToolBase.specify_args(self)

   156         self.parser.add_argument('src', metavar='source', type=str, help='Source database')

   157         self.parser.add_argument('dst', metavar='destination', type=str, help='Destination database')

   158         if self.allow_reverse:

   159             self.parser.add_argument('-r', '--reverse', action='store_true', help='Reverse operation. Swap source and destination.')

   160 

   161     def load_args(self, args=None, config_file=None):

   162         ToolBase.load_args(self, args, config_file)

   163         if self.is_reversed():

   164             self.args.src, self.args.dst = self.args.dst, self.args.src

   165 

   166     def setup(self, args=None):

   167         ToolBase.setup(self, args)

   168         self.prepare_conns(src=self.args.src, dst=self.args.dst)

   169 

   170     def is_reversed(self):

   171         return ('reverse' in self.args and self.args.reverse) or self.force_reverse

   172 

   173 

   174 class SrcDstTablesTool(SrcDstTool):

   175 

   176     def specify_args(self):

   177         SrcDstTool.specify_args(self)

   178         self.parser.add_argument('-t', '--src-table', metavar='source_table',

   179             dest='srctable', type=str, default='', help='Source table name.')

   180         self.parser.add_argument('-s', '--src-schema', metavar='source_schema',

   181             dest='srcschema', type=str, default='', help='Source schema name (default=public).')

   182         self.parser.add_argument('--dst-table', metavar='destination_table',

   183             dest='dsttable', type=str, default='', help='Destination table name (default=source_table).')

   184         self.parser.add_argument('--dst-schema', metavar='destination_schema',

   185             dest='dstschema', type=str, default='', help='Destination schema name (default=source_schema).')

   186         self.parser.add_argument('--regex', action='store_true', help="Use RE in schema or table name.")

   187 

   188     def load_args(self, args=None, config_file=None):

   189         SrcDstTool.load_args(self, args, config_file)

   190         self.load_table_names()

   191 

   192     def load_table_names(self):

   193         self.schema1 = self.args.srcschema

   194         self.table1 = self.args.srctable

   195         self.schema2 = self.args.dstschema

   196         self.table2 = self.args.dsttable

   197 

   198         # check regex - it applies to source name, dest name must not be specified

   199         # applies to only one - schema or table name

   200         if self.args.regex:

   201             if self.table2 or (self.schema2 and not self.table1):

   202                 raise BadArgsError('Cannot specify both --regex and --dst-schema, --dst-table.')

   203         # schema defaults to public

   204         if self.table1 and not self.schema1:

   205             self.schema1 = 'public'

   206         # dest defaults to source

   207         if not self.schema2:

   208             self.schema2 = self.schema1

   209         if not self.table2:

   210             self.table2 = self.table1

   211 

   212         # swap src, dst when in reverse mode

   213         if self.is_reversed():

   214             self.schema1, self.schema2 = self.schema2, self.schema1

   215             self.table1, self.table2 = self.table2, self.table1

   216 

   217     def tables(self):

   218         '''Generator. Yields schema1, table1, schema2, table2.'''

   219         srcconn = self.pgm.get_conn('src')

   220         try:

   221             srcbrowser = pgbrowser.PgBrowser(srcconn)

   222             if self.args.regex:

   223                 if not self.table1:

   224                     # all tables from schemas defined by regex

   225                     for item in self._iter_schemas_regex(srcbrowser, self.schema1):

   226                         yield item

   227                 else:

   228                     # all tables defined by regex

   229                     for item in self._iter_tables_regex(srcbrowser, self.schema1, self.schema2, self.table1):

   230                         yield item

   231             else:

   232                 if not self.table1:

   233                     if not self.schema1:

   234                         # all tables from all schemas

   235                         for item in self._iter_schemas_regex(srcbrowser, self.schema1):

   236                             yield item

   237                     else:

   238                         # all tables from specified schema

   239                         for item in self._iter_tables_regex(srcbrowser, self.schema1, self.schema2, self.table1):

   240                             yield item

   241                 else:

   242                     # one table

   243                     yield (self.schema1, self.table1, self.schema2, self.table2)

   244         finally:

   245             self.pgm.put_conn(srcconn, 'src')

   246 

   247     def _iter_schemas_regex(self, browser, regex):

   248         for schema in browser.list_schemas():

   249             if schema['system']:

   250                 continue

   251             schemaname = schema['name']

   252             if re.match(regex, schemaname):

   253                 for item in self._iter_tables_regex(browser, schemaname, schemaname, ''):

   254                     yield item

   255 

   256     def _iter_tables_regex(self, browser, schema1, schema2, regex):

   257         for table in browser.list_tables(schema1):

   258             tablename = table['name']

   259             if re.match(regex, tablename):

   260                 yield (schema1, tablename, schema2, tablename)

   261