Spring JdbcTemplate - how to prepend every query for achieving multitenancy?

Question

Setup

I have an app using Spring 4.3, JdbcTemplate, Hibernate 5 and MySQL 8. I have implemented multitenancy in hibernate per schema where i switch schemas using using hibernates multitenancy mechanism - MultiTenantConnectionProvider and doing there basically:

connection.createStatement().execute("USE " + databaseNamePrefix + tenantIdentifier); 

and this works.

Now the reporting part of my app uses JdbcTemplate for querying the DB. And now i want to similarily before every query executed by JdbcTemplate issue this USE tenantIdentifier statement.

Question

How can i prepend some SQL or a statement to every query executed by JdbcTemplate?

What i have tried

I looked into JdbcTemplate and only thing i found is setting a NativeJdbcExtractor. I have tried the code below but it's not even loggin that he is going through this methods.

@Bean
@DependsOn("dataSource")
public JdbcTemplate jdbcTemplate() {
  JdbcTemplate jdbcTemplate = new JdbcTemplate(dataSource());
  jdbcTemplate.setNativeJdbcExtractor(new SimpleNativeJdbcExtractor(){
     @Override
     public Connection getNativeConnection(Connection con) throws SQLException {
        LOGGER.info("getNativeConnection");
        System.out.println("aaa");
        return super.getNativeConnection(con);
     }

     @Override
     public Connection getNativeConnectionFromStatement(Statement stmt) throws SQLException {
        System.out.println("aaa");
        LOGGER.info("getNativeConnectionFromStatement");
        return super.getNativeConnectionFromStatement(stmt);
     }

  });
  return jdbcTemplate;
}

Added feature request to Spring: https://jira.spring.io/browse/SPR-17342

EDIT: i looked at Spring 5 and they removed the JdbcExtractor thing so this is definetly the wrong path.

Messiaen

Answer 1

There won't be an easy way to do this with JdbcTemplate as some methods are very generic e.g. execute(ConnectionCallback<T> action) methods allows direct access to java.sql.Connection object.

It would be easier to have a separate DataSource bean for every tenant and resolve this with qualified auto-wiring in Spring.

Having per-tenant java.sql.Connection will allow to execute USE tenantIdentifier statement when new database connection is opened (some pool libraries support this out the box). As MySQL USE statement docs this can be done once per session:

The USE db_name statement tells MySQL to use the db_name database as the default (current) database for subsequent statements. The database remains the default until the end of the session or another USE statement is issued.

Abject

Answer 2

Do not create jdbc template bean. Instead, you can use entity manager factory to create new instance of jdbc template, every time you need to execute a query. This approach worked for me.

public class JdbcQueryTemplate {

    public JdbcTemplate getJdbcTemplate(EntityManagerFactory emf) {
        EntityManagerFactoryInfo info = (EntityManagerFactoryInfo) emf;
        return new JdbcTemplate(info.getDataSource());
    }

    public NamedParameterJdbcTemplate getNamedJdbcTemplate(EntityManagerFactory emf) {
        EntityManagerFactoryInfo info = (EntityManagerFactoryInfo) emf;
        return new NamedParameterJdbcTemplate(info.getDataSource());
    }
}

And then use the class for querying.

public class Test{

  @Autowired
  private EntityManagerFactory entityManagerFactory;

  public List<Entity> executeQuery() {
      return new JdbcQueryTemplate().getNamedJdbcTemplate(entityManagerFactory)
              .query("query", new BeanPropertyRowMapper<>(Entity.class));
  }
}

Mutant

The solution works fine but it result to Too many connections error. – Pahl

Answer 3

It's late, but this is my solution, maybe can help someone. Not very elegant, but it works very well. I override JdbcTemplate, for spring 5:

import org.springframework.dao.DataAccessException;
import org.springframework.jdbc.core.*;
import org.springframework.jdbc.datasource.ConnectionHolder;
import org.springframework.jdbc.datasource.DataSourceUtils;
import org.springframework.jdbc.support.JdbcUtils;
import org.springframework.transaction.support.TransactionSynchronizationManager;
import org.springframework.util.Assert;

import javax.sql.DataSource;
import java.sql.*;

public class TenantJdbcTemplate extends JdbcTemplate {

    public TenantJdbcTemplate(DataSource dataSource) {
        super(dataSource);
    }

    @Override
    public <T> T execute(StatementCallback<T> action) throws DataAccessException {
        Assert.notNull(action, "Callback object must not be null");

        Connection con = DataSourceUtils.getConnection(obtainDataSource());
        boolean isTransactionalCon = isTransactionalConnection(con, getDataSource());
        Statement stmt = null;
        try {
            stmt = con.createStatement();
            applyStatementSettings(stmt);
            Statement stmtToUse = stmt;
            if (!isTransactionalCon) {
                setSchema(stmtToUse);
            }
            T result = action.doInStatement(stmtToUse);
            handleWarnings(stmt);
            return result;
        }
        catch (SQLException ex) {
            // Release Connection early, to avoid potential connection pool deadlock
            // in the case when the exception translator hasn't been initialized yet.
            String sql = getSql(action);
            JdbcUtils.closeStatement(stmt);
            stmt = null;
            DataSourceUtils.releaseConnection(con, getDataSource());
            con = null;
            throw translateException("StatementCallback", sql, ex);
        }
        finally {
            JdbcUtils.closeStatement(stmt);
            DataSourceUtils.releaseConnection(con, getDataSource());
        }
    }

    @Override
    public <T> T execute(PreparedStatementCreator psc, PreparedStatementCallback<T> action)
                    throws DataAccessException {

        Assert.notNull(psc, "PreparedStatementCreator must not be null");
        Assert.notNull(action, "Callback object must not be null");
        if (logger.isDebugEnabled()) {
            String sql = getSql(psc);
            logger.debug("Executing prepared SQL statement" + (sql != null ? " [" + sql + "]" : ""));
        }

        Connection con = DataSourceUtils.getConnection(obtainDataSource());
        boolean isTransactionalCon = isTransactionalConnection(con, getDataSource());

        PreparedStatement ps = null;
        try {
            ps = psc.createPreparedStatement(con);
            applyStatementSettings(ps);
            if (!isTransactionalCon) {
                try (Statement stmt = con.createStatement()) {
                    setSchema(stmt);
                }
            }
            T result = action.doInPreparedStatement(ps);
            handleWarnings(ps);
            return result;
        }
        catch (SQLException ex) {
            // Release Connection early, to avoid potential connection pool deadlock
            // in the case when the exception translator hasn't been initialized yet.
            if (psc instanceof ParameterDisposer) {
                ((ParameterDisposer) psc).cleanupParameters();
            }
            String sql = getSql(psc);
            psc = null;
            JdbcUtils.closeStatement(ps);
            ps = null;
            DataSourceUtils.releaseConnection(con, getDataSource());
            con = null;
            throw translateException("PreparedStatementCallback", sql, ex);
        }
        finally {
            if (psc instanceof ParameterDisposer) {
                ((ParameterDisposer) psc).cleanupParameters();
            }
            JdbcUtils.closeStatement(ps);
            DataSourceUtils.releaseConnection(con, getDataSource());
        }
    }

    @Override
    public <T> T execute(CallableStatementCreator csc, CallableStatementCallback<T> action)
                    throws DataAccessException {

        Assert.notNull(csc, "CallableStatementCreator must not be null");
        Assert.notNull(action, "Callback object must not be null");
        if (logger.isDebugEnabled()) {
            String sql = getSql(csc);
            logger.debug("Calling stored procedure" + (sql != null ? " [" + sql  + "]" : ""));
        }

        Connection con = DataSourceUtils.getConnection(obtainDataSource());
        boolean isTransactionalCon = isTransactionalConnection(con, getDataSource());
        CallableStatement cs = null;
        try {
            cs = csc.createCallableStatement(con);
            applyStatementSettings(cs);
            if (!isTransactionalCon) {
                try (Statement stmt = con.createStatement()) {
                    setSchema(stmt);
                }
            }
            T result = action.doInCallableStatement(cs);
            handleWarnings(cs);
            return result;
        }
        catch (SQLException ex) {
            // Release Connection early, to avoid potential connection pool deadlock
            // in the case when the exception translator hasn't been initialized yet.
            if (csc instanceof ParameterDisposer) {
                ((ParameterDisposer) csc).cleanupParameters();
            }
            String sql = getSql(csc);
            csc = null;
            JdbcUtils.closeStatement(cs);
            cs = null;
            DataSourceUtils.releaseConnection(con, getDataSource());
            con = null;
            throw translateException("CallableStatementCallback", sql, ex);
        }
        finally {
            if (csc instanceof ParameterDisposer) {
                ((ParameterDisposer) csc).cleanupParameters();
            }
            JdbcUtils.closeStatement(cs);
            DataSourceUtils.releaseConnection(con, getDataSource());
        }
    }

    private static void setSchema(Statement stmt) throws SQLException {
        stmt.execute("set search_path=\"" + TenantIdHolder.getTenantId() + "\";");
    }

    private static String getSql(Object sqlProvider) {
        if (sqlProvider instanceof SqlProvider) {
            return ((SqlProvider) sqlProvider).getSql();
        }
        else {
            return null;
        }
    }

    private static boolean isTransactionalConnection(Connection connection, DataSource dataSource) {
        ConnectionHolder conHolder = (ConnectionHolder) TransactionSynchronizationManager.getResource(dataSource);
        return conHolder != null && conHolder.getConnection() == connection;
    }

}

Greenhaw