From e676f3c1d369b7d1ebb82f71327176c15be21dfc Mon Sep 17 00:00:00 2001
From: telcobot <147593099+telcobot@users.noreply.github.com>
Date: Tue, 19 Mar 2024 06:10:34 -0300
Subject: [PATCH] chore: Update charm libraries (#135)

---
 lib/charms/loki_k8s/v1/loki_push_api.py | 52 ++++++++++++++++++-------
 1 file changed, 38 insertions(+), 14 deletions(-)

diff --git a/lib/charms/loki_k8s/v1/loki_push_api.py b/lib/charms/loki_k8s/v1/loki_push_api.py
index 6c8549c5..4bc4e22a 100644
--- a/lib/charms/loki_k8s/v1/loki_push_api.py
+++ b/lib/charms/loki_k8s/v1/loki_push_api.py
@@ -497,6 +497,7 @@ def _alert_rules_error(self, event):
 from ops.charm import (
     CharmBase,
     HookEvent,
+    PebbleReadyEvent,
     RelationBrokenEvent,
     RelationCreatedEvent,
     RelationDepartedEvent,
@@ -518,7 +519,7 @@ def _alert_rules_error(self, event):
 
 # Increment this PATCH version before using `charmcraft publish-lib` or reset
 # to 0 if you are raising the major API version
-LIBPATCH = 5
+LIBPATCH = 6
 
 logger = logging.getLogger(__name__)
 
@@ -2486,7 +2487,10 @@ def enable_endpoints(
 
 
 class LogForwarder(ConsumerBase):
-    """Forward the standard outputs of all workloads operated by a charm to one or multiple Loki endpoints."""
+    """Forward the standard outputs of all workloads operated by a charm to one or multiple Loki endpoints.
+
+    This class implements Pebble log forwarding. Juju >= 3.4 is needed.
+    """
 
     def __init__(
         self,
@@ -2510,27 +2514,47 @@ def __init__(
         self.framework.observe(on.relation_departed, self._update_logging)
         self.framework.observe(on.relation_broken, self._update_logging)
 
+        for container_name in self._charm.meta.containers.keys():
+            snake_case_container_name = container_name.replace("-", "_")
+            self.framework.observe(
+                getattr(self._charm.on, f"{snake_case_container_name}_pebble_ready"),
+                self._on_pebble_ready,
+            )
+
+    def _on_pebble_ready(self, event: PebbleReadyEvent):
+        if not (loki_endpoints := self._retrieve_endpoints_from_relation()):
+            logger.warning("No Loki endpoints available")
+            return
+
+        self._update_endpoints(event.workload, loki_endpoints)
+
     def _update_logging(self, _):
         """Update the log forwarding to match the active Loki endpoints."""
+        if not (loki_endpoints := self._retrieve_endpoints_from_relation()):
+            logger.warning("No Loki endpoints available")
+            return
+
+        for container in self._charm.unit.containers.values():
+            self._update_endpoints(container, loki_endpoints)
+
+    def _retrieve_endpoints_from_relation(self) -> dict:
         loki_endpoints = {}
 
         # Get the endpoints from relation data
         for relation in self._charm.model.relations[self._relation_name]:
             loki_endpoints.update(self._fetch_endpoints(relation))
 
-        if not loki_endpoints:
-            logger.warning("No Loki endpoints available")
-            return
+        return loki_endpoints
 
-        for container in self._charm.unit.containers.values():
-            _PebbleLogClient.disable_inactive_endpoints(
-                container=container,
-                active_endpoints=loki_endpoints,
-                topology=self.topology,
-            )
-            _PebbleLogClient.enable_endpoints(
-                container=container, active_endpoints=loki_endpoints, topology=self.topology
-            )
+    def _update_endpoints(self, container: Container, loki_endpoints: dict):
+        _PebbleLogClient.disable_inactive_endpoints(
+            container=container,
+            active_endpoints=loki_endpoints,
+            topology=self.topology,
+        )
+        _PebbleLogClient.enable_endpoints(
+            container=container, active_endpoints=loki_endpoints, topology=self.topology
+        )
 
     def is_ready(self, relation: Optional[Relation] = None):
         """Check if the relation is active and healthy."""