如何通过psycopg2为PostgreSQL所有查询自动设置UserId属性？

可行解决方案 #

1. 数据库行级安全（RLS）推荐 #

这是最安全且完全符合DRY原则的方案，无需修改应用层任何SQL：

1. 为需要控制的表启用行级安全：

ALTER TABLE Users ENABLE ROW LEVEL SECURITY;
ALTER TABLE Profiles ENABLE ROW LEVEL SECURITY;
-- 其他涉及用户隔离的表同理

2. 创建全局访问策略，限制用户只能访问自身数据：

CREATE POLICY user_row_access ON Users
FOR ALL USING (UserId = current_setting('app.user_id')::integer);

CREATE POLICY user_row_access ON Profiles
FOR ALL USING (UserId = current_setting('app.user_id')::integer);
-- 其他表同理

3. 应用层每次建立数据库连接后，设置会话级参数：

# 获取游标后执行该语句
curs.execute("SET app.user_id = %s", (pyUserId,))

此后所有对表的SELECT/INSERT/UPDATE/DELETE操作都会自动应用该策略，数据库会强制过滤掉不属于当前pyUserId的行。如果需要自动填充插入的UserId，可额外添加插入策略：

CREATE POLICY user_insert_access ON Profiles
FOR INSERT WITH CHECK (UserId = current_setting('app.user_id')::integer);

这样插入时即使不指定UserId，数据库也会自动填入当前会话的app.user_id值。

2. 自定义执行函数封装逻辑 #

如果无法使用RLS，可基于现有自定义执行函数，结合SQL解析库（如sqlparse）安全修改SQL语句：

1. 先安装sqlparse：pip install sqlparse
2. 修改自定义执行函数：

import sqlparse
from sqlparse.tokens import Keyword, DML

def safe_execute(curs, sql, params=None):
    params = params or []
    parsed = sqlparse.parse(sql)[0]
    
    # 识别SQL语句类型
    stmt_type = None
    for token in parsed.tokens:
        if token.ttype == DML:
            stmt_type = token.value.upper()
            break
    
    if stmt_type in ('SELECT', 'UPDATE', 'DELETE'):
        # 检查是否存在WHERE子句
        has_where = any(token.value.upper() == 'WHERE' for token in parsed.tokens if token.ttype == Keyword)
        if has_where:
            modified_sql = sql.rstrip().rstrip(';') + ' AND UserId = %s'
        else:
            modified_sql = sql.rstrip().rstrip(';') + ' WHERE UserId = %s'
        params.append(pyUserId)
        curs.execute(modified_sql, params)
    
    elif stmt_type == 'INSERT':
        # 解析插入的列和值部分
        col_list = None
        val_list = None
        for token in parsed.tokens:
            if token.value == '(':
                col_list = token.next_token
                # 定位VALUES关键字
                while token:
                    if token.value.upper() == 'VALUES':
                        val_list = token.next_token
                        break
                    token = token.next_token
                break
        
        if col_list and val_list:
            # 更新列列表，添加UserId
            cols = [t.value for t in col_list.tokens if t.value not in (',', ' ')]
            cols.append('UserId')
            new_col_part = '(' + ', '.join(cols) + ')'
            
            # 更新值列表，添加参数占位符
            vals = [t.value for t in val_list.tokens if t.value not in (',', ' ')]
            vals.append('%s')
            new_val_part = '(' + ', '.join(vals) + ')'
            
            # 重新拼接SQL
            modified_sql = str(parsed).replace(str(col_list), new_col_part).replace(str(val_list), new_val_part)
            params.append(pyUserId)
            curs.execute(modified_sql, params)
        else:
            # 无指定列的插入，直接追加值占位符
            modified_sql = sql.rstrip().rstrip(')') + ', %s)'
            params.append(pyUserId)
            curs.execute(modified_sql, params)
    
    else:
        # 其他类型语句直接执行
        curs.execute(sql, params)

这种方法通过解析SQL结构而非字符串拼接降低注入风险，但需要维护解析逻辑，可靠性不如RLS。

3. 视图+触发器方案 #

为每个表创建仅包含当前用户数据的视图，并通过触发器自动处理插入：

-- 创建用户专属视图
CREATE VIEW user_users AS
SELECT * FROM Users WHERE UserId = current_setting('app.user_id')::integer;

-- 创建插入触发器函数
CREATE OR REPLACE FUNCTION insert_user_data()
RETURNS TRIGGER AS $$
BEGIN
    NEW.UserId = current_setting('app.user_id')::integer;
    RETURN NEW;
END;
$$ LANGUAGE plpgsql;

-- 绑定触发器到视图
CREATE TRIGGER trigger_insert_users
INSTEAD OF INSERT ON user_users
FOR EACH ROW EXECUTE FUNCTION insert_user_data();

应用层直接访问user_users类视图即可，但需要为每个表创建视图和触发器，维护成本较高。

内容的提问来源于stack exchange，提问作者Paul Smith