✨ Refactor database connection mapping and introduce atomic transaction helpers

Author: Paillat-dev

src/database/config/__init__.py

@@ -8,7 +8,7 @@
 from tortoise import Tortoise
 
 from src.config import config
-from src.config.models import DbExtraApp, DbParams
+from src.config.models import DbParams
 
 logger = getLogger("bot").getChild("database")
 
@@ -25,44 +25,63 @@ def apply_params(uri: str, params: DbParams | None) -> str:
     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()
-}
+def get_app_url_mapping() -> dict[str, str]:
+    return {
+        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)
+def get_url_apps_mapping() -> dict[str, list[str]]:
+    app_url_mapping: dict[str, str] = get_app_url_mapping()
+    mapping: dict[str, list[str]] = defaultdict(list)
 
-app_connections_mapping: dict[str, str] = {}
-connection_url_mapping: dict[str, str] = {}
+    for app_name, app_url in app_url_mapping.items():
+        mapping[app_url].append(app_name)
 
-i: int = 0
+    return mapping
 
-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)
+def parse_url_apps_mapping(url_apps_mapping: dict[str, list[str]]) -> tuple[dict[str, str], dict[str, str]]:
+    app_connection: dict[str, str] = {}
+    connection_url: dict[str, str] = {}
 
-config.db.extra_apps["models"] = DbExtraApp(
-    models=["src.database.models", "aerich.models"],
-)
+    for i, (url, apps) in enumerate(url_apps_mapping.items()):
+        connection_name = f"connection_{i}"
+        connection_url[connection_name] = url
+        for app in apps:
+            app_connection[app] = connection_name
 
-TORTOISE_ORM = {
-    "connections": connection_url_mapping,
-    "apps": {
+    app_connection["models"] = "default"
+    connection_url["default"] = apply_params(config.db.url, config.db.params)
+
+    return app_connection, connection_url
+
+
+APP_CONNECTION_MAPPING: dict[str, str]
+CONNECTION_URL_MAPPING: dict[str, str]
+
+APP_CONNECTION_MAPPING, CONNECTION_URL_MAPPING = parse_url_apps_mapping(get_url_apps_mapping())  # pyright: ignore[reportConstantRedefinition]
+
+
+def get_apps() -> dict[str, dict[str, list[str] | str]]:
+    apps = {
         app_name: {
             "models": app.models,
-            "default_connection": app_connections_mapping[app_name],
+            "default_connection": APP_CONNECTION_MAPPING[app_name],
         }
         for app_name, app in config.db.extra_apps.items()
-    },
+    }
+    apps["models"] = {
+        "models": ["src.database.models", "aerich.models"],
+        "default_connection": "default",
+    }
+    return apps
+
+
+TORTOISE_ORM = {
+    "connections": CONNECTION_URL_MAPPING,
+    "apps": get_apps(),
 }
 
 
@@ -82,4 +101,4 @@ async def shutdown() -> None:
     await Tortoise.close_connections()
 
 
-__all__ = ["init", "shutdown"]
+__all__ = ["APP_CONNECTION_MAPPING", "init", "shutdown"]

src/database/utils/atomics.py

@@ -0,0 +1,22 @@
+# Copyright (c) NiceBots
+# SPDX-License-Identifier: MIT
+
+from collections.abc import Callable
+from typing import Any
+
+from tortoise.backends.base.client import TransactionContext
+from tortoise.transactions import atomic as t_atomic
+from tortoise.transactions import in_transaction as t_in_transaction
+
+from src.database.config import APP_CONNECTION_MAPPING
+
+
+def atomic[F: Callable[..., Any]](connection_name: str | None = None) -> Callable[[F], F]:
+    return t_atomic(APP_CONNECTION_MAPPING[connection_name or "models"])
+
+
+def in_transaction(app_name: str | None = None) -> TransactionContext:  # pyright: ignore[reportMissingTypeArgument, reportUnknownParameterType]
+    return t_in_transaction(APP_CONNECTION_MAPPING[app_name or "models"])  # pyright: ignore[reportUnknownVariableType]
+
+
+__all__ = ("atomic", "in_transaction")