generic coordinator and worker classes

.gitignore

@@ -129,3 +129,6 @@ dmypy.json
 .pyre/
 cos-tool-*
 .idea
+
+# charm libs used by the coordinator charm
+lib/

pyproject.toml

@@ -13,6 +13,7 @@ readme = "README.md"
 requires-python = ">=3.8"
 dependencies = [
   "ops",
+  "pydantic",
   "PyYAML",
   "typing-extensions"
 ]
@@ -49,31 +50,31 @@ target-version = ["py38"]
 [tool.isort]
 profile = "black"
 
-# Linting tools configuration
-[tool.flake8]
-max-line-length = 99
-max-doc-length = 99
-max-complexity = 10
-exclude = [".git", "__pycache__", ".tox", "build", "dist", "*.egg_info", "venv"]
-select = ["E", "W", "F", "C", "N", "R", "D", "H"]
-# Ignore W503, E501 because using black creates errors with this
-# Ignore D107 Missing docstring in __init__
-ignore = ["W503", "E501", "D107"]
-# D100, D101, D102, D103: Ignore missing docstrings in tests
-per-file-ignores = ["tests/*:D100,D101,D102,D103"]
-docstring-convention = "google"
-# Check for properly formatted copyright header in each file
-copyright-check = "True"
-copyright-author = "Canonical Ltd."
-copyright-regexp = "Copyright\\s\\d{4}([-,]\\d{4})*\\s+%(author)s"
-
 [tool.pyright]
+include = ["src"]
+extraPaths = ["lib", "src/cosl"]
 pythonPlatform = "All"
 typeCheckingMode = "strict"
 reportIncompatibleMethodOverride = false
 reportImportCycles = false
 reportTypeCommentUsage = false
 
+[tool.ruff]
+line-length = 99
+extend-exclude = ["__pycache__", "*.egg_info"]
+
+[tool.ruff.lint]
+select = ["E", "W", "F", "C", "N", "R", "D", "I001"]
+# Ignore E501 because using black creates errors with this
+# Ignore D107 Missing docstring in __init__
+# Ignore E402 because charm libraries can't be imported at the top
+ignore = ["E501", "D107", "RET504", "C901", "E402"]
+# D100, D101, D102, D103: Ignore missing docstrings in tests
+per-file-ignores = {"tests/*" = ["D100","D101","D102","D103"]}
+
+[tool.ruff.lint.pydocstyle]
+convention = "google"
+
 # Static analysis tools configuration
 [tool.mypy]
 pretty = true
@@ -92,4 +93,4 @@ no_implicit_optional = false
 
 [tool.codespell]
 skip = ".git,.tox,build,lib,venv*,.mypy_cache"
-ignore-words-list = "assertIn"
+ignore-words-list = "assertIn"

src/cosl/__init__.py

@@ -4,6 +4,9 @@
 """Utils for observability Juju charms."""
 
 from .cos_tool import CosTool
+from .distributed.cluster import ClusterProvider, ClusterRequirer
+from .distributed.coordinator import Coordinator
+from .distributed.worker import Worker
 from .grafana_dashboard import GrafanaDashboard
 from .juju_topology import JujuTopology
 from .mandatory_relation_pairs import MandatoryRelationPairs
@@ -16,4 +19,8 @@
     "AlertRules",
     "RecordingRules",
     "MandatoryRelationPairs",
+    "Coordinator",
+    "ClusterProvider",
+    "ClusterRequirer",
+    "Worker",
 ]

src/cosl/databag_model.py

@@ -0,0 +1,150 @@
+#!/usr/bin/env python3
+# Copyright 2024 Canonical
+# See LICENSE file for licensing details.
+
+"""DatabagModel implementation from traefik.v1.ingress charm lib."""
+
+import json
+import logging
+from typing import MutableMapping, Optional, cast
+
+import pydantic
+
+logger = logging.getLogger(__name__)
+
+
+_RawDatabag = MutableMapping[str, str]
+
+class DataValidationError(Exception):
+    """Raised when relation databag validation fails."""
+
+
+PYDANTIC_IS_V1 = int(pydantic.version.VERSION.split(".")[0]) < 2
+if PYDANTIC_IS_V1:
+
+    class DatabagModel(pydantic.BaseModel):  # type: ignore
+        """Base databag model."""
+
+        class Config:
+            """Pydantic config."""
+
+            allow_population_by_field_name = True
+            """Allow instantiating this class by field name (instead of forcing alias)."""
+
+        _NEST_UNDER = None
+
+        @classmethod
+        def load(cls, databag: _RawDatabag) -> "DatabagModel":
+            """Load this model from a Juju databag."""
+            if cls._NEST_UNDER:
+                return cast(DatabagModel, cls.parse_obj(json.loads(databag[cls._NEST_UNDER]))) # type: ignore
+
+            try:
+                data = {
+                    k: json.loads(v)
+                    for k, v in databag.items()
+                    # Don't attempt to parse model-external values
+                    if k in {f.alias for f in cls.__fields__.values()}  # type: ignore
+                }
+            except json.JSONDecodeError as e:
+                msg = f"invalid databag contents: expecting json. {databag}"
+                logger.error(msg)
+                raise DataValidationError(msg) from e
+
+            try:
+                return cls.parse_raw(json.dumps(data))  # type: ignore
+            except pydantic.ValidationError as e:
+                msg = f"failed to validate databag: {databag}"
+                logger.debug(msg, exc_info=True)
+                raise DataValidationError(msg) from e
+
+        def dump(self, databag: Optional[_RawDatabag] = None, clear: bool = True):
+            """Write the contents of this model to Juju databag.
+
+            :param databag: the databag to write the data to.
+            :param clear: ensure the databag is cleared before writing it.
+            """
+            if clear and databag:
+                databag.clear()
+
+            if databag is None:
+                databag = {}
+
+            if self._NEST_UNDER:
+                databag[self._NEST_UNDER] = self.json(by_alias=True, exclude_defaults=True) # type: ignore
+                return databag
+
+            for key, value in self.dict(by_alias=True, exclude_defaults=True).items():  # type: ignore
+                databag[key] = json.dumps(value)
+
+            return databag
+
+else:
+    from pydantic import ConfigDict
+
+    class DatabagModel(pydantic.BaseModel):
+        """Base databag model."""
+
+        model_config = ConfigDict(
+            # tolerate additional keys in databag
+            extra="ignore",
+            # Allow instantiating this class by field name (instead of forcing alias).
+            populate_by_name=True,
+            # Custom config key: whether to nest the whole datastructure (as json)
+            # under a field or spread it out at the toplevel.
+            _NEST_UNDER=None,
+            # TODO: Check if this is necessary / good to have
+            # Protected namespaces: will warn if keys starts with those
+            protected_namespaces=(),
+        )  # type: ignore
+        """Pydantic config."""
+
+        @classmethod
+        def load(cls, databag: _RawDatabag):
+            """Load this model from a Juju databag."""
+            nest_under = cls.model_config.get("_NEST_UNDER")
+            if nest_under:
+                return cls.model_validate(json.loads(databag[nest_under]))  # type: ignore
+
+            try:
+                data = {
+                    k: json.loads(v)
+                    for k, v in databag.items()
+                    # Don't attempt to parse model-external values
+                    if k in {(f.alias or n) for n, f in cls.__fields__.items()}  # type: ignore
+                }
+            except json.JSONDecodeError as e:
+                msg = f"invalid databag contents: expecting json. {databag}"
+                logger.error(msg)
+                raise DataValidationError(msg) from e
+
+            try:
+                return cls.model_validate_json(json.dumps(data))  # type: ignore
+            except pydantic.ValidationError as e:
+                msg = f"failed to validate databag: {databag}"
+                logger.debug(msg, exc_info=True)
+                raise DataValidationError(msg) from e
+
+        def dump(self, databag: Optional[_RawDatabag] = None, clear: bool = True):
+            """Write the contents of this model to Juju databag.
+
+            :param databag: the databag to write the data to.
+            :param clear: ensure the databag is cleared before writing it.
+            """
+            if clear and databag:
+                databag.clear()
+
+            if databag is None:
+                databag = {}
+            nest_under = self.model_config.get("_NEST_UNDER")
+            if nest_under:
+                databag[nest_under] = self.model_dump_json(  # type: ignore
+                    by_alias=True,
+                    # skip keys whose values are default
+                    exclude_defaults=True,
+                )
+                return databag
+
+            dct = self.model_dump(mode="json", by_alias=True, exclude_defaults=True)  # type: ignore
+            databag.update({k: json.dumps(v) for k, v in dct.items()})
+            return databag