diff --git a/rhodecode/model/__init__.py b/rhodecode/model/__init__.py --- a/rhodecode/model/__init__.py +++ b/rhodecode/model/__init__.py @@ -27,25 +27,32 @@ from rhodecode.lib.utils2 import obfusca 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): diff --git a/rhodecode/model/meta.py b/rhodecode/model/meta.py --- a/rhodecode/model/meta.py +++ b/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, - ) - ) +session_factory = sessionmaker( + expire_on_commit=True, + future=True +) -# pass empty regions so we can fetch it on-demand inside ORMCache +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) -# The declarative Base -Base = declarative_base() +def raw_query_executor(engine=None): + """ + + :param engine: + :return: + """ + if not engine: + engine = Session.bind + session = SA_Session(engine) + return session -def raw_query_executor(): - engine = Base.metadata.bind - session = SASession(engine) - return session +def get_engine(): + return Session.bind + + +def bind_engine_to_session(engine): + Session.remove() + Session.configure(bind=engine)