[DPE-6344] LDAP I: Create access groups (#881)

Author: sinclert-canonical

lib/charms/postgresql_k8s/v0/postgresql.py

@@ -35,7 +35,19 @@
 
 # Increment this PATCH version before using `charmcraft publish-lib` or reset
 # to 0 if you are raising the major API version
-LIBPATCH = 46
+LIBPATCH = 47
+
+# Groups to distinguish HBA access
+ACCESS_GROUP_IDENTITY = "identity_access"
+ACCESS_GROUP_INTERNAL = "internal_access"
+ACCESS_GROUP_RELATION = "relation_access"
+
+# List of access groups to filter role assignments by
+ACCESS_GROUPS = [
+    ACCESS_GROUP_IDENTITY,
+    ACCESS_GROUP_INTERNAL,
+    ACCESS_GROUP_RELATION,
+]
 
 # Groups to distinguish database permissions
 PERMISSIONS_GROUP_ADMIN = "admin"
@@ -57,10 +69,18 @@
 logger = logging.getLogger(__name__)
 
 
+class PostgreSQLAssignGroupError(Exception):
+    """Exception raised when assigning to a group fails."""
+
+
 class PostgreSQLCreateDatabaseError(Exception):
     """Exception raised when creating a database fails."""
 
 
+class PostgreSQLCreateGroupError(Exception):
+    """Exception raised when creating a group fails."""
+
+
 class PostgreSQLCreateUserError(Exception):
     """Exception raised when creating a user fails."""
 
@@ -93,6 +113,10 @@ class PostgreSQLGetPostgreSQLVersionError(Exception):
     """Exception raised when retrieving PostgreSQL version fails."""
 
 
+class PostgreSQLListGroupsError(Exception):
+    """Exception raised when retrieving PostgreSQL groups list fails."""
+
+
 class PostgreSQLListUsersError(Exception):
     """Exception raised when retrieving PostgreSQL users list fails."""
 
@@ -160,6 +184,24 @@ def _connect_to_database(
         connection.autocommit = True
         return connection
 
+    def create_access_groups(self) -> None:
+        """Create access groups to distinguish HBA authentication methods."""
+        connection = None
+        try:
+            with self._connect_to_database() as connection, connection.cursor() as cursor:
+                for group in ACCESS_GROUPS:
+                    cursor.execute(
+                        SQL("CREATE ROLE {} NOLOGIN;").format(
+                            Identifier(group),
+                        )
+                    )
+        except psycopg2.Error as e:
+            logger.error(f"Failed to create access groups: {e}")
+            raise PostgreSQLCreateGroupError() from e
+        finally:
+            if connection is not None:
+                connection.close()
+
     def create_database(
         self,
         database: str,
@@ -321,6 +363,50 @@ def delete_user(self, user: str) -> None:
             logger.error(f"Failed to delete user: {e}")
             raise PostgreSQLDeleteUserError() from e
 
+    def grant_internal_access_group_memberships(self) -> None:
+        """Grant membership to the internal access-group to existing internal users."""
+        connection = None
+        try:
+            with self._connect_to_database() as connection, connection.cursor() as cursor:
+                for user in self.system_users:
+                    cursor.execute(
+                        SQL("GRANT {} TO {};").format(
+                            Identifier(ACCESS_GROUP_INTERNAL),
+                            Identifier(user),
+                        )
+                    )
+        except psycopg2.Error as e:
+            logger.error(f"Failed to grant internal access group memberships: {e}")
+            raise PostgreSQLAssignGroupError() from e
+        finally:
+            if connection is not None:
+                connection.close()
+
+    def grant_relation_access_group_memberships(self) -> None:
+        """Grant membership to the relation access-group to existing relation users."""
+        rel_users = self.list_users_from_relation()
+        if not rel_users:
+            return
+
+        connection = None
+        try:
+            with self._connect_to_database() as connection, connection.cursor() as cursor:
+                rel_groups = SQL(",").join(Identifier(group) for group in [ACCESS_GROUP_RELATION])
+                rel_users = SQL(",").join(Identifier(user) for user in rel_users)
+
+                cursor.execute(
+                    SQL("GRANT {groups} TO {users};").format(
+                        groups=rel_groups,
+                        users=rel_users,
+                    )
+                )
+        except psycopg2.Error as e:
+            logger.error(f"Failed to grant relation access group memberships: {e}")
+            raise PostgreSQLAssignGroupError() from e
+        finally:
+            if connection is not None:
+                connection.close()
+
     def enable_disable_extensions(
         self, extensions: Dict[str, bool], database: Optional[str] = None
     ) -> None:
@@ -534,6 +620,26 @@ def is_tls_enabled(self, check_current_host: bool = False) -> bool:
             # Connection errors happen when PostgreSQL has not started yet.
             return False
 
+    def list_access_groups(self) -> Set[str]:
+        """Returns the list of PostgreSQL database access groups.
+
+        Returns:
+            List of PostgreSQL database access groups.
+        """
+        try:
+            with self._connect_to_database() as connection, connection.cursor() as cursor:
+                cursor.execute(
+                    "SELECT groname FROM pg_catalog.pg_group WHERE groname LIKE '%_access';"
+                )
+                access_groups = cursor.fetchall()
+                return {group[0] for group in access_groups}
+        except psycopg2.Error as e:
+            logger.error(f"Failed to list PostgreSQL database access groups: {e}")
+            raise PostgreSQLListGroupsError() from e
+        finally:
+            if connection is not None:
+                connection.close()
+
     def list_users(self) -> Set[str]:
         """Returns the list of PostgreSQL database users.
 
@@ -548,6 +654,29 @@ def list_users(self) -> Set[str]:
         except psycopg2.Error as e:
             logger.error(f"Failed to list PostgreSQL database users: {e}")
             raise PostgreSQLListUsersError() from e
+        finally:
+            if connection is not None:
+                connection.close()
+
+    def list_users_from_relation(self) -> Set[str]:
+        """Returns the list of PostgreSQL database users that were created by a relation.
+
+        Returns:
+            List of PostgreSQL database users.
+        """
+        try:
+            with self._connect_to_database() as connection, connection.cursor() as cursor:
+                cursor.execute(
+                    "SELECT usename FROM pg_catalog.pg_user WHERE usename LIKE 'relation_id_%';"
+                )
+                usernames = cursor.fetchall()
+                return {username[0] for username in usernames}
+        except psycopg2.Error as e:
+            logger.error(f"Failed to list PostgreSQL database users: {e}")
+            raise PostgreSQLListUsersError() from e
+        finally:
+            if connection is not None:
+                connection.close()
 
     def list_valid_privileges_and_roles(self) -> Tuple[Set[str], Set[str]]:
         """Returns two sets with valid privileges and roles.

src/charm.py

@@ -35,6 +35,7 @@
 from charms.grafana_k8s.v0.grafana_dashboard import GrafanaDashboardProvider
 from charms.loki_k8s.v1.loki_push_api import LogProxyConsumer
 from charms.postgresql_k8s.v0.postgresql import (
+    ACCESS_GROUPS,
     REQUIRED_PLUGINS,
     PostgreSQL,
     PostgreSQLEnableDisableExtensionError,
@@ -1098,6 +1099,11 @@ def _initialize_cluster(self, event: WorkloadEvent) -> bool:
 
         self.postgresql.set_up_database()
 
+        access_groups = self.postgresql.list_access_groups()
+        if access_groups != set(ACCESS_GROUPS):
+            self.postgresql.create_access_groups()
+            self.postgresql.grant_internal_access_group_memberships()
+
         # Mark the cluster as initialised.
         self._peers.data[self.app]["cluster_initialised"] = "True"
 

src/relations/db.py

@@ -7,6 +7,7 @@
 from typing import Iterable
 
 from charms.postgresql_k8s.v0.postgresql import (
+    ACCESS_GROUP_RELATION,
     PostgreSQLCreateDatabaseError,
     PostgreSQLCreateUserError,
     PostgreSQLDeleteUserError,
@@ -173,9 +174,11 @@ def set_up_relation(self, relation: Relation) -> bool:
             # created in a previous relation changed event.
             user = f"relation_id_{relation.id}"
             password = unit_relation_databag.get("password", new_password())
-            self.charm.postgresql.create_user(user, password, self.admin)
-            plugins = self.charm.get_plugins()
+            self.charm.postgresql.create_user(
+                user, password, self.admin, extra_user_roles=[ACCESS_GROUP_RELATION]
+            )
 
+            plugins = self.charm.get_plugins()
             self.charm.postgresql.create_database(
                 database, user, plugins=plugins, client_relations=self.charm.client_relations
             )

src/relations/postgresql_provider.py

@@ -10,6 +10,8 @@
     DatabaseRequestedEvent,
 )
 from charms.postgresql_k8s.v0.postgresql import (
+    ACCESS_GROUP_RELATION,
+    ACCESS_GROUPS,
     INVALID_EXTRA_USER_ROLE_BLOCKING_MESSAGE,
     PostgreSQLCreateDatabaseError,
     PostgreSQLCreateUserError,
@@ -71,7 +73,10 @@ def _sanitize_extra_roles(extra_roles: str | None) -> list[str]:
         if extra_roles is None:
             return []
 
-        return [role.lower() for role in extra_roles.split(",")]
+        # Make sure the access-groups are not in the list
+        extra_roles_list = [role.lower() for role in extra_roles.split(",")]
+        extra_roles_list = [role for role in extra_roles_list if role not in ACCESS_GROUPS]
+        return extra_roles_list
 
     def _on_database_requested(self, event: DatabaseRequestedEvent) -> None:
         """Handle the legacy postgresql-client relation changed event.
@@ -89,8 +94,9 @@ def _on_database_requested(self, event: DatabaseRequestedEvent) -> None:
         # Retrieve the database name and extra user roles using the charm library.
         database = event.database
 
-        # Make sure that certain groups are not in the list
+        # Make sure the relation access-group is added to the list
         extra_user_roles = self._sanitize_extra_roles(event.extra_user_roles)
+        extra_user_roles.append(ACCESS_GROUP_RELATION)
 
         try:
             # Creates the user and the database for this specific relation.

src/upgrade.py

@@ -12,6 +12,7 @@
     DependencyModel,
     KubernetesClientError,
 )
+from charms.postgresql_k8s.v0.postgresql import ACCESS_GROUPS
 from lightkube.core.client import Client
 from lightkube.core.exceptions import ApiError
 from lightkube.resources.apps_v1 import StatefulSet
@@ -121,6 +122,7 @@ def _on_postgresql_pebble_ready(self, event: WorkloadEvent) -> None:
                 event.defer()
                 return
             self._set_up_new_credentials_for_legacy()
+            self._set_up_new_access_roles_for_legacy()
 
         try:
             for attempt in Retrying(stop=stop_after_attempt(6), wait=wait_fixed(10)):
@@ -270,6 +272,16 @@ def _set_first_rolling_update_partition(self) -> None:
         except KubernetesClientError as e:
             raise ClusterNotReadyError(e.message, e.cause) from e
 
+    def _set_up_new_access_roles_for_legacy(self) -> None:
+        """Create missing access groups and their memberships."""
+        access_groups = self.charm.postgresql.list_access_groups()
+        if access_groups == set(ACCESS_GROUPS):
+            return
+
+        self.charm.postgresql.create_access_groups()
+        self.charm.postgresql.grant_internal_access_group_memberships()
+        self.charm.postgresql.grant_relation_access_group_memberships()
+
     def _set_up_new_credentials_for_legacy(self) -> None:
         """Create missing password and user."""
         for key in (MONITORING_PASSWORD_KEY, PATRONI_PASSWORD_KEY):

tests/unit/test_db.py

@@ -5,6 +5,7 @@
 
 import pytest
 from charms.postgresql_k8s.v0.postgresql import (
+    ACCESS_GROUP_RELATION,
     PostgreSQLCreateDatabaseError,
     PostgreSQLCreateUserError,
     PostgreSQLGetPostgreSQLVersionError,
@@ -230,7 +231,9 @@ def test_set_up_relation(harness):
             )
         assert harness.charm.legacy_db_relation.set_up_relation(relation)
         user = f"relation_id_{rel_id}"
-        postgresql_mock.create_user.assert_called_once_with(user, "test-password", False)
+        postgresql_mock.create_user.assert_called_once_with(
+            user, "test-password", False, extra_user_roles=[ACCESS_GROUP_RELATION]
+        )
         postgresql_mock.create_database.assert_called_once_with(
             DATABASE, user, plugins=["pgaudit"], client_relations=[relation]
         )
@@ -275,7 +278,9 @@ def test_set_up_relation(harness):
             )
             clear_relation_data(harness)
         assert harness.charm.legacy_db_relation.set_up_relation(relation)
-        postgresql_mock.create_user.assert_called_once_with(user, "test-password", False)
+        postgresql_mock.create_user.assert_called_once_with(
+            user, "test-password", False, extra_user_roles=[ACCESS_GROUP_RELATION]
+        )
         postgresql_mock.create_database.assert_called_once_with(
             DATABASE, user, plugins=["pgaudit"], client_relations=[relation]
         )

tests/unit/test_postgresql.py

@@ -5,6 +5,8 @@
 import psycopg2
 import pytest
 from charms.postgresql_k8s.v0.postgresql import (
+    ACCESS_GROUP_INTERNAL,
+    ACCESS_GROUPS,
     PERMISSIONS_GROUP_ADMIN,
     PostgreSQLCreateDatabaseError,
     PostgreSQLGetLastArchivedWALError,
@@ -19,6 +21,7 @@
     PEER,
     REPLICATION_USER,
     REWIND_USER,
+    SYSTEM_USERS,
     USER,
 )
 
@@ -35,6 +38,21 @@ def harness():
     harness.cleanup()
 
 
+def test_create_access_groups(harness):
+    with patch(
+        "charms.postgresql_k8s.v0.postgresql.PostgreSQL._connect_to_database"
+    ) as _connect_to_database:
+        execute = _connect_to_database.return_value.__enter__.return_value.cursor.return_value.__enter__.return_value.execute
+        harness.charm.postgresql.create_access_groups()
+
+        execute.assert_has_calls([
+            *(
+                call(SQL("CREATE ROLE {} NOLOGIN;").format(Identifier(group)))
+                for group in ACCESS_GROUPS
+            ),
+        ])
+
+
 def test_create_database(harness):
     with (
         patch(
@@ -165,6 +183,35 @@ def test_create_database(harness):
         _enable_disable_extensions.assert_not_called()
 
 
+def test_grant_internal_access_group_memberships(harness):
+    with patch(
+        "charms.postgresql_k8s.v0.postgresql.PostgreSQL._connect_to_database"
+    ) as _connect_to_database:
+        execute = _connect_to_database.return_value.__enter__.return_value.cursor.return_value.__enter__.return_value.execute
+        harness.charm.postgresql.grant_internal_access_group_memberships()
+
+        internal_group = Identifier(ACCESS_GROUP_INTERNAL)
+
+        execute.assert_has_calls([
+            *(
+                call(SQL("GRANT {} TO {};").format(internal_group, Identifier(user)))
+                for user in SYSTEM_USERS
+            ),
+        ])
+
+
+def test_grant_relation_access_group_memberships(harness):
+    with patch(
+        "charms.postgresql_k8s.v0.postgresql.PostgreSQL._connect_to_database"
+    ) as _connect_to_database:
+        execute = _connect_to_database.return_value.__enter__.return_value.cursor.return_value.__enter__.return_value.execute
+        harness.charm.postgresql.grant_relation_access_group_memberships()
+
+        execute.assert_has_calls([
+            call("SELECT usename FROM pg_catalog.pg_user WHERE usename LIKE 'relation_id_%';"),
+        ])
+
+
 def test_generate_database_privileges_statements(harness):
     # Test with only one established relation.
     assert harness.charm.postgresql._generate_database_privileges_statements(