bstein/ariadne
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

db: wrap pool config

Browse Source
This commit is contained in:
Brad Stein 2026-01-22 14:43:40 -03:00
parent 35fae20d36
commit 44c86e739d
3 changed files with 52 additions and 44 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

34
ariadne/app.py
View File

@ -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)

44
ariadne/db/database.py
View File

@ -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):

18
ariadne/migrate.py
View File

@ -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,
),
)