From 9eea733169b475013e15fb33ccab1e614428c8fa Mon Sep 17 00:00:00 2001 From: Moetez Fradi Date: Mon, 24 Aug 2026 11:03:54 +0100 Subject: [PATCH] Add Pydantic validation for coordinator Mimir config Type-check existing fields so invalid config fails at build time instead of sending broken data over relations. --- coordinator/src/mimir_config.py | 172 +++++++++++++++++++- coordinator/tests/unit/test_mimir_config.py | 67 +++++++- 2 files changed, 236 insertions(+), 3 deletions(-) diff --git a/coordinator/src/mimir_config.py b/coordinator/src/mimir_config.py index fd0c345..ed5ed0e 100755 --- a/coordinator/src/mimir_config.py +++ b/coordinator/src/mimir_config.py @@ -7,7 +7,7 @@ import logging from enum import Enum, unique from pathlib import Path -from typing import Any, Dict, List, Optional, Set +from typing import Any, Dict, List, Optional, Set, Union from urllib.parse import urlparse import yaml @@ -15,6 +15,7 @@ from coordinated_workers.interfaces.cluster import ClusterProvider from coordinated_workers.worker import CERT_FILE, CLIENT_CA_FILE, KEY_FILE from cosl import JujuTopology +from pydantic import BaseModel, ConfigDict, Field logger = logging.getLogger(__name__) @@ -101,6 +102,172 @@ class MimirRole(str, Enum): # Please visit https://grafana.com/docs/mimir/latest/manage/use-exemplars/store-exemplars/ for more info. EXEMPLARS_FLOOR = 100000 + +class ShardingRing(BaseModel): + """Ring replication settings.""" + + replication_factor: int + + +class Filesystem(BaseModel): + """Local filesystem storage.""" + + dir: str + + +class CommonStorage(BaseModel): + """Shared object-storage settings.""" + + backend: str + # Pass through coordinator S3 keys as-is (names differ by provider/track). + s3: Dict[str, Any] + + +class Common(BaseModel): + """Common Mimir section.""" + + storage: Optional[CommonStorage] = None + + +class Alertmanager(BaseModel): + """Alertmanager schema.""" + + data_dir: str + sharding_ring: ShardingRing + + +class Compactor(BaseModel): + """Compactor schema.""" + + data_dir: str + + +class Frontend(BaseModel): + """Query-frontend schema.""" + + scheduler_address: Optional[str] = None + + +class FrontendWorker(BaseModel): + """Querier frontend_worker schema.""" + + scheduler_address: Optional[str] = None + frontend_address: Optional[str] = None + + +class IngesterRing(BaseModel): + """Ingester ring schema.""" + + replication_factor: int + + +class Ingester(BaseModel): + """Ingester schema.""" + + ring: IngesterRing + + +class Ruler(BaseModel): + """Ruler schema.""" + + rule_path: str + alertmanager_url: str + + +class ComponentStorage(BaseModel): + """Filesystem or S3-prefixed storage for a Mimir component.""" + + filesystem: Optional[Filesystem] = None + storage_prefix: Optional[str] = None + + +class StoreGateway(BaseModel): + """Store-gateway schema.""" + + sharding_ring: ShardingRing + + +class BucketStore(BaseModel): + """Blocks bucket-store schema.""" + + sync_dir: str + + +class Tsdb(BaseModel): + """Ingester TSDB schema.""" + + dir: str + + +class BlocksStorage(BaseModel): + """Blocks storage schema.""" + + bucket_store: Optional[BucketStore] = None + filesystem: Optional[Filesystem] = None + tsdb: Optional[Tsdb] = None + storage_prefix: Optional[str] = None + + +class Memberlist(BaseModel): + """Memberlist schema.""" + + cluster_label: str + join_members: List[str] + + +class Limits(BaseModel): + """Limits schema.""" + + ruler_max_rules_per_rule_group: int + ruler_max_rule_groups_per_tenant: int + max_global_series_per_user: int + ingestion_rate: int + ingestion_burst_size: int + max_global_exemplars_per_user: int + # Charm emits 0 as int, or a duration string such as "1w". + compactor_blocks_retention_period: Union[int, str] + + +class HttpTlsConfig(BaseModel): + """HTTP TLS schema.""" + + cert_file: str + key_file: str + client_ca_file: str + client_auth_type: str + + +class Server(BaseModel): + """Server schema.""" + + http_tls_config: HttpTlsConfig + + +class MimirConfigModel(BaseModel): + """Minimal Mimir config schema for fields this charm emits. + + Unknown keys are ignored so upgrade paths (e.g. 2.17 -> dev) do not fail + when a newer track adds a field this model does not know about. + """ + + model_config = ConfigDict(extra="ignore") + + common: Common = Field(default_factory=Common) + alertmanager: Alertmanager + alertmanager_storage: ComponentStorage + compactor: Compactor + frontend: Frontend = Field(default_factory=Frontend) + frontend_worker: FrontendWorker = Field(default_factory=FrontendWorker) + ingester: Ingester + ruler: Ruler + ruler_storage: ComponentStorage + store_gateway: StoreGateway + blocks_storage: BlocksStorage + memberlist: Memberlist + limits: Limits + server: Optional[Server] = None + + class MimirConfig: """Config builder for the Mimir Coordinator.""" @@ -156,7 +323,8 @@ def config(self, coordinator: Coordinator) -> str: if coordinator.nginx.are_certificates_on_disk: mimir_config["server"] = self._build_tls_config() - return yaml.dump(mimir_config) + validated = MimirConfigModel.model_validate(mimir_config) + return yaml.dump(validated.model_dump(mode="json", by_alias=True, exclude_none=True)) def _build_tls_config(self) -> Dict[str, Any]: tls_config = { diff --git a/coordinator/tests/unit/test_mimir_config.py b/coordinator/tests/unit/test_mimir_config.py index 8b97987..7dc0f37 100644 --- a/coordinator/tests/unit/test_mimir_config.py +++ b/coordinator/tests/unit/test_mimir_config.py @@ -2,9 +2,11 @@ from unittest.mock import MagicMock import pytest +import yaml from deepdiff import DeepDiff +from pydantic import ValidationError -from src.mimir_config import MimirConfig +from src.mimir_config import Alertmanager, MimirConfig, MimirConfigModel @pytest.fixture(scope="module") @@ -364,5 +366,68 @@ def test_retention_period_logic(mimir_config, retention_period_config, expected_ # Assert that the value for compactor_blocks_retention_period matches the expected value assert limits_config["compactor_blocks_retention_period"] == expected_value + +def _load_built_config(mimir_config, coordinator): + return yaml.safe_load(mimir_config.config(coordinator)) + + +def test_config_round_trip_without_s3_or_tls(mimir_config, coordinator): + coordinator.s3_ready = False + coordinator.nginx.are_certificates_on_disk = False + + raw = _load_built_config(mimir_config, coordinator) + parsed = MimirConfigModel.model_validate(raw) + + assert parsed.common.storage is None + assert parsed.server is None + assert "server" not in raw + assert raw["common"] == {} + + +def test_config_model_ignores_unknown_fields(mimir_config, coordinator): + coordinator.s3_ready = False + coordinator.nginx.are_certificates_on_disk = False + + raw = _load_built_config(mimir_config, coordinator) + raw["not_a_real_mimir_key"] = 1 + raw["limits"]["also_unknown"] = True + + parsed = MimirConfigModel.model_validate(raw) + dumped = parsed.model_dump(exclude_none=True) + assert "not_a_real_mimir_key" not in dumped + assert "also_unknown" not in dumped["limits"] + + +def test_config_with_s3_and_tls_validates(mimir_config, coordinator): + coordinator.s3_ready = True + coordinator.nginx.are_certificates_on_disk = True + coordinator._s3_config = { + "endpoint": "s3.com:port", + "access_key_id": "your_access_key", + "secret_access_key": "your_secret_key", + "bucket_name": "your_bucket", + "region": "your_region", + "insecure": "true", + } + + raw = _load_built_config(mimir_config, coordinator) + parsed = MimirConfigModel.model_validate(raw) + + assert parsed.common.storage is not None + assert parsed.common.storage.backend == "s3" + assert parsed.server is not None + assert parsed.blocks_storage.storage_prefix == "blocks" + assert parsed.blocks_storage.filesystem is None + assert parsed.ruler_storage.storage_prefix == "rules" + assert parsed.alertmanager_storage.storage_prefix == "alerts" + + +def test_config_model_rejects_invalid_types(): + with pytest.raises(ValidationError): + Alertmanager.model_validate( + {"data_dir": "/data/data-alertmanager", "sharding_ring": {"replication_factor": "nope"}} + ) + + if __name__ == "__main__": unittest.main()