From 54f7b1d807827bed1e922a20519b021bbb04a5a2 Mon Sep 17 00:00:00 2001
From: Michael Dmitry <33381599+michaeldmitry@users.noreply.github.com>
Date: Fri, 13 Dec 2024 13:35:53 +0200
Subject: [PATCH] Implement `grafana_datasource_exchange` (#476)

* add datasource_exchange endpoint

* fetch lib

* fetch lib

* add raise_on_error=false

* PR comments
---
 lib/charms/grafana_k8s/v0/grafana_source.py   | 23 +++++-
 metadata.yaml                                 |  4 +
 requirements.txt                              |  2 +-
 src/charm.py                                  | 41 ++++++++++
 tests/integration/helpers.py                  |  2 +
 tests/interface/conftest.py                   | 81 +++++++++++++++++++
 .../test_grafana_datasource_exchange.py       | 13 +++
 tests/interface/test_grafana_source.py        | 11 +++
 tests/scenario/conftest.py                    | 19 ++++-
 tests/scenario/test_datasource_exchange.py    | 78 ++++++++++++++++++
 tox.ini                                       |  9 +++
 11 files changed, 277 insertions(+), 6 deletions(-)
 create mode 100644 tests/interface/conftest.py
 create mode 100644 tests/interface/test_grafana_datasource_exchange.py
 create mode 100644 tests/interface/test_grafana_source.py
 create mode 100644 tests/scenario/test_datasource_exchange.py

diff --git a/lib/charms/grafana_k8s/v0/grafana_source.py b/lib/charms/grafana_k8s/v0/grafana_source.py
index f3492be2d..6e09c5088 100644
--- a/lib/charms/grafana_k8s/v0/grafana_source.py
+++ b/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 = 22
+LIBPATCH = 24
 
 logger = logging.getLogger(__name__)
 
@@ -432,13 +432,22 @@ def update_source(self, source_url: Optional[str] = ""):
     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.
+        Returns a mapping from remote application UIDs 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"])
+            app_databag = rel.data[rel.app]
+            grafana_uid = app_databag.get("grafana_uid")
+            if not grafana_uid:
+                logger.warning(
+                    "remote end is using an old grafana_datasource interface: "
+                    "`grafana_uid` field not found."
+                )
+                continue
+
+            uids[grafana_uid] = json.loads(app_databag.get("datasource_uids", "{}"))
         return uids
 
     def _set_sources_from_event(self, event: RelationJoinedEvent) -> None:
@@ -568,6 +577,14 @@ def _publish_source_uids(self, rel: Relation, uids: Dict[str, str]):
 
         Assumes only leader unit will call this method
         """
+        unique_grafana_name = "juju_{}_{}_{}_{}".format(
+            self._charm.model.name,
+            self._charm.model.uuid,
+            self._charm.model.app.name,
+            self._charm.model.unit.name.split("/")[1],  # type: ignore
+        )
+
+        rel.data[self._charm.app]["grafana_uid"] = unique_grafana_name
         rel.data[self._charm.app]["datasource_uids"] = json.dumps(uids)
 
     def _get_source_config(self, rel: Relation):
diff --git a/metadata.yaml b/metadata.yaml
index de46282ee..bb374d66e 100644
--- a/metadata.yaml
+++ b/metadata.yaml
@@ -62,6 +62,10 @@ provides:
     interface: prometheus_scrape
   grafana-dashboard:
     interface: grafana_dashboard
+  send-datasource:
+    interface: grafana_datasource_exchange
+    description: |
+      Integration to share with other COS components this charm's grafana datasources, and receive theirs.
 
 requires:
   alertmanager:
diff --git a/requirements.txt b/requirements.txt
index efc80ea7b..3b1f28848 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -1,4 +1,4 @@
-cosl>=0.0.12
+cosl>=0.0.47
 # pinned to 2.16 as 2.17 breaks our unittests
 ops
 kubernetes
diff --git a/src/charm.py b/src/charm.py
index d22eb8de8..6070af247 100755
--- a/src/charm.py
+++ b/src/charm.py
@@ -47,6 +47,7 @@
 from charms.tempo_coordinator_k8s.v0.tracing import TracingEndpointRequirer, charm_tracing_config
 from charms.traefik_k8s.v1.ingress_per_unit import IngressPerUnitRequirer
 from cosl import JujuTopology
+from cosl.interfaces.datasource_exchange import DatasourceDict, DatasourceExchange
 from ops import CollectStatusEvent, StoredState
 from ops.charm import CharmBase
 from ops.main import main
@@ -231,6 +232,12 @@ def __init__(self, *args):
             self.charm_tracing, self._ca_cert_path
         )
 
+        self.datasource_exchange = DatasourceExchange(
+            self,
+            provider_endpoint="send-datasource",
+            requirer_endpoint=None,
+        )
+
         self.framework.observe(
             self.workload_tracing.on.endpoint_changed,  # type: ignore
             self._on_workload_tracing_endpoint_changed,
@@ -252,12 +259,29 @@ def __init__(self, *args):
             self._loki_push_api_alert_rules_changed,
         )
         self.framework.observe(self.on.logging_relation_changed, self._on_logging_relation_changed)
+
+        self.framework.observe(
+            self.on.send_datasource_relation_changed, self._on_grafana_source_changed
+        )
+        self.framework.observe(
+            self.on.send_datasource_relation_departed, self._on_grafana_source_changed
+        )
+        self.framework.observe(
+            self.on.grafana_source_relation_changed, self._on_grafana_source_changed
+        )
+        self.framework.observe(
+            self.on.grafana_source_relation_departed, self._on_grafana_source_changed
+        )
+
         self.framework.observe(self.on.collect_unit_status, self._on_collect_unit_status)
 
     ##############################################
     #           CHARM HOOKS HANDLERS             #
     ##############################################
 
+    def _on_grafana_source_changed(self, _):
+        self._update_datasource_exchange()
+
     def _on_collect_unit_status(self, event: CollectStatusEvent):
         # "Pull" statuses
         # TODO refactor _configure to turn the "rules" status into a "pull" status.
@@ -871,6 +895,23 @@ def _tsdb_versions_migration_dates(self) -> List[Dict[str, str]]:
         ret.append({"version": "v13", "date": tomorrow.strftime(date_format)})
         return ret
 
+    def _update_datasource_exchange(self) -> None:
+        """Update the grafana-datasource-exchange relations."""
+        if not self.unit.is_leader():
+            return
+
+        # we might have multiple grafana-source relations, this method collects them all and returns a mapping from
+        # the `grafana_uid` to the contents of the `datasource_uids` field
+        # for simplicity, we assume that we're sending the same data to different grafanas.
+        # read more in https://discourse.charmhub.io/t/tempo-ha-docs-correlating-traces-metrics-logs/16116
+        grafana_uids_to_units_to_uids = self.grafana_source_provider.get_source_uids()
+        raw_datasources: List[DatasourceDict] = []
+
+        for grafana_uid, ds_uids in grafana_uids_to_units_to_uids.items():
+            for _, ds_uid in ds_uids.items():
+                raw_datasources.append({"type": "loki", "uid": ds_uid, "grafana_uid": grafana_uid})
+        self.datasource_exchange.publish(datasources=raw_datasources)
+
 
 if __name__ == "__main__":
     main(LokiOperatorCharm)
diff --git a/tests/integration/helpers.py b/tests/integration/helpers.py
index 71c32ba54..58e90898b 100644
--- a/tests/integration/helpers.py
+++ b/tests/integration/helpers.py
@@ -539,6 +539,8 @@ async def deploy_tempo_cluster(ops_test: OpsTest):
             status="active",
             timeout=2000,
             idle_period=30,
+            # TODO: remove when https://github.com/canonical/tempo-coordinator-k8s-operator/issues/90 is fixed
+            raise_on_error=False,
         )
 
 
diff --git a/tests/interface/conftest.py b/tests/interface/conftest.py
new file mode 100644
index 000000000..e774f8d11
--- /dev/null
+++ b/tests/interface/conftest.py
@@ -0,0 +1,81 @@
+# Copyright 2024 Canonical Ltd.
+# See LICENSE file for licensing details.
+
+import json
+from contextlib import ExitStack
+from unittest.mock import MagicMock, patch
+
+import ops
+import pytest
+from charms.tempo_coordinator_k8s.v0.charm_tracing import charm_tracing_disabled
+from interface_tester import InterfaceTester
+from ops import ActiveStatus
+from scenario.state import Container, Exec, Relation, State
+
+from charm import LokiOperatorCharm
+
+
+@pytest.fixture(autouse=True, scope="module")
+def patch_all():
+    with ExitStack() as stack:
+        stack.enter_context(patch("lightkube.core.client.GenericSyncClient"))
+        stack.enter_context(
+            patch.multiple(
+                "charms.observability_libs.v0.kubernetes_compute_resources_patch.KubernetesComputeResourcesPatch",
+                _namespace="test-namespace",
+                _patch=lambda _: None,
+                is_ready=MagicMock(return_value=True),
+                get_status=lambda _: ActiveStatus(""),
+            )
+        )
+        stack.enter_context(charm_tracing_disabled())
+
+        yield
+
+
+loki_container = Container(
+    name="loki",
+    can_connect=True,
+    execs={Exec(["update-ca-certificates", "--fresh"], return_code=0)},
+    layers={"loki": ops.pebble.Layer({"services": {"loki": {}}})},
+    service_statuses={"loki": ops.pebble.ServiceStatus.ACTIVE},
+)
+
+grafana_source_relation = Relation(
+    "grafana-source",
+    remote_app_data={
+        "datasource_uids": json.dumps({"loki/0": "01234"}),
+        "grafana_uid": "5678",
+    },
+)
+
+grafana_datasource_exchange_relation = Relation(
+    "send-datasource",
+    remote_app_data={
+        "datasources": json.dumps([{"type": "loki", "uid": "01234", "grafana_uid": "5678"}])
+    },
+)
+
+
+@pytest.fixture
+def grafana_datasource_tester(interface_tester: InterfaceTester):
+    interface_tester.configure(
+        charm_type=LokiOperatorCharm,
+        state_template=State(
+            leader=True, containers=[loki_container], relations=[grafana_source_relation]
+        ),
+    )
+    yield interface_tester
+
+
+@pytest.fixture
+def grafana_datasource_exchange_tester(interface_tester: InterfaceTester):
+    interface_tester.configure(
+        charm_type=LokiOperatorCharm,
+        state_template=State(
+            leader=True,
+            containers=[loki_container],
+            relations=[grafana_source_relation, grafana_datasource_exchange_relation],
+        ),
+    )
+    yield interface_tester
diff --git a/tests/interface/test_grafana_datasource_exchange.py b/tests/interface/test_grafana_datasource_exchange.py
new file mode 100644
index 000000000..b2301deef
--- /dev/null
+++ b/tests/interface/test_grafana_datasource_exchange.py
@@ -0,0 +1,13 @@
+# Copyright 2024 Canonical Ltd.
+# See LICENSE file for licensing details.
+from interface_tester import InterfaceTester
+
+
+def test_grafana_datasource_exchange_v0_interface(
+    grafana_datasource_exchange_tester: InterfaceTester,
+):
+    grafana_datasource_exchange_tester.configure(
+        interface_name="grafana_datasource_exchange",
+        interface_version=0,
+    )
+    grafana_datasource_exchange_tester.run()
diff --git a/tests/interface/test_grafana_source.py b/tests/interface/test_grafana_source.py
new file mode 100644
index 000000000..be1f0b55f
--- /dev/null
+++ b/tests/interface/test_grafana_source.py
@@ -0,0 +1,11 @@
+# Copyright 2024 Canonical Ltd.
+# See LICENSE file for licensing details.
+from interface_tester import InterfaceTester
+
+
+def test_grafana_datasource_v0_interface(grafana_datasource_tester: InterfaceTester):
+    grafana_datasource_tester.configure(
+        interface_name="grafana_datasource",
+        interface_version=0,
+    )
+    grafana_datasource_tester.run()
diff --git a/tests/scenario/conftest.py b/tests/scenario/conftest.py
index e5757d8db..96951b572 100644
--- a/tests/scenario/conftest.py
+++ b/tests/scenario/conftest.py
@@ -1,7 +1,10 @@
 from unittest.mock import PropertyMock, patch
 
+import ops
 import pytest
+from charms.tempo_coordinator_k8s.v0.charm_tracing import charm_tracing_disabled
 from ops.testing import Context
+from scenario import Container, Exec
 
 from charm import LokiOperatorCharm
 
@@ -11,7 +14,7 @@ def tautology(*_, **__) -> bool:
 
 
 @pytest.fixture
-def loki_charm():
+def loki_charm(tmp_path):
     with patch.multiple(
         "charm.KubernetesComputeResourcesPatch",
         _namespace=PropertyMock("test-namespace"),
@@ -20,9 +23,21 @@ def loki_charm():
     ):
         with patch("socket.getfqdn", new=lambda *args: "fqdn"):
             with patch("lightkube.core.client.GenericSyncClient"):
-                yield LokiOperatorCharm
+                with charm_tracing_disabled():
+                    yield LokiOperatorCharm
 
 
 @pytest.fixture
 def context(loki_charm):
     return Context(loki_charm)
+
+
+@pytest.fixture(scope="function")
+def loki_container():
+    return Container(
+        "loki",
+        can_connect=True,
+        execs={Exec(["update-ca-certificates", "--fresh"], return_code=0)},
+        layers={"loki": ops.pebble.Layer({"services": {"loki": {}}})},
+        service_statuses={"loki": ops.pebble.ServiceStatus.INACTIVE},
+    )
diff --git a/tests/scenario/test_datasource_exchange.py b/tests/scenario/test_datasource_exchange.py
new file mode 100644
index 000000000..a14390210
--- /dev/null
+++ b/tests/scenario/test_datasource_exchange.py
@@ -0,0 +1,78 @@
+# Copyright 2024 Canonical Ltd.
+# See LICENSE file for licensing details.
+
+import json
+
+import pytest
+from cosl.interfaces.datasource_exchange import (
+    DatasourceExchange,
+    DSExchangeAppData,
+    GrafanaDatasource,
+)
+from scenario import Relation, State
+
+from charm import LokiOperatorCharm
+
+ds_tempo = [
+    {"type": "tempo", "uid": "3", "grafana_uid": "4"},
+]
+
+ds_mimir = [
+    {"type": "prometheus", "uid": "8", "grafana_uid": "9"},
+]
+
+mimir_dsx = Relation(
+    "send-datasource",
+    remote_app_data=DSExchangeAppData(datasources=json.dumps(ds_mimir)).dump(),
+)
+tempo_dsx = Relation(
+    "send-datasource",
+    remote_app_data=DSExchangeAppData(datasources=json.dumps(ds_tempo)).dump(),
+)
+
+ds = Relation(
+    "grafana-source",
+    remote_app_data={
+        "grafana_uid": "9",
+        "datasource_uids": json.dumps({"loki/0": "1234"}),
+    },
+)
+
+
+@pytest.mark.parametrize("event_type", ("changed", "created", "joined"))
+@pytest.mark.parametrize("relation_to_observe", (ds, mimir_dsx, tempo_dsx))
+def test_datasource_send(context, loki_container, relation_to_observe, event_type):
+
+    state_in = State(
+        relations=[
+            ds,
+            mimir_dsx,
+            tempo_dsx,
+        ],
+        containers=[loki_container],
+        leader=True,
+    )
+
+    # WHEN we receive a datasource-related event
+    with context(
+        getattr(context.on, f"relation_{event_type}")(relation_to_observe), state_in
+    ) as mgr:
+        charm: LokiOperatorCharm = mgr.charm
+        # THEN we can find all received datasource uids
+        dsx: DatasourceExchange = charm.datasource_exchange
+        received = dsx.received_datasources
+        assert received == (
+            GrafanaDatasource(type="tempo", uid="3", grafana_uid="4"),
+            GrafanaDatasource(type="prometheus", uid="8", grafana_uid="9"),
+        )
+        state_out = mgr.run()
+
+    # AND THEN we publish our own datasource information to mimir and tempo
+    published_dsx_mimir = state_out.get_relation(mimir_dsx.id).local_app_data
+    published_dsx_tempo = state_out.get_relation(tempo_dsx.id).local_app_data
+    assert published_dsx_tempo == published_dsx_mimir
+    assert json.loads(published_dsx_tempo["datasources"])[0] == {
+        "type": "loki",
+        "uid": "1234",
+        "grafana_uid": "9",
+    }
diff --git a/tox.ini b/tox.ini
index ce1cc8b15..c3c3b7b2f 100644
--- a/tox.ini
+++ b/tox.ini
@@ -98,3 +98,12 @@ deps =
     minio
 commands =
     pytest -v --tb native --log-cli-level=INFO --color=yes -s {posargs} {toxinidir}/tests/integration
+
+[testenv:interface]
+description = Run interface tests
+deps =
+    pytest
+    -r{toxinidir}/requirements.txt
+    pytest-interface-tester
+commands =
+    pytest -v --tb native --log-cli-level=INFO -s {posargs} {[vars]tst_path}/interface
\ No newline at end of file