diff --git a/pyproject.toml b/pyproject.toml index 9d89fe8..90a4138 100644 --- a/pyproject.toml +++ b/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 = [ diff --git a/single_kernel_postgresql/utils/postgresql.py b/single_kernel_postgresql/utils/postgresql.py index e72c11a..6b76184 100644 --- a/single_kernel_postgresql/utils/postgresql.py +++ b/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 diff --git a/tests/unit/test_postgresql.py b/tests/unit/test_postgresql.py index fb80e06..cfac0cc 100644 --- a/tests/unit/test_postgresql.py +++ b/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 diff --git a/uv.lock b/uv.lock index b562cc8..ea80c08 100644 --- a/uv.lock +++ b/uv.lock @@ -294,7 +294,7 @@ wheels = [ [[package]] name = "postgresql-charms-single-kernel" -version = "16.1.1" +version = "16.1.2" source = { virtual = "." } [package.dev-dependencies]