From 75b7b62f7f50049671d0b5288b96ab64867b6660 Mon Sep 17 00:00:00 2001
From: Judit Novak <judit.novak@canonical.com>
Date: Thu, 16 Nov 2023 16:46:01 +0100
Subject: [PATCH] Updates safe to switch from secret ID to label

---
 src/charm.py | 23 ++++++++++++++++++-----
 1 file changed, 18 insertions(+), 5 deletions(-)

diff --git a/src/charm.py b/src/charm.py
index cbf11f9abc..43fd2bdde4 100755
--- a/src/charm.py
+++ b/src/charm.py
@@ -19,7 +19,6 @@
     PostgreSQLEnableDisableExtensionError,
     PostgreSQLUpdateUserPasswordError,
 )
-from charms.postgresql_k8s.v0.postgresql_secrets import SecretCache, generate_secret_label
 from charms.postgresql_k8s.v0.postgresql_tls import PostgreSQLTLS
 from charms.prometheus_k8s.v0.prometheus_scrape import MetricsEndpointProvider
 from charms.rolling_ops.v0.rollingops import RollingOpsManager, RunWithLock
@@ -212,6 +211,20 @@ def _translate_field_to_secret_key(self, key: str) -> str:
         new_key = key.replace("_", "-")
         return new_key.strip("-")
 
+    def _safe_get_secret(self, scope: Scopes, label: str) -> SecretCache:
+        """Safety measure, for upgrades between versions based on secret URI usage to others with labels usage.
+
+        If the secret can't be retrieved by label, we search for the uri -- and if found, we "stick" the
+        label on the secret for further usage.
+        """
+        secret_uri = self._peer_data(scope).get(SECRET_INTERNAL_LABEL, None)
+        secret = self.secrets.get(label, secret_uri)
+
+        # Since now we switched to labels, the databag reference can be removed
+        if secret_uri and secret and scope == APP_SCOPE and self.unit.is_leader():
+            self._peer_data(scope).pop(SECRET_INTERNAL_LABEL, None)
+        return secret
+
     def get_secret(self, scope: Scopes, key: str) -> Optional[str]:
         """Get secret from the secret storage."""
         if scope not in get_args(Scopes):
@@ -221,15 +234,15 @@ def get_secret(self, scope: Scopes, key: str) -> Optional[str]:
             return value
 
         if JujuVersion.from_environ().has_secrets:
-            secret_key = self._translate_field_to_secret_key(key)
             label = generate_secret_label(self, scope)
             for attempt in Retrying(stop=stop_after_attempt(3), wait=wait_fixed(1), reraise=True):
                 with attempt:
-                    secret = self.secrets.get(label)
+                    secret = self._safe_get_secret(scope, label)
 
             if not secret:
                 return
 
+            secret_key = self._translate_field_to_secret_key(key)
             value = secret.get_content().get(secret_key)
             if value != SECRET_DELETED_LABEL:
                 return value
@@ -249,7 +262,7 @@ def set_secret(self, scope: Scopes, key: str, value: Optional[str]) -> Optional[
 
             secret_key = self._translate_field_to_secret_key(key)
             label = generate_secret_label(self, scope)
-            secret = self.secrets.get(label)
+            secret = self._safe_get_secret(scope, label)
             if not secret:
                 self.secrets.add(label, {secret_key: value}, scope)
             else:
@@ -268,7 +281,7 @@ def remove_secret(self, scope: Scopes, key: str) -> None:
         if JujuVersion.from_environ().has_secrets:
             secret_key = self._translate_field_to_secret_key(key)
             label = generate_secret_label(self, scope)
-            secret = self.secrets.get(label)
+            secret = self._safe_get_secret(scope, label)
 
             if not secret:
                 return