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/c1ef734fd749b360e14fdb5a24c55315cb0a9000.zip"}
 
 [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,21 @@
 from typing import Literal, get_args
 from urllib.parse import urlparse
 
+from single_kernel_postgresql.utils.postgresql import (
+    ACCESS_GROUP_IDENTITY,
+    ACCESS_GROUPS,
+    REQUIRED_PLUGINS,
+    PostgreSQL,
+    PostgreSQLCreatePredefinedRolesError,
+    PostgreSQLCreateUserError,
+    PostgreSQLEnableDisableExtensionError,
+    PostgreSQLGetCurrentTimelineError,
+    PostgreSQLGrantDatabasePrivilegesToUserError,
+    PostgreSQLListGroupsError,
+    PostgreSQLListUsersError,
+    PostgreSQLUpdateUserPasswordError,
+)
+
 from authorisation_rules_observer import (
     AuthorisationRulesChangeCharmEvents,
     AuthorisationRulesObserver,
@@ -42,16 +57,6 @@
 from charms.data_platform_libs.v0.data_models import TypedCharmBase
 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_GROUP_IDENTITY,
-    ACCESS_GROUPS,
-    REQUIRED_PLUGINS,
-    PostgreSQL,
-    PostgreSQLEnableDisableExtensionError,
-    PostgreSQLGetCurrentTimelineError,
-    PostgreSQLListGroupsError,
-    PostgreSQLUpdateUserPasswordError,
-)
 from charms.postgresql_k8s.v0.postgresql_tls import PostgreSQLTLS
 from charms.prometheus_k8s.v0.prometheus_scrape import MetricsEndpointProvider
 from charms.rolling_ops.v0.rollingops import RollingOpsManager, RunWithLock
@@ -1164,6 +1169,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 +1215,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 +1636,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

@@ -9,7 +9,10 @@
     DatabaseProvides,
     DatabaseRequestedEvent,
 )
-from charms.postgresql_k8s.v0.postgresql import (
+from ops.charm import CharmBase, RelationBrokenEvent, RelationDepartedEvent
+from ops.framework import Object
+from ops.model import ActiveStatus, BlockedStatus, Relation
+from single_kernel_postgresql.utils.postgresql import (
     ACCESS_GROUP_RELATION,
     ACCESS_GROUPS,
     INVALID_EXTRA_USER_ROLE_BLOCKING_MESSAGE,
@@ -18,9 +21,6 @@
     PostgreSQLDeleteUserError,
     PostgreSQLGetPostgreSQLVersionError,
 )
-from ops.charm import CharmBase, RelationBrokenEvent, RelationDepartedEvent
-from ops.framework import Object
-from ops.model import ActiveStatus, BlockedStatus, Relation
 
 from constants import DATABASE_PORT
 from utils import new_password
@@ -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.

src/upgrade.py

@@ -13,13 +13,13 @@
     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
 from ops.charm import UpgradeCharmEvent, WorkloadEvent
 from ops.model import BlockedStatus, MaintenanceStatus, RelationDataContent
 from pydantic import BaseModel
+from single_kernel_postgresql.utils.postgresql import ACCESS_GROUPS
 from tenacity import RetryError, Retrying, stop_after_attempt, wait_fixed
 
 from constants import APP_SCOPE, MONITORING_PASSWORD_KEY, MONITORING_USER, PATRONI_PASSWORD_KEY

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 %}
@@ -146,6 +154,13 @@ postgresql:
   pg_hba:
   - local all backup peer map=operator
   - local all monitoring password
+  - {{ 'hostssl' if enable_tls else 'host' }} all +charmed_stats 0.0.0.0/0 scram-sha-256
+  - {{ 'hostssl' if enable_tls else 'host' }} all +charmed_read 0.0.0.0/0 scram-sha-256
+  - {{ 'hostssl' if enable_tls else 'host' }} all +charmed_dml 0.0.0.0/0 scram-sha-256
+  - {{ 'hostssl' if enable_tls else 'host' }} all +charmed_backup 0.0.0.0/0 scram-sha-256
+  - {{ 'hostssl' if enable_tls else 'host' }} all +charmed_dba 0.0.0.0/0 scram-sha-256
+  - {{ 'hostssl' if enable_tls else 'host' }} all +charmed_admin 0.0.0.0/0 scram-sha-256
+  - {{ 'hostssl' if enable_tls else 'host' }} all +charmed_databases_owner 0.0.0.0/0 scram-sha-256
   {%- if not connectivity %}
   - {{ 'hostssl' if enable_tls else 'host' }} all all {{ endpoint }}.{{ namespace }}.svc.cluster.local md5
   - {{ 'hostssl' if enable_tls else 'host' }} all all 0.0.0.0/0 reject