models: new sqlalchemy logic and encryption setup for python3

rhodecode/model/__init__.py

@@ -27,25 +27,32 @@ from rhodecode.lib.utils2 import obfuscate_url_pw, get_encryption_key
 log = logging.getLogger(__name__)
 
 
-def init_model(engine, encryption_key=None):
+def init_model(engine, encryption_key: bytes = b''):
     """
     Initializes db session, bind the engine with the metadata,
     Call this before using any of the tables or classes in the model,
     preferably once in application start
 
     :param engine: engine to bind to
+    :param encryption_key: key used for encryption
     """
+
     engine_str = obfuscate_url_pw(str(engine.url))
     log.info("RhodeCode %s initializing db for %s", rhodecode.__version__, engine_str)
-    meta.Base.metadata.bind = engine
-    db.ENCRYPTION_KEY = encryption_key
+
+    meta.bind_engine_to_session(engine)
+    init_model_encryption(db, enc_key=encryption_key)
 
 
-def init_model_encryption(migration_models, config=None):
-    from pyramid.threadlocal import get_current_registry
-    config = config or get_current_registry().settings
-    migration_models.ENCRYPTION_KEY = get_encryption_key(config)
-    db.ENCRYPTION_KEY = get_encryption_key(config)
+def init_model_encryption(*db_models, enc_key: bytes = b'', config=None):
+    if not enc_key:
+        from pyramid.threadlocal import get_current_registry
+        config = config or get_current_registry().settings
+        enc_key = get_encryption_key(config)
+
+    for db_model in db_models:
+        log.debug('setting encryption key for model %s', db_model)
+        db_model.ENCRYPTION_KEY = enc_key
 
 
 class BaseModel(object):

rhodecode/model/meta.py

@@ -23,33 +23,53 @@ SQLAlchemy Metadata and Session object
 
 from sqlalchemy.orm import declarative_base
 from sqlalchemy.orm import scoped_session, sessionmaker
-from sqlalchemy.orm import Session as SASession
+from sqlalchemy.orm import Session as SA_Session
 from rhodecode.lib.caching_query import ORMCache
 
 
-__all__ = ['Base', 'Session', 'raw_query_executor']
+__all__ = [
+    'Base', 'Session', 'SA_Session',
+    'raw_query_executor',
+    'bind_engine_to_session',
+    'get_engine'
+]
 
 # scoped_session.  Apply our custom CachingQuery class to it,
 # using a callable that will associate the dictionary
 # of regions with the Query.
 # to use cache use this in query
 # .options(FromCache("sqlalchemy_cache_type", "cachekey"))
-Session = scoped_session(
-                sessionmaker(
-                    expire_on_commit=True,
-                )
-          )
-
-# pass empty regions so we can fetch it on-demand inside ORMCache
-cache = ORMCache(regions={})
-cache.listen_on_session(Session)
+session_factory = sessionmaker(
+    expire_on_commit=True,
+    future=True
+)
 
+Session = scoped_session(session_factory)
 
 # The declarative Base
 Base = declarative_base()
 
+# pass empty regions, so we can fetch it on-demand inside ORMCache
+cache = ORMCache(regions={})
+cache.listen_on_session(Session)
 
-def raw_query_executor():
-    engine = Base.metadata.bind
-    session = SASession(engine)
+
+def raw_query_executor(engine=None):
+    """
+
+    :param engine:
+    :return:
+    """
+    if not engine:
+        engine = Session.bind
+    session = SA_Session(engine)
     return session
+
+
+def get_engine():
+    return Session.bind
+
+
+def bind_engine_to_session(engine):
+    Session.remove()
+    Session.configure(bind=engine)