Move the SQLite backend into 'connector' module

src/vpoller/connector.py

@@ -31,6 +31,7 @@
 """
 
 import logging
+import sqlite3
 
 import pyVmomi
 import pyVim.connect
@@ -227,3 +228,107 @@ def _get_object_view(self, obj_type, container=None):
 
         return view_ref
 
+class VConnectorDatabase(object):
+    """
+    VConnectorDatabase class
+
+    Provides an SQLite database backend for storing information
+    about vSphere Agents, such as hostname, username, password, etc.
+    
+    Returns:
+        VConnectorDatabase object
+    
+    Raises:
+        VConnectorException
+
+    """
+    def __init__(self, db):
+        """
+        Initializes a new VConnectorDatabase object
+
+        Args:
+            db (str): Path to the SQLite database file
+
+        """
+        self.db = db
+        self.conn = sqlite3.connect(self.db)
+        self.cursor = self.conn.cursor()
+
+    def db_init(self):
+        """
+        Initializes the vConnector Database backend
+
+        """
+        logging.info('Initializing vConnector database at %s', self.db)
+
+        sql = """
+        CREATE TABLE hosts (
+            host TEXT UNIQUE,
+            user TEXT,
+            pwd  TEXT,
+            enabled INTEGER
+        )
+        """
+
+        try:
+            self.cursor.execute(sql)
+        except sqlite3.OperationalError as e:
+            raise VConnectorException, 'Cannot initialize database: %s' % e
+
+        self.conn.commit()
+        self.cursor.close()
+
+    def add_update_agent(self, host, user, pwd, enabled):
+        """
+        Add/update a vSphere Agent in the vConnector database
+
+        Args:
+            host    (str): Hostname of the vSphere host
+            user    (str): Username to use when connecting
+            pwd     (str): Password to use when connecting
+            enabled (int): Enable or disable the vSphere Agent
+
+        """
+        logging.info('Adding/updating vSphere Agent %s in database', host)
+
+        self.cursor.execute('INSERT OR REPLACE INTO hosts VALUES (?,?,?,?)', (host, user, pwd, enabled))
+        self.conn.commit()
+        self.cursor.close()
+
+    def remove_agent(self, host):
+        """
+        Remove a vSphere Agent from the vConnector database
+
+        Args:
+            host (str): Hostname of the vSphere Agent to remove
+        
+        """
+        logging.info('Removing vSphere Agent %s from database', host)
+
+        self.cursor.execute('DELETE FROM hosts WHERE host = ?', (host,))
+        self.conn.commit()
+        self.cursor.close()
+
+    def get_agents(self, only_enabled=False):
+        """
+        Get the vSphere Agents from the vConnector database
+
+        Args:
+            only_enabled (bool): If True return only the Agents which are enabled
+
+        """
+        logging.debug('Getting vSphere Agents from database')
+
+        self.conn.row_factory = sqlite3.Row
+
+        if only_enabled:
+            sql = 'SELECT * FROM hosts WHERE enabled = 1'
+        else:
+            sql = 'SELECT * FROM hosts'
+
+        self.cursor.execute(sql)
+        result = cursor.fetchall()
+        self.cursor.close()
+
+        return result
+

src/vpoller/worker.py

@@ -410,95 +410,3 @@ def worker_shutdown(self, msg):
         self.time_to_die = True
 
         return { 'success': 0, 'msg': 'vPoller Worker is shutting down' }
-
-    def worker_db_init(self):
-        """
-        Initializes the vPoller Worker database
-
-        The vPoller Worker database is used for storing
-        configuration about our vSphere Agents, such as
-        hostnames, usernames, passwords, etc.
-
-        """
-        logging.info('Initializing vPoller Worker database at %s', self.worker_db)
-
-        if os.path.exists(self.worker_db):
-            raise VPollerException, 'vPoller Worker database already exists at %s' % self.worker_db
-
-        conn = sqlite3.connect(self.worker_db)
-        cursor = conn.cursor()
-
-        sql = """
-        CREATE TABLE hosts (
-            host TEXT UNIQUE,
-            user TEXT,
-            pwd  TEXT,
-            enabled INTEGER
-        )
-        """
-
-        cursor.execute(sql)
-        cursor.close()
-
-    def worker_db_add_update_agent(self, host, user, pwd, enabled):
-        """
-        Add/update a vSphere Agent in the vPoller Worker database
-
-        Args:
-            host    (str): Hostname of the vSphere host
-            user    (str): Username to use when connecting
-            pwd     (str): Password to use when connecting
-            enabled (int): Enable or disable the vSphere Agent
-
-        """
-        logging.info('Adding/updating vSphere Agent %s in database', host)
-
-        conn = sqlite3.connect(self.worker_db)
-        cursor = conn.cursor()
-
-        cursor.execute('INSERT OR REPLACE INTO hosts VALUES (?,?,?,?)', (host, user, pwd, enabled))
-        conn.commit()
-        cursor.close()
-
-    def worker_db_remove_agent(self, host):
-        """
-        Remove a vSphere Agent from the vPoller Worker database
-
-        Args:
-            host (str): Hostname of the vSphere Agent to remove
-        
-        """
-        logging.info('Removing vSphere Agent %s from database', host)
-
-        conn = sqlite3.connect(self.worker_db)
-        cursor = conn.cursor()
-
-        cursor.execute('DELETE FROM hosts WHERE host = ?', (host,))
-        conn.commit()
-        cursor.close()
-
-    def worker_db_get_agents(self, only_enabled=False):
-        """
-        Get the vSphere Agents from the vPoller Worker database
-
-        Args:
-            only_enabled (bool): If True return only the Agents which are enabled
-
-        """
-        logging.debug('Getting vSphere Agents from database')
-
-        conn = sqlite3.connect(self.worker_db)
-        conn.row_factory = sqlite3.Row
-        cursor = conn.cursor()
-
-        if only_enabled:
-            sql = 'SELECT * FROM hosts WHERE enabled = 1'
-        else:
-            sql = 'SELECT * FROM hosts'
-
-        cursor.execute(sql)
-        result = cursor.fetchall()
-        cursor.close()
-
-        return result
-