[DPE-1765] Juju 3 peer secrets (#106)

* Initial secrets implementation

* Secrets unit tests

* Fixes

* Revert lib bump

* Code review improvements

src/charm.py

@@ -17,12 +17,20 @@
 from charms.operator_libs_linux.v2 import snap
 from charms.pgbouncer_k8s.v0 import pgb
 from jinja2 import Template
+from ops import JujuVersion
 from ops.charm import CharmBase
 from ops.framework import StoredState
 from ops.main import main
-from ops.model import ActiveStatus, BlockedStatus, MaintenanceStatus, WaitingStatus
+from ops.model import (
+    ActiveStatus,
+    BlockedStatus,
+    MaintenanceStatus,
+    SecretNotFoundError,
+    WaitingStatus,
+)
 
 from constants import (
+    APP_SCOPE,
     AUTH_FILE_NAME,
     CLIENT_RELATION_NAME,
     EXTENSIONS_BLOCKING_MESSAGE,
@@ -35,8 +43,14 @@
     PGB_LOG_DIR,
     PGBOUNCER_EXECUTABLE,
     POSTGRESQL_SNAP_NAME,
+    SECRET_CACHE_LABEL,
+    SECRET_DELETED_LABEL,
+    SECRET_INTERNAL_LABEL,
+    SECRET_KEY_OVERRIDES,
+    SECRET_LABEL,
     SNAP_PACKAGES,
     SNAP_TMP_DIR,
+    UNIT_SCOPE,
 )
 from relations.backend_database import BackendDatabaseRequires
 from relations.db import DbProvides
@@ -56,6 +70,8 @@ class PgBouncerCharm(CharmBase):
     def __init__(self, *args):
         super().__init__(*args)
 
+        self.secrets = {APP_SCOPE: {}, UNIT_SCOPE: {}}
+
         self.framework.observe(self.on.install, self._on_install)
         self.framework.observe(self.on.remove, self._on_remove)
         self.framework.observe(self.on.start, self._on_start)
@@ -197,6 +213,176 @@ def version(self) -> str:
             logger.exception("Unable to get Pgbouncer version")
             return ""
 
+    def _normalize_secret_key(self, key: str) -> str:
+        new_key = key.replace("_", "-")
+        new_key = new_key.strip("-")
+
+        return new_key
+
+    def _scope_obj(self, scope: str):
+        if scope == APP_SCOPE:
+            return self.framework.model.app
+        if scope == UNIT_SCOPE:
+            return self.framework.model.unit
+
+    def _juju_secrets_get(self, scope: str) -> Optional[bool]:
+        """Helper function to get Juju secret."""
+        if scope == UNIT_SCOPE:
+            peer_data = self.peers.unit_databag
+        else:
+            peer_data = self.peers.app_databag
+
+        if not peer_data.get(SECRET_INTERNAL_LABEL):
+            return
+
+        if SECRET_CACHE_LABEL not in self.secrets[scope]:
+            try:
+                # NOTE: Secret contents are not yet available!
+                secret = self.model.get_secret(id=peer_data[SECRET_INTERNAL_LABEL])
+            except SecretNotFoundError as e:
+                logging.debug(f"No secret found for ID {peer_data[SECRET_INTERNAL_LABEL]}, {e}")
+                return
+
+            logging.debug(f"Secret {peer_data[SECRET_INTERNAL_LABEL]} downloaded")
+
+            # We keep the secret object around -- needed when applying modifications
+            self.secrets[scope][SECRET_LABEL] = secret
+
+            # We retrieve and cache actual secret data for the lifetime of the event scope
+            self.secrets[scope][SECRET_CACHE_LABEL] = secret.get_content()
+
+        return bool(self.secrets[scope].get(SECRET_CACHE_LABEL))
+
+    def _juju_secret_get_key(self, scope: str, key: str) -> Optional[str]:
+        if not key:
+            return
+
+        key = SECRET_KEY_OVERRIDES.get(key, self._normalize_secret_key(key))
+
+        if self._juju_secrets_get(scope):
+            secret_cache = self.secrets[scope].get(SECRET_CACHE_LABEL)
+            if secret_cache:
+                secret_data = secret_cache.get(key)
+                if secret_data and secret_data != SECRET_DELETED_LABEL:
+                    logging.debug(f"Getting secret {scope}:{key}")
+                    return secret_data
+        logging.debug(f"No value found for secret {scope}:{key}")
+
+    def get_secret(self, scope: str, key: str) -> Optional[str]:
+        """Get secret from the secret storage."""
+        if scope not in [APP_SCOPE, UNIT_SCOPE]:
+            raise RuntimeError("Unknown secret scope.")
+
+        if scope == UNIT_SCOPE:
+            result = self.peers.unit_databag.get(key, None)
+        else:
+            result = self.peers.app_databag.get(key, None)
+
+        # TODO change upgrade to switch to secrets once minor version upgrades is done
+        if result:
+            return result
+
+        juju_version = JujuVersion.from_environ()
+        if juju_version.has_secrets:
+            return self._juju_secret_get_key(scope, key)
+
+    def _juju_secret_set(self, scope: str, key: str, value: str) -> Optional[str]:
+        """Helper function setting Juju secret."""
+        if scope == UNIT_SCOPE:
+            peer_data = self.peers.unit_databag
+        else:
+            peer_data = self.peers.app_databag
+        self._juju_secrets_get(scope)
+
+        key = SECRET_KEY_OVERRIDES.get(key, self._normalize_secret_key(key))
+
+        secret = self.secrets[scope].get(SECRET_LABEL)
+
+        # It's not the first secret for the scope, we can re-use the existing one
+        # that was fetched in the previous call
+        if secret:
+            secret_cache = self.secrets[scope][SECRET_CACHE_LABEL]
+
+            if secret_cache.get(key) == value:
+                logging.debug(f"Key {scope}:{key} has this value defined already")
+            else:
+                secret_cache[key] = value
+                try:
+                    secret.set_content(secret_cache)
+                except OSError as error:
+                    logging.error(
+                        f"Error in attempt to set {scope}:{key}. "
+                        f"Existing keys were: {list(secret_cache.keys())}. {error}"
+                    )
+                    return
+                logging.debug(f"Secret {scope}:{key} was {key} set")
+
+        # We need to create a brand-new secret for this scope
+        else:
+            scope_obj = self._scope_obj(scope)
+
+            secret = scope_obj.add_secret({key: value})
+            if not secret:
+                raise RuntimeError(f"Couldn't set secret {scope}:{key}")
+
+            self.secrets[scope][SECRET_LABEL] = secret
+            self.secrets[scope][SECRET_CACHE_LABEL] = {key: value}
+            logging.debug(f"Secret {scope}:{key} published (as first). ID: {secret.id}")
+            peer_data.update({SECRET_INTERNAL_LABEL: secret.id})
+
+        return self.secrets[scope][SECRET_LABEL].id
+
+    def set_secret(self, scope: str, key: str, value: Optional[str]) -> Optional[str]:
+        """Set secret from the secret storage."""
+        if scope not in [APP_SCOPE, UNIT_SCOPE]:
+            raise RuntimeError("Unknown secret scope.")
+
+        if not value:
+            return self.remove_secret(scope, key)
+
+        juju_version = JujuVersion.from_environ()
+
+        if juju_version.has_secrets:
+            self._juju_secret_set(scope, key, value)
+            return
+        if scope == UNIT_SCOPE:
+            self.peers.unit_databag.update({key: value})
+        else:
+            self.peers.app_databag.update({key: value})
+
+    def _juju_secret_remove(self, scope: str, key: str) -> None:
+        """Remove a Juju 3.x secret."""
+        self._juju_secrets_get(scope)
+
+        key = SECRET_KEY_OVERRIDES.get(key, self._normalize_secret_key(key))
+
+        secret = self.secrets[scope].get(SECRET_LABEL)
+        if not secret:
+            logging.error(f"Secret {scope}:{key} wasn't deleted: no secrets are available")
+            return
+
+        secret_cache = self.secrets[scope].get(SECRET_CACHE_LABEL)
+        if not secret_cache or key not in secret_cache:
+            logging.error(f"No secret {scope}:{key}")
+            return
+
+        secret_cache[key] = SECRET_DELETED_LABEL
+        secret.set_content(secret_cache)
+        logging.debug(f"Secret {scope}:{key}")
+
+    def remove_secret(self, scope: str, key: str) -> None:
+        """Removing a secret."""
+        if scope not in [APP_SCOPE, UNIT_SCOPE]:
+            raise RuntimeError("Unknown secret scope.")
+
+        juju_version = JujuVersion.from_environ()
+        if juju_version.has_secrets:
+            return self._juju_secret_remove(scope, key)
+        if scope == UNIT_SCOPE:
+            del self.peers.unit_databag[key]
+        else:
+            del self.peers.app_databag[key]
+
     def _on_start(self, _) -> None:
         """On Start hook.
 
@@ -401,7 +587,7 @@ def render_prometheus_service(self):
         rendered = template.render(
             stats_user=self.backend.stats_user,
             pgb_service=f"{PGB}-{self.app.name}",
-            stats_password=self.peers.get_secret("app", MONITORING_PASSWORD_KEY),
+            stats_password=self.get_secret(APP_SCOPE, MONITORING_PASSWORD_KEY),
             listen_port=self.config["listen_port"],
             metrics_port=self.config["metrics_port"],
         )

src/constants.py

@@ -35,3 +35,17 @@
 MONITORING_PASSWORD_KEY = "monitoring_password"
 
 EXTENSIONS_BLOCKING_MESSAGE = "bad relation request - remote app requested extensions, which are unsupported. Please remove this relation."
+
+SECRET_LABEL = "secret"
+SECRET_CACHE_LABEL = "cache"
+SECRET_INTERNAL_LABEL = "internal-secret"
+SECRET_DELETED_LABEL = "None"
+
+APP_SCOPE = "app"
+UNIT_SCOPE = "unit"
+
+SECRET_KEY_OVERRIDES = {
+    "cfg_file": "cfg-file",
+    "monitoring_password": "monitoring-password",
+    "auth_file": "auth-file",
+}

src/relations/backend_database.py

@@ -67,6 +67,7 @@
 )
 
 from constants import (
+    APP_SCOPE,
     AUTH_FILE_NAME,
     BACKEND_RELATION_NAME,
     MONITORING_PASSWORD_KEY,
@@ -139,11 +140,9 @@ def _on_database_created(self, event: DatabaseCreatedEvent) -> None:
         self.initialise_auth_function([self.database.database, PG])
 
         # Add the monitoring user.
-        if not (
-            monitoring_password := self.charm.peers.get_secret("app", MONITORING_PASSWORD_KEY)
-        ):
+        if not (monitoring_password := self.charm.get_secret(APP_SCOPE, MONITORING_PASSWORD_KEY)):
             monitoring_password = pgb.generate_password()
-            self.charm.peers.set_secret("app", MONITORING_PASSWORD_KEY, monitoring_password)
+            self.charm.set_secret("app", MONITORING_PASSWORD_KEY, monitoring_password)
         hashed_monitoring_password = pgb.get_hashed_password(self.stats_user, monitoring_password)
 
         self.charm.render_auth_file(

src/relations/db.py

@@ -111,7 +111,7 @@
     WaitingStatus,
 )
 
-from constants import EXTENSIONS_BLOCKING_MESSAGE, PG
+from constants import APP_SCOPE, EXTENSIONS_BLOCKING_MESSAGE, PG
 
 logger = logging.getLogger(__name__)
 
@@ -228,7 +228,7 @@ def _on_relation_joined(self, join_event: RelationJoinedEvent):
         if self.charm.unit.is_leader():
             password = pgb.generate_password()
         else:
-            if not (password := self.charm.peers.app_databag.get(user, None)):
+            if not (password := self.charm.get_secret(APP_SCOPE, user)):
                 join_event.defer()
 
         self.update_databags(