[DPE-7498] Create DBA role (#932)

* Implement instance level predefined roles

* Fix minor bug introduced while rebasing off of 16/edge

* Add integration test for charmed_read and charmed_dml roles

* Revert all major changes except introduction of predefined roles

* Sweep diff and minor bug fixes

* Avoid creating set_user extension

* Port Carl's fix for broken unit tests

* Create DBA role

Signed-off-by: Marcelo Henrique Neppel <[email protected]>

* Bump postgresql charm lib for 16/edge to v1 due to backwards incompatible changes

* Add DBA user test

Signed-off-by: Marcelo Henrique Neppel <[email protected]>

* Test DBA role in replica

Signed-off-by: Marcelo Henrique Neppel <[email protected]>

* Grant reset_user function to DBA role

Signed-off-by: Marcelo Henrique Neppel <[email protected]>

* Test set_user function for unprivileged users

Signed-off-by: Marcelo Henrique Neppel <[email protected]>

* Re-add mistakenly removed patch statements

* Reset connection to None before creating a new connection

Signed-off-by: Marcelo Henrique Neppel <[email protected]>

---------

Signed-off-by: Marcelo Henrique Neppel <[email protected]>
Co-authored-by: Shayan Patel <[email protected]>

Author: marceloneppel

lib/charms/postgresql_k8s/v1/postgresql.py

@@ -53,6 +53,7 @@
 ROLE_READ = "charmed_read"
 ROLE_DML = "charmed_dml"
 ROLE_BACKUP = "charmed_backup"
+ROLE_DBA = "charmed_dba"
 
 # Groups to distinguish database permissions
 PERMISSIONS_GROUP_ADMIN = "admin"
@@ -341,6 +342,17 @@ def create_user(
 
     def create_predefined_roles(self) -> None:
         """Create predefined roles."""
+        connection = None
+        try:
+            for database in ["postgres", "template1"]:
+                with self._connect_to_database(
+                    database=database,
+                ) as connection, connection.cursor() as cursor:
+                    cursor.execute(SQL("CREATE EXTENSION IF NOT EXISTS set_user;"))
+        finally:
+            if connection is not None:
+                connection.close()
+            connection = None
         role_to_queries = {
             ROLE_STATS: [
                 f"CREATE ROLE {ROLE_STATS} NOSUPERUSER NOCREATEDB NOCREATEROLE NOREPLICATION NOLOGIN IN ROLE pg_monitor",
@@ -359,6 +371,14 @@ def create_predefined_roles(self) -> None:
                 f"GRANT execute ON FUNCTION pg_create_restore_point TO {ROLE_BACKUP}",
                 f"GRANT execute ON FUNCTION pg_switch_wal TO {ROLE_BACKUP}",
             ],
+            ROLE_DBA: [
+                f"CREATE ROLE {ROLE_DBA} NOSUPERUSER CREATEDB NOCREATEROLE NOLOGIN NOREPLICATION;",
+                f"GRANT execute ON FUNCTION set_user(text) TO {ROLE_DBA};",
+                f"GRANT execute ON FUNCTION set_user(text, text) TO {ROLE_DBA};",
+                f"GRANT execute ON FUNCTION set_user_u(text) TO {ROLE_DBA};"
+                f"GRANT execute ON FUNCTION reset_user() TO {ROLE_DBA};"
+                f"GRANT execute ON FUNCTION reset_user(text) TO {ROLE_DBA};"
+            ]
         }
 
         _, existing_roles = self.list_valid_privileges_and_roles()
@@ -377,6 +397,9 @@ def create_predefined_roles(self) -> None:
         except psycopg2.Error as e:
             logger.error(f"Failed to create predefined roles: {e}")
             raise PostgreSQLCreatePredefinedRolesError() from e
+        finally:
+            if connection is not None:
+                connection.close()
 
     def grant_database_privileges_to_user(
         self, user: str, database: str, privileges: list[str]

templates/patroni.yml.j2

@@ -96,7 +96,10 @@ bootstrap:
         log_truncate_on_rotation: 'on'
         logging_collector: 'on'
         wal_level: logical
-        shared_preload_libraries: 'timescaledb,pgaudit'
+        shared_preload_libraries: 'timescaledb,pgaudit,set_user'
+        set_user.block_log_statement: 'on'
+        set_user.exit_on_error: 'on'
+        set_user.superuser_allowlist: '+charmed_dba'
 
   {%- if restoring_backup %}
   method: pgbackrest
@@ -132,7 +135,10 @@ postgresql:
   bin_dir: /snap/charmed-postgresql/current/usr/lib/postgresql/{{ version }}/bin
   data_dir: {{ data_path }}
   parameters:
-    shared_preload_libraries: 'timescaledb,pgaudit'
+    shared_preload_libraries: 'timescaledb,pgaudit,set_user'
+    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 {{ pgbackrest_configuration_file }} --stanza={{ stanza }} archive-push %p'
     {% else %}

tests/integration/helpers.py

@@ -37,6 +37,7 @@
 DATABASE_APP_NAME = METADATA["name"]
 STORAGE_PATH = METADATA["storage"]["data"]["location"]
 APPLICATION_NAME = "postgresql-test-app"
+DATA_INTEGRATOR_APP_NAME = "data-integrator"
 
 
 class SecretNotFoundError(Exception):
@@ -737,6 +738,23 @@ def get_unit_address(ops_test: OpsTest, unit_name: str, model: Model = None) ->
     return model.units.get(unit_name).public_address
 
 
+def check_connected_user(
+    cursor, session_user: str, current_user: str, primary: bool = True
+) -> None:
+    cursor.execute("SELECT session_user,current_user;")
+    result = cursor.fetchone()
+    if result is not None:
+        instance = "primary" if primary else "replica"
+        assert result[0] == session_user, (
+            f"The session user should be the {session_user} user in the {instance}"
+        )
+        assert result[1] == current_user, (
+            f"The current user should be the {current_user} user in the {instance}"
+        )
+    else:
+        assert False, "No result returned from the query"
+
+
 async def check_tls(ops_test: OpsTest, unit_name: str, enabled: bool) -> bool:
     """Returns whether TLS is enabled on the specific PostgreSQL instance.
 

tests/integration/test_predefined_dba_role.py

@@ -0,0 +1,93 @@
+#!/usr/bin/env python3
+# Copyright 2025 Canonical Ltd.
+# See LICENSE file for licensing details.
+
+import asyncio
+import logging
+
+import psycopg2
+import psycopg2.sql
+import pytest
+from pytest_operator.plugin import OpsTest
+
+from .helpers import (
+    CHARM_BASE,
+    DATA_INTEGRATOR_APP_NAME,
+    DATABASE_APP_NAME,
+    check_connected_user,
+)
+from .new_relations.helpers import build_connection_string
+
+logger = logging.getLogger(__name__)
+
+
+@pytest.mark.abort_on_fail
+async def test_deploy(ops_test: OpsTest, charm: str):
+    """Deploy the postgresql charm along with data integrator charm."""
+    async with ops_test.fast_forward("10s"):
+        await asyncio.gather(
+            ops_test.model.deploy(
+                charm,
+                application_name=DATABASE_APP_NAME,
+                num_units=2,
+                base=CHARM_BASE,
+                config={"profile": "testing"},
+            ),
+            ops_test.model.deploy(
+                DATA_INTEGRATOR_APP_NAME,
+                base=CHARM_BASE,
+            ),
+        )
+
+        await ops_test.model.wait_for_idle(apps=[DATABASE_APP_NAME], status="active")
+        assert ops_test.model.applications[DATABASE_APP_NAME].units[0].workload_status == "active"
+        await ops_test.model.wait_for_idle(apps=[DATA_INTEGRATOR_APP_NAME], status="blocked")
+
+
+@pytest.mark.abort_on_fail
+async def test_charmed_dba_role(ops_test: OpsTest):
+    """Test the DBA predefined role."""
+    await ops_test.model.applications[DATA_INTEGRATOR_APP_NAME].set_config({
+        "database-name": "charmed_dba_database",
+        "extra-user-roles": "charmed_dba",
+    })
+    await ops_test.model.add_relation(DATA_INTEGRATOR_APP_NAME, DATABASE_APP_NAME)
+    await ops_test.model.wait_for_idle(
+        apps=[DATA_INTEGRATOR_APP_NAME, DATABASE_APP_NAME], status="active"
+    )
+
+    action = await ops_test.model.units[f"{DATA_INTEGRATOR_APP_NAME}/0"].run_action(
+        action_name="get-credentials"
+    )
+    result = await action.wait()
+    data_integrator_credentials = result.results
+    username = data_integrator_credentials["postgresql"]["username"]
+
+    for read_write_endpoint in [True, False]:
+        connection_string = await build_connection_string(
+            ops_test,
+            DATA_INTEGRATOR_APP_NAME,
+            "postgresql",
+            database="charmed_dba_database",
+            read_only_endpoint=(not read_write_endpoint),
+        )
+        connection = psycopg2.connect(connection_string)
+        connection.autocommit = True
+        try:
+            with connection.cursor() as cursor:
+                instance = "primary" if read_write_endpoint else "replica"
+                logger.info(f"Testing escalation to the rewind user in the {instance}")
+                cursor.execute("SELECT set_user('rewind'::TEXT);")
+                check_connected_user(cursor, username, "rewind", primary=read_write_endpoint)
+                logger.info(f"Resetting the user to the {username} user in the {instance}")
+                cursor.execute("SELECT reset_user();")
+                check_connected_user(cursor, username, username, primary=read_write_endpoint)
+                logger.info(f"Testing escalation to the operator user in the {instance}")
+                cursor.execute("SELECT set_user_u('operator'::TEXT);")
+                check_connected_user(cursor, username, "operator", primary=read_write_endpoint)
+                logger.info(f"Resetting the user to the {username} user in the {instance}")
+                cursor.execute("SELECT reset_user();")
+                check_connected_user(cursor, username, username, primary=read_write_endpoint)
+        finally:
+            if connection is not None:
+                connection.close()

tests/integration/test_predefined_roles.py

@@ -12,6 +12,7 @@
 
 from .helpers import (
     CHARM_BASE,
+    DATA_INTEGRATOR_APP_NAME,
     DATABASE_APP_NAME,
     db_connect,
     get_password,
@@ -27,7 +28,6 @@
 logger = logging.getLogger(__name__)
 
 TIMEOUT = 15 * 60
-DATA_INTEGRATOR_APP_NAME = "data-integrator"
 
 # NOTE: We are unable to test set_user_u('operator') as dba user because psycopg2
 # runs every query in a transaction and running set_user() is not supported in transactions.

tests/spread/test_predefined_dba_role.py/task.yaml

@@ -0,0 +1,7 @@
+summary: test_predefined_dba_role.py
+environment:
+  TEST_MODULE: test_predefined_dba_role.py
+execute: |
+  tox run -e integration -- "tests/integration/$TEST_MODULE" --model testing --alluredir="$SPREAD_TASK/allure-results"
+artifacts:
+  - allure-results