Spring动态切换数据源

11

//定义数据源枚举
public enum DataSourceKey {
    master,
    slave,
}

22

/**
 * 数据源路由
 */
@Slf4j
public class DynamicRoutingDataSource  extends AbstractRoutingDataSource {
    @Override
    protected Object determineCurrentLookupKey() {
        log.error("Current DataSource is [{}]", DynamicDataSourceContextHolder.getDataSourceKey());
        return DynamicDataSourceContextHolder.getDataSourceKey();
    }

}

/**
 * 数据源路由holder
 */
@Slf4j
public class DynamicDataSourceContextHolder {
    private static int counter = 0;


    private static final ThreadLocal<String> CONTEXT_HOLDER = ThreadLocal.withInitial(DataSourceKey.master::name);

    /**
     * 所有数据源key
     */
    public static List<Object> dataSourceKeys = new ArrayList<>();

    /**
     * 从数据源key
     */
    public static List<Object> slaveDataSourceKeys = new ArrayList<>();

    /**
     *
     * set数据源
     */
    public static void setDataSourceKey(String key) {
        CONTEXT_HOLDER.set(key);
    }

    /**
     * get数据源
     */
    public static String getDataSourceKey() {
        return CONTEXT_HOLDER.get();
    }
    /**
     * 主数据源
     */
    public static void useMasterDataSource() {
        CONTEXT_HOLDER.set(DataSourceKey.master.name());
    }

    /**
     * 从数据源
     */
    public static void useSlaveDataSource() {
        try {
            int datasourceKeyIndex = counter % slaveDataSourceKeys.size();//负载均衡-轮询
            CONTEXT_HOLDER.set(String.valueOf(slaveDataSourceKeys.get(datasourceKeyIndex)));
            counter++;
        } catch (Exception e) {
            log.error("Switch slave datasource failed, error message is {}", e.getMessage());
            useMasterDataSource();
            e.printStackTrace();
        }
    }

    public static void clearDataSourceKey() {
        CONTEXT_HOLDER.remove();
    }

    public static boolean containDataSourceKey(String key) {
        return dataSourceKeys.contains(key);
    }

}

33

@Configuration
@MapperScan(basePackages = "com.buyi.mytransaction.dao.one", sqlSessionFactoryRef = "mybatisSqlSessionFactoryOne")
public class MyDynamicMybatisDataSource {

    @Bean(name = "master")
    @Primary
    @ConfigurationProperties(prefix = "spring.datasource.one.master", ignoreInvalidFields = true)
    public DataSource master() {
        return new DruidDataSource();
    }

    @Bean(name = "slave")
    @ConfigurationProperties(prefix = "spring.datasource.one.slave", ignoreInvalidFields = true)
    public DataSource slave() {
        return new DruidDataSource();
    }

    @Bean("dynamicDataSource")
    public DataSource dynamicDataSource() {
        DynamicRoutingDataSource dynamicRoutingDataSource = new DynamicRoutingDataSource();
        Map<Object, Object> dataSourceMap = new HashMap<>(4);
        dataSourceMap.put(DataSourceKey.master.name(), master());
        dataSourceMap.put(DataSourceKey.slave.name(), slave());

        dynamicRoutingDataSource.setDefaultTargetDataSource(master());//默认数据源
        dynamicRoutingDataSource.setTargetDataSources(dataSourceMap);//数据源

        DynamicDataSourceContextHolder.dataSourceKeys.addAll(dataSourceMap.keySet());//所有数据源
        DynamicDataSourceContextHolder.slaveDataSourceKeys.addAll(dataSourceMap.keySet());//设置从数据源
        DynamicDataSourceContextHolder.slaveDataSourceKeys.remove(DataSourceKey.master.name());
        return dynamicRoutingDataSource;
    }

    //事务管理器
    @Bean(name = "mybatisTransactionManagerOne")
    @Primary
    public DataSourceTransactionManager mybatisTransactionManager(@Qualifier("dynamicDataSource") DataSource dataSource) throws Exception {
        return new DataSourceTransactionManager(dataSource);
    }

    //会话工厂
    @Bean("mybatisSqlSessionFactoryOne")
    @Primary
    public SqlSessionFactory mybatisSqlSessionFactory(@Qualifier("dynamicDataSource") DataSource dataSource) throws Exception {
        SqlSessionFactoryBean sf = new SqlSessionFactoryBean();
        sf.setDataSource(dataSource);
        sf.setMapperLocations(new PathMatchingResourcePatternResolver().getResources("classpath:repository/one/**/*Mapper.xml"));

        //自动转驼峰   resultType="***"
        org.apache.ibatis.session.Configuration configuration = new org.apache.ibatis.session.Configuration();
        configuration.setMapUnderscoreToCamelCase(true);
        sf.setConfiguration(configuration);
        return sf.getObject();
    }

}

44

@Slf4j
@Aspect
@Component
public class DynamicDataSourceAspect {

    private final String[] QUERY_PREFIX = {"select","query"};


    @Pointcut("execution( * com.buyi.mytransaction.dao.one.*.*(..))")
    public void daoAspect() {
    }

    /**
     * 切换数据源
     */
    @Before("daoAspect()")
    public void switchDataSource(JoinPoint point) {
        Boolean isQueryMethod = isQueryMethod(point.getSignature().getName());
        if (isQueryMethod) {
            DynamicDataSourceContextHolder.useSlaveDataSource();
            log.debug("Switch DataSource to [{}] in Method [{}]",
                    DynamicDataSourceContextHolder.getDataSourceKey(), point.getSignature());
        }
    }

    /**
     * 重置数据源
     */
    @After("daoAspect())")
    public void restoreDataSource(JoinPoint point) {
        DynamicDataSourceContextHolder.clearDataSourceKey();
        log.debug("Restore DataSource to [{}] in Method [{}]",
                DynamicDataSourceContextHolder.getDataSourceKey(), point.getSignature());
    }


    private Boolean isQueryMethod(String methodName) {
        for (String prefix : QUERY_PREFIX) {
            if (methodName.startsWith(prefix)) {
                return true;
            }
        }
        return false;
    }
}

    @Test
    //测试切换数据源
    public void test3() {
        Company c = Company.builder().companyId("789").companyName("百度").memo("人工智能").build();
        oneCompanyservice.oneInsert(c);

        List<Company> companyList = oneCompanyservice.selectAll();
        for (Company company : companyList) {
            System.out.println(company);
            //Company(companyId=999, companyName=百度子公司, memo=111)
            //此条数据在从数据库中已存在
        }
    }