--- 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')