Implement instance level predefined roles

Author: shayancanonical

lib/charms/postgresql_k8s/v0/postgresql.py

@@ -35,7 +35,7 @@
 
 # Increment this PATCH version before using `charmcraft publish-lib` or reset
 # to 0 if you are raising the major API version
-LIBPATCH = 52
+LIBPATCH = 53
 
 # Groups to distinguish HBA access
 ACCESS_GROUP_IDENTITY = "identity_access"
@@ -49,6 +49,11 @@
     ACCESS_GROUP_RELATION,
 ]
 
+ROLE_STATS = "charmed_stats"
+ROLE_READ = "charmed_read"
+ROLE_DML = "charmed_dml"
+ROLE_BACKUP = "charmed_backup"
+
 # Groups to distinguish database permissions
 PERMISSIONS_GROUP_ADMIN = "admin"
 
@@ -125,6 +130,14 @@ class PostgreSQLUpdateUserPasswordError(Exception):
     """Exception raised when updating a user password fails."""
 
 
+class PostgreSQLCreatePredefinedRolesError(Exception):
+    """Exception raised when creating predefined roles."""
+
+
+class PostgreSQLGrantDatabasePrivilegesToUserError(Exception):
+    """Exception raised when granting database privileges to user."""
+
+
 class PostgreSQL:
     """Class to encapsulate all operations related to interacting with PostgreSQL instance."""
 
@@ -326,6 +339,60 @@ def create_user(
             logger.error(f"Failed to create user: {e}")
             raise PostgreSQLCreateUserError() from e
 
+    def create_predefined_roles(self) -> None:
+        """Create predefined roles."""
+        role_to_queries = {
+            ROLE_STATS: [
+                f"CREATE ROLE {ROLE_STATS} NOSUPERUSER NOCREATEDB NOCREATEROLE NOREPLICATION NOLOGIN IN ROLE pg_monitor",
+            ],
+            ROLE_READ: [
+                f"CREATE ROLE {ROLE_READ} NOSUPERUSER NOCREATEDB NOCREATEROLE NOREPLICATION NOLOGIN IN ROLE pg_read_all_data",
+            ],
+            ROLE_DML: [
+                f"CREATE ROLE {ROLE_DML} NOSUPERUSER NOCREATEDB NOCREATEROLE NOREPLICATION NOLOGIN IN ROLE pg_write_all_data",
+            ],
+            ROLE_BACKUP: [
+                f"CREATE ROLE {ROLE_BACKUP} NOSUPERUSER NOCREATEDB NOCREATEROLE NOREPLICATION NOLOGIN IN ROLE pg_checkpoint",
+                f"GRANT {ROLE_STATS} TO {ROLE_BACKUP}",
+                f"GRANT execute ON FUNCTION pg_backup_start TO {ROLE_BACKUP}",
+                f"GRANT execute ON FUNCTION pg_backup_stop TO {ROLE_BACKUP}",
+                f"GRANT execute ON FUNCTION pg_create_restore_point TO {ROLE_BACKUP}",
+                f"GRANT execute ON FUNCTION pg_switch_wal TO {ROLE_BACKUP}",
+            ],
+        }
+
+        _, existing_roles = self.list_valid_privileges_and_roles()
+
+        try:
+            with self._connect_to_database() as connection, connection.cursor() as cursor:
+                for role, queries in role_to_queries.items():
+                    if role in existing_roles:
+                        logger.debug(f"Role {role} already exists")
+                        continue
+
+                    logger.info(f"Creating predefined role {role}")
+
+                    for query in queries:
+                        cursor.execute(SQL(query))
+        except psycopg2.Error as e:
+            logger.error(f"Failed to create predefined roles: {e}")
+            raise PostgreSQLCreatePredefinedRolesError() from e
+
+    def grant_database_privileges_to_user(
+        self, user: str, database: str, privileges: list[str]
+    ) -> None:
+        """Grant the specified privileges on the provided database for the user."""
+        try:
+            with self._connect_to_database() as connection, connection.cursor() as cursor:
+                cursor.execute(
+                    SQL("GRANT {} ON DATABASE {} TO {};").format(
+                        Identifier(", ".join(privileges)), Identifier(database), Identifier(user)
+                    )
+                )
+        except psycopg2.Error as e:
+            logger.error(f"Failed to grant privileges to user: {e}")
+            raise PostgreSQLGrantDatabasePrivilegesToUserError() from e
+
     def delete_user(self, user: str) -> None:
         """Deletes a database user.
 
@@ -727,7 +794,7 @@ def set_up_database(self, temp_location: Optional[str] = None) -> None:
                     )
                 self.create_user(
                     PERMISSIONS_GROUP_ADMIN,
-                    extra_user_roles=["pg_read_all_data", "pg_write_all_data"],
+                    extra_user_roles=[ROLE_READ, ROLE_DML],
                 )
                 cursor.execute("GRANT CONNECT ON DATABASE postgres TO admin;")
         except psycopg2.Error as e:

src/backups.py

@@ -22,7 +22,6 @@
 from ops import HookEvent
 from ops.charm import ActionEvent
 from ops.framework import Object
-from ops.jujuversion import JujuVersion
 from ops.model import ActiveStatus, MaintenanceStatus
 from ops.pebble import ChangeError, ExecError
 from tenacity import RetryError, Retrying, stop_after_attempt, wait_fixed
@@ -793,12 +792,11 @@ def _on_create_backup_action(self, event) -> None:  # noqa: C901
 
         # Test uploading metadata to S3 to test credentials before backup.
         datetime_backup_requested = datetime.now().strftime("%Y-%m-%dT%H:%M:%SZ")
-        juju_version = JujuVersion.from_environ()
         metadata = f"""Date Backup Requested: {datetime_backup_requested}
 Model Name: {self.model.name}
 Application Name: {self.model.app.name}
 Unit Name: {self.charm.unit.name}
-Juju Version: {juju_version!s}
+Juju Version: {self.charm.model.juju_version!s}
 """
         if not self._upload_content_to_s3(
             metadata,

src/charm.py

@@ -39,7 +39,10 @@
     ACCESS_GROUP_IDENTITY,
     ACCESS_GROUPS,
     REQUIRED_PLUGINS,
+    ROLE_BACKUP,
+    ROLE_STATS,
     PostgreSQL,
+    PostgreSQLCreatePredefinedRolesError,
     PostgreSQLEnableDisableExtensionError,
     PostgreSQLGetCurrentTimelineError,
     PostgreSQLUpdateUserPasswordError,
@@ -1134,16 +1137,26 @@ def _initialize_cluster(self, event: WorkloadEvent) -> bool:
             event.defer()
             return False
 
+        try:
+            self.postgresql.create_predefined_roles()
+        except PostgreSQLCreatePredefinedRolesError as e:
+            logger.exception(e)
+            self.unit.status = BlockedStatus("Failed to create pre-defined roles")
+            return
+
         pg_users = self.postgresql.list_users()
         # Create the backup user.
         if BACKUP_USER not in pg_users:
-            self.postgresql.create_user(BACKUP_USER, new_password(), admin=True)
+            self.postgresql.create_user(
+                BACKUP_USER, new_password(), extra_user_roles=[ROLE_BACKUP]
+            )
+            self.postgresql.grant_database_privileges_to_user(BACKUP_USER, "postgres", ["connect"])
         # Create the monitoring user.
         if MONITORING_USER not in pg_users:
             self.postgresql.create_user(
                 MONITORING_USER,
                 self.get_secret(APP_SCOPE, MONITORING_PASSWORD_KEY),
-                extra_user_roles=["pg_monitor"],
+                extra_user_roles=[ROLE_STATS],
             )
 
         self.postgresql.set_up_database(temp_location="/var/lib/postgresql/temp")
@@ -1308,13 +1321,14 @@ def _update_admin_password(self, admin_secret_id: str) -> None:
             return
 
         try:
+            updateable_users = [*SYSTEM_USERS, BACKUP_USER]
             # get the secret content and check each user configured there
             # only SYSTEM_USERS with changed passwords are processed, all others ignored
             updated_passwords = self.get_secret_from_id(secret_id=admin_secret_id)
             for user, password in list(updated_passwords.items()):
-                if user not in SYSTEM_USERS:
+                if user not in updateable_users:
                     logger.error(
-                        f"Can only update system users: {', '.join(SYSTEM_USERS)} not {user}"
+                        f"Can only update system users: {', '.join(updateable_users)} not {user}"
                     )
                     updated_passwords.pop(user)
                     continue

src/constants.py

@@ -26,7 +26,7 @@
     "/var/log/postgresql/postgresql*.log",
 ]
 # List of system usernames needed for correct work of the charm/workload.
-SYSTEM_USERS = [BACKUP_USER, REPLICATION_USER, REWIND_USER, USER, MONITORING_USER]
+SYSTEM_USERS = [REPLICATION_USER, REWIND_USER, USER, MONITORING_USER]
 
 # Labels are not confidential
 REPLICATION_PASSWORD_KEY = "replication-password"  # noqa: S105

src/relations/postgresql_provider.py

@@ -156,7 +156,7 @@ def _on_database_requested(self, event: DatabaseRequestedEvent) -> None:
             self.charm.unit.status = BlockedStatus(
                 e.message
                 if issubclass(type(e), PostgreSQLCreateUserError) and e.message is not None
-                else f"Failed to initialize {self.relation_name} relation"
+                else f"Failed to initialize relation {self.relation_name}"
             )
 
     def _on_relation_departed(self, event: RelationDepartedEvent) -> None:
@@ -291,6 +291,11 @@ def _update_unit_status(self, relation: Relation) -> None:
             and not self.check_for_invalid_extra_user_roles(relation.id)
         ):
             self.charm.unit.status = ActiveStatus()
+        if (
+            self.charm.is_blocked
+            and "Failed to initialize relation" in self.charm.unit.status.message
+        ):
+            self.charm.unit.status = ActiveStatus()
 
         self._update_unit_status_on_blocking_endpoint_simultaneously()
 

src/upgrade.py

@@ -12,7 +12,11 @@
     DependencyModel,
     KubernetesClientError,
 )
-from charms.postgresql_k8s.v0.postgresql import ACCESS_GROUPS
+from charms.postgresql_k8s.v0.postgresql import (
+    ACCESS_GROUPS,
+    ROLE_STATS,
+    PostgreSQLCreatePredefinedRolesError,
+)
 from lightkube.core.client import Client
 from lightkube.core.exceptions import ApiError
 from lightkube.resources.apps_v1 import StatefulSet
@@ -292,9 +296,16 @@ def _set_up_new_credentials_for_legacy(self) -> None:
             self.charm.postgresql.create_user(
                 MONITORING_USER,
                 self.charm.get_secret(APP_SCOPE, MONITORING_PASSWORD_KEY),
-                extra_user_roles="pg_monitor",
+                extra_user_roles=[ROLE_STATS],
             )
 
+        try:
+            self.postgresql.create_predefined_roles()
+        except PostgreSQLCreatePredefinedRolesError as e:
+            logger.exception(e)
+            self.unit.status = BlockedStatus("Failed to create pre-defined roles")
+            return
+
     @property
     def unit_upgrade_data(self) -> RelationDataContent:
         """Return the application upgrade data."""

tests/unit/test_backups.py

@@ -1197,7 +1197,6 @@ def test_on_create_backup_action(harness):
         ) as _is_primary,
         patch("charm.PostgreSQLBackups._upload_content_to_s3") as _upload_content_to_s3,
         patch("backups.datetime") as _datetime,
-        patch("ops.JujuVersion.from_environ") as _from_environ,
         patch("charm.PostgreSQLBackups._retrieve_s3_parameters") as _retrieve_s3_parameters,
         patch("charm.PostgreSQLBackups._can_unit_perform_backup") as _can_unit_perform_backup,
     ):
@@ -1233,13 +1232,12 @@ def test_on_create_backup_action(harness):
             [],
         )
         _datetime.now.return_value.strftime.return_value = "2023-01-01T09:00:00Z"
-        _from_environ.return_value = "test-juju-version"
         _upload_content_to_s3.return_value = False
         expected_metadata = f"""Date Backup Requested: 2023-01-01T09:00:00Z
 Model Name: {harness.charm.model.name}
 Application Name: {harness.charm.model.app.name}
 Unit Name: {harness.charm.unit.name}
-Juju Version: test-juju-version
+Juju Version: 0.0.0
 """
         harness.charm.backup._on_create_backup_action(mock_event)
         _upload_content_to_s3.assert_called_once_with(

tests/unit/test_postgresql.py

@@ -80,6 +80,13 @@ def test_create_database(harness):
         harness.charm.postgresql.create_database(database, user, plugins, client_relations)
         execute = _connect_to_database.return_value.cursor.return_value.execute
         execute.assert_has_calls([
+            call(
+                Composed([
+                    SQL("SELECT datname FROM pg_database WHERE datname="),
+                    Literal(database),
+                    SQL(";"),
+                ]),
+            ),
             call(
                 Composed([
                     SQL("REVOKE ALL PRIVILEGES ON DATABASE "),
@@ -105,15 +112,6 @@ def test_create_database(harness):
                     SQL(";"),
                 ])
             ),
-            call(
-                Composed([
-                    SQL("GRANT ALL PRIVILEGES ON DATABASE "),
-                    Identifier(database),
-                    SQL(" TO "),
-                    Identifier(BACKUP_USER),
-                    SQL(";"),
-                ])
-            ),
             call(
                 Composed([
                     SQL("GRANT ALL PRIVILEGES ON DATABASE "),