Druid多数据源 + mybatis动态分库映射

一、需求背景

具体的项目不便多说，简单来说是外部将结构化数据发送至 kafka 的数十个 topic。
本项目需要消费所有这些数据，并根据topic，存入不同的table。每条结构化数据都拥有公共字段 uid。
根据uid 和 table 将映射到某个具体的数据库 dataSource。

本文主要是实现根据已知的消息，自动将dao层的sql调用映射到不同的数据库连接，至于kafka 等数据来源是什么，并不重要。
因此，下文的实现，仅涉及到一些 Druid 和 aop层面的东西。

二、配置过程

2.1 添加Druid依赖

<dependency>
    <groupId>com.alibaba</groupId>
    <artifactId>druid-spring-boot-starter</artifactId>
    <version>1.1.10</version>
</dependency>

2.2 数据源配置

spring:
  datasource:
    type: com.alibaba.druid.pool.DruidDataSource
    druid:
      default-db:
        driverClassName: org.postgresql.Driver
        url: jdbc:postgresql://localhost:5432/defaultDb
        username: postgres
        password: Gjm6sCZV4h6zl3zSqgRoMMN5ckfKiCbJgKeiA...
      db1:
        driverClassName: org.postgresql.Driver
        url: jdbc:postgresql://localhost:5432/db1
        username: postgres
        password: Gjm6sCZV4h6zl3zSqgRoMMN5ckfKiCbJgKeiA...  

• 此处可以配置需要的任意多个数据库。
• 密码采用密文的时候，需要在创建 DataSource Bean 之前解析出明文。
• 配置DataSource不一定需要采用上述代码的配置方式，完全可以更改属性字段或者自定义配置文件。

2.3 java中获取配置好的DataSource 配置

采用其他的单独的置文件的可以诸多java自带的工具类获取，不赘述

采用 application.yml 方式获取配置非常简单。 如下构建 DruidProperties.java类即可一次性获取所有DB 配置：

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

    private Map<String, Map<String, String>> druid;

    private String type;

    public String getType() {
        return type;
    }

    public DruidProperties setType(String type) {
        this.type = type;
        return this;
    }

    public Map<String, DbProperties> getDruid() {
        Map<String, DbProperties> map = new HashMap<>();
        for(Map.Entry<String, Map<String, String>> entry : druid.entrySet()) {
            map.put(entry.getKey(), JSON.parseObject(JSON.toJSONString(entry.getValue()), DbProperties.class));
        }
        return map;
    }

    public DruidProperties setDruid(Map<String, Map<String, String>> druid) {
        this.druid = druid;
        return this;
    }

    protected static class DbProperties {
        private String driverClassName;
        private String url;
        private String username;
        private String password;

        public String getDriverClassName() {
            return driverClassName;
        }

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

        public String getUrl() {
            return url;
        }

        public DbProperties setUrl(String url) {
            this.url = url;
            return this;
        }

        public String getUsername() {
            return username;
        }

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

        public String getPassword() {
            return CommonUtils.decryptRsa(password);
        }

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

如上，在DruidProperties 类中，定义了一个 内部类 DbProperties，方便其他代码读取DataSource属性。

2.4 配置动态数据源

Druid动态数据源需要继承 AbstractRoutingDataSource。
构建动态数据源Bean，提供给 SqlSessionFactory，这样，当调用sql时，将自动获取sqlSessionFatory配置的动态数据源实例
通过调用此类中重写的 determineCurrentLookupKey() 获取指定的 DataSource。
如下, 一些介绍参照注释

public class DynamicRouteDataSource extends AbstractRoutingDataSource {

    // DynamicRouteDataSource 是一个单例Bean, 而同一时间不同的线程需要获取不同的数据源连接，
    // 因此使用线程本地变量 ThreadLoacal<String> 类用以控制当前线程的数据源连接。
    private static final ThreadLocal<String> contextHolder = ThreadLocal.withInitial(() -> "default-db");

    private static List<String> dataSourceNameList= new ArrayList<>();
    
    public DynamicRouteDataSource(DataSource defaultTargetDataSource, Map<Object, Object> targetDataSources) {
        super.setDefaultTargetDataSource(defaultTargetDataSource);
        super.setTargetDataSources(targetDataSources);
        DynamicRouteDataSource.dataSourceNameList = targetDataSources.keySet().stream().map(a -> (String)a).collect(Collectors.toList());
        super.afterPropertiesSet();
    }


    // 核心方法
    // 此方法确定当前线程选择的数据库，此方法是一个无参方法，因此动态数据源切换时，需要通过其他的方法修改某变量来控制。
    @Override
    protected Object determineCurrentLookupKey() {
        return getDataSource();
    }

    
    // 如果需要动态修改数据源，可重写此方法。
    @Override
    public void setTargetDataSources(Map<Object, Object> targetDataSources) {
        super.setTargetDataSources(targetDataSources);
        DynamicRouteDataSource.dataSourceNameList = targetDataSources.keySet().stream().map(a -> (String)a).collect(Collectors.toList());
        super.afterPropertiesSet();
    }

    public static void setDataSource(String dataSource) {
        contextHolder.set(dataSource);
    }

    public static String getDataSource() {
        return contextHolder.get();
    }

    public static void clearDataSource() {
        contextHolder.remove();
    }
    
    // 外部获取动态数据库数据，此处注意保护信息。不需要可删除
    protected List<String> getDataSourceNameList() {
        return dataSourceNameList;
    }
}

配置SqlSessionFactory，没什么可说的，无非是根据配置构建动态数据源Bean，以此Bean 设置 SqlSessionFactory。

@Configuration
@EnableConfigurationProperties(MybatisProperties.class)
public class DataSourceConfig {

    @Resource
    private MybatisProperties mybatisProperties;
    
    // 2.3 节读取的配置参数。
    @Resource
    public DruidProperties druidProperties;


    @Bean("dynamicDataSource")
    @Primary
    public DataSource dataSource() {
        Map<String, DruidProperties.DbProperties> dbPropertiesMap = druidProperties.getDruid();

        Map<Object, Object> targetDataSources = new HashMap<>();
        for(Map.Entry<String, DruidProperties.DbProperties> entry : dbPropertiesMap.entrySet()) {
            DataSource dataSource = DataSourceBuilder.create()
                    .driverClassName(entry.getValue().getDriverClassName())
                    .url(entry.getValue().getUrl())
                    .username(entry.getValue().getUsername())
                    .password(entry.getValue().getPassword())
                    .build();

            targetDataSources.put(entry.getKey(), dataSource);
        }

        return new DynamicRouteDataSource((DataSource) targetDataSources.get("default-db"), targetDataSources);
    }


    @Bean
    public SqlSessionFactory sqlSessionFactory(DataSource dynamicDataSource) throws Exception {
        SqlSessionFactoryBean bean = new SqlSessionFactoryBean();
        bean.setDataSource(dynamicDataSource);
        bean.setMapperLocations(mybatisProperties.resolveMapperLocations());
        bean.setTypeAliasesPackage(mybatisProperties.getTypeAliasesPackage());
        bean.setConfiguration(mybatisProperties.getConfiguration());
        return bean.getObject();
    }

    @Bean
    public PlatformTransactionManager transactionManager(DataSource dynamicDataSource) {
        return new DataSourceTransactionManager(dynamicDataSource);
    }

}

配置完成动态数据源后，即可在代码中通过注入 动态数据源类 DynamicRouteDataSource 实例，然后调用 setDataSource 方法实现手动选择数据源了。
如果想要达到业务代码层面无感自动切换，可以选择通过下面的AOP方式实现。

三、AOP自动切换动态数据源

先说说我的思路：

• 第一种分库方式：不同的 table 存储在不同数据库中，需要自动完成 table -> dataSource的映射。
  • 1、dao 层不同的方法对应唯一的table，维护好此对应关系，可直接通过 aop 获取方法名映射到具体的dataSource。
  • 2、从入参中可以获取table, 则抓取入参的table。
• 第二种分库方式：不同的数据库中存在相同的表，根据数据的主键、table等hash自动映射到数据库，需要实现 hash -> dataSource的映射
  • hash 往往是通过参数中的主键 uid、执行表名 table等计算得出的，从入参中提取出来即可。
• 第三种分库方式：dao层某方法具体对应唯一的dataSource，可直接通过注解的方式声明dataSource。唯一的麻烦是如果这种形式较多，维护将较为麻烦，此时可以选择第一种思路。

3.1 注解形式提前创建自定义注解

@Target(ElementType.METHOD)
@Retention(RetentionPolicy.RUNTIME)
@Documented
public @interface SelectDataSource {
    String name() default "";
}

3.2 实现 aop 切面

@Aspect
@Configuration
public class AutoMapDataSourceAspect {

    @Resource(name = "dynamicDataSource")
    DynamicRouteDataSource dynamicRouteDataSource;

    // 维护dao层接口到table的映射，table到dataSource的映射。（完全可以合二为一，属于个人代码风格问题)
    // 简单的思路就是维护一份固定的映射关系，也可以通过数据库、配置文件的形式，按需搞
    private static final HashMap<String, String> daoFunctionMapToTable = new HashMap<>();
    private static final HashMap<String, String> tableMapToDataSource = new HashMap<>();

    static {
        daoFunctionMapToTable.put("batchInsertToTableA", "tableA");
        daoFunctionMapToTable.put("batchInsertToTableB", "tableB");
       
        tableMapToDataSource.put("tableA", "default-db");
        tableMapToDataSource.put("tableB", "db1");

    }


    // 切点，此处选择了 DynamicDao 中所有的方法
    @Pointcut("execution(* com.test.db.dao.DynamicDao.*(..))")
    public void opePointCut() {

    }

    // 选择围绕的方式，因为需要在执行前设置具体的数据源，执行完毕之后清空该设置。
    @Around("opePointCut()")
    public Object around(ProceedingJoinPoint point) throws Throwable {
        MethodSignature signature = (MethodSignature) point.getSignature();
        Method method = signature.getMethod();


        // 检验是否有注解的方式，如果采用注解的方式，优先根据注解参数选择数据源
        SelectDataSource dataSource = method.getAnnotation(SelectDataSource.class);
        if(dataSource == null){
            // 获取tableName, 首先尝试通过方法名获取tableName, 若失败，尝试从参数中抓取。
            String tableName = getTableByMethodName(method.getName());
            if(StringUtils.isEmpty(tableName)) {
                Object temp = getArg(method.getParameters(), point.getArgs(), "table");
                tableName = temp == null ? null : (String)temp;
            }

            // 尝试根据tableName 直接获取动态数据源
            String dataSourceName = mapToDataSourceByTable(tableName);
            
            // 尝试获取UID，通过UID、tableName映射到动态数据源
            if(StringUtils.isEmpty(dataSourceName)) {
                // 尝试提取 uid 参数
                Object temp = getArg(method.getParameters(), point.getArgs(), "uid");
                if(temp != null) {
                    dataSourceName = mapToDataSourceByTableAndUid(tableName, (String)temp);
                } else {
                    // 尝试从 collection 类型的参数中提取uid，
                    // 提取后，将包含uid的 collection 参数，进行拆分。
                    Object[] args = point.getArgs();
                    Map<String, Collection> uidArgMap = new HashMap<>();
                    int index = Integer.MIN_VALUE;
                    for(int i = 0; i < args.length; i++) {
                        if(args[i] instanceof Collection) {
                            Collection collection = (Collection) args[i];
                            if(!collection.isEmpty()) {
                                Object obj = collection.toArray()[0];
                                try {
                                    Method tempMethod = obj.getClass().getMethod("getUid");
                                    for(Object tempObj : collection) {
                                        String uid = (String) tempMethod.invoke(tempObj);
                                        if(!uidArgMap.containsKey(uid)) {
                                            uidArgMap.put(uid, new ArrayList<>());
                                        }
                                        uidArgMap.get(uid).add(tempObj);
                                    }
                                    index = i;
                                    break;
                                } catch (NoSuchMethodException ignore){}
                            }
                        }
                    }

                    // 将需要分别选择数据源的入参数据拆分至不同的数据源
                    Map<String, Collection> dataSourceArgsMap = new HashMap<>();
                    for(Map.Entry<String, Collection> entry : uidArgMap.entrySet()) {
                        dataSourceName = mapToDataSourceByTableAndUid(tableName, entry.getKey());
                        if(!dataSourceArgsMap.containsKey(dataSourceName)) {
                            dataSourceArgsMap.put(dataSourceName, new ArrayList());
                        }
                        dataSourceArgsMap.get(dataSourceName).addAll(entry.getValue());
                    }
                    
                    // 依次选择数据源，调用方法
                    try {
                        for(Map.Entry<String, Collection> entry : dataSourceArgsMap.entrySet()) {
                            DynamicRouteDataSource.setDataSource(StringUtils.isNotEmpty(entry.getKey()) ? entry.getKey() : "default-db");
                            args[index] = entry.getValue();
                            point.proceed(args);
                        }
                    } finally {
                        DynamicRouteDataSource.clearDataSource();
                    }

                    return new Object();
                }
            }

            DynamicRouteDataSource.setDataSource(StringUtils.isNotEmpty(dataSourceName) ? dataSourceName : "default-db");
        }else {
            DynamicRouteDataSource.setDataSource(dataSource.name());
        }

        // 非 从collection中提取UID映射到不同数据的，从此处执行原方法。
        try {
            return point.proceed();
        } finally {
            DynamicRouteDataSource.clearDataSource();
        }
    }


    // 匹配获取参数列表中的某参数
    private Object getArg(Parameter[] parameters, Object[] args, String argName) {
        int index = Integer.MIN_VALUE;
        for(int i = 0; i < parameters.length; i++) {
            if(parameters[i].getName().toLowerCase().contains(argName.toLowerCase())) {
                index = i;
            }
        }

        if(index >= 0) {
            return args[index];
        }
        return null;
    }

    private String mapToDataSourceByTable(String tableName) {
        return tableMapToDataSource.get(tableName);
    }

    private String getTableByMethodName(String name){
        return daoFunctionMapToTable.get(name);
    }

    // 根据表名，UID自动映射到数据源。按需修改
    private String mapToDataSourceByTableAndUid(String tableName, String uid) {
        int h = (tableName + uid).hashCode();
        int hash = h ^ (h >>> 16);
        int size = dynamicRouteDataSource.getDataSourceNameList().size();
        return dynamicRouteDataSource.getDataSourceNameList().get(hash % size);
    }

四、结束

整体代码还不够优雅，尤其是根据UID、table映射到数据源的方式，仍然需要依赖于dao 层接口的编写，不够智能，字段名称也是以常量形式存在代码中
也可以考虑通过配置的形式设置aop 代码中的常量字段 “table”, "uid"等。

在入参到数据库这一套映射关系中，可以根据需求，灵活的编写。

over