Merge pull request #134 from ecmwf-projects/configurable_engine

more configurable engine on calling

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

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