Copy over safe secret get

canonical · Nov 23, 2023 · d62d95c · d62d95c
1 parent 30fd8e6
commit d62d95c
Show file tree

Hide file tree

Showing 4 changed files with 30 additions and 14 deletions.
diff --git a/...s/postgresql_k8s/v0/postgresql_secrets.py → ...rms/data_platform_libs/v0/data_secrets.py b/...s/postgresql_k8s/v0/postgresql_secrets.py → ...rms/data_platform_libs/v0/data_secrets.py
@@ -9,7 +9,7 @@
 from ops.model import SecretNotFoundError
 
 # The unique Charmhub library identifier, never change it
-LIBID = "cd223b249c294092b9ac703760e64b3f"
+LIBID = "d77fb3d01aba41ed88e837d0beab6be5"
 
 # Increment this major API version when introducing breaking changes
 LIBAPI = 0
@@ -21,14 +21,14 @@
 
 APP_SCOPE = "app"
 UNIT_SCOPE = "unit"
-Scopes = Literal[APP_SCOPE, UNIT_SCOPE]
+Scopes = Literal["app", "unit"]
 
 
-class PostgresSQLSecretsError(Exception):
+class DataSecretsError(Exception):
     """A secret that we want to create already exists."""
 
 
-class SecretAlreadyExistsError(PostgresSQLSecretsError):
+class SecretAlreadyExistsError(DataSecretsError):
     """A secret that we want to create already exists."""
 
 

diff --git a/src/charm.py b/src/charm.py
@@ -11,6 +11,7 @@
 from typing import Dict, List, Literal, Optional, Set, get_args
 
 from charms.data_platform_libs.v0.data_models import TypedCharmBase
+from charms.data_platform_libs.v0.data_secrets import SecretCache, generate_secret_label
 from charms.grafana_agent.v0.cos_agent import COSAgentProvider
 from charms.operator_libs_linux.v2 import snap
 from charms.postgresql_k8s.v0.postgresql import (
@@ -19,7 +20,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.rolling_ops.v0.rollingops import RollingOpsManager, RunWithLock
 from ops import JujuVersion
@@ -69,6 +69,7 @@
     REPLICATION_PASSWORD_KEY,
     REWIND_PASSWORD_KEY,
     SECRET_DELETED_LABEL,
+    SECRET_INTERNAL_LABEL,
     SECRET_KEY_OVERRIDES,
     SNAP_PACKAGES,
     SYSTEM_USERS,
@@ -210,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):
@@ -219,13 +234,13 @@ 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)
-            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
@@ -245,7 +260,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:

diff --git a/src/constants.py b/src/constants.py
@@ -55,6 +55,7 @@
 
 METRICS_PORT = "9187"
 
+SECRET_INTERNAL_LABEL = "internal-secret"
 SECRET_DELETED_LABEL = "None"
 
 APP_SCOPE = "app"

diff --git a/tests/unit/test_charm.py b/tests/unit/test_charm.py
@@ -919,15 +919,15 @@ def test_get_secret_juju(self, _, __):
             # Test application scope.
             _secret_cache.get.return_value = None
             assert self.charm.get_secret("app", "password") is None
-            _secret_cache.get.assert_called_once_with("postgresql.app")
+            _secret_cache.get.assert_called_once_with("postgresql.app", None)
             _secret_cache.reset_mock()
 
             _secret_cache.get.return_value = Mock()
             _secret_cache.get.return_value.get_content.return_value.get.return_value = (
                 sentinel.test_password
             )
             assert self.charm.get_secret("app", "password") == sentinel.test_password
-            _secret_cache.get.assert_called_once_with("postgresql.app")
+            _secret_cache.get.assert_called_once_with("postgresql.app", None)
             _secret_cache.get.return_value.get_content.return_value.get.assert_called_once_with(
                 "password"
             )
@@ -936,15 +936,15 @@ def test_get_secret_juju(self, _, __):
             # Test unit scope.
             _secret_cache.get.return_value = None
             assert self.charm.get_secret("unit", "password") is None
-            _secret_cache.get.assert_called_once_with("postgresql.unit")
+            _secret_cache.get.assert_called_once_with("postgresql.unit", None)
             _secret_cache.reset_mock()
 
             _secret_cache.get.return_value = Mock()
             _secret_cache.get.return_value.get_content.return_value.get.return_value = (
                 sentinel.test_password
             )
             assert self.charm.get_secret("unit", "password") == sentinel.test_password
-            _secret_cache.get.assert_called_once_with("postgresql.unit")
+            _secret_cache.get.assert_called_once_with("postgresql.unit", None)
             _secret_cache.get.return_value.get_content.return_value.get.assert_called_once_with(
                 "password"
             )
@@ -989,7 +989,7 @@ def test_set_secret_juju(self, _, __):
         with patch.object(self.charm, "secrets") as _secret_cache:
             # Test application scope.
             self.charm.set_secret("app", "password", "test-password")
-            _secret_cache.get.assert_called_once_with("postgresql.app")
+            _secret_cache.get.assert_called_once_with("postgresql.app", None)
             _secret_cache.get().get_content().update.assert_called_once_with(
                 {"password": "test-password"}
             )
@@ -1004,7 +1004,7 @@ def test_set_secret_juju(self, _, __):
 
             # Test unit scope.
             self.charm.set_secret("unit", "password", "test-password")
-            _secret_cache.get.assert_called_once_with("postgresql.unit")
+            _secret_cache.get.assert_called_once_with("postgresql.unit", None)
             _secret_cache.get().get_content().update.assert_called_once_with(
                 {"password": "test-password"}
             )