PostgreSQL:如何在匿名代码块(Do Block)内切换角色并确保退出后自动恢复原角色
Great question! Let's break this down clearly for PostgreSQL.
First, the straight answer: there's no built-in way to make a DO block automatically revert your session role after execution. That's because SET ROLE modifies your session-level identity, and DO blocks don't roll back session-level changes (they only handle transactional DML/DDL rollbacks on failure). So your example code will indeed leave your session as user1 unless you explicitly switch back.
But there are solid ways to either avoid forgetting the revert entirely, or add safeguards to ensure it happens even if things go wrong:
1. Use a Wrapper Function (Cleanest, Reusable Approach)
Instead of writing raw DO blocks every time, create a reusable function that handles role switching and cleanup automatically. This eliminates the risk of forgetting to revert the role:
CREATE OR REPLACE FUNCTION run_as_target_role(p_target_role text, p_sql_to_execute text) RETURNS void AS $$ DECLARE v_original_role text := current_role; BEGIN -- Switch to the target role SET ROLE p_target_role; -- Execute your desired SQL (use EXECUTE for dynamic commands) EXECUTE p_sql_to_execute; -- Revert to original role on success SET ROLE v_original_role; EXCEPTION WHEN OTHERS THEN -- Even if an error occurs, revert the role first before re-raising the error SET ROLE v_original_role; RAISE; END $$ LANGUAGE plpgsql SECURITY DEFINER;
To use it:
SHOW role; -- Output: postgres -- Run your operations as user1 SELECT run_as_target_role('user1', 'SELECT * FROM your_table; INSERT INTO another_table VALUES (1);'); SHOW role; -- Output: postgres (back to original, no manual revert needed)
2. Explicit Revert with Safety Nets (For One-Off DO Blocks)
If you prefer to stick with raw DO blocks, you can add safeguards to ensure the role is reverted even if an error occurs mid-execution. Since DO block variables are scoped to the block, use a session-level setting to store the original role:
SHOW role; -- Output: postgres -- Store original role in a custom session variable SET my_session.original_role = current_role; DO $$ BEGIN SET ROLE user1; -- Your operations go here -- Example: SELECT * FROM restricted_table; EXCEPTION WHEN OTHERS THEN -- Revert role on error before re-raising SET ROLE current_setting('my_session.original_role'); RAISE; END $$; -- Final revert and clean up the session variable SET ROLE current_setting('my_session.original_role'); RESET my_session.original_role; SHOW role; -- Output: postgres
Why Your Original Code Leaves the Role Changed
Just to clarify: SET ROLE is a session-wide change, not limited to the DO block. Unlike transactional changes (like inserting data), session-level settings aren't automatically rolled back when the DO block finishes—even if the block hits an error. That's why the explicit revert is mandatory unless you use a wrapper to handle it for you.
内容的提问来源于stack exchange,提问作者user1409708

