利用 AOP 实现动态数据源切换

PiPiD 2020-06-29

790

1. 数据库配置文件

首先新建两个数据库，master 和 slave。

建表 SQL 如下：

SET FOREIGN_KEY_CHECKS=0;

-- ----------------------------
-- Table structure for users
-- ----------------------------
DROP TABLE IF EXISTS `users`;
CREATE TABLE `users` (
  `id` int(11) NOT NULL AUTO_INCREMENT,
  `name` varchar(255) DEFAULT NULL,
  `sex` int(11) DEFAULT NULL,
  `address` varchar(255) DEFAULT NULL,
  PRIMARY KEY (`id`)
) ENGINE=InnoDB AUTO_INCREMENT=3 DEFAULT CHARSET=utf8mb4;

-- ----------------------------
-- Records of users
-- ----------------------------
INSERT INTO `users` VALUES ('1', 'zhang_san', '1', '四川省成都市');
INSERT INTO `users` VALUES ('2', 'lis_si', '1', '重庆市');

resources 目录下 application.yml 文件如下：

spring:
  datasource:
    master:
      url: jdbc:mysql://localhost/master?useUnicode=true&characterEncoding=utf-8&autoReconnect=true
      username: root
      password: root
    slave:
      url: jdbc:mysql://localhost/slave?useUnicode=true&characterEncoding=utf-8&autoReconnect=true
      username: root
      password: root

2. datasource 配置类

DbConfig 类，读取配置文件。


@Configuration
public class DbConfig {

    @Bean(name= "master")
    @Primary
    @ConfigurationProperties(prefix = "spring.datasource.master")
    public DataSource master(){
        return DruidDataSourceBuilder.create().build();
    }

    @Bean(name= "slave")
    @ConfigurationProperties(prefix = "spring.datasource.slave")
    public DataSource slave(){
        return DruidDataSourceBuilder.create().build();
    }
}

3. DataSourceEnum 类

DataSource 枚举类，定义所有的数据源名字，目前有 master 和 slave 两个数据源。需要新增数据源时，在此类中新增即可。

@Getter
public enum DataSourceEnum {
    MASTER("master"),
    SLAVE("slave"),
    ;

    protected String dataSourceName;

    DataSourceEnum(String dataSourceName) {
        this.dataSourceName = dataSourceName;
    }

    public String getDataSourceName() {
        return dataSourceName;
    }
}

4. DynamicDataSourceContextHolder 类

根据 contextHolder 的值来决定当前使用的数据库。初始值为 master。


public class DynamicDataSourceContextHolder {


    private static final ThreadLocal<String> contextHolder = new ThreadLocal<String>() {
        @Override
        protected String initialValue() {
            return DataSourceEnum.MASTER.getDataSourceName();
        }
    };


    static List<Object> dataSourceKeys = new ArrayList<Object>();

    static void setDataSourceKeys(String key) {
        contextHolder.set(key);
    }

    static String getDataSourceKey() {
        return contextHolder.get();
    }

    static void clearDataSourceKey() {
        contextHolder.remove();
    }

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

    // 切换数据源
    static void switchDataSourceKey() {
        String currentKey = getDataSourceKey();

        for (Object dataSourceKey : dataSourceKeys) {
            if (!dataSourceKey.equals(currentKey)) {
                setDataSourceKeys((String) dataSourceKey);
                return;
            }
        }
    }

}

5. DynamicRoutingDataSource 类

继承自 AbstractRoutingDataSource 类，用于动态确定使用的数据库。

@Log4j2
public class DynamicRoutingDataSource extends AbstractRoutingDataSource {

    @Override
    protected Object determineCurrentLookupKey() {
        log.info("determineCurrentLookupKey, {}", DynamicDataSourceContextHolder.getDataSourceKey());

        return DynamicDataSourceContextHolder.getDataSourceKey();
    }
}

6. 业务代码

// mapper 类
@Mapper
public interface UserMapper {

    @Select("SELECT id, name, sex, address FROM users")
    List<User> listAll();

}

// 实体类
@Data
public class User {

    private Long id;

    private String name;

    private Integer sex;

    private String address;

}

// 接口类
public interface UserService {

    JSONObject listUsers();
}

// 实现类
@Service
public class UserServiceImpl implements UserService {

    @Autowired
    private UserMapper userMapper;

    @Override
    public JSONObject listUsers() {
        JSONObject jsonObject = new JSONObject();

        jsonObject.put("data", userMapper.listAll());
        return jsonObject;
    }
}

7. 切面 Aspect

@Aspect
@Component
@Log4j2
@Order(1)
public class DynamicDataSourceAspect {


    @Value("${system.current.datasource:master}")
    private String currentDataSource;

    @Pointcut("execution(* com.seven.dynamicdatasource.mapper.*.*(..))")
    public void pointcut() {};

    private boolean retry = false;

    /**
     * 在 pointcut 之前执行，根据 {currentDataSource} 的值来设置调用的数据源
     * 如果 {currentDataSource} 设置错误，则使用默认的数据源
     * @param point
     */
    @Before("pointcut()")
    public void switchDataSource(JoinPoint point) {
        if (!retry) {
            if (!DynamicDataSourceContextHolder.containDataSourceKey(currentDataSource)) {
                log.error("[Before Advice] datasource [{}] doesn't exist, use default datasource：{}",
                        currentDataSource, DynamicDataSourceContextHolder.getDataSourceKey());
            } else {
                DynamicDataSourceContextHolder.setDataSourceKeys(currentDataSource);
            }
        }

        log.info("[Before Advice] pointcut: {}, current datasource is [{}]",
                point.getSignature().getName(),
                DynamicDataSourceContextHolder.getDataSourceKey());
    }

    /**
     * 在 pointcut 执行时执行，首先执行 joinPoint.proceed()
     * 如果执行成功，返回执行的结果
     * 如果执行失败，切换数据源，再次执行 joinPoint.proceed()
     * @param joinPoint
     * @return
     * @throws Throwable
     */
    @Around("pointcut()")
    public Object around(ProceedingJoinPoint joinPoint) throws Throwable {
        try {
            return joinPoint.proceed();
        } catch (Exception e) {
            log.error("[Around Advice] execute {} error: {}", joinPoint.getSignature(), e.getStackTrace());
            setRetry(true);
            DynamicDataSourceContextHolder.switchDataSourceKey();
            return joinPoint.proceed();
        }
    }

    /**
     * 在 poingcut 执行后执行，设置 retry 参数为 false，并且清除当前设置的数据源
     * @param point
     */
    @After("pointcut()")
    public void restoreDataSource(JoinPoint point) {
        setRetry(false);
        DynamicDataSourceContextHolder.clearDataSourceKey();
    }

    private void setRetry(boolean retry) {
        this.retry = retry;
    }
}

8. 项目启动类

启动类需要加上 @ComponentScan(basePackages = "com.seven")
注解。

@SpringBootApplication
@ComponentScan(basePackages = "com.seven")
public class DynamicDatasourceApplication {

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

}