db: wrap pool config

ariadne/app.py

@@ -11,7 +11,7 @@ from fastapi.responses import JSONResponse, Response
 from prometheus_client import CONTENT_TYPE_LATEST, generate_latest
 
 from .auth.keycloak import AuthContext, authenticator
-from .db.database import Database
+from .db.database import Database, DatabaseConfig
 from .db.storage import Storage, TaskRunRecord
 from .manager.provisioning import ProvisioningManager
 from .metrics.metrics import record_task_run
@@ -62,6 +62,7 @@ class PasswordResetRequest:
 
 portal_db = Database(
     settings.portal_database_url,
+    DatabaseConfig(
         pool_min=settings.ariadne_db_pool_min,
         pool_max=settings.ariadne_db_pool_max,
         connect_timeout_sec=settings.ariadne_db_connect_timeout_sec,
@@ -69,9 +70,11 @@ portal_db = Database(
         statement_timeout_sec=settings.ariadne_db_statement_timeout_sec,
         idle_in_tx_timeout_sec=settings.ariadne_db_idle_in_tx_timeout_sec,
         application_name="ariadne_portal",
+    ),
 )
 ariadne_db = Database(
     settings.ariadne_database_url,
+    DatabaseConfig(
         pool_min=settings.ariadne_db_pool_min,
         pool_max=settings.ariadne_db_pool_max,
         connect_timeout_sec=settings.ariadne_db_connect_timeout_sec,
@@ -79,6 +82,7 @@ ariadne_db = Database(
         statement_timeout_sec=settings.ariadne_db_statement_timeout_sec,
         idle_in_tx_timeout_sec=settings.ariadne_db_idle_in_tx_timeout_sec,
         application_name="ariadne",
+    ),
 )
 storage = Storage(ariadne_db, portal_db)
 provisioning = ProvisioningManager(portal_db, storage)

ariadne/db/database.py

@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 from contextlib import contextmanager
+from dataclasses import dataclass
 import logging
 from typing import Any, Iterable
 
@@ -12,38 +13,39 @@ from .schema import ARIADNE_ACCESS_REQUEST_ALTER_SQL, ARIADNE_TABLES_SQL
 logger = logging.getLogger(__name__)
 
 
+@dataclass(frozen=True)
+class DatabaseConfig:
+    pool_min: int = 0
+    pool_max: int = 5
+    connect_timeout_sec: int = 5
+    lock_timeout_sec: int = 5
+    statement_timeout_sec: int = 30
+    idle_in_tx_timeout_sec: int = 10
+    application_name: str = "ariadne"
+
+
 class Database:
-    def __init__(
+    def __init__(self, dsn: str, config: DatabaseConfig | None = None) -> None:
-        self,
-        dsn: str,
-        *,
-        pool_min: int = 0,
-        pool_max: int = 5,
-        connect_timeout_sec: int = 5,
-        lock_timeout_sec: int = 5,
-        statement_timeout_sec: int = 30,
-        idle_in_tx_timeout_sec: int = 10,
-        application_name: str = "ariadne",
-    ) -> None:
         if not dsn:
             raise RuntimeError("database URL is required")
+        config = config or DatabaseConfig()
         options = (
-            f"-c lock_timeout={lock_timeout_sec}s "
+            f"-c lock_timeout={config.lock_timeout_sec}s "
-            f"-c statement_timeout={statement_timeout_sec}s "
+            f"-c statement_timeout={config.statement_timeout_sec}s "
-            f"-c idle_in_transaction_session_timeout={idle_in_tx_timeout_sec}s"
+            f"-c idle_in_transaction_session_timeout={config.idle_in_tx_timeout_sec}s"
         )
         self._pool = ConnectionPool(
             conninfo=dsn,
-            min_size=pool_min,
+            min_size=config.pool_min,
-            max_size=pool_max,
+            max_size=config.pool_max,
             kwargs={
-                "connect_timeout": connect_timeout_sec,
+                "connect_timeout": config.connect_timeout_sec,
-                "application_name": application_name,
+                "application_name": config.application_name,
                 "options": options,
             },
         )
-        self._lock_timeout_sec = lock_timeout_sec
+        self._lock_timeout_sec = config.lock_timeout_sec
-        self._statement_timeout_sec = statement_timeout_sec
+        self._statement_timeout_sec = config.statement_timeout_sec
 
     @contextmanager
     def connection(self):

ariadne/migrate.py

@@ -1,6 +1,6 @@
 from __future__ import annotations
 
-from .db.database import Database
+from .db.database import Database, DatabaseConfig
 from .settings import settings
 
 
@@ -11,6 +11,7 @@ ARIADNE_MIGRATION_LOCK_ID = 982732
 def _build_db(dsn: str, application_name: str) -> Database:
     return Database(
         dsn,
+        DatabaseConfig(
             pool_min=settings.ariadne_db_pool_min,
             pool_max=settings.ariadne_db_pool_max,
             connect_timeout_sec=settings.ariadne_db_connect_timeout_sec,
@@ -18,6 +19,7 @@ def _build_db(dsn: str, application_name: str) -> Database:
             statement_timeout_sec=settings.ariadne_db_statement_timeout_sec,
             idle_in_tx_timeout_sec=settings.ariadne_db_idle_in_tx_timeout_sec,
             application_name=application_name,
+        ),
     )