pgtoolkit/toolbase.py
changeset 100 d6088dba8fea
parent 84 3b5dd9efba35
child 101 2a2d0d5df03b
equal deleted inserted replaced
99:245646538743 100:d6088dba8fea
       
     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 
     1 import argparse
     7 import argparse
     2 import logging
     8 import logging
     3 import re
     9 import re
     4 
    10 import textwrap
     5 from pycolib.configparser import ConfigParser
       
     6 from pycolib.coloredformatter import ColoredFormatter
       
     7 from pycolib.ansicolor import highlight
       
     8 
       
     9 from pgtoolkit import pgmanager, pgbrowser
       
    10 
    11 
    11 
    12 
    12 class ConnectionInfoNotFound(Exception):
    13 class ConnectionInfoNotFound(Exception):
    13     pass
    14     pass
    14 
    15 
    15 
    16 
    16 class BadArgsError(Exception):
    17 class BadArgsError(Exception):
    17     pass
    18     pass
    18 
    19 
    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 
    20 class ToolBase:
    28 class ToolBase:
    21     def __init__(self, name, desc, **kwargs):
    29     def __init__(self, name, desc, **kwargs):
    22         self.parser = argparse.ArgumentParser(description=desc)
    30         self.parser = argparse.ArgumentParser(prog=name, description=desc,
       
    31             formatter_class=ToolDescriptionFormatter)
    23         self.parser.add_argument('-d', dest='debug', action='store_true',
    32         self.parser.add_argument('-d', dest='debug', action='store_true',
    24             help='Debug mode - print database queries.')
    33             help='Debug mode - print database queries.')
    25 
    34 
    26         self.config = ConfigParser()
    35         self.config = ConfigParser()
    27         self.config.add_option('databases', dict)
    36         self.config.add_option('databases', dict)
    29         self.config.add_option('meta_query')
    38         self.config.add_option('meta_query')
    30 
    39 
    31         self.pgm = pgmanager.get_instance()
    40         self.pgm = pgmanager.get_instance()
    32         self.target_isolation_level = None
    41         self.target_isolation_level = None
    33 
    42 
    34     def init(self):
    43     def init(self, args=None):
    35         self.config.load('pgtoolkit.conf')
    44         self.config.load('pgtoolkit.conf')
    36         self.args = self.parser.parse_args()
    45         self.args = self.parser.parse_args(args)
    37         self.init_logging()
    46         self.init_logging()
    38 
    47 
    39     def init_logging(self):
    48     def init_logging(self):
    40         # logging
    49         # logging
    41         format = ColoredFormatter(highlight(1,7,0)+'%(asctime)s %(levelname)-5s'+highlight(0)+' %(message)s', '%H:%M:%S')
    50         format = ColoredFormatter(highlight(1,7,0)+'%(asctime)s %(levelname)-5s'+highlight(0)+' %(message)s', '%H:%M:%S')
   108 class SimpleTool(ToolBase):
   117 class SimpleTool(ToolBase):
   109     def __init__(self, name, desc, **kwargs):
   118     def __init__(self, name, desc, **kwargs):
   110         ToolBase.__init__(self, name, desc, **kwargs)
   119         ToolBase.__init__(self, name, desc, **kwargs)
   111         self.parser.add_argument('target', metavar='target', type=str, help='Target database')
   120         self.parser.add_argument('target', metavar='target', type=str, help='Target database')
   112 
   121 
   113     def init(self):
   122     def init(self, args=None):
   114         ToolBase.init(self)
   123         ToolBase.init(self, args)
   115         self.prepare_conns(target=self.args.target)
   124         self.prepare_conns(target=self.args.target)
   116 
   125 
   117 
   126 
   118 class SrcDstTool(ToolBase):
   127 class SrcDstTool(ToolBase):
   119     def __init__(self, name, desc, **kwargs):
   128     def __init__(self, name, desc, **kwargs):
   121         self.parser.add_argument('src', metavar='source', type=str, help='Source database')
   130         self.parser.add_argument('src', metavar='source', type=str, help='Source database')
   122         self.parser.add_argument('dst', metavar='destination', type=str, help='Destination database')
   131         self.parser.add_argument('dst', metavar='destination', type=str, help='Destination database')
   123         if 'allow_reverse' in kwargs and kwargs['allow_reverse']:
   132         if 'allow_reverse' in kwargs and kwargs['allow_reverse']:
   124             self.parser.add_argument('-r', '--reverse', action='store_true', help='Reverse operation. Swap source and destination.')
   133             self.parser.add_argument('-r', '--reverse', action='store_true', help='Reverse operation. Swap source and destination.')
   125 
   134 
   126     def init(self):
   135     def init(self, args=None):
   127         ToolBase.init(self)
   136         ToolBase.init(self, args)
   128         if self.is_reversed():
   137         if self.is_reversed():
   129             self.args.src, self.args.dst = self.args.dst, self.args.src
   138             self.args.src, self.args.dst = self.args.dst, self.args.src
   130         self.prepare_conns(src=self.args.src, dst=self.args.dst)
   139         self.prepare_conns(src=self.args.src, dst=self.args.dst)
   131 
   140 
   132     def is_reversed(self):
   141     def is_reversed(self):
   144             dest='dsttable', type=str, default='', help='Destination table name (default=source_table).')
   153             dest='dsttable', type=str, default='', help='Destination table name (default=source_table).')
   145         self.parser.add_argument('--dst-schema', metavar='destination_schema',
   154         self.parser.add_argument('--dst-schema', metavar='destination_schema',
   146             dest='dstschema', type=str, default='', help='Destination schema name (default=source_schema).')
   155             dest='dstschema', type=str, default='', help='Destination schema name (default=source_schema).')
   147         self.parser.add_argument('--regex', action='store_true', help="Use RE in schema or table name.")
   156         self.parser.add_argument('--regex', action='store_true', help="Use RE in schema or table name.")
   148 
   157 
   149     def init(self):
   158     def init(self, args=None):
   150         SrcDstTool.init(self)
   159         SrcDstTool.init(self, args)
   151 
   160 
   152         self.schema1 = self.args.srcschema
   161         self.schema1 = self.args.srcschema
   153         self.table1 = self.args.srctable
   162         self.table1 = self.args.srctable
   154         self.schema2 = self.args.dstschema
   163         self.schema2 = self.args.dstschema
   155         self.table2 = self.args.dsttable
   164         self.table2 = self.args.dsttable