Replace borg connection manager with thread safe db pool

src/translators/base_translator.py

@@ -17,7 +17,7 @@ def __init__(self, host, port, db_name):
         self.port = port
         self.db_name = db_name
 
-    def __enter__(self):
+    def __enter__(self, connection, query):
         self.setup()
         return self
 

src/translators/crate.py

@@ -43,15 +43,15 @@
 class CrateTranslator(sql_translator.SQLTranslator):
     NGSI_TO_SQL = NGSI_TO_SQL
 
-    def __init__(self, host, port=4200, db_name="ngsi-tsdb"):
-        super(CrateTranslator, self).__init__(host, port, db_name)
+    def __init__(self, connection, query , host, port=4200, db_name="ngsi-tsdb"):
+        super(CrateTranslator, self).__init__(host, connection, query , port, db_name)
         self.logger = logging.getLogger(__name__)
         self.dbCacheName = 'crate'
         self.ccm = None
         self.connection = None
         self.cursor = None
 
-    def setup(self):
+    def setup(self, connection, query):
         url = "{}:{}".format(self.host, self.port)
         self.ccm = ConnectionManager()
         self.connection = self.ccm.get_connection('crate')

src/translators/sql_translator.py

@@ -15,7 +15,8 @@
 
 from cache.factory import get_cache, is_cache_available
 from translators.insert_splitter import to_insert_batches
-from utils.connection_manager import Borg
+from utils.connection_manager import ConnectionManager
+
 # NGSI TYPES
 # Based on Orion output because official docs don't say much about these :(
 NGSI_DATETIME = 'DateTime'
@@ -117,9 +118,9 @@ class SQLTranslator(base_translator.BaseTranslator):
 
     start_time = None
 
-    def __init__(self, host, port, db_name):
+    def __init__(self, host, port, db_name, connection, query):
         super(SQLTranslator, self).__init__(host, port, db_name)
-        qcm = QueryCacheManager()
+        qcm = QueryCacheManager(connection, query)
         self.cache = qcm.get_query_cache()
         self.default_ttl = None
         if self.cache:
@@ -1781,11 +1782,11 @@ def _remove_from_cache(self, tenant_name, key):
                                     exc_info=True)
 
 
-class QueryCacheManager(Borg):
+class QueryCacheManager(ConnectionManager):
     cache = None
 
-    def __init__(self):
-        super(QueryCacheManager, self).__init__()
+    def __init__(self, connection, query):
+        super(QueryCacheManager, self ).__init__( connection, query)
         if is_cache_available() and self.cache is None:
             try:
                 self.cache = get_cache()

src/utils/connection_manager.py

@@ -1,25 +1,67 @@
-class Borg:
-    _shared_state = {}
+from crate import client
+from crate.client.sqlalchemy.dialect import CrateDialect
+import sqlalchemy.pool as pool
+import argparse
+from sqlalchemy import create_engine
 
-    def __init__(self):
-        self.__dict__ = self._shared_state
+class ConnectionManager():
+    """
+    Invoke queries to database in parallel.
+    """
 
+    def __init__(self, connection, query):
+        self.connection = connection
+        self.query = query
 
-class ConnectionManager(Borg):
+    def querydb_dbapi(connection, query):
+        """
+        Submit query to database.
+        """
+        cursor = connection.cursor()
+        cursor.execute(query)
+        result = cursor.fetchone()
+        cursor.close()
+        return result
 
-    connection = {}
-
-    def __init__(self):
-        Borg.__init__(self)
+    def querydb_sqlalchemy(connection, query):
+        """
+        Submit query to database.
+        """
+        result = connection.execute(query).fetchone()
+        return result
 
     def set_connection(self, db, connection):
         self.connection[db] = connection
 
-    def get_connection(self, db):
-        try:
-            return self.connection[db]
-        except KeyError as e:
-            return None
+    def get_connection(connection):
+        parser = argparse.ArgumentParser()
+        parser.add_argument('--driver')
+        parser.add_argument('--pool', action='store_true')
+        args = parser.parse_args()
+
+        if args.driver == "dbapi":
+        # Use DBAPI driver.
+            if args.pool:
+                query = querydb_dbapi
+                # Use a connection pool matching the number of workers.
+                engine  = create_engine("crate://localhost:4200", connect_args={"pool_size": 10})
+            else:
+                # Don't use a pool.
+                engine = create_engine("crate://localhost:4200")
+        elif args.driver == "sqlalchemy":
+            if args.pool:
+                query = querydb_sqlalchemy
+                # Use a connection pool matching the number of workers.
+                engine = create_engine("crate://localhost:4200", connect_args={"pool_size": 10})
+            else:
+                # Don't use a pool.
+                engine = create_engine("crate://localhost:4200")
+            connection = engine.connect()
+        else:
+            raise ValueError("Unknown value for --driver: Use 'dbapi' or 'sqlalchemy'.")
+
+    # Invoke some database queries.
+    query = 'SELECT 1;'
 
-    def reset_connection(self, db):
-        self.connection[db] = None
+if __name__ == '__main__':
+    main()