[Java知识库]使用注解通过AOP实现Mybatis多数据源动态切换

最近有个公共的项目，需要连接多个系统的数据库读取数据生成报表。框架使用的是Springboot+Mybatis。下面展示如何使用注解通过AOP实现多数据源的自动切换。

一、首先需要定义好数据源的类型。

package com.test.db;

public enum DatabaseType {

    DB_ORDER, // 订单库
    DB_MEMBER, // 会员库
    DB_STOCK; // 库存库

    /**
     * 通过名字获取类型
     * @param name
     * @return
     */
    public static DatabaseType getByName(String name) {
        for (DatabaseType dbType : DatabaseType.values()) {
            if (dbType.name().equals(name)) {
                return dbType;
            }
        }
        // 默认返回一个 比如：DB_ORDER
        return DatabaseType.DB_ORDER;
    }

}

二、定义一个容器来保存每个线程当前使用的数据库类型。

package com.test.db;

public class DatabaseContextHolder {

    private static final ThreadLocal<DatabaseType> DATABASE_TYPE_HOLDER = new ThreadLocal<>();

    public static void setDatabaseType(DatabaseType dbType) {
        DATABASE_TYPE_HOLDER.set(dbType);
    }

    public static void clearDatabaseType() {
        DATABASE_TYPE_HOLDER.remove();
    }

    public static DatabaseType getDatabaseType() {
        return DATABASE_TYPE_HOLDER.get();
    }

}

三、自定义动态数据源继承AbstractRoutingDataSource。

AbstractRoutingDataSource里面有个成员变量targetDataSources如下：
@Nullable private Map<Object, Object> targetDataSources;

这个Map就用来保存多个数据源。KEY可以保存类型，VALUE就是数据源对象。

package com.test.db;

import org.springframework.jdbc.datasource.lookup.AbstractRoutingDataSource;

public class DynamicDataSource extends AbstractRoutingDataSource {

    @Override
    protected Object determineCurrentLookupKey() {
        // 根据返回的key从targetDataSources中获取该key对应的value(即具体的数据源对象)
        return DatabaseContextHolder.getDatabaseType();
    }

}

四、配置数据源，将所有的数据源放置到targetDataSources中。

package com.test.config;

import com.alibaba.druid.pool.DruidDataSourceFactory;
import com.test.db.DatabaseType;
import com.test.db.DynamicDataSource;
import org.apache.ibatis.session.SqlSessionFactory;
import org.mybatis.spring.SqlSessionFactoryBean;
import org.mybatis.spring.annotation.MapperScan;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.beans.factory.annotation.Qualifier;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;
import org.springframework.context.annotation.Primary;
import org.springframework.core.io.support.PathMatchingResourcePatternResolver;
import org.springframework.jdbc.datasource.DataSourceTransactionManager;

import javax.sql.DataSource;
import java.util.HashMap;
import java.util.Map;
import java.util.Properties;

/**
 * springboot集成mybatis的基本入口 1）创建数据源(如果采用的是默认的tomcat-jdbc数据源，则不需要)
 * 2）创建SqlSessionFactory 3）配置事务管理器，除非需要使用事务，否则不用配置
 */
@Configuration
@MapperScan("com.test.mapper")
public class DynamicDataSourceConfig {

    private static final String MAPPER_LOCATION = "classpath:mapper/*.xml";

    @Autowired
    private OrderDataSourceProperties orderDataSourceProperties;

    @Autowired
    private MemberDataSourceProperties memberDataSourceProperties;

    @Autowired
    private StockDataSourceProperties stockDataSourceProperties;
   

    @Bean
    public DataSource orderDataSource() throws Exception {
        Properties props = new Properties();
        props.put("driverClassName", orderDataSourceProperties.getDriverClassName());
        props.put("url", orderDataSourceProperties.getJdbcUrl());
        props.put("username", orderDataSourceProperties.getUsername());
        props.put("password", orderDataSourceProperties.getPassword());
        return DruidDataSourceFactory.createDataSource(props);
    }

    @Bean
    public DataSource memberDataSource() throws Exception {
        Properties props = new Properties();
        props.put("driverClassName", memberDataSourceProperties.getDriverClassName());
        props.put("url", memberDataSourceProperties.getJdbcUrl());
        props.put("username", memberDataSourceProperties.getUsername());
        props.put("password", memberDataSourceProperties.getPassword());
        return DruidDataSourceFactory.createDataSource(props);
    }
    

    @Bean
    public DataSource stockDataSource() throws Exception {
        Properties props = new Properties();
        props.put("driverClassName", stockDataSourceProperties.getDriverClassName());
        props.put("url", stockDataSourceProperties.getJdbcUrl());
        props.put("username", stockDataSourceProperties.getUsername());
        props.put("password", stockDataSourceProperties.getPassword());
        return DruidDataSourceFactory.createDataSource(props);
    }

    /**
     * @Primary 该注解表示在同一个接口有多个实现类可以注入的时候，默认选择哪一个，而不是让@autowire注解报错
     * @Qualifier 根据名称进行注入，通常是在具有相同的多个类型的实例的一个注入（例如有多个DataSource类型的实例）
     */
    @Primary
    @Bean("dynamicDataSource")
    public DynamicDataSource dataSource(@Qualifier("orderDataSource") DataSource orderDataSource,
                                        @Qualifier("memberDataSource") DataSource memberDataSource,
                                        @Qualifier("stockDataSource") DataSource stockDataSource) {
        Map<Object, Object> targetDataSources = new HashMap<>();
        targetDataSources.put(DatabaseType.DB_ORDER, orderDataSource);
        targetDataSources.put(DatabaseType.DB_MEMBER, memberDataSource);
        targetDataSources.put(DatabaseType.DB_STOCK, stockDataSource);
        DynamicDataSource dataSource = new DynamicDataSource();
        dataSource.setTargetDataSources(targetDataSources);
        dataSource.setDefaultTargetDataSource(orderDataSource);
        return dataSource;
    }

    /**
     * 根据数据源创建SqlSessionFactory
     */
    @Bean
    public SqlSessionFactory sqlSessionFactory(@Qualifier("dynamicDataSource") DynamicDataSource ds) throws Exception {
        SqlSessionFactoryBean fb = new SqlSessionFactoryBean();
        fb.setDataSource(ds);
        /**
         * !!!这里我不指定的话，会报错：org.apache.ibatis.binding.BindingException: Invalid bound statement (not found)
         */  
        fb.setTypeAliasesPackage("com.test.entity");
        fb.setMapperLocations(new PathMatchingResourcePatternResolver().getResources(DynamicDataSourceConfig.MAPPER_LOCATION));
        return fb.getObject();
    }

    /**
     * 配置事务管理器
     */
    @Bean
    public DataSourceTransactionManager transactionManager(@Qualifier("dynamicDataSource") DynamicDataSource ds) throws Exception {
        return new DataSourceTransactionManager(ds);
    }
}

package com.test.config;

import org.springframework.boot.context.properties.ConfigurationProperties;
import org.springframework.context.annotation.Configuration;

@Configuration
@ConfigurationProperties(prefix = "spring.datasource.order")
public class OrderDataSourceProperties {

    private String username;
    private String password;
    private String driverClassName;
    private String jdbcUrl;

    public String getUsername() {
        return username;
    }

    public void setUsername(String username) {
        this.username = username;
    }

    public String getPassword() {
        return password;
    }

    public void setPassword(String password) {
        this.password = password;
    }

    public String getDriverClassName() {
        return driverClassName;
    }

    public void setDriverClassName(String driverClassName) {
        this.driverClassName = driverClassName;
    }

    public String getJdbcUrl() {
        return jdbcUrl;
    }

    public void setJdbcUrl(String jdbcUrl) {
        this.jdbcUrl = jdbcUrl;
    }
}

package com.test.config;

import org.springframework.boot.context.properties.ConfigurationProperties;
import org.springframework.context.annotation.Configuration;

@Configuration
@ConfigurationProperties(prefix = "spring.datasource.member")
public class MemberDataSourceProperties {

    private String username;
    private String password;
    private String driverClassName;
    private String jdbcUrl;

    public String getUsername() {
        return username;
    }

    public void setUsername(String username) {
        this.username = username;
    }

    public String getPassword() {
        return password;
    }

    public void setPassword(String password) {
        this.password = password;
    }

    public String getDriverClassName() {
        return driverClassName;
    }

    public void setDriverClassName(String driverClassName) {
        this.driverClassName = driverClassName;
    }

    public String getJdbcUrl() {
        return jdbcUrl;
    }

    public void setJdbcUrl(String jdbcUrl) {
        this.jdbcUrl = jdbcUrl;
    }
}

package com.test.config;

import org.springframework.boot.context.properties.ConfigurationProperties;
import org.springframework.context.annotation.Configuration;

@Configuration
@ConfigurationProperties(prefix = "spring.datasource.stock")
public class StockDataSourceProperties {

    private String username;
    private String password;
    private String driverClassName;
    private String jdbcUrl;

    public String getUsername() {
        return username;
    }

    public void setUsername(String username) {
        this.username = username;
    }

    public String getPassword() {
        return password;
    }

    public void setPassword(String password) {
        this.password = password;
    }

    public String getDriverClassName() {
        return driverClassName;
    }

    public void setDriverClassName(String driverClassName) {
        this.driverClassName = driverClassName;
    }

    public String getJdbcUrl() {
        return jdbcUrl;
    }

    public void setJdbcUrl(String jdbcUrl) {
        this.jdbcUrl = jdbcUrl;
    }
}

# Datasource
spring.datasource.order.type=com.alibaba.druid.pool.DruidDataSource
spring.datasource.order.username=order
spring.datasource.order.password=111111
spring.datasource.order.driver-class-name=com.microsoft.sqlserver.jdbc.SQLServerDriver
spring.datasource.order.jdbc-url=jdbc:sqlserver://127.0.0.1:1433;database=order;loginTimeout=30;

spring.datasource.member.type=com.alibaba.druid.pool.DruidDataSource
spring.datasource.member.username=member
spring.datasource.member.password=111111
spring.datasource.member.driver-class-name=com.microsoft.sqlserver.jdbc.SQLServerDriver
spring.datasource.member.jdbc-url=jdbc:sqlserver://127.0.0.1:1433;database=member;loginTimeout=30;

spring.datasource.stock.type=com.alibaba.druid.pool.DruidDataSource
spring.datasource.stock.username=stock
spring.datasource.stock.password=111111
spring.datasource.stock.driver-class-name=com.microsoft.sqlserver.jdbc.SQLServerDriver
spring.datasource.stock.jdbc-url=jdbc:sqlserver://127.0.0.1:1433;database=stock;loginTimeout=30;

# mybatis
mybatis-plus.mapper-locations=classpath:mapper/*.xml
mybatis-plus.type-aliases-package=com.test.entity

五、启动类的配置。

package com.test;

import org.mybatis.spring.annotation.MapperScan;
import org.springframework.boot.SpringApplication;
import org.springframework.boot.autoconfigure.SpringBootApplication;
import org.springframework.boot.autoconfigure.jdbc.DataSourceAutoConfiguration;

@SpringBootApplication(exclude = DataSourceAutoConfiguration.class)
@MapperScan(basePackages = "com.test.mapper")
public class DemoTestApplication {

	public static void main(String[] args) {
		SpringApplication.run(DemoTestApplication.class, args);
	}

}

这里如果不加(exclude = DataSourceAutoConfiguration.class)配置，会报循环依赖的错误。

      ↓
   sqlSessionFactory defined in class
┌─────┐
|  dynamicDataSource defined in class
↑     ↓
|  orderDataSource defined in class
↑     ↓
|  dataSourceInitializer
└─────┘

至此，我们的多数据源就搭建好了。下面开始定义注解，通过AOP实现自动切换。

六、自定义注解。

package com.test.db;

import org.springframework.stereotype.Component;

import java.lang.annotation.ElementType;
import java.lang.annotation.Retention;
import java.lang.annotation.RetentionPolicy;
import java.lang.annotation.Target;

@Target(ElementType.METHOD)
@Retention(RetentionPolicy.RUNTIME)
@Component
public @interface DataSource {

    String value() default "";

}

七、通过AOP实现自动切换。

package com.test.db;

import org.aspectj.lang.JoinPoint;
import org.aspectj.lang.annotation.AfterReturning;
import org.aspectj.lang.annotation.Aspect;
import org.aspectj.lang.annotation.Before;
import org.aspectj.lang.annotation.Pointcut;
import org.aspectj.lang.reflect.MethodSignature;
import org.springframework.stereotype.Component;

import java.lang.reflect.Method;

@Aspect
@Component
public class DataSourceAspect {


    @Pointcut("execution(* com.test.mapper..*.*(..))")
    private void anyMethod() {
    }

    @AfterReturning(value = "anyMethod()", returning = "result")
    public void afterReturning(JoinPoint joinPoint, Object result){
        DatabaseContextHolder.clearDatabaseType();
    }

    @Before(value="anyMethod()")
    public void before(JoinPoint joinPoint) throws Throwable {
        MethodSignature methodSignature = (MethodSignature) joinPoint.getSignature();
        Method method = methodSignature.getMethod();
        // 如果方法体上使用了DataSource注解
        if (method.isAnnotationPresent(DataSource.class)) {
            // 获取该方法上的注解名
            DataSource datasource = method.getAnnotation(DataSource.class);
            // 将方法体上的注解的值赋予给DataSourceHolder数据源持有类
            DatabaseContextHolder.setDatabaseType(DatabaseType.getByName(datasource.value()));
        }else{
            // 默认使用DB_ORDER
            DatabaseContextHolder.setDatabaseType(DatabaseType.DB_ORDER);
        }
    }
}

八、注解的使用。

package com.test.mapper;

import com.test.entity.Member;

public interface MemberMapper {
   
    @DataSource("DB_MEMBER")
    Member selectByPrimaryKey(String memberId);

}

package com.test.service;

import com.test.entity.Member;

public interface IMemberService {

    Member selectByPrimaryKey(String id);

}

package com.test.service.impl;

import com.test.entity.Member;
import com.test.mapper.MemberMapper;
import com.test.service.IMemberService;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Service;

@Service
public class MemberServiceImpl implements IMemberService {

    @Autowired
    private MemberMapper memberMapper;

    @Override
    public Member selectByPrimaryKey(String id) {
        return memberMapper.selectByPrimaryKey(id);
    }
}

package com.test.controller;

import com.test.entity.Member;
import com.test.mapper.MemberMapper;
import com.test.service.IMemberService;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.web.bind.annotation.GetMapping;
import org.springframework.web.bind.annotation.RequestMapping;
import org.springframework.web.bind.annotation.RequestParam;
import org.springframework.web.bind.annotation.RestController;

@RestController
@RequestMapping("/test")
public class TestController {

    @Autowired
    private IMemberService memberService;

    @GetMapping("/query")
    public String testQuery(@RequestParam("id") String id) {

        Member member = memberService.selectByPrimaryKey(id);

        return member.toString();
    }

}

package com.test.entity;

public class Member {
    
    private String id;
    
    private String name;
    
    private Integer age;
    
    private String phone;

    public String getId() {
        return id;
    }

    public void setId(String id) {
        this.id = id;
    }

    public String getName() {
        return name;
    }

    public void setName(String name) {
        this.name = name;
    }

    public Integer getAge() {
        return age;
    }

    public void setAge(Integer age) {
        this.age = age;
    }

    public String getPhone() {
        return phone;
    }

    public void setPhone(String phone) {
        this.phone = phone;
    }

    @Override
    public String toString() {
        return "Member{" +
                "id='" + id + '\'' +
                ", name='" + name + '\'' +
                ", age=" + age +
                ", phone='" + phone + '\'' +
                '}';
    }
}