✨ Add SSL support for database connections and refactor connection configuration handling

Author: Paillat-dev

src/database/config/__init__.py

@@ -1,8 +1,10 @@
 # Copyright (c) NiceBots
 # SPDX-License-Identifier: MIT
 
+import ssl
 from collections import defaultdict
 from logging import getLogger
+from pathlib import Path
 from typing import Any
 
 import aerich
@@ -13,12 +15,40 @@
 logger = getLogger("bot").getChild("database")
 
 
+def create_ssl_context() -> ssl.SSLContext:
+    """Create SSL context for Database connection."""
+    # Try to find the certificate in common locations
+    cert_paths = [
+        Path(__file__).parent.parent.parent / ".postgres" / "root.crt",  # Docker/production
+        Path.cwd() / ".postgres" / "root.crt",  # Local dev
+    ]
+
+    cert_path = None
+    for path in cert_paths:
+        if path.exists():
+            cert_path = path
+            break
+
+    if cert_path is None:
+        logger.warning("Database certificate not found, SSL verification will fail")
+        return True  # fallback to basic SSL
+
+    logger.info(f"Using SSL certificate: {cert_path}")
+    ssl_context = ssl.create_default_context(cafile=str(cert_path))
+    ssl_context.check_hostname = False
+    ssl_context.verify_mode = ssl.CERT_REQUIRED
+    return ssl_context
+
+
 def apply_params(uri: str, params: dict[str, Any] | None) -> str:
     if params is None:
         return uri
 
     first: bool = True
     for param, value in params.items():
+        # Skip 'ssl' param as we'll handle it separately
+        if param == "ssl":
+            continue
         if value is not None:
             uri += f"{'?' if first else '&'}{param}={value}"
             first = False
@@ -42,26 +72,35 @@ def get_url_apps_mapping() -> dict[str, list[str]]:
     return mapping
 
 
-def parse_url_apps_mapping(url_apps_mapping: dict[str, list[str]]) -> tuple[dict[str, str], dict[str, str]]:
+def parse_url_apps_mapping(url_apps_mapping: dict[str, list[str]]) -> tuple[dict[str, str], dict[str, dict[str, Any]]]:
     app_connection: dict[str, str] = {}
-    connection_url: dict[str, str] = {}
+    connection_config: dict[str, dict[str, Any]] = {}
+
+    # Create SSL context once if needed
+    ssl_context = create_ssl_context() if config.db.params and config.db.params.get("ssl") else None
 
     for i, (url, apps) in enumerate(url_apps_mapping.items()):
         connection_name = f"connection_{i}"
-        connection_url[connection_name] = url
+        connection_config[connection_name] = {"engine": "tortoise.backends.asyncpg", "credentials": {"dsn": url}}
+        if ssl_context:
+            connection_config[connection_name]["credentials"]["ssl"] = ssl_context
+
         for app in apps:
             app_connection[app] = connection_name
 
     app_connection["models"] = "default"
-    connection_url["default"] = apply_params(config.db.url, config.db.params)
+    default_url = apply_params(config.db.url, config.db.params)
+    connection_config["default"] = {"engine": "tortoise.backends.asyncpg", "credentials": {"dsn": default_url}}
+    if ssl_context:
+        connection_config["default"]["credentials"]["ssl"] = ssl_context
 
-    return app_connection, connection_url
+    return app_connection, connection_config
 
 
 APP_CONNECTION_MAPPING: dict[str, str]
-CONNECTION_URL_MAPPING: dict[str, str]
+CONNECTION_CONFIG_MAPPING: dict[str, dict[str, Any]]
 
-APP_CONNECTION_MAPPING, CONNECTION_URL_MAPPING = parse_url_apps_mapping(get_url_apps_mapping())  # pyright: ignore[reportConstantRedefinition]
+APP_CONNECTION_MAPPING, CONNECTION_CONFIG_MAPPING = parse_url_apps_mapping(get_url_apps_mapping())
 
 
 def get_apps() -> dict[str, dict[str, list[str] | str]]:
@@ -80,7 +119,7 @@ def get_apps() -> dict[str, dict[str, list[str] | str]]:
 
 
 TORTOISE_ORM = {
-    "connections": CONNECTION_URL_MAPPING,
+    "connections": CONNECTION_CONFIG_MAPPING,
     "apps": get_apps(),
 }
 
@@ -93,7 +132,7 @@ async def init() -> None:
     )
     await command.init()
     migrated = await command.upgrade(run_in_transaction=True)
-    logger.success(f"Successfully migrated {migrated} migrations")  # pyright: ignore [reportAttributeAccessIssue]
+    logger.success(f"Successfully migrated {migrated} migrations")
     await Tortoise.init(config=TORTOISE_ORM)