✨ Enhance database configuration with dynamic connection mapping and introduce DbParams and DbExtraApp models

Author: Paillat-dev

src/config/models.py

@@ -3,7 +3,7 @@
 
 from typing import Any, Literal, overload, override
 
-from pydantic import BaseModel
+from pydantic import BaseModel, ConfigDict, Field
 
 type Extension = dict[str, Any]
 
@@ -68,9 +68,29 @@ class UseConfig(BaseModel):
  backend: bool = False
 
 
+class DbParams(BaseModel):
+ model_config = ConfigDict(serialize_by_alias=True)
+ minsize: int | None = None
+ maxsize: int | None = None
+ max_queries: int | None = None
+ max_inactive_connection_lifetime: float | None = None
+ schema_name: str | None = Field(
+ alias="schema", validation_alias="schema", serialization_alias="schema", default=None
+ )
+ ssl: bool | None = None
+
+
+class DbExtraApp(BaseModel):
+ url: str | None = None
+ params: DbParams | None = None
+ models: list[str] = []
+
+
 class DbConfig(BaseModel):
  url: str
  enabled: bool = True
+ params: DbParams | None = None
+ extra_apps: dict[str, DbExtraApp] = {}
 
 
 class Config(BaseModel):

src/database/config/__init__.py

@@ -1,22 +1,67 @@
 # Copyright (c) NiceBots
 # SPDX-License-Identifier: MIT
 
+from collections import defaultdict
 from logging import getLogger
 
 import aerich
 from tortoise import Tortoise
 
 from src.config import config
+from src.config.models import DbExtraApp, DbParams
 
 logger = getLogger("bot").getChild("database")
 
+
+def apply_params(uri: str, params: DbParams | None) -> str:
+ if params is None:
+ return uri
+
+ first: bool = True
+ for param, value in params.model_dump().items():
+ if value is not None:
+ uri += f"{'?' if first else '&'}{param}={value}"
+ first = False
+ return uri
+
+
+app_url_mapping: dict[str, str] = {
+ app_name: apply_params(app_config.url or config.db.url, app_config.params)
+ for app_name, app_config in config.db.extra_apps.items()
+}
+
+url_apps_mapping: dict[str, list[str]] = defaultdict(list)
+
+for app_name, app_url in app_url_mapping.items():
+ url_apps_mapping[app_url].append(app_name)
+
+app_connections_mapping: dict[str, str] = {}
+connection_url_mapping: dict[str, str] = {}
+
+i: int = 0
+
+for url, apps in url_apps_mapping.items():
+ connection_name = f"connection_{i}"
+ connection_url_mapping[connection_name] = url
+ for app in apps:
+ app_connections_mapping[app] = connection_name
+ i += 1 # noqa: SIM113 # Incompatible with .items()
+
+app_connections_mapping["models"] = "default"
+connection_url_mapping["default"] = apply_params(config.db.url, config.db.params)
+
+config.db.extra_apps["models"] = DbExtraApp(
+ models=["src.database.models", "aerich.models"],
+)
+
 TORTOISE_ORM = {
- "connections": {"default": config.db.url},
+ "connections": connection_url_mapping,
  "apps": {
- "models": {
- "models": ["src.database.models", "aerich.models"],
- "default_connection": "default",
+ app_name: {
+ "models": app.models,
+ "default_connection": app_connections_mapping[app_name],
  }
+ for app_name, app in config.db.extra_apps.items()
  },
 }