Migrate config to pydantic-settings

Author: ChrisLovering

.env.example

@@ -0,0 +1,7 @@
+loki_api_url=http://localhost:3100/
+loki_jobs=["default/bot"]
+
+discord_webhook_url=some_discord_webhook
+
+service_interval_minutes=5
+service_tokens=[{"token":"ERROR","color":"#ff5f5f"},{"token":"WARN","color":"#ffe24d"},{"token":"INFO"}]

olli/__main__.py

@@ -6,14 +6,14 @@
 from loguru import logger
 
 from olli.alert import run
-from olli.config import CONFIG
+from olli.config import SERVICE_CONFIG
 
 
 @logger.catch
 def start() -> None:
     """Start the Olli process."""
     logger.info("Starting Olli")
-    schedule.every(CONFIG.olli.interval_minutes).minutes.do(run)
+    schedule.every(SERVICE_CONFIG.interval_minutes).minutes.do(run)
 
     run()
 

olli/alert.py

@@ -4,31 +4,27 @@
 
 from olli import webhook
 from olli.api import LokiHTTPClient
-from olli.config import CONFIG, TokenConfig
+from olli.config import SERVICE_CONFIG, TokenConfig
 from olli.structures import TokenMatch
 
 api_client = LokiHTTPClient()
 
 
 def get_match(token: TokenConfig) -> TokenMatch:
     """Search the configured service logs for a given token."""
+    logger.debug(f"Searching for token {token.token}")
     try:
-        logger.debug(f"Searching for token {token.token}")
         svc_logs = api_client.get_token_logs(token)
     except httpx.ConnectError:
         logger.error("Could not connect to Loki")
         return webhook.send_olli_error("Loki refused to connect.")
     except httpx.HTTPStatusError as e:
         logger.error(f"Loki returned error status code: {e.response.status_code}")
-        return webhook.send_olli_error(
-            f"Loki returned a bad status code: `{e.response.status_code}`"
-        )
+        return webhook.send_olli_error(f"Loki returned a bad status code: `{e.response.status_code}`")
 
     if svc_logs["status"] != "success":
         logger.error(f"Received an error response from Loki for token {token.token}")
-        return webhook.send_olli_error(
-            "Loki returned an error, check your service names are correct."
-        )
+        return webhook.send_olli_error("Loki returned an error, check your service names are correct.")
 
     match = TokenMatch(token, {})
 
@@ -57,7 +53,7 @@ def run() -> None:
     logger.info("Running Olli search")
     matches = []
 
-    for token in CONFIG.olli.tokens:
+    for token in SERVICE_CONFIG.tokens:
         if match := get_match(token):
             matches.append(match)
 

olli/api.py

@@ -4,7 +4,7 @@
 
 import httpx
 
-from olli.config import CONFIG, TokenConfig
+from olli.config import LOKI_CONFIG, SERVICE_CONFIG, TokenConfig
 
 
 class LokiHTTPClient:
@@ -13,7 +13,7 @@ class LokiHTTPClient:
     @staticmethod
     def route(path: str) -> str:
         """Generate the Loki API route for a given path."""
-        return CONFIG.loki.api_url + "/loki/api/v1/" + path
+        return LOKI_CONFIG.api_url + "/loki/api/v1/" + path
 
     def get_token_logs(self, token: TokenConfig) -> dict[str, Any]:
         """
@@ -22,18 +22,18 @@ def get_token_logs(self, token: TokenConfig) -> dict[str, Any]:
         The term is searched case-insensitively in the logs for the interval
         configured in the config.toml file.
         """
-        td = datetime.timedelta(minutes=CONFIG.olli.interval_minutes)
+        td = datetime.timedelta(minutes=SERVICE_CONFIG.interval_minutes)
         start_time = datetime.datetime.now() - td
         start_ts = start_time.timestamp() * 1_000_000_000
 
-        job_regex = "|".join(CONFIG.loki.jobs)
+        job_regex = "|".join(LOKI_CONFIG.jobs)
 
         case_filter = '(?i)' if not token.case_sensitive else ''
 
         resp = httpx.get(self.route("query_range"), params={
             "query": f'{{job=~"({job_regex})"}} |~ "{case_filter}{token.token}"',
             "start": f"{start_ts:0.0f}",
-            "limit": CONFIG.loki.max_logs
+            "limit": LOKI_CONFIG.max_logs,
         })
 
         resp.raise_for_status()

olli/config.py

@@ -1,65 +1,51 @@
 """Configuration loading and validation."""
-import os
-from pathlib import Path
+
 from typing import Optional
 
-import tomllib
-from dotenv import load_dotenv
 from loguru import logger
-from pydantic import BaseModel, validator
-
-load_dotenv()
+from pydantic import BaseModel
+from pydantic_settings import BaseSettings
 
 
-# All the locations where we might find a config file
-CONFIG_PATHS = [
-    "./olli.toml",
-    "/config/olli.toml",
-    "/olli/config.toml",
-    "/etc/olli/config.toml"
-]
+class EnvConfig(
+    BaseSettings,
+    env_file=".env",
+    env_file_encoding="utf-8",
+    env_nested_delimiter="__",
+    extra="ignore",
+):
+    """Our default configuration for models that should load from .env files."""
 
 
-class TokenConfig(BaseModel):
-    """Class representing a token config entry."""
-
-    token: str
-    color: Optional[str] = "#7289DA"
-    case_sensitive: Optional[bool] = False
-
-
-class LokiConfig(BaseModel):
+class _LokiConfig(EnvConfig, env_prefix="loki_"):
     """Loki specific configuration."""
 
     api_url: str
     jobs: list[str]
     max_logs: Optional[int] = 5_000
 
 
-class DiscordConfig(BaseModel):
+LOKI_CONFIG = _LokiConfig()
+
+
+class _DiscordConfig(EnvConfig, env_prefix="discord_"):
     """Configuration for Discord alerting."""
 
     webhook_url: Optional[str]
 
-    @validator("webhook_url", always=True)
-    def env_provided_webhook(cls, value: Optional[str]) -> str:
-        """
-        If no webhook is specified in the config, try fetch from environment.
 
-        If not found in the environment either then raise a validation error.
-        """
-        if value:
-            return value
+DISCORD_CONFIG = _DiscordConfig()
 
-        if webhook := os.environ.get("WEBHOOK_URL"):
-            return webhook
 
-        raise ValueError(
-            "Must specify webhook_url under [discord] or WEBHOOK_URL env var"
-        )
+class TokenConfig(BaseModel):
+    """Class representing a token config entry."""
 
+    token: str
+    color: Optional[str] = "#7289DA"
+    case_sensitive: Optional[bool] = False
 
-class ServiceConfig(BaseModel):
+
+class _ServiceConfig(EnvConfig, env_prefix="service_"):
     """Configuration of the Olli status."""
 
     interval_minutes: int
@@ -81,40 +67,10 @@ def warn_above_ten(cls, value: list[TokenConfig]) -> list[TokenConfig]:
         This is because we cannot handle more than 10 token triggers at once until we
         batch triggers into groups of 10 to distribute to the webhook.
         """
-        if len(value) > 10:
-            logger.warning(
-                "More than 10 token triggers in one period cannot be handled, be careful."
-            )
+        if len(value) > 10:  # noqa: PLR2004
+            logger.warning("More than 10 token triggers in one period cannot be handled, be careful.")
 
         return value
 
 
-class OlliConfig(BaseModel):
-    """Class representing root Olli config."""
-
-    loki: LokiConfig
-    olli: ServiceConfig
-    discord: DiscordConfig = DiscordConfig()
-
-
-def get_config() -> OlliConfig:
-    """Open the config file, parse the TOML and convert to Pydantic objects."""
-    logger.info("Searching for config file")
-
-    path = None
-
-    for file_path in CONFIG_PATHS:
-        if (config_file := Path(file_path)).exists():
-            path = config_file
-            logger.info(f"Found config at {file_path}")
-            break
-
-    if not path:
-        logger.critical("Could not find a config file. Please refer to the documentation.")
-        raise SystemExit(1)
-
-    with open(config_file, "rb") as conf_file:
-        return OlliConfig(**tomllib.load(conf_file))
-
-
-CONFIG = get_config()
+SERVICE_CONFIG = _ServiceConfig()

olli/webhook.py

@@ -5,7 +5,7 @@
 import httpx
 from loguru import logger
 
-from olli.config import CONFIG
+from olli.config import DISCORD_CONFIG, SERVICE_CONFIG
 from olli.structures import TokenMatch
 
 
@@ -26,12 +26,12 @@ def send_with_backoff(url: str, json: dict[str, Any], n: int = 5) -> None:
 def send_olli_error(error: str) -> None:
     """Send an error embed containing the passed error message to Discord."""
     logger.info("Sending error payload to Discord")
-    send_with_backoff(CONFIG.discord.webhook_url, {
+    send_with_backoff(DISCORD_CONFIG.webhook_url, {
         "embeds": [{
             "title": "Olli Error",
             "color": 0xff5f5f,
             "description": f"Olli encountered an error: {error}",
-            "timestamp": datetime.utcnow().isoformat()
+            "timestamp": datetime.utcnow().isoformat(),
         }]
     })
 
@@ -56,7 +56,7 @@ def send_token_matches(matches: list[TokenMatch]) -> None:
                 "name": "Olli"
             },
             "footer": {
-                "text": f"Last {CONFIG.olli.interval_minutes} minutes"
+                "text": f"Last {SERVICE_CONFIG.interval_minutes} minutes",
             },
             "timestamp": datetime.utcnow().isoformat(),
             "fields": []
@@ -73,7 +73,7 @@ def send_token_matches(matches: list[TokenMatch]) -> None:
 
     if len(embeds) > 0:
         logger.info("Sending alerts payload to Discord")
-        send_with_backoff(CONFIG.discord.webhook_url, {
+        send_with_backoff(DISCORD_CONFIG.webhook_url, {
             "embeds": embeds
         })
     else: