Adapt postgresql provider class

canonical · Mar 4, 2025 · de8a7ac · de8a7ac
1 parent e87e366
commit de8a7ac
Show file tree

Hide file tree

Showing 2 changed files with 29 additions and 6 deletions.
diff --git a/src/charm.py b/src/charm.py
@@ -22,6 +22,7 @@
 from charms.grafana_agent.v0.cos_agent import COSAgentProvider, charm_tracing_config
 from charms.operator_libs_linux.v1 import systemd
 from charms.operator_libs_linux.v2 import snap
+from charms.postgresql_k8s.v0.postgresql import PERMISSIONS_GROUP_ADMIN
 from charms.postgresql_k8s.v0.postgresql_tls import PostgreSQLTLS
 from charms.tempo_coordinator_k8s.v0.charm_tracing import trace_charm
 from jinja2 import Template
@@ -697,16 +698,23 @@ def generate_relation_databases(self) -> Dict[str, Dict[str, Union[str, bool]]]:
             fields=["database", "extra-user-roles"]
         ).items():
             database = data.get("database")
-            roles = data.get("extra-user-roles", "").lower().split(",")
+            extra_user_roles = data.get("extra-user-roles")
+            extra_user_roles = self.client_relation.sanitize_extra_roles(extra_user_roles)
             if database:
                 databases[str(rel_id)] = {
                     "name": database,
                     "legacy": False,
                 }
-            if "admin" in roles or "superuser" in roles or "createdb" in roles:
+            if (
+                PERMISSIONS_GROUP_ADMIN in extra_user_roles
+                or "superuser" in extra_user_roles
+                or "createdb" in extra_user_roles
+            ):
                 add_wildcard = True
+
         if add_wildcard:
             databases["*"] = {"name": "*", "auth_dbname": database, "legacy": False}
+
         self.set_relation_databases(databases)
         return databases
 

diff --git a/src/relations/pgbouncer_provider.py b/src/relations/pgbouncer_provider.py
@@ -45,6 +45,7 @@
 )
 from charms.pgbouncer_k8s.v0 import pgb
 from charms.postgresql_k8s.v0.postgresql import (
+    PERMISSIONS_GROUP_ADMIN,
     PostgreSQLCreateDatabaseError,
     PostgreSQLCreateUserError,
     PostgreSQLDeleteUserError,
@@ -96,6 +97,14 @@ def __init__(self, charm: CharmBase, relation_name: str = CLIENT_RELATION_NAME)
             charm.on[self.relation_name].relation_broken, self._on_relation_broken
         )
 
+    @staticmethod
+    def sanitize_extra_roles(extra_roles: str | None) -> list[str]:
+        """Standardize and sanitize user extra-roles."""
+        if extra_roles is None:
+            return []
+
+        return [role.lower() for role in extra_roles.split(",")]
+
     def _depart_flag(self, relation):
         return f"{self.relation_name}_{relation.id}_departing"
 
@@ -128,13 +137,19 @@ def _on_database_requested(self, event: DatabaseRequestedEvent) -> None:
 
         # Retrieve the database name and extra user roles using the charm library.
         database = event.database
-        extra_user_roles = event.extra_user_roles or ""
+
+        # Make sure that certain groups are not in the list
+        extra_user_roles = self.sanitize_extra_roles(event.extra_user_roles)
 
         dbs = self.charm.generate_relation_databases()
-        dbs[str(event.relation.id)] = {"name": database, "legacy": False}
-        roles = extra_user_roles.lower().split(",")
-        if "admin" in roles or "superuser" in roles or "createdb" in roles:
+        dbs[str(rel_id)] = {"name": database, "legacy": False}
+        if (
+            PERMISSIONS_GROUP_ADMIN in extra_user_roles
+            or "superuser" in extra_user_roles
+            or "createdb" in extra_user_roles
+        ):
             dbs["*"] = {"name": "*", "auth_dbname": database}
+
         self.charm.set_relation_databases(dbs)
 
         pgb_dbs_hash = shake_128(