pgtoolkit/toolbase.py
changeset 20 73f0d53fef6b
parent 17 f768a3529ee7
child 29 8636719a30f6
--- a/pgtoolkit/toolbase.py	Sun Nov 27 17:02:42 2011 +0100
+++ b/pgtoolkit/toolbase.py	Tue Dec 06 16:20:53 2011 +0100
@@ -6,6 +6,10 @@
 from pgtoolkit.highlight import highlight
 
 
+class ConnectionInfoNotFound(Exception):    
+    pass
+
+
 class ToolBase:
     def __init__(self, name, desc):
         self.parser = argparse.ArgumentParser(description=desc)
@@ -19,11 +23,12 @@
         
         self.pgm = pgmanager.get_instance()
         
-    def init(self):
+    def init(self, *args):
         self.config.load('pgtoolkit.conf')
         self.args = self.parser.parse_args()
         self.init_logging()
-        
+        self.prepare_conns(*args)
+    
     def init_logging(self):
         # logging
         handler = logging.StreamHandler()
@@ -41,32 +46,43 @@
             logger = logging.getLogger('pgmanager')
             logger.addHandler(handler)
 
-    def buildconn(self, name, lookup_name):
+    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.
+        
+        '''
         with self.pgm.cursor('meta') as curs:
             curs.execute(self.config.meta_query, [lookup_name])
             row = curs.fetchone_dict()
             curs.connection.commit()
-            
-            if not row:
-                raise Exception('Unknown database "%s"' % lookup_name)
-            
-            self.pgm.create_conn(name=name, **row)
+            if row:
+                self.pgm.create_conn(name=name, **row)
+                return True
 
-    def prepareconns(self, *pgm_names):
-        # use meta db
+    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, dsn=dsn)
+                return True
+
+    def prepare_conns(self, *pgm_names):
         if self.config.meta_db:
             self.pgm.create_conn(name='meta', dsn=self.config.meta_db)
-            for name in pgm_names:
-                lookup_name = self.args.__dict__[name]
-                self.buildconn(name, lookup_name)
-            self.pgm.close_conn('meta')
+            
+        for name in pgm_names:
+            lookup_name = self.args.__dict__[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)
         
-        # use config dict 'databases'
-        if self.config.databases:
-            for name in pgm_names:
-                lookup_name = self.args.__dict__[name]
-                dsn = self.config.databases[lookup_name]
-                self.pgm.create_conn(name=name, dsn=dsn)
+        if self.config.meta_db:
+            self.pgm.close_conn('meta')
 
 
 class SimpleTool(ToolBase):
@@ -75,8 +91,7 @@
         self.parser.add_argument('target', metavar='target', type=str, help='Target database')
       
     def init(self):
-        ToolBase.init(self)
-        self.prepareconns('target')
+        ToolBase.init(self, 'target')
 
 
 class SrcDstTool(ToolBase):
@@ -86,6 +101,5 @@
         self.parser.add_argument('dst', metavar='destination', type=str, help='Destination database')
       
     def init(self):
-        ToolBase.init(self)
-        self.prepareconns('src', 'dst')
+        ToolBase.init(self, 'src', 'dst')