Merge pull request #308 from omegaml/improve-sql-cnx-caching

Improve sql cnx caching

Author: omegaml

omegaml/backends/sqlalchemy.py

@@ -1,15 +1,15 @@
-import warnings
-
+import logging
+import os
 import string
+import warnings
 from getpass import getuser
+from hashlib import sha256
 from logging import warning
 
-import logging
-import os
 import pandas as pd
 import sqlalchemy
-
 from omegaml.backends.basedata import BaseDataBackend
+from omegaml.util import ProcessLocal, KeepMissing
 
 try:
     import snowflake
@@ -19,8 +19,15 @@
 except:
     pass
 
-CNX_CACHE = {}
-ALWAYS_CACHE = False
+#: override by setting om.defaults.SQLALCHEMY_ALWAYS_CACHE
+ALWAYS_CACHE = True
+# -- enabled by default as this is the least-surprised option
+# -- consistent with sqlalchemy connection pooling defaults
+#: kwargs for create_engine()
+ENGINE_KWARGS = dict(echo=False, pool_pre_ping=True, pool_recycle=3600)
+# -- echo=False - do not log to stdout
+# -- pool_pre_ping=True - always check, re-establish connection if no longer working
+# -- pool_recylce=N - do not reuse connections older than N seconds
 
 logger = logging.getLogger(__name__)
 
@@ -135,26 +142,42 @@ class SQLAlchemyBackend(BaseDataBackend):
 
     """
     KIND = 'sqlalchemy.conx'
+    #: sqlalchemy.Engine cache to enable pooled connections
+    __CNX_CACHE = ProcessLocal()
+
+    # -- https://docs.sqlalchemy.org/en/14/core/pooling.html#module-sqlalchemy.pool
+    # -- create_engine() must be called per-process, hence using ProcessLocal
+    # -- meaning when using a multiprocessing.Pool or other fork()-ed processes,
+    #    the cache will be cleared in child processes, forcing the engine to be
+    #    recreated automatically in _get_connection
 
     @classmethod
     def supports(cls, obj, name, insert=False, data_store=None, model_store=None, *args, **kwargs):
         valid = cls._is_valid_url(cls, obj)
         support_via = cls._supports_via(cls, data_store, name, obj)
         return valid or support_via
 
-    def drop(self, name, **kwargs):
-        if name in CNX_CACHE:
-            del CNX_CACHE[name]
+    def drop(self, name, secrets=None, **kwargs):
+        # ensure cache is cleared
+        clear_cache = True if secrets is None else False
+        try:
+            self.get(name, secrets=secrets, raw=True, keep=False)
+        except KeyError as e:
+            warnings.warn(f'Connection cache was cleared, however secret {e} was missing.')
+            clear_cache = True
+        if clear_cache:
+            self.__CNX_CACHE.clear()
         return super().drop(name, **kwargs)
 
     def get(self, name, sql=None, chunksize=None, raw=False, sqlvars=None,
-            secrets=None, index=True, keep=False, lazy=False, table=None, *args, **kwargs):
+            secrets=None, index=True, keep=None, lazy=False, table=None, *args, **kwargs):
         """ retrieve a stored connection or query data from connection
 
         Args:
             name (str): the name of the connection
             secrets (dict): dict to resolve variables in the connection string
-            keep (bool): if True connection is kept open.
+            keep (bool): if True connection is kept open, defaults to True (change
+              default as om.defaults.SQLALCHEMY_ALWAYS_CACHE)
             table (str): the name of the table, will be prefixed with the
                store's bucket name unless the table is specified as ':name'
 
@@ -195,8 +218,10 @@ def get(self, name, sql=None, chunksize=None, raw=False, sqlvars=None,
         if not raw and not valid_sql(sql):
             sql = f'select * from {table}'
         chunksize = chunksize or meta.kind_meta.get('chunksize')
-        keep = getattr(self.data_store.defaults, 'SQLALCHEMY_ALWAYS_CACHE',
-                       ALWAYS_CACHE) or keep
+        _default_keep = getattr(self.data_store.defaults,
+                                'SQLALCHEMY_ALWAYS_CACHE',
+                                ALWAYS_CACHE)
+        keep = keep if keep is not None else _default_keep
         if connection_str:
             secrets = self._get_secrets(meta, secrets)
             connection = self._get_connection(name, connection_str, secrets=secrets, keep=keep)
@@ -352,24 +377,30 @@ def _get_connection(self, name, connection_str, secrets=None, keep=False):
         from sqlalchemy import create_engine
 
         connection = None
+        cache_key = None
         try:
+            # SECDEV: the cache key is a secret in order to avoid privilege escalation
+            # -- if it is not secret, user A could create the connection (=> cache)
+            # -- user B could reuse the connection by retrieving the dataset without secrets
+            # -- this way the user needs to have the same secrets in order to reuse the connection
             connection_str = connection_str.format(**(secrets or {}))
-            engine = create_engine(connection_str, echo=False)
-            connection = CNX_CACHE.get(name) or engine.connect()
+            cache_key = sha256(f'{name}:{connection_str}'.encode('utf8')).hexdigest()
+            engine = self.__CNX_CACHE.get(cache_key) or create_engine(connection_str, **ENGINE_KWARGS)
+            connection = engine.connect()
         except KeyError as e:
             msg = ('{e}, ensure secrets are specified for connection '
                    '>{connection_str}<'.format(**locals()))
             raise KeyError(msg)
         except Exception as e:
             if connection is not None:
                 connection.close()
+                self.__CNX_CACHE.pop(cache_key, None)
             raise
         else:
             if keep:
-                CNX_CACHE[name] = connection
+                self.__CNX_CACHE[cache_key] = engine
             else:
-                if name in CNX_CACHE:
-                    del CNX_CACHE[name]
+                self.__CNX_CACHE.pop(cache_key, None)
         return connection
 
     def copy_from_sql(self, sql, connstr, name, chunksize=10000,
@@ -445,7 +476,7 @@ def _supports_via(self, data_store, name, obj):
 
     def _get_secrets(self, meta, secrets):
         secrets_specs = meta.kind_meta.get('secrets')
-        values = dict(os.environ)
+        values = dict(os.environ) if self.data_store.defaults.OMEGA_ALLOW_ENV_CONFIG else dict()
         values.update(**self.data_store.defaults)
         if not secrets and secrets_specs:
             dsname = secrets_specs['dsname']
@@ -500,7 +531,7 @@ def _sanitize_statement(self, sql, sqlvars):
             # -- sqlvars is not used in constructing sql text
             v = sqlvars[k]
             if isinstance(v, (list, tuple)):
-                bind_vars = { f'{k}_{i}': lv for i, lv in enumerate(v)}
+                bind_vars = {f'{k}_{i}': lv for i, lv in enumerate(v)}
                 placeholders = ','.join(f':{bk}' for bk in bind_vars)
                 sql = sql.replace(f':{k}', f'({placeholders})')
                 sqlvars.update(bind_vars)
@@ -565,7 +596,7 @@ def _format_dict(d, replace=None, **kwargs):
         if replace:
             del d[k]
             k = k.replace(*replace) if replace else k
-        d[k] = v.format(**kwargs) if isinstance(v, str) else v
+        d[k] = v.format_map(KeepMissing(kwargs)) if isinstance(v, str) else v
     return d
 
 

omegaml/tests/core/test_sqlalchemy.py

@@ -51,6 +51,27 @@ def test_put_connection_with_secrets(self):
         conn = om.datasets.get('testsqlite', raw=True)
         self.assertIsInstance(conn, Connection)
 
+    def test_connection_cache(self):
+        """ test connection caching
+        """
+        from omegaml.backends import sqlalchemy
+        om = self.om
+        cnx = 'sqlite:///{user}.db'
+        om.datasets.put(cnx, 'testsqlite', kind=SQLAlchemyBackend.KIND)
+        conn = om.datasets.get('testsqlite', raw=True, secrets=dict(user='user'), keep=True)
+        conn_ = om.datasets.get('testsqlite', raw=True, secrets=dict(user='user'), keep=True)
+        self.assertEqual(conn.engine, conn_.engine)
+        # drop should clear cache
+        om.datasets.drop('testsqlite', secrets=dict(user='user'))
+        conn_ = om.datasets.get('testsqlite', raw=True, secrets=dict(user='user'), keep=True)
+        self.assertIsNone(conn_)
+        self.assertTrue(len(sqlalchemy.SQLAlchemyBackend._SQLAlchemyBackend__CNX_CACHE) == 0)
+        # even if we drop without secrets, cache is cleared
+        om.datasets.put(cnx, 'testsqlite', kind=SQLAlchemyBackend.KIND)
+        om.datasets.get('testsqlite', raw=True, secrets=dict(user='user'), keep=True)
+        om.datasets.drop('testsqlite')
+        self.assertTrue(len(sqlalchemy.SQLAlchemyBackend._SQLAlchemyBackend__CNX_CACHE) == 0)
+
     def test_put_connection_with_sql(self):
         """
         store generic sqlalchemy connection with sql, same principle as a view

omegaml/util.py

@@ -1088,3 +1088,35 @@ def __str__(self):
         # on posix systems, this is a noop
         path = super().__str__()
         return path if os.name != 'nt' else path.replace('\\', '/')
+
+
+class ProcessLocal(dict):
+    def __init__(self, *args, **kwargs):
+        self._pid = os.getpid()
+        super().__init__(*args, **kwargs)
+
+    def _check_pid(self):
+        if self._pid != os.getpid():
+            self.clear()
+            self._pid = os.getpid()
+
+    def __getitem__(self, k):
+        self._check_pid()
+        return super().__getitem__(k)
+
+    def keys(self):
+        self._check_pid()
+        return super().keys()
+
+    def __contains__(self, item):
+        self._check_pid()
+        return super().__contains__(item)
+
+
+class KeepMissing(dict):
+    # a missing '{key}' is replaced by '{key}'
+    # in order to avoid raising KeyError
+    # see str.format_map
+    def __missing__(self, key):
+        return '{' + key + '}'
+