1 import argparse

     2 

     3 from tools import pgmanager

     4 from common import config

     5 

     6 

     7 class ToolBase:

     8     def __init__(self, name, desc):

     9         self.parser = argparse.ArgumentParser(description=desc)

    10         

    11         self.config = config.ConfigParser()

    12         self.config.add_argument('databases', type=dict)

    13         self.config.add_argument('meta_db')

    14         self.config.add_argument('meta_query')

    15         

    16         self.pgm = pgmanager.get_instance()

    17         

    18     def init(self):

    19         self.config.load('pgtoolkit.conf')

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

    21 

    22     def buildconn(self, name, targetname):

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

    24             curs.execute(self.config.meta_query, [targetname])

    25             row = curs.fetchone_dict()

    26             curs.connection.commit()

    27             

    28             if not row:

    29                 raise Exception('Unknown database "%s"' % targetname)

    30             

    31             self.pgm.create_conn(name=name, **row)

    32 

    33     def prepareconns(self, *args):        

    34         if self.config.meta_db:

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

    36             for name in args:

    37                 self.buildconn(name, self.args.__dict__[name])

    38             self.pgm.close_conn('meta')

    39 

    40 

    41 class SimpleTool(ToolBase):

    42     def __init__(self, name, desc):

    43         ToolBase.__init__(self, name, desc)        

    44         self.parser.add_argument('target', metavar='target', type=str, help='Target database')

    45       

    46     def init(self):

    47         ToolBase.init(self)

    48         self.prepareconns('target')

    49 

    50 

    51 class SrcDstTool(ToolBase):

    52     def __init__(self, name, desc):

    53         ToolBase.__init__(self, name, desc)        

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

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

    56       

    57     def init(self):

    58         ToolBase.init(self)

    59         self.prepareconns('src', 'dst')

    60