[DPE-6345] LDAP I: Create access groups

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,12 +620,34 @@ 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.
+        """
+        connection = None
+        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.
 
         Returns:
             List of PostgreSQL database users.
         """
+        connection = None
         try:
             with self._connect_to_database() as connection, connection.cursor() as cursor:
                 cursor.execute("SELECT usename FROM pg_catalog.pg_user;")
@@ -548,6 +656,30 @@ 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.
+        """
+        connection = None
+        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

@@ -23,6 +23,7 @@
 from charms.grafana_agent.v0.cos_agent import COSAgentProvider, charm_tracing_config
 from charms.operator_libs_linux.v2 import snap
 from charms.postgresql_k8s.v0.postgresql import (
+    ACCESS_GROUPS,
     REQUIRED_PLUGINS,
     PostgreSQL,
     PostgreSQLCreateUserError,
@@ -1360,6 +1361,11 @@ def _start_primary(self, event: StartEvent) -> None:
 
         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()
+
         self.postgresql_client_relation.oversee_users()
 
         # Set the flag to enable the replicas to start the Patroni service.

src/relations/db.py

@@ -7,6 +7,7 @@
 from collections.abc import Iterable
 
 from charms.postgresql_k8s.v0.postgresql import (
+    ACCESS_GROUP_RELATION,
     PostgreSQLCreateDatabaseError,
     PostgreSQLCreateUserError,
     PostgreSQLGetPostgreSQLVersionError,
@@ -198,10 +199,11 @@ def set_up_relation(self, relation: Relation) -> bool:
             # non-leader units when the cluster topology changes.
             self.charm.set_secret(APP_SCOPE, user, password)
             self.charm.set_secret(APP_SCOPE, f"{user}-database", database)
+            self.charm.postgresql.create_user(
+                user, password, self.admin, extra_user_roles=[ACCESS_GROUP_RELATION]
+            )
 
-            self.charm.postgresql.create_user(user, password, self.admin)
             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:
         """Generate password and handle user and database creation for the related application."""
@@ -93,8 +98,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,
     UpgradeGrantedEvent,
 )
+from charms.postgresql_k8s.v0.postgresql import ACCESS_GROUPS
 from ops.model import MaintenanceStatus, RelationDataContent, WaitingStatus
 from pydantic import BaseModel
 from tenacity import RetryError, Retrying, stop_after_attempt, wait_fixed
@@ -247,6 +248,17 @@ def _prepare_upgrade_from_legacy(self) -> None:
                 extra_user_roles="pg_monitor",
             )
         self.charm.postgresql.set_up_database()
+        self._set_up_new_access_roles_for_legacy()
+
+    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()
 
     @property
     def unit_upgrade_data(self) -> RelationDataContent:

tests/unit/test_db.py

@@ -5,6 +5,7 @@
 
 import pytest
 from charms.postgresql_k8s.v0.postgresql import (
+    ACCESS_GROUP_RELATION,
     PostgreSQLCreateDatabaseError,
     PostgreSQLCreateUserError,
     PostgreSQLGetPostgreSQLVersionError,
@@ -226,7 +227,9 @@ def test_set_up_relation(harness):
             )
         assert harness.charm.legacy_db_relation.set_up_relation(relation)
         user = f"relation-{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]
         )
@@ -253,7 +256,9 @@ def test_set_up_relation(harness):
                 {"database": DATABASE},
             )
         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]
         )
@@ -274,7 +279,9 @@ def test_set_up_relation(harness):
                 {"database": ""},
             )
         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(
             "test_database", user, plugins=["pgaudit"], client_relations=[relation]
         )

tests/unit/test_postgresql_provider.py

@@ -5,6 +5,7 @@
 
 import pytest
 from charms.postgresql_k8s.v0.postgresql import (
+    ACCESS_GROUP_RELATION,
     PostgreSQLCreateDatabaseError,
     PostgreSQLCreateUserError,
     PostgreSQLGetPostgreSQLVersionError,
@@ -124,10 +125,12 @@ def test_on_database_requested(harness):
 
         # Assert that the correct calls were made.
         user = f"relation-{rel_id}"
+        expected_user_roles = [role.lower() for role in EXTRA_USER_ROLES.split(",")]
+        expected_user_roles.append(ACCESS_GROUP_RELATION)
         postgresql_mock.create_user.assert_called_once_with(
             user,
             "test-password",
-            extra_user_roles=[role.lower() for role in EXTRA_USER_ROLES.split(",")],
+            extra_user_roles=expected_user_roles,
         )
         database_relation = harness.model.get_relation(RELATION_NAME)
         client_relations = [database_relation]