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