Port predefined roles changes from VM charm and use the single kernel library

pyproject.toml

@@ -17,6 +17,7 @@ jinja2 = "^3.1.6"
 lightkube = "^0.17.2"
 lightkube-models = "^1.28.1.4"
 psycopg2 = "^2.9.10"
+postgresql-charms-single-kernel = {url = "https://github.com/marceloneppel/postgresql-single-kernel-library/archive/refs/tags/test1.tar.gz"}
 
 [tool.poetry.group.charm-libs.dependencies]
 # data_platform_libs/v0/data_interfaces.py
@@ -69,6 +70,7 @@ psycopg2-binary = "^2.9.10"
 boto3 = "*"
 tenacity = "^9.1.2"
 allure-pytest = "^2.15.0"
+jubilant = "^1.3.0"
 
 [build-system]
 requires = ["poetry-core>=1.0.0"]

src/charm.py

@@ -18,6 +18,14 @@
 from typing import Literal, get_args
 from urllib.parse import urlparse
 
+from single_kernel_postgresql.utils.postgresql_connection import (
+    PostgreSQLConnection,
+    PostgreSQLCreatePredefinedRolesError,
+    PostgreSQLCreateUserError,
+    PostgreSQLGrantDatabasePrivilegesToUserError,
+    PostgreSQLListUsersError,
+)
+
 from authorisation_rules_observer import (
     AuthorisationRulesChangeCharmEvents,
     AuthorisationRulesObserver,
@@ -450,7 +458,7 @@ def is_unit_stopped(self) -> bool:
     @property
     def postgresql(self) -> PostgreSQL:
         """Returns an instance of the object used to interact with the database."""
-        return PostgreSQL(
+        return PostgreSQLConnection(
             primary_host=self.primary_endpoint,
             current_host=self.endpoint,
             user=USER,
@@ -1164,6 +1172,33 @@ def _initialize_cluster(self, event: WorkloadEvent) -> bool:
             event.defer()
             return False
 
+        try:
+            self._setup_users()
+        except PostgreSQLCreatePredefinedRolesError as e:
+            logger.exception(e)
+            self.unit.status = BlockedStatus("Failed to create pre-defined roles")
+            return False
+        except PostgreSQLGrantDatabasePrivilegesToUserError as e:
+            logger.exception(e)
+            self.unit.status = BlockedStatus("Failed to grant database privileges to user")
+            return False
+        except PostgreSQLCreateUserError as e:
+            logger.exception(e)
+            self.set_unit_status(BlockedStatus("Failed to create postgres user"))
+            return False
+        except PostgreSQLListUsersError:
+            logger.warning("Deferring on_start: Unable to list users")
+            event.defer()
+            return False
+
+        # Mark the cluster as initialised.
+        self._peers.data[self.app]["cluster_initialised"] = "True"
+
+        return True
+
+    def _setup_users(self) -> None:
+        self.postgresql.create_predefined_instance_roles()
+
         pg_users = self.postgresql.list_users()
         # Create the backup user.
         if BACKUP_USER not in pg_users:
@@ -1183,11 +1218,6 @@ def _initialize_cluster(self, event: WorkloadEvent) -> bool:
             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"
-
-        return True
-
     @property
     def is_blocked(self) -> bool:
         """Returns whether the unit is in a blocked state."""
@@ -1609,6 +1639,12 @@ def _was_restore_successful(self, container: Container, service: ServiceInfo) ->
             logger.debug("Restore check early exit: Patroni has not started yet")
             return False
 
+        try:
+            self._setup_users()
+        except Exception as e:
+            logger.exception(e)
+            return False
+
         restoring_backup = self.app_peer_data.get("restoring-backup")
         restore_timeline = self.app_peer_data.get("restore-timeline")
         restore_to_time = self.app_peer_data.get("restore-to-time")

src/relations/postgresql_provider.py

@@ -107,12 +107,13 @@ def _on_database_requested(self, event: DatabaseRequestedEvent) -> None:
         try:
             # Creates the user and the database for this specific relation.
             user = f"relation_id_{event.relation.id}"
-            password = new_password()
-            self.charm.postgresql.create_user(user, password, extra_user_roles=extra_user_roles)
             plugins = self.charm.get_plugins()
 
-            self.charm.postgresql.create_database(
-                database, user, plugins=plugins, client_relations=self.charm.client_relations
+            self.charm.postgresql.create_database(database, plugins=plugins)
+
+            password = new_password()
+            self.charm.postgresql.create_user(
+                user, password, extra_user_roles=extra_user_roles, database=database
             )
 
             # Share the credentials with the application.

templates/patroni.yml.j2

@@ -46,7 +46,11 @@ bootstrap:
         log_truncate_on_rotation: 'on'
         logging_collector: 'on'
         wal_level: logical
-        shared_preload_libraries: 'timescaledb,pgaudit'
+        shared_preload_libraries: 'timescaledb,pgaudit,set_user'
+        session_preload_libraries: 'login_hook'
+        set_user.block_log_statement: 'on'
+        set_user.exit_on_error: 'on'
+        set_user.superuser_allowlist: '+charmed_dba'
     {%- if slots %}
     slots:
     {%- for slot, database in slots.items() %}
@@ -124,7 +128,11 @@ postgresql:
   bin_dir: /usr/lib/postgresql/{{ version }}/bin
   listen: 0.0.0.0:5432
   parameters:
-    shared_preload_libraries: 'timescaledb,pgaudit'
+    shared_preload_libraries: 'timescaledb,pgaudit,set_user'
+    session_preload_libraries: 'login_hook'
+    set_user.block_log_statement: 'on'
+    set_user.exit_on_error: 'on'
+    set_user.superuser_allowlist: '+charmed_dba'
     {%- if enable_pgbackrest_archiving %}
     archive_command: 'pgbackrest --stanza={{ stanza }} archive-push %p'
     {% else %}