From 44c86e739d8d9971838223df91381eb3317d0c84 Mon Sep 17 00:00:00 2001 From: Brad Stein Date: Thu, 22 Jan 2026 14:43:40 -0300 Subject: [PATCH] db: wrap pool config --- ariadne/app.py | 34 ++++++++++++++++++-------------- ariadne/db/database.py | 44 ++++++++++++++++++++++-------------------- ariadne/migrate.py | 18 +++++++++-------- 3 files changed, 52 insertions(+), 44 deletions(-) diff --git a/ariadne/app.py b/ariadne/app.py index 2413544..d625ddf 100644 --- a/ariadne/app.py +++ b/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,23 +62,27 @@ class PasswordResetRequest: portal_db = Database( settings.portal_database_url, - pool_min=settings.ariadne_db_pool_min, - pool_max=settings.ariadne_db_pool_max, - connect_timeout_sec=settings.ariadne_db_connect_timeout_sec, - lock_timeout_sec=settings.ariadne_db_lock_timeout_sec, - 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", + DatabaseConfig( + pool_min=settings.ariadne_db_pool_min, + pool_max=settings.ariadne_db_pool_max, + connect_timeout_sec=settings.ariadne_db_connect_timeout_sec, + lock_timeout_sec=settings.ariadne_db_lock_timeout_sec, + 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, - pool_min=settings.ariadne_db_pool_min, - pool_max=settings.ariadne_db_pool_max, - connect_timeout_sec=settings.ariadne_db_connect_timeout_sec, - lock_timeout_sec=settings.ariadne_db_lock_timeout_sec, - 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", + DatabaseConfig( + pool_min=settings.ariadne_db_pool_min, + pool_max=settings.ariadne_db_pool_max, + connect_timeout_sec=settings.ariadne_db_connect_timeout_sec, + lock_timeout_sec=settings.ariadne_db_lock_timeout_sec, + 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) diff --git a/ariadne/db/database.py b/ariadne/db/database.py index 581693b..961910d 100644 --- a/ariadne/db/database.py +++ b/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__( - 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: + def __init__(self, dsn: str, config: DatabaseConfig | None = None) -> 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 statement_timeout={statement_timeout_sec}s " - f"-c idle_in_transaction_session_timeout={idle_in_tx_timeout_sec}s" + f"-c lock_timeout={config.lock_timeout_sec}s " + f"-c statement_timeout={config.statement_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, - max_size=pool_max, + min_size=config.pool_min, + max_size=config.pool_max, kwargs={ - "connect_timeout": connect_timeout_sec, - "application_name": application_name, + "connect_timeout": config.connect_timeout_sec, + "application_name": config.application_name, "options": options, }, ) - self._lock_timeout_sec = lock_timeout_sec - self._statement_timeout_sec = statement_timeout_sec + self._lock_timeout_sec = config.lock_timeout_sec + self._statement_timeout_sec = config.statement_timeout_sec @contextmanager def connection(self): diff --git a/ariadne/migrate.py b/ariadne/migrate.py index e3f4435..2877149 100644 --- a/ariadne/migrate.py +++ b/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,13 +11,15 @@ ARIADNE_MIGRATION_LOCK_ID = 982732 def _build_db(dsn: str, application_name: str) -> Database: return Database( dsn, - pool_min=settings.ariadne_db_pool_min, - pool_max=settings.ariadne_db_pool_max, - connect_timeout_sec=settings.ariadne_db_connect_timeout_sec, - lock_timeout_sec=settings.ariadne_db_lock_timeout_sec, - 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, + DatabaseConfig( + pool_min=settings.ariadne_db_pool_min, + pool_max=settings.ariadne_db_pool_max, + connect_timeout_sec=settings.ariadne_db_connect_timeout_sec, + lock_timeout_sec=settings.ariadne_db_lock_timeout_sec, + 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, + ), )