Spring R2DBC Postgres Row Level Security

Question

I'm trying to implement Postgres Row Level Security on my app that uses R2DBC.

I found this AWS post that implements this but uses a non-reactive approach.

I'm having problems converting this to a reactive approach since I can't find a class equivalent to the AbstractRoutingDataSource:

public class TenantAwareDataSource extends AbstractRoutingDataSource {

    private static final Logger LOGGER = LoggerFactory.getLogger(TenantAwareDataSource.class);
    
    @Override
    protected Object determineCurrentLookupKey() {
        Object key = null;
        // Pull the currently authenticated tenant from the security context
        // of the HTTP request and use it as the key in the map that points
        // to the connection pool (data source) for each tenant.
        Authentication authentication = SecurityContextHolder.getContext().getAuthentication();
        try {
            if (!(authentication instanceof AnonymousAuthenticationToken)) {
                Tenant currentTenant = (Tenant) authentication.getPrincipal();
                key = currentTenant.getId();
            }
        } catch (Exception e) {
            LOGGER.error("Failed to get current tenant for data source lookup", e);
            throw new RuntimeException(e);
        }
        return key;
    }
    
    @Override
    public Connection getConnection() throws SQLException {
        // Every time the app asks the data source for a connection
        // set the PostgreSQL session variable to the current tenant
        // to enforce data isolation.
        Connection connection = super.getConnection();
        try (Statement sql = connection.createStatement()) {
            LOGGER.info("Setting PostgreSQL session variable app.current_tenant = '{}' on {}", determineCurrentLookupKey().toString(), this);
            sql.execute("SET SESSION app.current_tenant = '" + determineCurrentLookupKey().toString() + "'");
        } catch (Exception e) {
            LOGGER.error("Failed to execute: SET SESSION app.current_tenant = '{}'", determineCurrentLookupKey().toString(), e);
        }
        return connection;
    }

    @Override
    public String toString() {
        return determineTargetDataSource().toString();
    }
}

What would be the equivalent on R2DBC to AbstractRoutingDataSource?

Thanks

Full source code here.

Spring R2DBC Postgres Row Level Security

0 Answers0