Springboot 2.3.1 dynamically update Jdbc template's schema in Multi-tenant environment

Question

My Project is on spring-boot-starter-parent - "1.5.9.RELEASE" and I'm migrating it to spring-boot-starter-parent - "2.3.1.RELEASE".

This is multi-tenant env application, where one database will have multiple schemas, and based on the tenant-id, execution switches between schemas.

I had achieved this schema switching using SimpleNativeJdbcExtractor but in the latest Springboot version NativeJdbcExtractor is no longer available.

Code snippet for the existing implementation:

 @Bean
@Scope(
        value = ConfigurableBeanFactory.SCOPE_PROTOTYPE,
        proxyMode = ScopedProxyMode.TARGET_CLASS)
public JdbcTemplate jdbcTemplate() {
    JdbcTemplate jdbcTemplate = new JdbcTemplate(dataSource);
    SimpleNativeJdbcExtractor simpleNativeJdbcExtractor = new SimpleNativeJdbcExtractor() {
        @Override
        public Connection getNativeConnection(Connection con) throws SQLException {
            LOGGER.debug("Set schema for getNativeConnection "+Utilities.getTenantId());
            con.setSchema(Utilities.getTenantId());
            return super.getNativeConnection(con);
        }

        @Override
        public Connection getNativeConnectionFromStatement(Statement stmt) throws SQLException {
            LOGGER.debug("Set schema for getNativeConnectionFromStatement "+Utilities.getTenantId());
            Connection nativeConnectionFromStatement = super.getNativeConnectionFromStatement(stmt);
            nativeConnectionFromStatement.setSchema(Utilities.getTenantId());
            return nativeConnectionFromStatement;
        }
    };

    simpleNativeJdbcExtractor.setNativeConnectionNecessaryForNativeStatements(true);
    simpleNativeJdbcExtractor.setNativeConnectionNecessaryForNativePreparedStatements(true);

    jdbcTemplate.setNativeJdbcExtractor(simpleNativeJdbcExtractor);
    return jdbcTemplate;
}

Here Utilities.getTenantId() ( Stored value in ThreadLocal) would give the schema name based on the REST request.

Questions:

• What are the alternates to NativeJdbcExtractor so that schema can be dynamically changed for JdbcTemplate?
• Is there any other way, where while creating the JdbcTemplate bean I can set the schema based on the request.

Any help, code snippet, or guidance to solve this issue is deeply appreciated.

Thanks.

Answer 1

When I was running the application in debug mode I saw Spring was selecting Hikari Datasource.

I had to intercept getConnection call and update schema.

So I did something like below,

Created a Custom class which extends HikariDataSource

public class CustomHikariDataSource extends HikariDataSource {
@Override
public Connection getConnection() throws SQLException {

    Connection connection =  super.getConnection();
    connection.setSchema(Utilities.getTenantId());
    return connection;
}
}

Then in the config class, I created bean for my CustomHikariDataSource class.

 @Bean
public DataSource customDataSource(DataSourceProperties properties) {

    final CustomHikariDataSource dataSource = (CustomHikariDataSource) properties
            .initializeDataSourceBuilder().type(CustomHikariDataSource.class).build();
    if (properties.getName() != null) {
        dataSource.setPoolName(properties.getName());
    }
    return dataSource;
}

Which will be used by the JdbcTemplate bean.

 @Bean
@Scope(
        value = ConfigurableBeanFactory.SCOPE_PROTOTYPE,
        proxyMode = ScopedProxyMode.TARGET_CLASS)
public JdbcTemplate jdbcTemplate() throws SQLException {
    JdbcTemplate jdbcTemplate = new JdbcTemplate(dataSource);
    return jdbcTemplate;
}

With this approach, I will have DataSource bean created only once and for every JdbcTemplate access, the proper schema will be updated during runtime.