fixed github issues/3225 主键生成器支持多数据源

mybatis-plus-annotation/src/main/java/com/baomidou/mybatisplus/annotation/KeySequence.java

@@ -34,4 +34,9 @@ public @interface KeySequence {
      * 序列名
      */
     String value() default "";
+
+    /**
+     * 数据库类型，未配置默认使用注入 IKeyGenerator 实现，多个实现必须指定
+     */
+    DbType dbType() default DbType.OTHER;
 }

mybatis-plus-core/src/main/java/com/baomidou/mybatisplus/core/config/GlobalConfig.java

@@ -28,6 +28,7 @@ import lombok.experimental.Accessors;
 import org.apache.ibatis.session.SqlSessionFactory;
 
 import java.io.Serializable;
+import java.util.List;
 import java.util.Set;
 import java.util.concurrent.ConcurrentSkipListSet;
 
@@ -159,7 +160,7 @@ public class GlobalConfig implements Serializable {
         /**
          * 表主键生成器
          */
-        private IKeyGenerator keyGenerator;
+        private List<IKeyGenerator> keyGenerators;
         /**
          * 逻辑删除全局属性名
          */

mybatis-plus-core/src/main/java/com/baomidou/mybatisplus/core/incrementer/IKeyGenerator.java

@@ -16,6 +16,7 @@
 package com.baomidou.mybatisplus.core.incrementer;
 
 
+import com.baomidou.mybatisplus.annotation.DbType;
 import com.baomidou.mybatisplus.annotation.KeySequence;
 
 /**
@@ -33,4 +34,9 @@ public interface IKeyGenerator {
      * @return sql
      */
     String executeSql(String incrementerName);
+
+    /**
+     * 数据库类型
+     */
+    DbType dbType();
 }

mybatis-plus-core/src/main/java/com/baomidou/mybatisplus/core/metadata/TableInfoHelper.java

@@ -232,7 +232,7 @@ public class TableInfoHelper {
         tableInfo.setTableName(targetTableName);
 
         /* 开启了自定义 KEY 生成器 */
-        if (null != dbConfig.getKeyGenerator()) {
+        if (CollectionUtils.isNotEmpty(dbConfig.getKeyGenerators())) {
             tableInfo.setKeySequence(clazz.getAnnotation(KeySequence.class));
         }
         return excludeProperty;
@@ -313,12 +313,12 @@ public class TableInfoHelper {
 
             /* 有 @TableField 注解的字段初始化 */
             if (tableField != null) {
-                fieldList.add(new TableFieldInfo(dbConfig, tableInfo, field, tableField, reflector, existTableLogic,existOrderBy));
+                fieldList.add(new TableFieldInfo(dbConfig, tableInfo, field, tableField, reflector, existTableLogic, existOrderBy));
                 continue;
             }
 
             /* 无 @TableField  注解的字段初始化 */
-            fieldList.add(new TableFieldInfo(dbConfig, tableInfo, field, reflector, existTableLogic,existOrderBy));
+            fieldList.add(new TableFieldInfo(dbConfig, tableInfo, field, reflector, existTableLogic, existOrderBy));
         }
 
         /* 字段列表 */
@@ -362,7 +362,7 @@ public class TableInfoHelper {
      * @param list 字段列表
      * @return true 为存在 @TableId 注解;
      */
-    public static boolean isExistOrderBy(List<Field> list){
+    public static boolean isExistOrderBy(List<Field> list) {
         return list.stream().anyMatch(field -> field.isAnnotationPresent(OrderBy.class));
     }
 
@@ -498,10 +498,22 @@ public class TableInfoHelper {
     }
 
     public static KeyGenerator genKeyGenerator(String baseStatementId, TableInfo tableInfo, MapperBuilderAssistant builderAssistant) {
-        IKeyGenerator keyGenerator = GlobalConfigUtils.getKeyGenerator(builderAssistant.getConfiguration());
-        if (null == keyGenerator) {
+        List<IKeyGenerator> keyGenerators = GlobalConfigUtils.getKeyGenerators(builderAssistant.getConfiguration());
+        if (CollectionUtils.isEmpty(keyGenerators)) {
             throw new IllegalArgumentException("not configure IKeyGenerator implementation class.");
         }
+        IKeyGenerator keyGenerator = null;
+        if (keyGenerators.size() > 1) {
+            // 多个主键生成器
+            KeySequence keySequence = tableInfo.getKeySequence();
+            if (null != keySequence && DbType.OTHER != keySequence.dbType()) {
+                keyGenerator = keyGenerators.stream().filter(k -> k.dbType() == keySequence.dbType()).findFirst().get();
+            }
+        }
+        // 无法找到注解指定生成器，默认使用第一个生成器
+        if (null == keyGenerator) {
+            keyGenerator = keyGenerators.get(0);
+        }
         Configuration configuration = builderAssistant.getConfiguration();
         String id = builderAssistant.getCurrentNamespace() + StringPool.DOT + baseStatementId + SelectKeyGenerator.SELECT_KEY_SUFFIX;
         ResultMap resultMap = new ResultMap.Builder(builderAssistant.getConfiguration(), id, tableInfo.getKeyType(), new ArrayList<>()).build();

mybatis-plus-core/src/main/java/com/baomidou/mybatisplus/core/toolkit/GlobalConfigUtils.java

@@ -25,6 +25,7 @@ import com.baomidou.mybatisplus.core.metadata.TableInfoHelper;
 import org.apache.ibatis.session.Configuration;
 import org.apache.ibatis.session.SqlSessionFactory;
 
+import java.util.List;
 import java.util.Map;
 import java.util.Optional;
 import java.util.Set;
@@ -87,8 +88,8 @@ public class GlobalConfigUtils {
         return CollectionUtils.computeIfAbsent(GLOBAL_CONFIG, key, k -> defaults());
     }
 
-    public static IKeyGenerator getKeyGenerator(Configuration configuration) {
-        return getGlobalConfig(configuration).getDbConfig().getKeyGenerator();
+    public static List<IKeyGenerator> getKeyGenerators(Configuration configuration) {
+        return getGlobalConfig(configuration).getDbConfig().getKeyGenerators();
     }
 
     public static IdType getIdType(Configuration configuration) {

mybatis-plus-extension/src/main/java/com/baomidou/mybatisplus/extension/incrementer/DB2KeyGenerator.java

@@ -15,6 +15,7 @@
  */
 package com.baomidou.mybatisplus.extension.incrementer;
 
+import com.baomidou.mybatisplus.annotation.DbType;
 import com.baomidou.mybatisplus.core.incrementer.IKeyGenerator;
 
 /**
@@ -29,4 +30,9 @@ public class DB2KeyGenerator implements IKeyGenerator {
     public String executeSql(String incrementerName) {
         return "values nextval for " + incrementerName;
     }
+
+    @Override
+    public DbType dbType() {
+        return DbType.DB2;
+    }
 }

mybatis-plus-extension/src/main/java/com/baomidou/mybatisplus/extension/incrementer/H2KeyGenerator.java

@@ -15,6 +15,7 @@
  */
 package com.baomidou.mybatisplus.extension.incrementer;
 
+import com.baomidou.mybatisplus.annotation.DbType;
 import com.baomidou.mybatisplus.core.incrementer.IKeyGenerator;
 
 /**
@@ -29,4 +30,9 @@ public class H2KeyGenerator implements IKeyGenerator {
     public String executeSql(String incrementerName) {
         return "select " + incrementerName + ".nextval";
     }
+
+    @Override
+    public DbType dbType() {
+        return DbType.H2;
+    }
 }

mybatis-plus-extension/src/main/java/com/baomidou/mybatisplus/extension/incrementer/KingbaseKeyGenerator.java

@@ -15,6 +15,7 @@
  */
 package com.baomidou.mybatisplus.extension.incrementer;
 
+import com.baomidou.mybatisplus.annotation.DbType;
 import com.baomidou.mybatisplus.core.incrementer.IKeyGenerator;
 
 /**
@@ -29,4 +30,9 @@ public class KingbaseKeyGenerator implements IKeyGenerator {
     public String executeSql(String incrementerName) {
         return "select nextval('" + incrementerName + "')";
     }
+
+    @Override
+    public DbType dbType() {
+        return DbType.KINGBASE_ES;
+    }
 }

mybatis-plus-extension/src/main/java/com/baomidou/mybatisplus/extension/incrementer/OracleKeyGenerator.java

@@ -15,6 +15,7 @@
  */
 package com.baomidou.mybatisplus.extension.incrementer;
 
+import com.baomidou.mybatisplus.annotation.DbType;
 import com.baomidou.mybatisplus.core.incrementer.IKeyGenerator;
 
 /**
@@ -29,4 +30,9 @@ public class OracleKeyGenerator implements IKeyGenerator {
     public String executeSql(String incrementerName) {
         return "SELECT " + incrementerName + ".NEXTVAL FROM DUAL";
     }
+
+    @Override
+    public DbType dbType() {
+        return DbType.ORACLE;
+    }
 }

mybatis-plus-extension/src/main/java/com/baomidou/mybatisplus/extension/incrementer/PostgreKeyGenerator.java

@@ -15,6 +15,7 @@
  */
 package com.baomidou.mybatisplus.extension.incrementer;
 
+import com.baomidou.mybatisplus.annotation.DbType;
 import com.baomidou.mybatisplus.core.incrementer.IKeyGenerator;
 
 /**
@@ -29,4 +30,9 @@ public class PostgreKeyGenerator implements IKeyGenerator {
     public String executeSql(String incrementerName) {
         return "select nextval('" + incrementerName + "')";
     }
+
+    @Override
+    public DbType dbType() {
+        return DbType.POSTGRE_SQL;
+    }
 }

mybatis-plus/src/test/java/com/baomidou/mybatisplus/test/h2/keygenerator/KeyGeneratorConfig.java

@@ -15,8 +15,10 @@
  */
 package com.baomidou.mybatisplus.test.h2.keygenerator;
 
+import com.baomidou.mybatisplus.annotation.DbType;
 import com.baomidou.mybatisplus.core.MybatisConfiguration;
 import com.baomidou.mybatisplus.core.config.GlobalConfig;
+import com.baomidou.mybatisplus.core.incrementer.IKeyGenerator;
 import com.baomidou.mybatisplus.extension.incrementer.H2KeyGenerator;
 import com.baomidou.mybatisplus.extension.plugins.MybatisPlusInterceptor;
 import com.baomidou.mybatisplus.extension.plugins.inner.BlockAttackInnerInterceptor;
@@ -31,6 +33,7 @@ import org.springframework.context.annotation.Bean;
 import org.springframework.context.annotation.Configuration;
 
 import javax.sql.DataSource;
+import java.util.Arrays;
 
 @Configuration
 @MapperScan("com.baomidou.mybatisplus.test.h2.keygenerator.mapper")
@@ -59,7 +62,21 @@ public class KeyGeneratorConfig {
     @Bean
     public GlobalConfig globalConfiguration() {
         GlobalConfig conf = new GlobalConfig();
-        conf.setDbConfig(new GlobalConfig.DbConfig().setKeyGenerator(new H2KeyGenerator()));
+        conf.setDbConfig(new GlobalConfig.DbConfig().setKeyGenerators(Arrays.asList(
+            new H2KeyGenerator(),
+            new IKeyGenerator() {
+
+                @Override
+                public String executeSql(String incrementerName) {
+                    return "select " + incrementerName + ".nextval";
+                }
+
+                @Override
+                public DbType dbType() {
+                    return DbType.POSTGRE_SQL;
+                }
+            }
+        )));
         return conf;
     }
 }

mybatis-plus/src/test/java/com/baomidou/mybatisplus/test/h2/keygenerator/model/LongKeyGeneratorModel.java

@@ -15,15 +15,12 @@
  */
 package com.baomidou.mybatisplus.test.h2.keygenerator.model;
 
-import com.baomidou.mybatisplus.annotation.IdType;
-import com.baomidou.mybatisplus.annotation.KeySequence;
-import com.baomidou.mybatisplus.annotation.TableId;
-import com.baomidou.mybatisplus.annotation.TableName;
+import com.baomidou.mybatisplus.annotation.*;
 import lombok.Data;
 
 @Data
 @TableName(value = "key_generator_model")
-@KeySequence(value = "key_generator_model_seq")
+@KeySequence(value = "key_generator_model_seq", dbType = DbType.POSTGRE_SQL)
 public class LongKeyGeneratorModel {
 
     @TableId(type = IdType.INPUT)