1 import argparse

     2 

     3 from pgtookit import pgmanager, config

     4 

     5 

     6 class ToolBase:

     7     def __init__(self, name, desc):

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

     9         

    10         self.config = config.ConfigParser()

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

    12         self.config.add_argument('meta_db')

    13         self.config.add_argument('meta_query')

    14         

    15         self.pgm = pgmanager.get_instance()

    16         

    17     def init(self):

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

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

    20 

    21     def buildconn(self, name, targetname):

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

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

    24             row = curs.fetchone_dict()

    25             curs.connection.commit()

    26             

    27             if not row:

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

    29             

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

    31 

    32     def prepareconns(self, *args):        

    33         if self.config.meta_db:

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

    35             for name in args:

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

    37             self.pgm.close_conn('meta')

    38 

    39 

    40 class SimpleTool(ToolBase):

    41     def __init__(self, name, desc):

    42         ToolBase.__init__(self, name, desc)        

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

    44       

    45     def init(self):

    46         ToolBase.init(self)

    47         self.prepareconns('target')

    48 

    49 

    50 class SrcDstTool(ToolBase):

    51     def __init__(self, name, desc):

    52         ToolBase.__init__(self, name, desc)        

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

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

    55       

    56     def init(self):

    57         ToolBase.init(self)

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

    59