[DPE-7302] Remove user from database (#31)

* Remove user from database

* Revoke predefined roles

Author: dragomirp

pyproject.toml

@@ -4,7 +4,7 @@
 [project]
 name = "postgresql-charms-single-kernel"
 description = "Shared and reusable code for PostgreSQL-related charms"
-version = "16.1.1"
+version = "16.1.2"
 readme = "README.md"
 license = "Apache-2.0"
 authors = [

single_kernel_postgresql/utils/postgresql.py

@@ -1885,7 +1885,7 @@ def list_databases(self, prefix: Optional[str] = None) -> List[str]:
     def add_user_to_databases(
         self, user: str, databases: List[str], extra_user_roles: Optional[List[str]] = None
     ) -> None:
-        """Grant user access to database."""
+        """Grant user access to a database."""
         try:
             roles, _ = self._process_extra_user_roles(user, extra_user_roles)
             connect_stmt = []
@@ -1907,5 +1907,34 @@ def add_user_to_databases(
                 for statement in connect_stmt:
                     cursor.execute(statement)
         except psycopg2.Error as e:
-            logger.error(f"Failed to create user: {e}")
+            logger.error(f"Failed to add user: {e}")
+            raise PostgreSQLUpdateUserError() from e
+
+    def remove_user_from_databases(self, user: str, databases: List[str]) -> None:
+        """Revoke user access to a database."""
+        try:
+            for database in databases:
+                with self._connect_to_database() as connection, connection.cursor() as cursor:
+                    cursor.execute(
+                        SQL("REVOKE CONNECT ON DATABASE {} FROM {};").format(
+                            Identifier(database), Identifier(user)
+                        )
+                    )
+                    cursor.execute(
+                        SQL("REVOKE {} FROM {};").format(
+                            Identifier(f"charmed_{database}_owner"), Identifier(user)
+                        )
+                    )
+                    cursor.execute(
+                        SQL("REVOKE {} FROM {};").format(
+                            Identifier(f"charmed_{database}_admin"), Identifier(user)
+                        )
+                    )
+                    cursor.execute(
+                        SQL("REVOKE {} FROM {};").format(
+                            Identifier(f"charmed_{database}_dml"), Identifier(user)
+                        )
+                    )
+        except psycopg2.Error as e:
+            logger.error(f"Failed to remove user: {e}")
             raise PostgreSQLUpdateUserError() from e

tests/unit/test_postgresql.py

@@ -922,3 +922,78 @@ def test_add_user_to_databases():
         with pytest.raises(PostgreSQLUpdateUserError):
             pg.add_user_to_databases("test-user", ["db1", "db2"])
             assert False
+
+
+def test_remove_user_from_databases():
+    with (
+        patch(
+            "single_kernel_postgresql.utils.postgresql.PostgreSQL._connect_to_database"
+        ) as _connect_to_database,
+    ):
+        pg = PostgreSQL(
+            Substrates.VM, "primary", "current", "operator", "password", "postgres", None
+        )
+        execute = _connect_to_database.return_value.__enter__.return_value.cursor.return_value.__enter__.return_value.execute
+
+        pg.remove_user_from_databases("test-user", ["db1", "db2"])
+        assert execute.call_count == 8
+        execute.assert_any_call(
+            Composed([
+                SQL("REVOKE CONNECT ON DATABASE "),
+                Identifier("db1"),
+                SQL(" FROM "),
+                Identifier("test-user"),
+                SQL(";"),
+            ])
+        )
+        execute.assert_any_call(
+            Composed([
+                SQL("REVOKE CONNECT ON DATABASE "),
+                Identifier("db2"),
+                SQL(" FROM "),
+                Identifier("test-user"),
+                SQL(";"),
+            ])
+        )
+        execute.assert_any_call(
+            Composed([
+                SQL("REVOKE "),
+                Identifier("charmed_db1_admin"),
+                SQL(" FROM "),
+                Identifier("test-user"),
+                SQL(";"),
+            ])
+        )
+        execute.assert_any_call(
+            Composed([
+                SQL("REVOKE "),
+                Identifier("charmed_db1_dml"),
+                SQL(" FROM "),
+                Identifier("test-user"),
+                SQL(";"),
+            ])
+        )
+        execute.assert_any_call(
+            Composed([
+                SQL("REVOKE "),
+                Identifier("charmed_db2_admin"),
+                SQL(" FROM "),
+                Identifier("test-user"),
+                SQL(";"),
+            ])
+        )
+        execute.assert_any_call(
+            Composed([
+                SQL("REVOKE "),
+                Identifier("charmed_db2_dml"),
+                SQL(" FROM "),
+                Identifier("test-user"),
+                SQL(";"),
+            ])
+        )
+
+        # Exception
+        execute.side_effect = psycopg2.Error
+        with pytest.raises(PostgreSQLUpdateUserError):
+            pg.remove_user_from_databases("test-user", ["db1", "db2"])
+            assert False