diff --git a/cads_catalogue/config.py b/cads_catalogue/config.py index aae9d86..91920e6 100644 --- a/cads_catalogue/config.py +++ b/cads_catalogue/config.py @@ -45,6 +45,9 @@ class SqlalchemySettings: catalogue_db_host_read: str | None = None catalogue_db_name: str | None = None pool_recycle: int = 60 + pool_size: int = 5 + pool_timeout: float = 30 + max_overflow: int = 15 def __init__(self, **kwargs): self.match_args = kwargs diff --git a/cads_catalogue/database.py b/cads_catalogue/database.py index 309e386..28c69ca 100644 --- a/cads_catalogue/database.py +++ b/cads_catalogue/database.py @@ -392,7 +392,24 @@ class Licence(BaseModel): ) -def ensure_session_obj(read_only: bool = False) -> sa.orm.sessionmaker: +def ensure_engine(connection_string, **kwargs) -> sa.engine.Engine: + """Return a sqlalchemy engine.""" + settings = config.ensure_settings(config.dbsettings) + engine_kwargs = { + "pool_recycle": settings.pool_recycle, + "pool_size": settings.pool_size, + "pool_timeout": settings.pool_timeout, + "max_overflow": settings.max_overflow, + } + engine_kwargs.update(kwargs) + if engine_kwargs["pool_size"] == -1: + engine = sa.create_engine(connection_string, poolclass=sa.pool.NullPool) + else: + engine = sa.create_engine(connection_string, **engine_kwargs) + return engine + + +def ensure_session_obj(read_only: bool = False, **kwargs) -> sa.orm.sessionmaker: """Create a new session object bound to the catalogue database. Parameters @@ -406,13 +423,10 @@ def ensure_session_obj(read_only: bool = False) -> sa.orm.sessionmaker: """ settings = config.ensure_settings(config.dbsettings) if read_only: - engine = sa.create_engine( - settings.connection_string_read, pool_recycle=settings.pool_recycle - ) + connection_string = settings.connection_string_read else: - engine = sa.create_engine( - settings.connection_string, pool_recycle=settings.pool_recycle - ) + connection_string = settings.connection_string + engine = ensure_engine(connection_string, **kwargs) session_obj = sa.orm.sessionmaker(engine) return session_obj