mybatisplus 自定义主键生成器逻辑

一、自动配置类

一、自动配置是springboot的一个特点，mybatisplus自动配置类是MybatisPlusAutoConfiguration.java，下面一段源码中可以看到如果从IOC容器中可以获取IdentifierGenerator类型的组件则，设置到globalConfig中，否则globalConfig中的IdentifierGenerator是空。如果不设置会使用默认的DefaultIdentifierGenerator.java

        GlobalConfig globalConfig = this.properties.getGlobalConfig();
        this.getBeanThen(MetaObjectHandler.class, globalConfig::setMetaObjectHandler);
        this.getBeanThen(AnnotationHandler.class, globalConfig::setAnnotationHandler);
        this.getBeanThen(PostInitTableInfoHandler.class, globalConfig::setPostInitTableInfoHandler);
        this.getBeansThen(IKeyGenerator.class, i -> globalConfig.getDbConfig().setKeyGenerators(i));
        this.getBeanThen(ISqlInjector.class, globalConfig::setSqlInjector);
        //这里是从IOC容器中获取IdentifierGenerator类型组件放入到GlobalConfig 
        this.getBeanThen(IdentifierGenerator.class, globalConfig::setIdentifierGenerator);
        factory.setGlobalConfig(globalConfig);
        return factory.getObject();
    }

    /**
     * 检查spring容器里是否有对应的bean,有则进行消费
     *
     * @param clazz    class
     * @param consumer 消费
     * @param <T>      泛型
     */
    private <T> void getBeanThen(Class<T> clazz, Consumer<T> consumer) {
        if (this.applicationContext.getBeanNamesForType(clazz, false, false).length > 0) {
            consumer.accept(this.applicationContext.getBean(clazz));
        }
    }

2、继续跟踪代码MybatisSqlSessionFactoryBean执行getObject()方法，
->afterPropertiesSet()->buildSqlSessionFactory()中执行 靠后可以跟踪到 final SqlSessionFactory sqlSessionFactory = this.sqlSessionFactoryBuilder.build(targetConfiguration);
继续看build方法进入MybatisSqlSessionFactoryBuilder中bulid如下：

 public SqlSessionFactory build(Configuration configuration) {
        GlobalConfig globalConfig = GlobalConfigUtils.getGlobalConfig(configuration);

        final IdentifierGenerator identifierGenerator;
        //如果没有设置自定义主键生成器，会生成DefaultIdentifierGenerator放入到GlobalConfig
        if (null == globalConfig.getIdentifierGenerator()) {
            GlobalConfig.Sequence sequence = globalConfig.getSequence();
            if (sequence.getWorkerId() != null && sequence.getDatacenterId() != null) {
                identifierGenerator = new DefaultIdentifierGenerator(sequence.getWorkerId(), sequence.getDatacenterId());
            } else {
                NetUtils.NetProperties netProperties = new NetUtils.NetProperties(sequence.getPreferredNetworks(), sequence.getIgnoredInterfaces());
                InetAddress inetAddress = new NetUtils(netProperties).findFirstNonLoopbackAddress();
                identifierGenerator = new DefaultIdentifierGenerator(inetAddress);
            }
            globalConfig.setIdentifierGenerator(identifierGenerator);
        } else {
            identifierGenerator = globalConfig.getIdentifierGenerator();
        }
        IdWorker.setIdentifierGenerator(identifierGenerator);

        if (globalConfig.isEnableSqlRunner()) {
            new SqlRunnerInjector().inject(configuration);
        }

        SqlSessionFactory sqlSessionFactory = super.build(configuration);

        // 缓存 sqlSessionFactory
        globalConfig.setSqlSessionFactory(sqlSessionFactory);

        return sqlSessionFactory;
    }

如果没有设置自定义主键生成器，会生成DefaultIdentifierGenerator放入到GlobalConfig中，所以不设置会使用DefaultIdentifierGenerator中的主键生成方法。
策略类型与生成器映射

二、自定义主键生成器调用

Mybatisplus在启动，解析mapper.xml时会往mappedStatement中放入insert、delete等方法，可以供开发者直接使用。mybatisplus中把mapper通过jdk代理生成MybatisMapperProxy.java,使用mybatisplus拓展的增删改查方法时，会调用MybatisMapperProxy中的invoke,最终调用MybatisMapperMethod.execute->org.mybatis.spring.SqlSessionTemplate.insert->org.apache.ibatis.session.defaults.DefaultSqlSession.insert跟踪到org.apache.ibatis.executor.SimpleExecutor.doUpdate

  @Override
  public int doUpdate(MappedStatement ms, Object parameter) throws SQLException {
    Statement stmt = null;
    try {
      Configuration configuration = ms.getConfiguration();
      StatementHandler handler = configuration.newStatementHandler(this, ms, parameter, RowBounds.DEFAULT, null, null);
      stmt = prepareStatement(handler, ms.getStatementLog());
      return handler.update(stmt);
    } finally {
      closeStatement(stmt);
    }
  }

在configuration.newStatementHandler(this, ms, parameter, RowBounds.DEFAULT, null, null)，可以看到调用栈

定位到MybatisParameterHandler构造器里面的processParameter(parameter);当执行insert和update时会进入 extractParameters(parameter).forEach(this::process);

    public void processParameter(Object parameter) {
        /* 只处理插入或更新操作 */
        if (parameter != null && !SimpleTypeRegistry.isSimpleType(parameter.getClass())) {
            if (SqlCommandType.INSERT == this.sqlCommandType || SqlCommandType.UPDATE == this.sqlCommandType) {
                extractParameters(parameter).forEach(this::process);
            }
        }
    }```
可以看到每个参数会循环执行process,其中populateKeys就是设置主键的方法。

```java
 private void process(Object parameter) {
        if (parameter != null) {
            TableInfo tableInfo = null;
            Object entity = parameter;
            if (parameter instanceof Map) {
                // 处理单参数使用注解标记的时候，尝试提取et来获取实体参数
                Map<?, ?> map = (Map<?, ?>) parameter;
                Object et = null;
                if(map.containsKey(Constants.ENTITY)){
                    et = map.get(Constants.ENTITY);

                } else if(map.containsKey(Constants.MP_FILL_ET)){
                    et = map.get(Constants.MP_FILL_ET);
                }
                if (et != null) {
                    entity = et;
                    tableInfo = TableInfoHelper.getTableInfo(entity.getClass());
                }
            } else {
                tableInfo = TableInfoHelper.getTableInfo(parameter.getClass());
            }
            if (tableInfo != null) {
                //到这里就应该转换到实体参数对象了,因为填充和ID处理都是针对实体对象处理的,不用传递原参数对象下去.
                MetaObject metaObject = this.configuration.newMetaObject(entity);
                if (SqlCommandType.INSERT == this.sqlCommandType) {
                    populateKeys(tableInfo, metaObject, entity);
                    insertFill(metaObject, tableInfo);
                } else {
                    updateFill(metaObject, tableInfo);
                }
            }
        }
    }

这里还可以看到insertFill和updateFill方法实现了MetaObjectHandler接口，并添加到IOC容器的组件可以在这执行insertFill和updateFill，一般用作统一设置创建时间，创建人或修改时间，修改人的操作。现在重点看下当执行insert的时候populateKeys方法。总体来看就是在给sql参数赋值时，设置参数值。

 protected void populateKeys(TableInfo tableInfo, MetaObject metaObject, Object entity) {
        final IdType idType = tableInfo.getIdType();
        final String keyProperty = tableInfo.getKeyProperty();
        if (StringUtils.isNotBlank(keyProperty) && null != idType && idType.getKey() >= 3) {
            final IdentifierGenerator identifierGenerator = GlobalConfigUtils.getGlobalConfig(this.configuration).getIdentifierGenerator();
            Object idValue = metaObject.getValue(keyProperty);
            if (identifierGenerator.assignId(idValue)) {
                if (idType.getKey() == IdType.ASSIGN_ID.getKey()) {
                    Number id = identifierGenerator.nextId(entity);
                    metaObject.setValue(keyProperty, OgnlOps.convertValue(id, tableInfo.getKeyType()));
                } else if (idType.getKey() == IdType.ASSIGN_UUID.getKey()) {
                    if(String.class.equals(tableInfo.getKeyType())) {
                        metaObject.setValue(keyProperty, identifierGenerator.nextUUID(entity));
                    } else {
                        log.warn("The current ID generation strategy does not support: " + tableInfo.getKeyType());
                    }
                }
            }
        }
    }

可以看到StringUtils.isNotBlank(keyProperty) && null != idType && idType.getKey() >= 3时会从globalConfig中获取主键生成器，如果有自定义的就从这获取，没有的话就是DefaultIdentifierGenerator.java下面是IdType的代码。

public enum IdType {
    AUTO(0),
    NONE(1),
    INPUT(2),
    ASSIGN_ID(3),
    ASSIGN_UUID(4);

    private final int key;

    private IdType(int key) {
        this.key = key;
    }

    public int getKey() {
        return this.key;
    }
}

会判断 实体类中主键@TableId注解中的配置，如何IdType.ASSIGN_ID会调用主键生成器的nextId, IdType.ASSIGN_UUID则调用nextUUID，给需要执行的sql参数赋值。

 if (idType.getKey() == IdType.ASSIGN_ID.getKey()) {
                    Number id = identifierGenerator.nextId(entity);
                    metaObject.setValue(keyProperty, OgnlOps.convertValue(id, tableInfo.getKeyType()));
                } else if (idType.getKey() == IdType.ASSIGN_UUID.getKey()) {
                    if(String.class.equals(tableInfo.getKeyType())) {
                        metaObject.setValue(keyProperty, identifierGenerator.nextUUID(entity));
                    } else {
                        log.warn("The current ID generation strategy does not support: " + tableInfo.getKeyType());
                    }
                }