DataSourcePoolService.java

package pwc.taxtech.atms.service.impl;

import com.alibaba.druid.pool.DruidDataSource;
import com.google.common.collect.Maps;
import org.apache.commons.lang3.StringUtils;
import org.apache.commons.pool2.impl.GenericKeyedObjectPool;
import org.apache.commons.pool2.impl.GenericKeyedObjectPoolConfig;
import org.apache.ibatis.session.SqlSessionFactory;
import org.mybatis.spring.SqlSessionFactoryBean;
import org.mybatis.spring.SqlSessionTemplate;
import org.springframework.beans.factory.annotation.Qualifier;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.context.annotation.Bean;
import org.springframework.core.io.support.PathMatchingResourcePatternResolver;
import org.springframework.jdbc.core.JdbcTemplate;
import org.springframework.jdbc.datasource.lookup.AbstractRoutingDataSource;
import org.springframework.stereotype.Service;
import pwc.taxtech.atms.common.datasource.DataSourceConfig;
import pwc.taxtech.atms.common.datasource.DataSourceFactory;
import pwc.taxtech.atms.common.datasource.ShardingContextHolder;
import pwc.taxtech.atms.exception.ParameterException;

import javax.annotation.PostConstruct;
import javax.sql.DataSource;

@Service("DynamicDataSource")
public class DataSourcePoolService extends AbstractRoutingDataSource {
    private GenericKeyedObjectPool<String, DruidDataSource> dataSourcePool;

    @Value("${jdbc_url_demo}")
    private String demoUrl;
    @Value("${jdbc_user}")
    private String jdbcUser;
    @Value("${jdbc_password}")
    private String jdbcPwd;

    @PostConstruct
    public void init() throws Exception {
        DataSourceConfig config = new DataSourceConfig();
        config.setUrl(demoUrl);
        config.setUsername(jdbcUser);
        config.setPassword(jdbcPwd);
        GenericKeyedObjectPoolConfig poolConfig = new GenericKeyedObjectPoolConfig();
        poolConfig.setTestOnBorrow(true);
        poolConfig.setTestOnReturn(true);
        poolConfig.setMaxTotalPerKey(1);
        poolConfig.setMaxIdlePerKey(1);
        dataSourcePool = new GenericKeyedObjectPool<>(new DataSourceFactory(config), poolConfig);
        this.setTargetDataSources(Maps.newHashMap());
    }

    @Override
    protected Object determineCurrentLookupKey() {
        return ShardingContextHolder.getDataSourceKey();
    }

    @Override
    protected synchronized DataSource determineTargetDataSource() {
        try {
            String key = ShardingContextHolder.getDataSourceKey();
            DruidDataSource dataSource = dataSourcePool.borrowObject(key);
            dataSourcePool.returnObject(key, dataSource);
            return dataSource;
        } catch (Exception e) {
            throw new RuntimeException("determineTargetDataSource error.", e);
        }
    }

    private void createDataSource(String key) {
        if (StringUtils.isBlank(key)) {
            throw new ParameterException("DataSource name is empty.");
        }
        try {
            dataSourcePool.addObject(key);
        } catch (Exception e) {
            throw new RuntimeException("createDataSource error.", e);
        }
    }

    @Bean
    public SqlSessionFactory dynamicSqlSessionFactory(
            @Qualifier("DynamicDataSource") DataSource dynamicDataSource)
            throws Exception {
        SqlSessionFactoryBean bean = new SqlSessionFactoryBean();
        bean.setDataSource(dynamicDataSource);
        bean.setMapperLocations(new PathMatchingResourcePatternResolver()
                .getResources("classpath*:pwc/taxtech/atms/vat/dao/*Mapper.xml"));
        return bean.getObject();

    }

    @Bean
    public SqlSessionTemplate dynamicSqlSessionTemplate(
            @Qualifier("dynamicSqlSessionFactory") SqlSessionFactory sqlSessionFactory)
            throws Exception {
        return new SqlSessionTemplate(sqlSessionFactory);
    }

    @Bean
    public JdbcTemplate dynamicJdbcTemplate(@Qualifier("DynamicDataSource")DataSource dynamicDataSource) {
        return new JdbcTemplate(dynamicDataSource);
    }

}