DBScript.add_connect_string_profile: single place for profile logic

2 files changed, 14 insertions, 15 deletions

diff --git a/python/pgq/cascade/admin.py b/python/pgq/cascade/admin.py
index 9eefcdfc..c5685667 100644
--- a/python/pgq/cascade/admin.py
+++ b/python/pgq/cascade/admin.py
@@ -433,14 +433,10 @@ class CascadeAdmin(skytools.AdminScript):
         """Show set status."""
         self.load_local_info()
 
-        cstr_extra = self.cf.get('remote_extra_connstr', '')
-
-        # prepare structs for workers
+        # prepare data for workers
         members = Queue.Queue()
         for m in self.queue_info.member_map.itervalues():
-            cstr = m.location
-            if cstr_extra:
-                cstr += ' ' + cstr_extra
+            cstr = self.add_connect_string_profile(m.location, 'remote')
             members.put( (m.name, cstr) )
         nodes = Queue.Queue()
 
diff --git a/python/skytools/scripting.py b/python/skytools/scripting.py
index df0a2ea2..a5e82663 100644
--- a/python/skytools/scripting.py
+++ b/python/skytools/scripting.py
@@ -696,6 +696,15 @@ class DBScript(BaseScript):
     def set_database_defaults(self, dbname, **kwargs):
         self._db_defaults[dbname] = kwargs
 
+    def add_connect_string_profile(self, connstr, profile):
+        """Add extra profile info to connect string.
+        """
+        if profile:
+            extra = self.cf.get("%s_extra_connstr" % profile, '')
+            if extra:
+                connstr += ' ' + extra
+        return connstr
+
     def get_database(self, dbname, autocommit = 0, isolation_level = -1,
                      cache = None, connstr = None, profile = None):
         """Load cached database connection.
@@ -725,22 +734,16 @@ class DBScript(BaseScript):
             params['max_age'] = max_age
 
         if cache in self.db_cache:
+            dbc = self.db_cache[cache]
             if connstr is None:
                 connstr = self.cf.get(dbname, '')
-            if profile:
-                extra = self.cf.get("%s_extra_connstr" % profile, '')
-                if extra:
-                    connstr += ' ' + extra
-            dbc = self.db_cache[cache]
             if connstr:
+                connstr = self.add_connect_string_profile(connstr, profile)
                 dbc.check_connstr(connstr)
         else:
             if not connstr:
                 connstr = self.cf.get(dbname)
-            if profile:
-                extra = self.cf.get("%s_extra_connstr" % profile, '')
-                if extra:
-                    connstr += ' ' + extra
+            connstr = self.add_connect_string_profile(connstr, profile)
 
             # connstr might contain password, it is not a good idea to log it
             filtered_connstr = connstr