Rework secrets (#27)

* Rework secrets management

* Add unit tests

* Add constants

src/charm.py

@@ -5,7 +5,7 @@
 """Charmed Kubernetes Operator for the PostgreSQL database."""
 import json
 import logging
-from typing import List
+from typing import Dict, List, Optional
 
 from charms.postgresql_k8s.v0.postgresql import PostgreSQL
 from lightkube import ApiError, Client, codecs
@@ -30,7 +30,7 @@
 from requests import ConnectionError
 from tenacity import RetryError
 
-from constants import PEER, USER
+from constants import PEER, REPLICATION_PASSWORD_KEY, USER, USER_PASSWORD_KEY
 from patroni import NotReadyError, Patroni
 from relations.db import DbProvides
 from relations.postgresql_provider import PostgreSQLProvider
@@ -69,6 +69,48 @@ def __init__(self, *args):
         self.legacy_db_relation = DbProvides(self, admin=False)
         self.legacy_db_admin_relation = DbProvides(self, admin=True)
 
+    @property
+    def app_peer_data(self) -> Dict:
+        """Application peer relation data object."""
+        relation = self.model.get_relation(PEER)
+        if relation is None:
+            return {}
+
+        return relation.data[self.app]
+
+    @property
+    def unit_peer_data(self) -> Dict:
+        """Unit peer relation data object."""
+        relation = self.model.get_relation(PEER)
+        if relation is None:
+            return {}
+
+        return relation.data[self.unit]
+
+    def _get_secret(self, scope: str, key: str) -> Optional[str]:
+        """Get secret from the secret storage."""
+        if scope == "unit":
+            return self.unit_peer_data.get(key, None)
+        elif scope == "app":
+            return self.app_peer_data.get(key, None)
+        else:
+            raise RuntimeError("Unknown secret scope.")
+
+    def _set_secret(self, scope: str, key: str, value: Optional[str]) -> None:
+        """Get secret from the secret storage."""
+        if scope == "unit":
+            if not value:
+                del self.unit_peer_data[key]
+                return
+            self.unit_peer_data.update({key: value})
+        elif scope == "app":
+            if not value:
+                del self.app_peer_data[key]
+                return
+            self.app_peer_data.update({key: value})
+        else:
+            raise RuntimeError("Unknown secret scope.")
+
     @property
     def postgresql(self) -> PostgreSQL:
         """Returns an instance of the object used to interact with the database."""
@@ -249,15 +291,11 @@ def _get_hostname_from_unit(self, member: str) -> str:
 
     def _on_leader_elected(self, event: LeaderElectedEvent) -> None:
         """Handle the leader-elected event."""
-        data = self._peers.data[self.app]
-        operator_password = data.get("operator-password", None)
-        replication_password = data.get("replication-password", None)
-
-        if operator_password is None:
-            self._peers.data[self.app]["operator-password"] = new_password()
+        if self._get_secret("app", USER_PASSWORD_KEY) is None:
+            self._set_secret("app", USER_PASSWORD_KEY, new_password())
 
-        if replication_password is None:
-            self._peers.data[self.app]["replication-password"] = new_password()
+        if self._get_secret("app", REPLICATION_PASSWORD_KEY) is None:
+            self._set_secret("app", REPLICATION_PASSWORD_KEY, new_password())
 
         # Create resources and add labels needed for replication.
         self._create_resources()
@@ -389,7 +427,7 @@ def _create_resources(self) -> None:
 
     def _on_get_operator_password(self, event: ActionEvent) -> None:
         """Returns the password for the operator user as an action response."""
-        event.set_results({"operator-password": self._get_operator_password()})
+        event.set_results({USER_PASSWORD_KEY: self._get_operator_password()})
 
     def _on_get_primary(self, event: ActionEvent) -> None:
         """Get primary instance."""
@@ -498,14 +536,12 @@ def _peers(self) -> Relation:
 
     def _get_operator_password(self) -> str:
         """Get operator user password."""
-        data = self._peers.data[self.app]
-        return data.get("operator-password", None)
+        return self._get_secret("app", USER_PASSWORD_KEY)
 
     @property
     def _replication_password(self) -> str:
         """Get replication user password."""
-        data = self._peers.data[self.app]
-        return data.get("replication-password", None)
+        return self._get_secret("app", REPLICATION_PASSWORD_KEY)
 
     def _unit_name_to_pod_name(self, unit_name: str) -> str:
         """Converts unit name to pod name.

src/constants.py

@@ -5,4 +5,6 @@
 
 DATABASE_PORT = "5432"
 PEER = "database-peers"
+REPLICATION_PASSWORD_KEY = "replication-password"
 USER = "operator"
+USER_PASSWORD_KEY = "operator-password"

tests/unit/test_charm.py

@@ -31,6 +31,8 @@ def setUp(self):
             "app_name": self.harness.model.app.name,
         }
 
+        self.rel_id = self.harness.add_relation(self._peer_relation, self.charm.app.name)
+
     @patch_network_get(private_address="1.1.1.1")
     @patch("charm.Patroni.render_postgresql_conf_file")
     def test_on_install(
@@ -46,7 +48,6 @@ def test_on_install(
     @patch("charm.PostgresqlOperatorCharm._create_resources")
     def test_on_leader_elected(self, _, __, _render_postgresql_conf_file, ___):
         # Assert that there is no password in the peer relation.
-        self.harness.add_relation(self._peer_relation, self.charm.app.name)
         self.assertIsNone(self.charm._peers.data[self.charm.app].get("postgres-password", None))
         self.assertIsNone(self.charm._peers.data[self.charm.app].get("replication-password", None))
 
@@ -86,7 +87,6 @@ def test_on_postgresql_pebble_ready(
         # Check that the initial plan is empty.
         plan = self.harness.get_container_pebble_plan(self._postgresql_container)
         self.assertEqual(plan.to_dict(), {})
-        self.harness.add_relation(self._peer_relation, self.charm.app.name)
 
         # Get the current and the expected layer from the pebble plan and the _postgresql_layer
         # method, respectively.
@@ -202,7 +202,6 @@ def test_patch_pod_labels(self, _client):
     @patch("charm.PostgresqlOperatorCharm._create_resources")
     def test_postgresql_layer(self, _, __, ___, ____):
         # Test with the already generated password.
-        self.harness.add_relation(self._peer_relation, self.charm.app.name)
         self.harness.set_leader()
         plan = self.charm._postgresql_layer().to_dict()
         expected = {
@@ -238,11 +237,52 @@ def test_postgresql_layer(self, _, __, ___, ____):
     @patch("charm.PostgresqlOperatorCharm._create_resources")
     def test_get_operator_password(self, _, __, ___, ____):
         # Test for a None password.
-        self.harness.add_relation(self._peer_relation, self.charm.app.name)
         self.assertIsNone(self.charm._get_operator_password())
 
         # Then test for a non empty password after leader election and peer data set.
         self.harness.set_leader()
         password = self.charm._get_operator_password()
         self.assertIsNotNone(password)
         self.assertNotEqual(password, "")
+
+    @patch("charm.Patroni.reload_patroni_configuration")
+    @patch("charm.Patroni.render_postgresql_conf_file")
+    @patch("charm.PostgresqlOperatorCharm._create_resources")
+    def test_get_secret(self, _, __, ___):
+        self.harness.set_leader()
+
+        # Test application scope.
+        assert self.charm._get_secret("app", "password") is None
+        self.harness.update_relation_data(
+            self.rel_id, self.charm.app.name, {"password": "test-password"}
+        )
+        assert self.charm._get_secret("app", "password") == "test-password"
+
+        # Test unit scope.
+        assert self.charm._get_secret("unit", "password") is None
+        self.harness.update_relation_data(
+            self.rel_id, self.charm.unit.name, {"password": "test-password"}
+        )
+        assert self.charm._get_secret("unit", "password") == "test-password"
+
+    @patch("charm.Patroni.reload_patroni_configuration")
+    @patch("charm.Patroni.render_postgresql_conf_file")
+    @patch("charm.PostgresqlOperatorCharm._create_resources")
+    def test_set_secret(self, _, __, ___):
+        self.harness.set_leader()
+
+        # Test application scope.
+        assert "password" not in self.harness.get_relation_data(self.rel_id, self.charm.app.name)
+        self.charm._set_secret("app", "password", "test-password")
+        assert (
+            self.harness.get_relation_data(self.rel_id, self.charm.app.name)["password"]
+            == "test-password"
+        )
+
+        # Test unit scope.
+        assert "password" not in self.harness.get_relation_data(self.rel_id, self.charm.unit.name)
+        self.charm._set_secret("unit", "password", "test-password")
+        assert (
+            self.harness.get_relation_data(self.rel_id, self.charm.unit.name)["password"]
+            == "test-password"
+        )