Spring Boot如何在每次PostgreSQL连接前后设置/清除会话变量（含非事务查询）

针对你的场景——Spring Boot结合Spring Data、PostgreSQL和HikariCP，需要在**所有数据库连接（包括非事务查询）**获取前设置user_id会话变量，归还时清除，以适配Employee表的行级安全策略，可通过以下方案实现：

核心思路 #

通过代理HikariCP的数据源和数据库连接，实现每次从连接池取出连接时自动注入当前用户的user_id，连接归还到池前清除该变量，确保无论事务内还是非事务的查询都能触发RLS规则。

具体实现 #

1. 代理数据库连接，处理会话变量的设置与清除 #

实现java.sql.Connection接口，包装原始连接，在构造时设置会话变量，重写close()方法在归还连接前清除变量：

import java.sql.*;
import your.package.UserService;

public class RlsAwareConnectionProxy implements Connection {

    private final Connection targetConnection;
    private final UserService userService;

    public RlsAwareConnectionProxy(Connection targetConnection, UserService userService) {
        this.targetConnection = targetConnection;
        this.userService = userService;
        setSessionUserId();
    }

    // 设置RLS会话变量，用PreparedStatement避免SQL注入
    private void setSessionUserId() {
        Integer userId = userService.getCurrentUserId();
        if (userId == null) {
            return;
        }
        try (PreparedStatement stmt = targetConnection.prepareStatement("SET user_id = ?::integer")) {
            stmt.setInt(1, userId);
            stmt.executeUpdate();
        } catch (SQLException e) {
            throw new RuntimeException("Failed to set RLS session variable 'user_id'", e);
        }
    }

    // 清除会话变量，出错仅记录不抛出，避免影响连接归还
    private void clearSessionUserId() {
        try (Statement stmt = targetConnection.createStatement()) {
            stmt.executeUpdate("RESET user_id");
        } catch (SQLException e) {
            // 替换为你的日志组件，如log.error("Failed to clear RLS session variable", e);
            System.err.println("Failed to clear RLS session variable: " + e.getMessage());
        }
    }

    @Override
    public void close() throws SQLException {
        clearSessionUserId();
        targetConnection.close();
    }

    // 以下为Connection接口所有方法的委托实现，直接调用原始连接的对应方法
    @Override
    public Statement createStatement() throws SQLException {
        return targetConnection.createStatement();
    }

    @Override
    public PreparedStatement prepareStatement(String sql) throws SQLException {
        return targetConnection.prepareStatement(sql);
    }

    @Override
    public CallableStatement prepareCall(String sql) throws SQLException {
        return targetConnection.prepareCall(sql);
    }

    @Override
    public String nativeSQL(String sql) throws SQLException {
        return targetConnection.nativeSQL(sql);
    }

    @Override
    public void setAutoCommit(boolean autoCommit) throws SQLException {
        targetConnection.setAutoCommit(autoCommit);
    }

    @Override
    public boolean getAutoCommit() throws SQLException {
        return targetConnection.getAutoCommit();
    }

    @Override
    public void commit() throws SQLException {
        targetConnection.commit();
    }

    @Override
    public void rollback() throws SQLException {
        targetConnection.rollback();
    }

    @Override
    public boolean isClosed() throws SQLException {
        return targetConnection.isClosed();
    }

    @Override
    public DatabaseMetaData getMetaData() throws SQLException {
        return targetConnection.getMetaData();
    }

    @Override
    public void setReadOnly(boolean readOnly) throws SQLException {
        targetConnection.setReadOnly(readOnly);
    }

    @Override
    public boolean isReadOnly() throws SQLException {
        return targetConnection.isReadOnly();
    }

    @Override
    public void setCatalog(String catalog) throws SQLException {
        targetConnection.setCatalog(catalog);
    }

    @Override
    public String getCatalog() throws SQLException {
        return targetConnection.getCatalog();
    }

    @Override
    public void setTransactionIsolation(int level) throws SQLException {
        targetConnection.setTransactionIsolation(level);
    }

    @Override
    public int getTransactionIsolation() throws SQLException {
        return targetConnection.getTransactionIsolation();
    }

    @Override
    public SQLWarning getWarnings() throws SQLException {
        return targetConnection.getWarnings();
    }

    @Override
    public void clearWarnings() throws SQLException {
        targetConnection.clearWarnings();
    }

    @Override
    public Statement createStatement(int resultSetType, int resultSetConcurrency) throws SQLException {
        return targetConnection.createStatement(resultSetType, resultSetConcurrency);
    }

    @Override
    public PreparedStatement prepareStatement(String sql, int resultSetType, int resultSetConcurrency) throws SQLException {
        return targetConnection.prepareStatement(sql, resultSetType, resultSetConcurrency);
    }

    @Override
    public CallableStatement prepareCall(String sql, int resultSetType, int resultSetConcurrency) throws SQLException {
        return targetConnection.prepareCall(sql, resultSetType, resultSetConcurrency);
    }

    @Override
    public Map<String, Class<?>> getTypeMap() throws SQLException {
        return targetConnection.getTypeMap();
    }

    @Override
    public void setTypeMap(Map<String, Class<?>> map) throws SQLException {
        targetConnection.setTypeMap(map);
    }

    @Override
    public void setHoldability(int holdability) throws SQLException {
        targetConnection.setHoldability(holdability);
    }

    @Override
    public int getHoldability() throws SQLException {
        return targetConnection.getHoldability();
    }

    @Override
    public Savepoint setSavepoint() throws SQLException {
        return targetConnection.setSavepoint();
    }

    @Override
    public Savepoint setSavepoint(String name) throws SQLException {
        return targetConnection.setSavepoint(name);
    }

    @Override
    public void rollback(Savepoint savepoint) throws SQLException {
        targetConnection.rollback(savepoint);
    }

    @Override
    public void releaseSavepoint(Savepoint savepoint) throws SQLException {
        targetConnection.releaseSavepoint(savepoint);
    }

    @Override
    public Statement createStatement(int resultSetType, int resultSetConcurrency, int resultSetHoldability) throws SQLException {
        return targetConnection.createStatement(resultSetType, resultSetConcurrency, resultSetHoldability);
    }

    @Override
    public PreparedStatement prepareStatement(String sql, int resultSetType, int resultSetConcurrency, int resultSetHoldability) throws SQLException {
        return targetConnection.prepareStatement(sql, resultSetType, resultSetConcurrency, resultSetHoldability);
    }

    @Override
    public CallableStatement prepareCall(String sql, int resultSetType, int resultSetConcurrency, int resultSetHoldability) throws SQLException {
        return targetConnection.prepareCall(sql, resultSetType, resultSetConcurrency, resultSetHoldability);
    }

    @Override
    public PreparedStatement prepareStatement(String sql, int autoGeneratedKeys) throws SQLException {
        return targetConnection.prepareStatement(sql, autoGeneratedKeys);
    }

    @Override
    public PreparedStatement prepareStatement(String sql, int[] columnIndexes) throws SQLException {
        return targetConnection.prepareStatement(sql, columnIndexes);
    }

    @Override
    public PreparedStatement prepareStatement(String sql, String[] columnNames) throws SQLException {
        return targetConnection.prepareStatement(sql, columnNames);
    }

    @Override
    public Clob createClob() throws SQLException {
        return targetConnection.createClob();
    }

    @Override
    public Blob createBlob() throws SQLException {
        return targetConnection.createBlob();
    }

    @Override
    public NClob createNClob() throws SQLException {
        return targetConnection.createNClob();
    }

    @Override
    public SQLXML createSQLXML() throws SQLException {
        return targetConnection.createSQLXML();
    }

    @Override
    public boolean isValid(int timeout) throws SQLException {
        return targetConnection.isValid(timeout);
    }

    @Override
    public void setClientInfo(String name, String value) throws SQLClientInfoException {
        targetConnection.setClientInfo(name, value);
    }

    @Override
    public void setClientInfo(Properties properties) throws SQLClientInfoException {
        targetConnection.setClientInfo(properties);
    }

    @Override
    public String getClientInfo(String name) throws SQLException {
        return targetConnection.getClientInfo(name);
    }

    @Override
    public Properties getClientInfo() throws SQLException {
        return targetConnection.getClientInfo();
    }

    @Override
    public Array createArrayOf(String typeName, Object[] elements) throws SQLException {
        return targetConnection.createArrayOf(typeName, elements);
    }

    @Override
    public Struct createStruct(String typeName, Object[] attributes) throws SQLException {
        return targetConnection.createStruct(typeName, attributes);
    }

    @Override
    public void setSchema(String schema) throws SQLException {
        targetConnection.setSchema(schema);
    }

    @Override
    public String getSchema() throws SQLException {
        return targetConnection.getSchema();
    }

    @Override
    public void abort(Executor executor) throws SQLException {
        targetConnection.abort(executor);
    }

    @Override
    public void setNetworkTimeout(Executor executor, int milliseconds) throws SQLException {
        targetConnection.setNetworkTimeout(executor, milliseconds);
    }

    @Override
    public int getNetworkTimeout() throws SQLException {
        return targetConnection.getNetworkTimeout();
    }

    @Override
    public <T> T unwrap(Class<T> iface) throws SQLException {
        return targetConnection.unwrap(iface);
    }

    @Override
    public boolean isWrapperFor(Class<?> iface) throws SQLException {
        return targetConnection.isWrapperFor(iface);
    }
}

2. 代理Hikari数据源，返回包装后的连接 #

创建代理数据源类，注入原始Hikari数据源和UserService，每次获取连接时返回上面定义的代理连接：

import javax.sql.DataSource;
import java.sql.Connection;
import java.sql.SQLException;
import java.sql.SQLFeatureNotSupportedException;
import java.util.logging.Logger;
import org.springframework.beans.factory.annotation.Qualifier;
import org.springframework.stereotype.Component;
import your.package.UserService;

@Component
public class RlsAwareDataSourceProxy implements DataSource {

    private final DataSource hikariDataSource;
    private final UserService userService;

    public RlsAwareDataSourceProxy(
            @Qualifier("dataSource") DataSource hikariDataSource,
            UserService userService) {
        this.hikariDataSource = hikariDataSource;
        this.userService = userService;
    }

    @Override
    public Connection getConnection() throws SQLException {
        Connection originalConn = hikariDataSource.getConnection();
        return new RlsAwareConnectionProxy(originalConn, userService);
    }

    @Override
    public Connection getConnection(String username, String password) throws SQLException {
        // HikariCP已配置全局用户名密码，直接复用getConnection()逻辑
        return getConnection();
    }

    // 以下为DataSource接口方法的委托实现，直接调用原始数据源的对应方法
    @Override
    public <T> T unwrap(Class<T> iface) throws SQLException {
        return hikariDataSource.unwrap(iface);
    }

    @Override
    public boolean isWrapperFor(Class<?> iface) throws SQLException {
        return hikariDataSource.isWrapperFor(iface);
    }

    @Override
    public java.io.PrintWriter getLogWriter() throws SQLException {
        return hikariDataSource.getLogWriter();
    }

    @Override
    public void setLogWriter(java.io.PrintWriter out) throws SQLException {
        hikariDataSource.setLogWriter(out);
    }

    @Override
    public void setLoginTimeout(int seconds) throws SQLException {
        hikariDataSource.setLoginTimeout(seconds);
    }

    @Override
    public int getLoginTimeout() throws SQLException {
        return hikariDataSource.getLoginTimeout();
    }

    @Override
    public Logger getParentLogger() throws SQLFeatureNotSupportedException {
        return hikariDataSource.getParentLogger();
    }
}

3. 配置代理数据源为默认数据源 #

在配置类中标记代理数据源为@Primary，确保Spring Data的Repository使用该数据源：

import javax.sql.DataSource;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;
import org.springframework.context.annotation.Primary;
import your.package.UserService;

@Configuration
public class DataSourceConfig {

    @Bean
    @Primary
    public DataSource primaryDataSource(
            @Qualifier("dataSource") DataSource hikariDataSource,
            UserService userService) {
        return new RlsAwareDataSourceProxy(hikariDataSource, userService);
    }
}

关键注意事项 #

• 用户ID传递：UserService.getCurrentUserId()需通过ThreadLocal存储当前请求的用户ID（如登录后将用户ID存入ThreadLocal），确保在获取连接时能拿到正确值；异步任务需配合TaskDecorator传递ThreadLocal上下文。
• SQL注入防护：使用PreparedStatement设置会话变量，避免字符串拼接带来的注入风险。
• 异常处理：清除会话变量时若出错，仅记录日志不抛出异常，防止连接无法正常归还到池。
• 连接复用：该方案适配HikariCP的连接复用机制，每次取出连接都会重新设置当前用户的user_id，不会出现用户数据串扰。

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