Not able to switch database after defining Spring AOP

Question

I'm working on a multitenant architecture and we've to deal with a dynamic creation and switching of databases.

The problem I'm facing is that I wrote the entire code of switching database at the AOP level and then at the controller or service level I'm not able to make another switch.

AOP class

/**
 * The type Api util.
 */
@Aspect
@Component
@Order(2)
public class APIUtil {
    private static final Logger log = LoggerFactory.getLogger(APIUtil.class);
    private final SchoolMasterService schoolMasterService;

    /**
     * Instantiates a new Api util.
     *
     * @param schoolMasterService the school master service
     * @param helperService
     */
    public APIUtil(SchoolMasterService schoolMasterService) {
        this.schoolMasterService = schoolMasterService;
    }

    /**
     * Around controller methods object.
     *
     * @param proceedingJoinPoint the proceeding join point
     * @return the object
     * @throws Throwable the throwable
     */
    @Around("execution(* com.example.board.controller.*.*(..))")
    public Object aroundControllerMethods(ProceedingJoinPoint proceedingJoinPoint) throws Throwable {
        log.info("All schools loaded!");
        TenantContext.setCurrentTenant(DEFAULT_TENANT_ID);
        schoolMasterService.findAllMasters();
        HttpServletRequest request = ((ServletRequestAttributes) RequestContextHolder.currentRequestAttributes()).getRequest();
        String tenantID = request.getHeader(XTENANTID).trim();
        return filterByHeader(tenantID, proceedingJoinPoint);
    }

   private Object filterByHeader(String tenantID, ProceedingJoinPoint joinPoint) throws Throwable {
        SchoolMaster schoolMaster = schoolMasterService.findBySchoolId(Long.parseLong(tenantID));
        log.info(format("Current school is %s", schoolMaster.getDataSourceKey()));
        TenantContext.setCurrentTenant(schoolMaster.getDataSourceKey());
        return joinPoint.proceed();
    }
}

CurrentTenantIdentifierResolverImpl class

@Component
public class CurrentTenantIdentifierResolverImpl implements CurrentTenantIdentifierResolver {
    private final Logger logger = LoggerFactory.getLogger(getClass());

    @Override
    public String resolveCurrentTenantIdentifier() {
        String currentDataSourceKey = TenantContext.getCurrentTenant();
        if (Objects.isNull(currentDataSourceKey)) {
            currentDataSourceKey = DEFAULT_TENANT_ID;
        }
        logger.debug("currentDataSourceKey {}", currentDataSourceKey);
        return currentDataSourceKey;
    }

    @Override
    public boolean validateExistingCurrentSessions() {
        return true;
    }
}

MultiTenantConnectionProviderImpl class

@Component(MultiTenantConnectionProviderImpl.BEAN_ID)
public class MultiTenantConnectionProviderImpl extends AbstractDataSourceBasedMultiTenantConnectionProviderImpl {
    /**
     * The constant BEAN_ID.
     */
    public static final String BEAN_ID = "multiTenantConnectionProvider";
    private static final long serialVersionUID = 7395318315512114572L;

    @Autowired
    private TenantDataSource tenantDataSource;

    private Logger log;

    /**
     * Instantiates a new Multi tenant connection provider.
     */
    public MultiTenantConnectionProviderImpl() {
        log = LoggerFactory.getLogger(getClass());
    }

    @Override
    protected DataSource selectAnyDataSource() {
        log.debug("selectAnyDataSource , returning dafault tenantid");
        return tenantDataSource.getDataSource(DEFAULT_TENANT_ID);
    }

    @Override
    protected DataSource selectDataSource(String tenantIdentifier) {
        log.debug("selected Datasource {} ", tenantIdentifier);
        return tenantDataSource.getDataSource(tenantIdentifier);
    }
}

TenantDataSource class

@Component
public class TenantDataSource {
    private final Map<Object, Object> tenantDataSourcesMap = new HashMap<>();
    private final Logger log = LoggerFactory.getLogger(TenantDataSource.class);

    private final DataSourceProperties dataSourceProperties;
    private final DataSource dataSource;
    private final JdbcTemplate jdbcTemplateObject;

    /**
     * Instantiates a new Tenant data source.
     *
     * @param properties the properties
     * @param source     the source
     * @param object     the object
     */
    public TenantDataSource(DataSourceProperties properties, DataSource source, JdbcTemplate object) {
        this.dataSourceProperties = properties;
        this.dataSource = source;
        this.jdbcTemplateObject = object;
    }

    /**
     * Add default datasource to map.
     */
    @PostConstruct
    void addDefaultDatasourceToMap() {
        tenantDataSourcesMap.put(DEFAULT_TENANT_ID, dataSource);
    }

    /**
     * Gets data source.
     *
     * @param dataSourceName the data source name
     * @return the data source
     */
    public DataSource getDataSource(String dataSourceName) {
        DataSource currentDatasource = null;
        log.debug("getDataSource().dataSourceName {}", dataSourceName);
        if (tenantDataSourcesMap.containsKey(dataSourceName)) {
            currentDatasource = (DataSource) tenantDataSourcesMap.get(dataSourceName);
        }
        return currentDatasource;
    }

    /**
     * Load tenant boolean.
     *
     * @param tenantDatasource the tenant datasource
     * @return the boolean
     */
    public boolean loadTenant(SchoolMaster tenantDatasource) {
        try {
            if (!verifyPort(tenantDatasource))
                return false;
            DataSource temp = createDataSource(tenantDatasource);
            boolean result = verifyConnection(temp);
            if (result) {
                tenantDataSourcesMap.putIfAbsent(tenantDatasource.getDataSourceKey(), temp);
            }
            return result;
        } catch (Exception h) {
            return false;
        }
    }

    /**
     * Load all tenants.
     *
     * @param tenantDatasourcesList the tenant datasources list
     */
    public void loadAllTenants(List<SchoolMaster> tenantDatasourcesList) {
        tenantDatasourcesList.forEach(tenant -> tenantDataSourcesMap.putIfAbsent(tenant.getDataSourceKey(), createDataSource(tenant)));
    }

    /**
     * Create data source data source.
     *
     * @param tenantDatasource the tenant datasource
     * @return the data source
     */
    public DataSource createDataSource(SchoolMaster tenantDatasource) {
        HikariDataSource hikariDataSource = null;
        if (Objects.nonNull(tenantDatasource)) {
            String url = JDBCMYSQL + tenantDatasource.getSchoolIP().trim() + ":" + tenantDatasource.getDataSourcePort().trim() + SLASH + tenantDatasource.getDataSourceKey().trim() + "?createDatabaseIfNotExist=true&useSSL=true";
            hikariDataSource = (HikariDataSource) DataSourceBuilder.create()
                    .driverClassName(dataSourceProperties.getDriverClassName())
                    .username(tenantDatasource.getDataSourceUserName()).password(tenantDatasource.getDataSourcePassword())
                    .url(url)
                    .build();
            setConnectionPooling(hikariDataSource);
        }
        return hikariDataSource;
    }

    /**
     * Create schema.
     *
     * @param dataSourceName the data source name
     * @throws SQLException the sql exception
     */
    public void createSchema(String dataSourceName) throws SQLException {
        if (tenantDataSourcesMap.containsKey(dataSourceName)) {
            jdbcTemplateObject.execute(CREATE_SCHEMA + " " + dataSourceName);
            jdbcTemplateObject.execute(USE_SCHEMA + " " + dataSourceName);
            DataSource currentDataSource = (DataSource) tenantDataSourcesMap.get(dataSourceName);
            ClassPathResource resource = new ClassPathResource("dbscripts/schema.sql");
            try (Connection connection = currentDataSource.getConnection()) {
                ScriptUtils.executeSqlScript(connection, new EncodedResource(resource, "UTF-8"));
            }
            jdbcTemplateObject.execute(USE_SCHEMA + " " + DEFAULT_TENANT_ID);
        }
    }

    /**
     * Drop schema.
     *
     * @param dataSourceName the data source name
     */
    public void dropSchema(String dataSourceName) {
        if (tenantDataSourcesMap.containsKey(dataSourceName)) {
            DataSource currentDataSource = (DataSource) tenantDataSourcesMap.get(dataSourceName);
            JdbcTemplate template = new JdbcTemplate(currentDataSource);
            template.execute(DROP_SCHEMA + " " + dataSourceName);
        }
    }

    /**
     * Sets connection pooling.
     *
     * @param hikariDataSource the hikari data source
     */
    private void setConnectionPooling(HikariDataSource hikariDataSource) {
        hikariDataSource.setMinimumIdle(2);
        hikariDataSource.setMaximumPoolSize(5);
        hikariDataSource.setIdleTimeout(100000);
        hikariDataSource.setMaxLifetime(3000000);
        hikariDataSource.setConnectionTimeout(200000);
        hikariDataSource.setLeakDetectionThreshold(2100);
        hikariDataSource.setConnectionTestQuery("SELECT 1 FROM DUAL");
        hikariDataSource.setAutoCommit(false);
    }


    /**
     * Verify connection boolean.
     *
     * @param currentDatasource the current datasource
     * @return the boolean
     */
    private boolean verifyConnection(DataSource currentDatasource) {
        try (Connection ignored = currentDatasource.getConnection()) {
            return true;
        } catch (Exception e) {
            return false;
        }
    }

    /**
     * Verify port boolean.
     *
     * @param tenantDataSource the tenant data source
     * @return the boolean
     */
    private boolean verifyPort(SchoolMaster tenantDataSource) {
        return tenantDataSource.getDataSourcePort().trim().chars().allMatch(Character::isDigit);
    }
}

TenantContext class

public final class TenantContext {
    private static final ThreadLocal<String> currentTenant = ThreadLocal.withInitial(() -> DEFAULT_TENANT_ID);

    private TenantContext() {
    }

    /**
     * Gets current tenant.
     *
     * @return the current tenant
     */
    public static String getCurrentTenant() {
        return currentTenant.get();
    }

    /**
     * Sets current tenant.
     *
     * @param tenant the tenant
     */
    public static void setCurrentTenant(String tenant) {
        currentTenant.set(tenant);
    }

    /**
     * Clear.
     */
    public static void clear() {
        currentTenant.remove();
    }
}

At the controller/service level, I'm not able to switch the database using TenantContext.setCurrentTenant(String identifier) yet I'm able to do the same thing on AOP.

What is the reason for this? Any way how to fix this?

I would be really grateful for any help.

I've added the example of the Service Code

Where I need to switch database twice which is not being possible.

 @Transactional(readOnly = true)
    @Override
    public List<SyllabusListingResponseDto> getSyllabusByBoardId(Long rackId, Long languageId) {
        logger.info("getSyllabusByBoardId Method called in BoardManagementServiceImpl");
        ResourceRackModel resourceRackModel = resourceRackService.getByRackIdAndStatus(rackId, ACTIVE_STATUS);
        if (Objects.nonNull(resourceRackModel)) {
            TenantContext.setCurrentTenant(DEFAULT_TENANT_ID);
            List<Long> rackIds = resourceRackService.findAllRackIdsByBoardId(rackId);
            rackIds.add(rackId);
            ResourceRackModel boardModel = resourceRackModel;
            if (!boardModel.getParentPath().isEmpty()) {
                String[] ids = resourceRackModel.getParentPath().split(",", 2);
                boardModel = resourceRackService.getByRackIdAndStatus(Long.parseLong(ids[INT_ZERO]), ACTIVE_STATUS);
            }
            TenantContext.setCurrentTenant("S_" + 1);
            BoardVersionModel activeVersionModel = boardVersionRepository.findByBoardIdAndStatusAndVersionStatus(boardModel.getRackId(), ACTIVE_STATUS, ACTIVE_STATUS);
            ContentCountDto contentCountDto = new ContentCountDto().setStatus(true).setRackIds(rackIds).setActiveBoardVersionId(activeVersionModel.getVersionId().toString());
            ResponseModel responseModel = nemrSTCManagementClient.findContentCount(Common.getTenantId(), contentCountDto).getBody();
            if (Objects.nonNull(responseModel)) {
                Map<String, List<String>> lookup = (Map<String, List<String>>) responseModel.getObject();
                String languageCode = languageMasterService.findByIdAndStatus(languageId, ACTIVE_STATUS).getLanguageCode();
                String defaultLanguageCode = languageMasterService.findByIdAndStatus(resourceRackModel.getDefaultLanguageId(), ACTIVE_STATUS).getLanguageCode();
                List<ResourceRackModel> resourceRackModels = resourceRackService.findByParentIdAndStatus(rackId, ACTIVE_STATUS);
                if (resourceRackModels.isEmpty()) {
                    return Collections.emptyList();
                }
                Map<Integer, String> rackTypes = new HashMap<>();
                return getResult(languageId, lookup, new SyllabusListingResponseDto(), boardModel, languageCode, defaultLanguageCode, resourceRackModels, rackTypes);
            }
        }
        return Collections.emptyList();
    }

Answer 1

1. As I have not used MultiTenantConnectionProvider so not sure at what point the datasource is chosen by the framework.

2. But I have a great suspicion it is done by the interceptor auto created by @Transactional, and the framework never reads the TenantContext inside the method. I.e one datasource for one transactional method. So may be worth identifying if my suspicion is true.

   • You can put a breakpoint at the following :

       logger.info("getSyllabusByBoardId Method called in BoardManagement...");
   

   • Then clear the console logs when the break point is there.

   • Then let the method execute and see if any of the following logs are printed when you step through your service method lines.

       log.debug("selectAnyDataSource , returning dafault tenantid");
   

      log.debug("selected Datasource {} ", tenantIdentifier);
   

3. If my suspicion in step 2 is correct, then either you have to remove the @Transactional annotation or split the methods in two and set the correct TenantContext.setCurrentTenant for each method in the controller before calling the service method.