Add support to validate and load configuration from a schema

When passing a `dataclass` as a schema, the receiver will load the *raw*
data from the dict and validate it. If the validation passed, the
`dataclass` will be constructed with the validated data (using
`load_config()`, so using `marshmallow`).

If a configuration is invalid, an error will be logged and the update
will be skipped (not sent to the receiver).

Signed-off-by: Leandro Lucarella <[email protected]>

Author: llucax

src/frequenz/sdk/config/_manager.py

@@ -8,12 +8,14 @@
 import pathlib
 from collections.abc import Mapping, MutableMapping, Sequence, Set
 from datetime import timedelta
-from typing import Any, assert_never, overload
+from typing import Any, TypeGuard, assert_never, cast, overload
 
 from frequenz.channels import Broadcast, Receiver
 from frequenz.channels.file_watcher import EventType
+from marshmallow import Schema, ValidationError
 
 from ._actor import ConfigManagingActor
+from ._util import DataclassT, load_config
 
 _logger = logging.getLogger(__name__)
 
@@ -128,6 +130,21 @@ async def new_receiver(
         skip_unchanged: bool = True,
     ) -> Receiver[Mapping[str, Any]]: ...
 
+    @overload
+    async def new_receiver(  # pylint: disable=too-many-arguments
+        self,
+        *,
+        wait_for_first: bool = True,
+        skip_unchanged: bool = True,
+        # We need to specify the key here because we have kwargs, so if it is not
+        # present is not considered None as the only possible value, as any value can be
+        # accepted as part of the kwargs.
+        key: None = None,
+        schema: type[DataclassT],
+        base_schema: type[Schema] | None = None,
+        **marshmallow_load_kwargs: Any,
+    ) -> Receiver[DataclassT]: ...
+
     @overload
     async def new_receiver(
         self,
@@ -137,14 +154,29 @@ async def new_receiver(
         key: str,
     ) -> Receiver[Mapping[str, Any] | None]: ...
 
+    @overload
+    async def new_receiver(  # pylint: disable=too-many-arguments
+        self,
+        *,
+        wait_for_first: bool = True,
+        skip_unchanged: bool = True,
+        key: str,
+        schema: type[DataclassT],
+        base_schema: type[Schema] | None,
+        **marshmallow_load_kwargs: Any,
+    ) -> Receiver[DataclassT | None]: ...
+
     # The noqa DOC502 is needed because we raise TimeoutError indirectly.
     async def new_receiver(  # pylint: disable=too-many-arguments # noqa: DOC502
         self,
         *,
         wait_for_first: bool = False,
         skip_unchanged: bool = True,
         key: str | None = None,
-    ) -> Receiver[Mapping[str, Any] | None]:
+        schema: type[DataclassT] | None = None,
+        base_schema: type[Schema] | None = None,
+        **marshmallow_load_kwargs: Any,
+    ) -> Receiver[Mapping[str, Any] | DataclassT | None]:
         """Create a new receiver for the configuration.
 
         This method has a lot of features and functionalities to make it easier to
@@ -168,6 +200,26 @@ async def new_receiver(  # pylint: disable=too-many-arguments # noqa: DOC502
         configuration under the specified key is received, or `None` if the key is not
         found.
 
+        ### Schema validation
+
+        The configuration is received as a dictionary unless a `schema` is provided. In
+        this case, the configuration will be validated against the schema and received
+        as an instance of the configuration class.
+
+        The configuration `schema` class is expected to be
+        a [`dataclasses.dataclass`][], which is used to create
+        a [`marshmallow.Schema`][] schema to validate the configuration dictionary.
+
+        To customize the schema derived from the configuration dataclass, you can
+        use [`marshmallow_dataclass.dataclass`][] to specify extra metadata.
+
+        Configurations that don't pass the validation will be ignored and not sent to
+        the receiver, but an error will be logged. Errors other than `ValidationError`
+        will not be handled and will be raised.
+
+        Additional arguments can be passed to [`marshmallow.Schema.load`][] using keyword
+        arguments.
+
         ### Waiting for the first configuration
 
         If `wait_for_first` is `True`, the receiver will wait for the first
@@ -196,6 +248,13 @@ async def new_receiver(  # pylint: disable=too-many-arguments # noqa: DOC502
             key: The key to filter the configuration. A nested key can be specified by
                 using a dot (`.`) as separator. For example "key.subkey" will get only
                 `config[key][subkey]` If `None`, all configurations will be received.
+            schema: The type of the configuration. If provided, the configuration
+                will be validated against this type.
+            base_schema: An optional class to be used as a base schema for the
+                configuration class. This allow using custom fields for example. Will be
+                passed to [`marshmallow_dataclass.class_schema`][].
+            **marshmallow_load_kwargs: Additional arguments to be passed to
+                [`marshmallow.Schema.load`][].
 
         Returns:
             The receiver for the configuration.
@@ -214,8 +273,36 @@ async def new_receiver(  # pylint: disable=too-many-arguments # noqa: DOC502
                 lambda config: self._config_has_changed(config, key)
             )
 
-        if key is not None:
-            receiver = receiver.map(lambda config: config.get(key))
+        match (key, schema):
+            case (None, None):
+                pass
+            case (None, schema):
+                # This cast can be removed after frequenz-channels is upgraded to v1.3.0
+                receiver = cast(
+                    Receiver[DataclassT],
+                    receiver.map(
+                        lambda config: _load_config_with_logging(
+                            config,
+                            schema,
+                            base_schema=base_schema,
+                            **marshmallow_load_kwargs,
+                        )
+                    ).filter(_is_valid_and_not_none),
+                )
+            case (key, None):
+                receiver = receiver.map(lambda config: config.get(key))
+            case (key, schema):
+                # This cast can be removed after frequenz-channels is upgraded to v1.3.0
+                receiver = cast(
+                    Receiver[DataclassT | None],
+                    receiver.map(
+                        lambda config: load_config(
+                            schema, config.get(key, {}), **marshmallow_load_kwargs
+                        )
+                    ).filter(_is_valid_or_none),
+                )
+            case unexpected:
+                assert_never(unexpected)
 
         if wait_for_first:
             async with asyncio.timeout(self.wait_for_first_timeout.total_seconds()):
@@ -254,3 +341,78 @@ def _config_has_changed(
             _logger.debug("Old configuration%s being kept: %r", key_str, old_config)
 
         return has_changed
+
+
+class _InvalidConfig:
+    """A sentinel to represent an invalid configuration."""
+
+
+_INVALID_CONFIG = _InvalidConfig()
+"""A sentinel singleton instance to represent an invalid configuration."""
+
+
+@overload
+def _load_config_with_logging(
+    config: Mapping[str, Any],
+    schema: type[DataclassT],
+    key: str,
+    base_schema: type[Schema] | None = None,
+    **marshmallow_load_kwargs: Any,
+) -> DataclassT | None | _InvalidConfig: ...
+
+
+@overload
+def _load_config_with_logging(
+    config: Mapping[str, Any],
+    schema: type[DataclassT],
+    key: None = None,
+    base_schema: type[Schema] | None = None,
+    **marshmallow_load_kwargs: Any,
+) -> DataclassT | _InvalidConfig: ...
+
+
+def _load_config_with_logging(
+    config: Mapping[str, Any],
+    schema: type[DataclassT],
+    key: str | None = None,
+    base_schema: type[Schema] | None = None,
+    **marshmallow_load_kwargs: Any,
+) -> DataclassT | None | _InvalidConfig:
+    """Try to load a configuration and log any validation errors."""
+    if key is not None:
+        maybe_config = config.get(key, None)
+        if maybe_config is None:
+            _logger.debug(
+                "Configuration key %s not found, sending None: config=%r", key, config
+            )
+            return None
+        config = maybe_config
+
+    try:
+        return load_config(
+            schema, config, base_schema=base_schema, **marshmallow_load_kwargs
+        )
+    except ValidationError as err:
+        key_str = ""
+        if key:
+            key_str = f" for key '{key}'"
+        _logger.error(
+            "The configuration%s is invalid, the configuration update will be skipped: %s",
+            key_str,
+            err,
+        )
+        return _INVALID_CONFIG
+
+
+def _is_valid_or_none(
+    config: DataclassT | _InvalidConfig | None,
+) -> TypeGuard[DataclassT | None]:
+    """Return whether the configuration is valid or `None`."""
+    return config is not _INVALID_CONFIG
+
+
+def _is_valid_and_not_none(
+    config: DataclassT | _InvalidConfig | None,
+) -> TypeGuard[DataclassT]:
+    """Return whether the configuration is valid and not `None`."""
+    return config is not None and config is not _INVALID_CONFIG