[DPE-2167] Added admin extra user role (#201)

* Added admin extra user role

* Fixed admin role permissions and added integration test

* Removed flag from poetry exportD

* Fixed comment and removed restriction to superuser

* Simplify admin role privileges

Author: marceloneppel

lib/charms/postgresql_k8s/v0/postgresql.py

@@ -32,7 +32,7 @@
 
 # Increment this PATCH version before using `charmcraft publish-lib` or reset
 # to 0 if you are raising the major API version
-LIBPATCH = 10
+LIBPATCH = 11
 
 
 logger = logging.getLogger(__name__)
@@ -129,7 +129,7 @@ def create_database(self, database: str, user: str) -> None:
                     sql.Identifier(database)
                 )
             )
-            for user_to_grant_access in [user] + self.system_users:
+            for user_to_grant_access in [user, "admin"] + self.system_users:
                 cursor.execute(
                     sql.SQL("GRANT ALL PRIVILEGES ON DATABASE {} TO {};").format(
                         sql.Identifier(database), sql.Identifier(user_to_grant_access)
@@ -165,7 +165,7 @@ def create_database(self, database: str, user: str) -> None:
             raise PostgreSQLCreateDatabaseError()
 
     def create_user(
-        self, user: str, password: str, admin: bool = False, extra_user_roles: str = None
+        self, user: str, password: str = None, admin: bool = False, extra_user_roles: str = None
     ) -> None:
         """Creates a database user.
 
@@ -178,28 +178,28 @@ def create_user(
         try:
             with self._connect_to_database() as connection, connection.cursor() as cursor:
                 # Separate roles and privileges from the provided extra user roles.
+                admin_role = False
                 roles = privileges = None
                 if extra_user_roles:
                     extra_user_roles = tuple(extra_user_roles.lower().split(","))
                     cursor.execute(
                         "SELECT rolname FROM pg_roles WHERE rolname IN %s;", (extra_user_roles,)
                     )
-                    roles = [role[0] for role in cursor.fetchall()]
-                    privileges = [
+                    admin_role = "admin" in extra_user_roles
+                    roles = [role[0] for role in cursor.fetchall() if role[0] != "admin"]
+                    privileges = {
                         extra_user_role
                         for extra_user_role in extra_user_roles
-                        if extra_user_role not in roles
-                    ]
+                        if extra_user_role not in roles and extra_user_role != "admin"
+                    }
 
                 # Create or update the user.
                 cursor.execute(f"SELECT TRUE FROM pg_roles WHERE rolname='{user}';")
                 if cursor.fetchone() is not None:
                     user_definition = "ALTER ROLE {}"
                 else:
                     user_definition = "CREATE ROLE {}"
-                user_definition += (
-                    f"WITH LOGIN{' SUPERUSER' if admin else ''} ENCRYPTED PASSWORD '{password}'"
-                )
+                user_definition += f"WITH {'NOLOGIN' if user == 'admin' else 'LOGIN'}{' SUPERUSER' if admin else ''} ENCRYPTED PASSWORD '{password}'{'IN ROLE admin CREATEDB' if admin_role else ''}"
                 if privileges:
                     user_definition += f' {" ".join(privileges)}'
                 cursor.execute(sql.SQL(f"{user_definition};").format(sql.Identifier(user)))
@@ -347,15 +347,21 @@ def set_up_database(self) -> None:
         """Set up postgres database with the right permissions."""
         connection = None
         try:
+            self.create_user(
+                "admin",
+                extra_user_roles="pg_read_all_data,pg_write_all_data",
+            )
             with self._connect_to_database() as connection, connection.cursor() as cursor:
                 # Allow access to the postgres database only to the system users.
                 cursor.execute("REVOKE ALL PRIVILEGES ON DATABASE postgres FROM PUBLIC;")
+                cursor.execute("REVOKE CREATE ON SCHEMA public FROM PUBLIC;")
                 for user in self.system_users:
                     cursor.execute(
                         sql.SQL("GRANT ALL PRIVILEGES ON DATABASE postgres TO {};").format(
                             sql.Identifier(user)
                         )
                     )
+                cursor.execute("GRANT CONNECT ON DATABASE postgres TO admin;")
         except psycopg2.Error as e:
             logger.error(f"Failed to set up databases: {e}")
             raise PostgreSQLDatabasesSetupError()

tests/integration/new_relations/test_new_relations.py

@@ -3,6 +3,8 @@
 # See LICENSE file for licensing details.
 import asyncio
 import logging
+import secrets
+import string
 from pathlib import Path
 
 import psycopg2
@@ -27,6 +29,7 @@
 APPLICATION_APP_NAME = "application"
 DATABASE_APP_NAME = "database"
 ANOTHER_DATABASE_APP_NAME = "another-database"
+DATA_INTEGRATOR_APP_NAME = "data-integrator"
 APP_NAMES = [APPLICATION_APP_NAME, DATABASE_APP_NAME, ANOTHER_DATABASE_APP_NAME]
 DATABASE_APP_METADATA = yaml.safe_load(Path("./metadata.yaml").read_text())
 FIRST_DATABASE_RELATION_NAME = "first-database"
@@ -397,3 +400,87 @@ async def test_relation_with_no_database_name(ops_test: OpsTest):
             f"{DATABASE_APP_NAME}", f"{APPLICATION_APP_NAME}:{NO_DATABASE_RELATION_NAME}"
         )
         await ops_test.model.wait_for_idle(apps=APP_NAMES, status="active", raise_on_blocked=True)
+
+
+async def test_admin_role(ops_test: OpsTest):
+    """Test that the admin role gives access to all the databases."""
+    all_app_names = [DATA_INTEGRATOR_APP_NAME]
+    all_app_names.extend(APP_NAMES)
+    async with ops_test.fast_forward():
+        await ops_test.model.deploy(DATA_INTEGRATOR_APP_NAME)
+        await ops_test.model.wait_for_idle(apps=[DATA_INTEGRATOR_APP_NAME], status="blocked")
+        await ops_test.model.applications[DATA_INTEGRATOR_APP_NAME].set_config(
+            {
+                "database-name": DATA_INTEGRATOR_APP_NAME.replace("-", "_"),
+                "extra-user-roles": "admin",
+            }
+        )
+        await ops_test.model.wait_for_idle(apps=[DATA_INTEGRATOR_APP_NAME], status="blocked")
+        await ops_test.model.add_relation(DATA_INTEGRATOR_APP_NAME, DATABASE_APP_NAME)
+        await ops_test.model.wait_for_idle(apps=all_app_names, status="active")
+
+    # Check that the user can access all the databases.
+    for database in [
+        "postgres",
+        "application_first_database",
+        "another_application_first_database",
+    ]:
+        logger.info(f"connecting to the following database: {database}")
+        connection_string = await build_connection_string(
+            ops_test, DATA_INTEGRATOR_APP_NAME, "postgresql", database=database
+        )
+        connection = None
+        should_fail = False
+        try:
+            with psycopg2.connect(connection_string) as connection, connection.cursor() as cursor:
+                # Check the version that the application received is the same on the
+                # database server.
+                cursor.execute("SELECT version();")
+                data = cursor.fetchone()[0].split(" ")[1]
+
+                # Get the version of the database and compare with the information that
+                # was retrieved directly from the database.
+                version = await get_application_relation_data(
+                    ops_test, DATA_INTEGRATOR_APP_NAME, "postgresql", "version"
+                )
+                assert version == data
+
+                # Write some data (it should fail in the "postgres" database).
+                random_name = (
+                    f"test_{''.join(secrets.choice(string.ascii_lowercase) for _ in range(10))}"
+                )
+                should_fail = database == "postgres"
+                cursor.execute(f"CREATE TABLE {random_name}(data TEXT);")
+                if should_fail:
+                    assert (
+                        False
+                    ), f"failed to run a statement in the following database: {database}"
+        except psycopg2.errors.InsufficientPrivilege as e:
+            if not should_fail:
+                logger.exception(e)
+                assert (
+                    False
+                ), f"failed to connect to or run a statement in the following database: {database}"
+        finally:
+            if connection is not None:
+                connection.close()
+
+    # Test the creation and deletion of databases.
+    connection_string = await build_connection_string(
+        ops_test, DATA_INTEGRATOR_APP_NAME, "postgresql", database="postgres"
+    )
+    connection = psycopg2.connect(connection_string)
+    connection.autocommit = True
+    cursor = connection.cursor()
+    random_name = f"test_{''.join(secrets.choice(string.ascii_lowercase) for _ in range(10))}"
+    cursor.execute(f"CREATE DATABASE {random_name};")
+    cursor.execute(f"DROP DATABASE {random_name};")
+    try:
+        cursor.execute("DROP DATABASE postgres;")
+        assert False, "the admin extra user role was able to drop the `postgres` system database"
+    except psycopg2.errors.InsufficientPrivilege:
+        # Ignore the error, as the admin extra user role mustn't be able to drop
+        # the "postgres" system database.
+        pass
+    finally:
+        connection.close()