grafana source uid added to grafana-source interface (#359)

* type fixes and grafana source uid added to grafana-source interface

* added test for provider

* fmt

* updated return object type

* lint

* update ds names on relation removal

* fix

* better comment

---------

Co-authored-by: michael <[email protected]>

lib/charms/grafana_k8s/v0/grafana_source.py

@@ -162,7 +162,7 @@ def __init__(self, *args):
 
 # Increment this PATCH version before using `charmcraft publish-lib` or reset
 # to 0 if you are raising the major API version
-LIBPATCH = 21
+LIBPATCH = 22
 
 logger = logging.getLogger(__name__)
 
@@ -429,6 +429,18 @@ def update_source(self, source_url: Optional[str] = ""):
                 continue
             self._set_sources(rel)
 
+    def get_source_uids(self) -> Dict[str, Dict[str, str]]:
+        """Get the datasource UID(s) assigned by the remote end(s) to this datasource.
+
+        Returns a mapping from remote application names to unit names to datasource uids.
+        """
+        uids = {}
+        for rel in self._charm.model.relations.get(self._relation_name, []):
+            if not rel:
+                continue
+            uids[rel.app.name] = json.loads(rel.data[rel.app]["datasource_uids"])
+        return uids
+
     def _set_sources_from_event(self, event: RelationJoinedEvent) -> None:
         """Get a `Relation` object from the event to pass on."""
         self._set_sources(event.relation)
@@ -551,6 +563,13 @@ def _on_grafana_peer_changed(self, _: RelationChangedEvent) -> None:
         self.on.sources_changed.emit()  # pyright: ignore
         self.on.sources_to_delete_changed.emit()  # pyright: ignore
 
+    def _publish_source_uids(self, rel: Relation, uids: Dict[str, str]):
+        """Share the datasource UIDs back to the datasources.
+
+        Assumes only leader unit will call this method
+        """
+        rel.data[self._charm.app]["datasource_uids"] = json.dumps(uids)
+
     def _get_source_config(self, rel: Relation):
         """Generate configuration from data stored in relation data by providers."""
         source_data = json.loads(rel.data[rel.app].get("grafana_source_data", "{}"))  # type: ignore
@@ -588,6 +607,10 @@ def _get_source_config(self, rel: Relation):
                 sources_to_delete.remove(host_data["source_name"])
 
             data.append(host_data)
+
+        # share the unique source names back to the datasource units
+        self._publish_source_uids(rel, {ds["unit"]: ds["source_name"] for ds in data})
+
         self.set_peer_data("sources_to_delete", list(sources_to_delete))
         return data
 
@@ -653,6 +676,12 @@ def _remove_source_from_datastore(self, event: RelationDepartedEvent) -> bool:
                     self._remove_source(host["source_name"])
 
             self.set_peer_data("sources", stored_sources)
+
+            # update this relation's shared datasource names after removing this unit/source
+            self._publish_source_uids(
+                event.relation, {ds["unit"]: ds["source_name"] for ds in removed_source}
+            )
+
             return True
         return False
 

src/charm.py

@@ -76,7 +76,7 @@
 from charms.tempo_coordinator_k8s.v0.tracing import TracingEndpointRequirer, charm_tracing_config
 from ops.framework import StoredState
 from ops.main import main
-from ops.model import ActiveStatus, BlockedStatus, MaintenanceStatus, OpenedPort
+from ops.model import ActiveStatus, BlockedStatus, MaintenanceStatus, Port
 
 from ops.pebble import (
     APIError,
@@ -669,7 +669,7 @@ def _update_dashboards(self, event) -> None:
 
     def set_ports(self):
         """Open necessary (and close no longer needed) workload ports."""
-        planned_ports = {OpenedPort("tcp", PORT)} if self.unit.is_leader() else set()
+        planned_ports = {Port(protocol="tcp", port=PORT)} if self.unit.is_leader() else set()
         actual_ports = self.unit.opened_ports()
 
         # Ports may change across an upgrade, so need to sync
@@ -957,12 +957,7 @@ def restart_grafana(self) -> None:
 
             if self._poll_container(self.containers["workload"].can_connect):
                 # We should also make sure sqlite is in WAL mode for replication
-                self.containers["workload"].push(
-                    "/usr/local/bin/sqlite3",
-                    Path("sqlite-static").read_bytes(),
-                    permissions=0o755,
-                    make_dirs=True,
-                )
+                self._push_sqlite_static()
 
                 pragma = self.containers["workload"].exec(
                     [
@@ -1321,7 +1316,9 @@ def _get_admin_password(self) -> str:
 
         return self._stored.admin_password  # type: ignore
 
-    def _poll_container(self, func: Callable, timeout: float = 2.0, delay: float = 0.1) -> bool:
+    def _poll_container(
+        self, func: Callable[[], bool], timeout: float = 2.0, delay: float = 0.1
+    ) -> bool:
         """Try to poll the container to work around Container.is_connect() being point-in-time.
 
         Args:
@@ -1582,6 +1579,15 @@ def _on_oauth_info_changed(self, event: OAuthInfoChangedEvent) -> None:
         """Event handler for the oauth_info_changed event."""
         self._configure()
 
+    def _push_sqlite_static(self):
+        # for ease of mocking in unittests, this is a standalone function
+        self.containers["workload"].push(
+            "/usr/local/bin/sqlite3",
+            Path("sqlite-static").read_bytes(),
+            permissions=0o755,
+            make_dirs=True,
+        )
+
 
 if __name__ == "__main__":
     main(GrafanaCharm, use_juju_for_storage=True)

tests/scenario/conftest.py

@@ -13,6 +13,7 @@ def tautology(*_, **__) -> bool:
 @pytest.fixture
 def ctx():
     patches = (
+        patch("charm.GrafanaCharm._push_sqlite_static", new=lambda _: None),
         patch("lightkube.core.client.GenericSyncClient"),
         patch("socket.getfqdn", new=lambda *args: "grafana-k8s-0.testmodel.svc.cluster.local"),
         patch("socket.gethostbyname", new=lambda *args: "1.2.3.4"),

tests/scenario/test_datasources.py

@@ -0,0 +1,83 @@
+import json
+
+from ops import CharmBase, Framework
+from ops.testing import Container, State
+from scenario import Relation, PeerRelation, Context
+
+from charms.grafana_k8s.v0.grafana_source import GrafanaSourceProvider
+
+containers = [
+    Container(name="grafana", can_connect=True),
+    Container(name="litestream", can_connect=True),
+]
+
+
+def test_datasource_sharing(ctx):
+    # GIVEN a datasource relation with two remote units
+    datasource = Relation(
+        "grafana-source",
+        remote_app_name="remote_host",
+        remote_units_data={
+            0: {"grafana_source_host": "remote_host.0"},
+            1: {"grafana_source_host": "remote_host.1"},
+        },
+        remote_app_data={
+            "grafana_source_data": json.dumps(
+                {"model": "foo", "model_uuid": "bar", "application": "baz", "type": "tempo"}
+            )
+        },
+    )
+    state = State(
+        leader=True, containers=containers, relations={datasource, PeerRelation("grafana")}
+    )
+
+    # WHEN relation-changed fires for a datasource relation
+    out = ctx.run(ctx.on.relation_changed(datasource), state)
+
+    # THEN grafana shares back over the same relation a mapping of datasource uids
+    datasource_out = out.get_relation(datasource.id)
+    local_app_data = datasource_out.local_app_data
+    ds_uids = json.loads(local_app_data["datasource_uids"])
+    assert ds_uids == {
+        "remote_host/0": "juju_foo_bar_baz_0",
+        "remote_host/1": "juju_foo_bar_baz_1",
+    }
+
+
+def test_datasource_get():
+    # GIVEN a datasource relation with two remote units
+    local_ds_uids = {
+        "prometheus/0": "some-datasource-uid",
+        "prometheus/1": "some-datasource-uid",
+    }
+    datasource = Relation(
+        "grafana-source",
+        remote_app_name="remote_host",
+        local_unit_data={"grafana_source_host": "somehost:80"},
+        local_app_data={
+            "grafana_source_data": json.dumps(
+                {"model": "foo", "model_uuid": "bar", "application": "baz", "type": "tempo"}
+            )
+        },
+        remote_app_data={"datasource_uids": json.dumps(local_ds_uids)},
+    )
+    state = State(leader=True, relations={datasource})
+
+    # WHEN relation-changed fires for a datasource relation
+    class MyProviderCharm(CharmBase):
+        META = {
+            "name": "edgar",
+            "provides": {"grafana-source": {"interface": "grafana_datasource"}},
+        }
+
+        def __init__(self, framework: Framework):
+            super().__init__(framework)
+            self.source_provider = GrafanaSourceProvider(
+                self, "tempo", source_url="somehost", source_port="80"
+            )
+
+    ctx = Context(MyProviderCharm, MyProviderCharm.META)
+    with ctx(ctx.on.relation_changed(datasource), state) as mgr:
+        charm = mgr.charm
+        # THEN we can see our datasource uids via the provider
+        assert list(charm.source_provider.get_source_uids().values())[0] == local_ds_uids

tox.ini

@@ -50,7 +50,7 @@ commands =
 [testenv:static-{charm,lib}]
 description = Run static analysis checks
 deps =
-    pyright==1.1.316
+    pyright==1.1.389
     charm: -r{toxinidir}/requirements.txt
     lib: ops
     lib: jinja2