Add remote write endpoints to relation data (#64)

* remote write endpoints

* refactor comment

* add test

* reorder

* add TypedDict

pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "hatchling.build"
 
 [project]
 name = "cosl"
-version = "0.0.25"
+version = "0.0.26"
 authors = [
   { name="sed-i", email="82407168+sed-i@users.noreply.github.com" },
 ]

src/cosl/coordinated_workers/coordinator.py

@@ -29,7 +29,7 @@
 import yaml
 
 import cosl
-from cosl.coordinated_workers.interface import ClusterProvider
+from cosl.coordinated_workers.interface import ClusterProvider, RemoteWriteEndpoint
 from cosl.coordinated_workers.nginx import (
     Nginx,
     NginxMappingOverrides,
@@ -175,6 +175,7 @@ def __init__(
         resources_limit_options: Optional[_ResourceLimitOptionsMapping] = None,
         resources_requests: Optional[Callable[["Coordinator"], Dict[str, str]]] = None,
         container_name: Optional[str] = None,
+        remote_write_endpoints: Optional[Callable[[], List[RemoteWriteEndpoint]]] = None,
     ):
         """Constructor for a Coordinator object.
 
@@ -200,6 +201,8 @@ def __init__(
                 their respective config options in config.yaml.
             container_name: The container for which to apply the resources requests & limits.
                 Required if `resources_requests` is provided.
+            remote_write_endpoints: A function generating endpoints to which the workload
+                and the worker charm can push metrics to.
 
         Raises:
         ValueError:
@@ -231,6 +234,7 @@ def __init__(
         )
         self._container_name = container_name
         self._resources_limit_options = resources_limit_options or {}
+        self.remote_write_endpoints_getter = remote_write_endpoints
 
         self.nginx = Nginx(
             self._charm,
@@ -601,6 +605,9 @@ def _on_config_changed(self, _: ops.ConfigChangedEvent):
     def _on_collect_unit_status(self, e: ops.CollectStatusEvent):
         # todo add [nginx.workload] statuses
 
+        if self.resources_patch and self.resources_patch.get_status().name != "active":
+            e.add_status(self.resources_patch.get_status())
+
         if not self.cluster.has_workers:
             e.add_status(ops.BlockedStatus("[consistency] Missing any worker relation."))
         if not self.is_coherent:
@@ -613,9 +620,6 @@ def _on_collect_unit_status(self, e: ops.CollectStatusEvent):
         else:
             e.add_status(ops.ActiveStatus())
 
-        if self.resources_patch:
-            e.add_status(self.resources_patch.get_status())
-
     ###################
     # UTILITY METHODS #
     ###################
@@ -681,6 +685,11 @@ def update_cluster(self):
             tracing_receivers=(
                 self._tracing_receivers_getter() if self._tracing_receivers_getter else None
             ),
+            remote_write_endpoints=(
+                self.remote_write_endpoints_getter()
+                if self.remote_write_endpoints_getter
+                else None
+            ),
         )
 
     def _render_workers_alert_rules(self):

src/cosl/coordinated_workers/interface.py

@@ -31,6 +31,7 @@
 import yaml
 from ops import EventSource, Object, ObjectEvents, RelationCreatedEvent
 from pydantic import ConfigDict
+from typing_extensions import TypedDict
 
 import cosl
 
@@ -122,6 +123,12 @@ def dump(self, databag: Optional[_RawDatabag] = None, clear: bool = True) -> _Ra
 # =============
 
 
+class RemoteWriteEndpoint(TypedDict):
+    """Type of the remote write endpoints to be passed to the worker through cluster relation data."""
+
+    url: str
+
+
 class ConfigReceivedEvent(ops.EventBase):
     """Event emitted when the "-cluster" provider has shared a new  config."""
 
@@ -188,6 +195,8 @@ class ClusterProviderAppData(DatabagModel):
     """Endpoints to which the workload (and the worker charm) can push logs to."""
     tracing_receivers: Optional[Dict[str, str]] = None
     """Endpoints to which the workload (and the worker charm) can push traces to."""
+    remote_write_endpoints: Optional[List[RemoteWriteEndpoint]] = None
+    """Endpoints to which the workload (and the worker charm) can push metrics to."""
 
     ### TLS stuff
     ca_cert: Optional[str] = None
@@ -275,6 +284,7 @@ def publish_data(
         privkey_secret_id: Optional[str] = None,
         loki_endpoints: Optional[Dict[str, str]] = None,
         tracing_receivers: Optional[Dict[str, str]] = None,
+        remote_write_endpoints: Optional[List[RemoteWriteEndpoint]] = None,
     ) -> None:
         """Publish the config to all related worker clusters."""
         for relation in self._relations:
@@ -286,6 +296,7 @@ def publish_data(
                     server_cert=server_cert,
                     privkey_secret_id=privkey_secret_id,
                     tracing_receivers=tracing_receivers,
+                    remote_write_endpoints=remote_write_endpoints,
                 )
                 local_app_databag.dump(relation.data[self.model.app])
 
@@ -540,3 +551,10 @@ def get_tracing_receivers(self) -> Optional[Dict[str, str]]:
         if data:
             return data.tracing_receivers or {}
         return {}
+
+    def get_remote_write_endpoints(self) -> List[RemoteWriteEndpoint]:
+        """Fetch the remote write endpoints from the coordinator databag."""
+        data = self._get_data_from_coordinator()
+        if data:
+            return data.remote_write_endpoints or []
+        return []

src/cosl/coordinated_workers/worker.py

@@ -287,6 +287,9 @@ def status(self) -> ServiceEndpointStatus:
             return ServiceEndpointStatus.down
 
     def _on_collect_status(self, e: ops.CollectStatusEvent):
+        if self.resources_patch and self.resources_patch.get_status().name != "active":
+            e.add_status(self.resources_patch.get_status())
+
         if not self._container.can_connect():
             e.add_status(WaitingStatus(f"Waiting for `{self._name}` container"))
         if not self.model.get_relation(self._endpoints["cluster"]):
@@ -317,9 +320,6 @@ def _on_collect_status(self, e: ops.CollectStatusEvent):
             )
         )
 
-        if self.resources_patch:
-            e.add_status(self.resources_patch.get_status())
-
     # Utility functions
     @property
     def roles(self) -> List[str]:

tests/test_coordinated_workers/test_worker.py

@@ -1,11 +1,12 @@
+import json
 from unittest.mock import MagicMock, patch
 
 import ops
 import pytest
 import tenacity
 from ops import Framework
 from ops.pebble import Layer, ServiceStatus
-from scenario import Container, Context, Mount, State
+from scenario import Container, Context, Mount, Relation, State
 from scenario.runtime import UncaughtCharmError
 
 from cosl.coordinated_workers.worker import CONFIG_FILE, Worker
@@ -274,3 +275,51 @@ def raise_change_error(*args):
                 # so we don't have to wait for minutes:
                 mgr.charm.worker.SERVICE_START_RETRY_WAIT = tenacity.wait_none()
                 mgr.charm.worker.SERVICE_START_RETRY_STOP = tenacity.stop_after_delay(2)
+
+
+@pytest.mark.parametrize(
+    "remote_databag, expected",
+    (
+        (
+            {
+                "remote_write_endpoints": json.dumps([{"url": "test-url.com"}]),
+                "worker_config": json.dumps("test"),
+            },
+            [{"url": "test-url.com"}],
+        ),
+        ({"remote_write_endpoints": json.dumps(None), "worker_config": json.dumps("test")}, []),
+        (
+            {
+                "remote_write_endpoints": json.dumps(
+                    [{"url": "test-url.com"}, {"url": "test2-url.com"}]
+                ),
+                "worker_config": json.dumps("test"),
+            },
+            [{"url": "test-url.com"}, {"url": "test2-url.com"}],
+        ),
+    ),
+)
+def test_get_remote_write_endpoints(remote_databag, expected):
+    ctx = Context(
+        MyCharm,
+        meta={
+            "name": "foo",
+            "requires": {"cluster": {"interface": "cluster"}},
+            "containers": {"foo": {"type": "oci-image"}},
+        },
+        config={"options": {"role-all": {"type": "boolean", "default": True}}},
+    )
+    container = Container(
+        "foo",
+        can_connect=True,
+    )
+    relation = Relation(
+        "cluster",
+        remote_app_data=remote_databag,
+    )
+    with ctx.manager(
+        relation.changed_event, State(containers=[container], relations=[relation])
+    ) as mgr:
+        charm = mgr.charm
+        mgr.run()
+        assert charm.worker.cluster.get_remote_write_endpoints() == expected