基于R2DBC的Postgres行级安全：AbstractRoutingDataSource替代方案问询

R2DBC生态中没有直接对应JDBC里AbstractRoutingDataSource的类，因为两者核心模型不同（JDBC依赖DataSource，R2DBC核心是ConnectionFactory）。针对Postgres行级安全（RLS）场景，通常无需切换数据源，而是在创建数据库连接时动态设置租户上下文变量，具体实现方案如下：

1. 用Reactor Context传递租户ID #

响应式场景下不能用ThreadLocal存储租户信息，需借助Reactor的Context传递：

public class TenantContext {
    public static final String TENANT_ID_KEY = "tenantId";

    public static Mono<String> getTenantId() {
        return Mono.deferContextual(ctx -> Mono.justOrEmpty(ctx.getOrDefault(TENANT_ID_KEY, null)));
    }
}

2. 实现自定义ConnectionFactory包装类 #

包装原生Postgres ConnectionFactory，在创建连接时自动设置RLS租户变量：

public class TenantAwareConnectionFactory implements ConnectionFactory {
    private final ConnectionFactory targetConnectionFactory;

    public TenantAwareConnectionFactory(ConnectionFactory targetConnectionFactory) {
        this.targetConnectionFactory = targetConnectionFactory;
    }

    @Override
    public Publisher<? extends Connection> create() {
        return TenantContext.getTenantId()
                .flatMap(tenantId -> targetConnectionFactory.create()
                        .flatMap(connection -> {
                            // 参数化设置RLS租户变量，避免SQL注入
                            return connection.createStatement("SET app.current_tenant = $1;")
                                    .bind(0, tenantId)
                                    .execute()
                                    .then(Mono.just(connection));
                        }));
    }

    @Override
    public ConnectionFactoryMetadata getMetadata() {
        return targetConnectionFactory.getMetadata();
    }
}

注：app.current_tenant是你在Postgres中为RLS定义的会话变量，需提前在数据库中配置对应的RLS策略

3. 配置自定义ConnectionFactory #

在Spring Boot中替换默认的ConnectionFactory：

@Configuration
public class R2dbcConfig {
    @Bean
    public ConnectionFactory connectionFactory(ConnectionFactoryProperties properties) {
        PostgresConnectionConfiguration config = PostgresConnectionConfiguration.builder()
                .host(properties.getHost())
                .port(properties.getPort())
                .database(properties.getDatabase())
                .username(properties.getUsername())
                .password(properties.getPassword())
                .build();
        PostgresConnectionFactory originalFactory = new PostgresConnectionFactory(config);
        return new TenantAwareConnectionFactory(originalFactory);
    }
}

4. 在请求链路中注入租户ID #

通过WebFilter将请求中的租户ID（比如从请求头X-Tenant-ID获取）存入Reactor Context：

@Component
public class TenantWebFilter implements WebFilter {
    @Override
    public Mono<Void> filter(ServerWebExchange exchange, WebFilterChain chain) {
        String tenantId = exchange.getRequest().getHeaders().getFirst("X-Tenant-ID");
        return chain.filter(exchange)
                .contextWrite(ctx -> ctx.put(TenantContext.TENANT_ID_KEY, tenantId));
    }
}

扩展：多租户独立数据库场景 #

如果你的场景是每个租户对应独立数据库，需要动态切换ConnectionFactory，可以修改自定义ConnectionFactory实现：

public class MultiTenantRoutingConnectionFactory implements ConnectionFactory {
    private final Map<String, ConnectionFactory> tenantConnectionFactories;
    private final ConnectionFactory defaultConnectionFactory;

    public MultiTenantRoutingConnectionFactory(Map<String, ConnectionFactory> tenantConnectionFactories, 
                                               ConnectionFactory defaultConnectionFactory) {
        this.tenantConnectionFactories = tenantConnectionFactories;
        this.defaultConnectionFactory = defaultConnectionFactory;
    }

    @Override
    public Publisher<? extends Connection> create() {
        return TenantContext.getTenantId()
                .map(tenantId -> tenantConnectionFactories.getOrDefault(tenantId, defaultConnectionFactory))
                .flatMap(ConnectionFactory::create);
    }

    @Override
    public ConnectionFactoryMetadata getMetadata() {
        return defaultConnectionFactory.getMetadata();
    }
}

内容的提问来源于stack exchange，提问作者kylie.zoltan