moved admin_command method into main DB class

Author: Tim Vaillancourt

MongoBackup/DB.py

@@ -1,16 +1,19 @@
 import logging
 
 from pymongo import MongoClient
+from time import sleep
 
 
 class DB:
-    def __init__(self, host='localhost', port=27017, username=None, password=None, authdb="admin", conn_timeout=5000):
-        self.host = host
-        self.port = port
-        self.username = username
-        self.password = password
-        self.authdb = authdb
+    def __init__(self, host='localhost', port=27017, username=None, password=None, authdb="admin", conn_timeout=5000, retries=5):
+        self.host         = host
+        self.port         = port
+        self.username     = username
+        self.password     = password
+        self.authdb       = authdb
         self.conn_timeout = conn_timeout
+        self.retries      = retries
+
         self._conn = None
         self.connect()
         self.auth_if_required()
@@ -41,6 +44,22 @@ def auth_if_required(self):
         else:
             pass
 
+    def admin_command(self, admin_command):
+        tries  = 0
+        status = None
+        while not status and tries < self.retries:
+            try:
+                status = self._conn['admin'].command(admin_command)
+                if not status:
+                    raise e
+            except Exception, e:
+                logging.error("Error running admin command '%s': %s" % (admin_command, e))
+                tries += 1
+                sleep(1)
+        if not status:
+            raise Exception, "Could not get output from command: '%s' after %i retries!" % (admin_command, self.retries), None
+        return status
+
     def connection(self):
         return self._conn
 

MongoBackup/ReplsetHandler.py

@@ -1,52 +1,36 @@
 import logging
 
 from math import ceil
-from time import mktime, sleep
+from time import mktime
 
 from DB import DB
 from ShardingHandler import ShardingHandler
 
 
 class ReplsetHandler:
-    def __init__(self, host, port, user, password, authdb, max_lag_secs, retries=5):
+    def __init__(self, host, port, user, password, authdb, max_lag_secs):
         self.host         = host
         self.port         = port
         self.user         = user
         self.password     = password
         self.authdb       = authdb
         self.max_lag_secs = max_lag_secs
-        self.retries      = retries
 
         try:
-            self.connection = DB(self.host, self.port, self.user, self.password, self.authdb).connection()
+            self.db         = DB(self.host, self.port, self.user, self.password, self.authdb)
+            self.connection = self.db.connection()
         except Exception, e:
             logging.fatal("Could not get DB connection! Error: %s" % e)
             raise e
 
     def close(self):
         return self.connection.close()
 
-    def admin_command(self, admin_command):
-        tries  = 0
-        status = None
-        while not status and tries < self.retries:
-            try:
-                status = self.connection['admin'].command(admin_command)
-                if not status:
-                    raise e
-            except Exception, e:
-                logging.error("Error running command '%s': %s" % (admin_command, e))
-                tries += 1
-                sleep(1)
-        if not status:
-            raise Exception, "Could not get output from command: '%s' after %i retries!" % (admin_command, self.retries), None
-        return status
-
     def get_rs_status(self):
-        return self.admin_command('replSetGetStatus')
+        return self.db.admin_command('replSetGetStatus')
 
     def get_rs_config(self):
-        return self.admin_command('replSetGetConfig')
+        return self.db.admin_command('replSetGetConfig')
 
     def find_desirable_secondary(self):
         rs_status    = self.get_rs_status()