provide config overrides for db engine pool params

CHANGELOG.md

@@ -16,6 +16,9 @@ The types of changes are:
 
 ## [Unreleased](https://github.com/ethyca/fides/compare/2.6.1...main)
 
+### Added
+* Added config properties to override database Engine parameters [#2511](https://github.com/ethyca/fides/pull/2511)
+
 ## [2.6.1](https://github.com/ethyca/fides/compare/2.6.0...2.6.1)
 
 ## [2.6.0](https://github.com/ethyca/fides/compare/2.5.1...2.6.0)

src/fides/api/ops/api/deps.py

@@ -30,7 +30,11 @@ def get_api_session() -> Session:
     """Gets the shared database session to use for API functionality"""
     global _engine  # pylint: disable=W0603
     if not _engine:
-        _engine = get_db_engine(config=CONFIG)
+        _engine = get_db_engine(
+            config=CONFIG,
+            pool_size=CONFIG.database.api_engine_pool_size,
+            max_overflow=CONFIG.database.api_engine_max_overflow,
+        )
     SessionLocal = get_db_session(CONFIG, engine=_engine)
     db = SessionLocal()
     return db

src/fides/api/ops/tasks/__init__.py

@@ -25,7 +25,11 @@ def get_new_session(self) -> ContextManager[Session]:
         # only one engine will be instantiated in a given task scope, i.e
         # once per celery process.
         if self._task_engine is None:
-            _task_engine = get_db_engine(config=CONFIG)
+            _task_engine = get_db_engine(
+                config=CONFIG,
+                pool_size=CONFIG.database.task_engine_pool_size,
+                max_overflow=CONFIG.database.task_engine_max_overflow,
+            )
 
         # same for the sessionmaker
         if self._sessionmaker is None:

src/fides/core/config/database_settings.py

@@ -23,6 +23,11 @@ class DatabaseSettings(FidesSettings):
     db: str = "default_db"
     test_db: str = "default_test_db"
 
+    api_engine_pool_size: int = 5
+    api_engine_max_overflow: int = 10
+    task_engine_pool_size: int = 5
+    task_engine_max_overflow: int = 10
+
     sqlalchemy_database_uri: Optional[str] = None
     sqlalchemy_test_database_uri: Optional[str] = None
 

src/fides/core/config/utils.py

@@ -24,6 +24,10 @@ def get_dev_mode() -> bool:
         "port",
         "db",
         "test_db",
+        "api_engine_pool_size",
+        "api_engine_max_overflow",
+        "task_engine_pool_size",
+        "task_engine_max_overflow",
     ],
     "notifications": [
         "send_request_completion_notification",

src/fides/lib/db/session.py

@@ -14,6 +14,8 @@ def get_db_engine(
     *,
     config: FidesConfig | None = None,
     database_uri: str | URL | None = None,
+    pool_size: int = 5,
+    max_overflow: int = 10,
 ) -> Engine:
     """Return a database engine.
 
@@ -29,7 +31,9 @@ def get_db_engine(
             database_uri = config.database.sqlalchemy_test_database_uri
         else:
             database_uri = config.database.sqlalchemy_database_uri
-    return create_engine(database_uri, pool_pre_ping=True)
+    return create_engine(
+        database_uri, pool_pre_ping=True, pool_size=pool_size, max_overflow=max_overflow
+    )
 
 
 def get_db_session(

tests/ops/api/test_deps.py

@@ -1,6 +1,12 @@
+# pylint: disable=protected-access
+
 import pytest
+from sqlalchemy.engine import Engine
+from sqlalchemy.orm import Session
+from sqlalchemy.pool import QueuePool
 
-from fides.api.ops.api.deps import get_cache
+import fides.api.ops.api.deps
+from fides.api.ops.api.deps import get_api_session, get_cache
 from fides.api.ops.common_exceptions import FunctionalityNotConfigured
 from fides.core.config import get_config
 
@@ -15,7 +21,36 @@ def mock_config():
     CONFIG.redis.enabled = redis_enabled
 
 
+@pytest.fixture
+def mock_config_changed_db_engine_settings():
+    pool_size = CONFIG.database.api_engine_pool_size
+    CONFIG.database.api_engine_pool_size = pool_size + 5
+    max_overflow = CONFIG.database.api_engine_max_overflow
+    CONFIG.database.api_engine_max_overflow = max_overflow + 5
+    yield
+    CONFIG.database.api_engine_pool_size = pool_size
+    CONFIG.database.api_engine_max_overflow = max_overflow
+
+
 @pytest.mark.usefixtures("mock_config")
 def test_get_cache_not_enabled():
     with pytest.raises(FunctionalityNotConfigured):
         next(get_cache())
+
+
+@pytest.mark.parametrize(
+    "config_fixture", [None, "mock_config_changed_db_engine_settings"]
+)
+def test_get_api_session(config_fixture, request):
+    if config_fixture is not None:
+        request.getfixturevalue(
+            config_fixture
+        )  # used to invoke config fixture if provided
+    fides.api.ops.api.deps._engine = None
+    pool_size = CONFIG.database.api_engine_pool_size
+    max_overflow = CONFIG.database.api_engine_max_overflow
+    session: Session = get_api_session()
+    engine: Engine = session.get_bind()
+    pool: QueuePool = engine.pool
+    assert pool.size() == pool_size
+    assert pool._max_overflow == max_overflow

tests/ops/tasks/test_celery.py

@@ -1,9 +1,26 @@
-from fides.api.ops.tasks import _create_celery
+# pylint: disable=protected-access
+import pytest
+from sqlalchemy.engine import Engine
+from sqlalchemy.orm import Session
+from sqlalchemy.pool import QueuePool
+
+from fides.api.ops.tasks import DatabaseTask, _create_celery
 from fides.core.config import get_config
 
 CONFIG = get_config()
 
 
+@pytest.fixture
+def mock_config_changed_db_engine_settings():
+    pool_size = CONFIG.database.task_engine_pool_size
+    CONFIG.database.task_engine_pool_size = pool_size + 5
+    max_overflow = CONFIG.database.task_engine_max_overflow
+    CONFIG.database.task_engine_max_overflow = max_overflow + 5
+    yield
+    CONFIG.database.task_engine_pool_size = pool_size
+    CONFIG.database.task_engine_max_overflow = max_overflow
+
+
 def test_create_task(celery_session_app, celery_session_worker):
     @celery_session_app.task
     def multiply(x, y):
@@ -32,3 +49,21 @@ def test_celery_config_override() -> None:
     celery_app = _create_celery(config=config)
     assert celery_app.conf["event_queue_prefix"] == "overridden_fides_worker"
     assert celery_app.conf["task_default_queue"] == "overridden_fides"
+
+
+@pytest.mark.parametrize(
+    "config_fixture", [None, "mock_config_changed_db_engine_settings"]
+)
+def test_get_task_session(config_fixture, request):
+    if config_fixture is not None:
+        request.getfixturevalue(
+            config_fixture
+        )  # used to invoke config fixture if provided
+    pool_size = CONFIG.database.task_engine_pool_size
+    max_overflow = CONFIG.database.task_engine_max_overflow
+    t = DatabaseTask()
+    session: Session = t.get_new_session()
+    engine: Engine = session.get_bind()
+    pool: QueuePool = engine.pool
+    assert pool.size() == pool_size
+    assert pool._max_overflow == max_overflow