Peer data magic

src/relations/db.py

@@ -84,6 +84,7 @@
 ------------------------------------------------------------------------------------------------------------
 """  # noqa: W505
 
+import json
 import logging
 from typing import Dict, Iterable, List
 
@@ -238,14 +239,13 @@ def _on_relation_joined(self, join_event: RelationJoinedEvent):
         dbs[str(join_event.relation.id)] = {"name": database, "legacy": True}
         self.charm.set_relation_databases(dbs)
 
-        self.update_databags(
-            join_event.relation,
-            {
-                "user": user,
-                "password": password,
-                "database": database,
-            },
-        )
+        creds = {
+            "user": user,
+            "password": password,
+            "database": database,
+        }
+        self.charm.peers.app_databag[user] = json.dumps(creds)
+        self.update_databags(join_event.relation, creds)
 
         # Create user and database in backend postgresql database
         try:
@@ -292,7 +292,8 @@ def _on_relation_changed(self, change_event: RelationChangedEvent):
         # No backup values because if databag isn't populated, this relation isn't initialised.
         # This means that the database and user requested in this relation haven't been created,
         # so we defer this event until the databag is populated.
-        databag = self.get_databags(change_event.relation)[0]
+        user = self._generate_username(change_event.relation)
+        databag = json.loads(self.charm.peers.app_databag.get(user, "{}"))
         database = databag.get("database")
         user = databag.get("user")
         password = databag.get("password")
@@ -304,22 +305,28 @@ def _on_relation_changed(self, change_event: RelationChangedEvent):
             change_event.defer()
             return
 
+        dbs = self.charm.get_relation_databases()
+        if str(change_event.relation.id) not in dbs:
+            pass
+            logger.debug("relation not fully initialised - database mapping not yet set")
+            change_event.defer()
+            return
+
         self.charm.render_pgb_config(reload_pgbouncer=True)
-        if self.charm.unit.is_leader():
-            self.update_connection_info(change_event.relation, self.charm.config["listen_port"])
-            self.update_databags(
-                change_event.relation,
-                {
-                    "allowed-subnets": self.get_allowed_subnets(change_event.relation),
-                    "allowed-units": self.get_allowed_units(change_event.relation),
-                    "version": self.charm.backend.postgres.get_postgresql_version(),
-                    "host": "localhost",
-                    "user": user,
-                    "password": password,
-                    "database": database,
-                    "state": self._get_state(),
-                },
-            )
+        self.update_databags(
+            change_event.relation,
+            {
+                "allowed-subnets": self.get_allowed_subnets(change_event.relation),
+                "allowed-units": self.get_allowed_units(change_event.relation),
+                "version": self.charm.backend.postgres.get_postgresql_version(),
+                "host": "localhost",
+                "user": user,
+                "password": password,
+                "database": database,
+                "state": "master",
+            },
+        )
+        self.update_connection_info(change_event.relation, self.charm.config["listen_port"])
 
     def update_connection_info(self, relation: Relation, port: str = None):
         """Updates databag connection information."""
@@ -456,8 +463,7 @@ def update_databags(self, relation, updates: Dict[str, str]):
         for key, item in updates.items():
             updates[key] = str(item)
 
-        databag = relation.data[self.charm.unit]
-        databag.update(updates)
+        relation.data[self.charm.unit].update(updates)
 
     def _generate_username(self, relation):
         """Generates a unique username for this relation."""
@@ -466,20 +472,6 @@ def _generate_username(self, relation):
         model_name = self.model.name
         return f"{app_name}_user_{relation_id}_{model_name}".replace("-", "_")
 
-    def _get_state(self) -> str:
-        """Gets the given state for this unit.
-
-        Args:
-            standbys: the comma-separated list of postgres standbys
-
-        Returns:
-            The described state of this unit. Can be 'master' or 'standby'.
-        """
-        if self.charm.unit.is_leader():
-            return "master"
-        else:
-            return "standby"
-
     def get_allowed_subnets(self, relation: Relation) -> str:
         """Gets the allowed subnets from this relation."""
 

tests/unit/relations/test_db.py

@@ -1,6 +1,7 @@
 # Copyright 2022 Canonical Ltd.
 # See LICENSE file for licensing details.
 
+import json
 import unittest
 from unittest.mock import Mock, PropertyMock, patch
 
@@ -124,17 +125,17 @@ def test_on_relation_joined(
     @patch(
         "relations.backend_database.BackendDatabaseRequires.postgres", new_callable=PropertyMock
     )
-    @patch("relations.db.DbProvides.get_databags", return_value=[{}])
+    @patch("relations.peers.Peers.app_databag", new_callable=PropertyMock, return_value={})
     @patch("relations.db.DbProvides.update_connection_info")
     @patch("relations.db.DbProvides.update_databags")
     @patch("relations.db.DbProvides.get_allowed_units")
     @patch("relations.db.DbProvides.get_allowed_subnets")
-    @patch("relations.db.DbProvides._get_state")
+    @patch("charm.PgBouncerCharm.get_relation_databases", return_value={"1": {"some": "data"}})
     @patch("charm.PgBouncerCharm.render_pgb_config")
     def test_on_relation_changed(
         self,
         _render_pgb_config,
-        _get_state,
+        _,
         _allowed_subnets,
         _allowed_units,
         _update_databags,
@@ -143,19 +144,24 @@ def test_on_relation_changed(
         _backend_postgres,
         _check_backend,
     ):
-        self.harness.set_leader(True)
+        with self.harness.hooks_disabled():
+            self.harness.set_leader(True)
+
+        event = Mock()
+        event.relation.id = 1
 
         database = "test_db"
-        user = "test_user"
+        user = "pgbouncer_user_1_None"
         password = "test_pw"
-        _get_databags.return_value[0] = {
-            "database": database,
-            "user": user,
-            "password": password,
+        _get_databags.return_value = {
+            user: json.dumps({
+                "database": database,
+                "user": user,
+                "password": password,
+            })
         }
 
         # Call the function
-        event = Mock()
         self.db_relation._on_relation_changed(event)
 
         _update_connection_info.assert_called_with(
@@ -171,7 +177,7 @@ def test_on_relation_changed(
                 "user": user,
                 "password": password,
                 "database": database,
-                "state": _get_state.return_value,
+                "state": "master",
             },
         )
         _render_pgb_config.assert_called_once_with(reload_pgbouncer=True)