接上篇MyBatis-Plus 源码阅读（四）逻辑删除原理深度剖析，本文我们讲解一下 MP 中主键自动生成的实现原理，并给出几个在 ASSIGN_ID 策略下防止 id 重复的方案。
环境
- 核心依赖：mybatis-plus-spring-boot3-starter 3.5.14
- 基础框架：Spring Boot 3.5.6
- JDK 版本：17
- 开发工具：IntelliJ IDEA
- 数据库：MySQL 8.0.31
前言
在插入数据时，我们通常都需要生成记录的主键。虽然 mybatis 支持通过数据库生成递增主键，但用起来不太方便（需要使用 xml 配置），也不支持雪花算法这样的 id 生成方式。mp 的出现让主键生成变得简单和灵活，我们只需在实体类的主键属性上添加 @TableId 就可以实现各种主键生成策略。
源码解析
public enum IdType {
AUTO(0), // 数据库自增策略
NONE(1), // 不设置主键类型，默认使用全局配置中的主键策略
INPUT(2), // 手动设置主键值
ASSIGN_ID(3), // 雪花算法策略
ASSIGN_UUID(4); // UUID策略
mp 支持上面五种主键生成策略，其中 AUTO 和 ASSIGN_ID 使用最为广泛，也是我们重点分析的对象。
数据库自增策略 AUTO
适用于数据库支持自增主键的情况（如 MySQL 的 AUTO_INCREMENT）。
mp 相关逻辑
此策略实际是由 mybatis 框架原生实现的，mp 只是省略了 xml 配置的编写，在 sql 注入时指定数据库主键生成需要的各种配置。（keyGenerator、keyProperty 和 keyColumn）
sql 注入的代码见
com.baomidou.mybatisplus.core.injector.methods.Insert#injectMappedStatement
@Override
public MappedStatement injectMappedStatement(Class<?> mapperClass, Class<?> modelClass, TableInfo tableInfo) {
// ...
if (StringUtils.isNotBlank(tableInfo.getKeyProperty())) {
// AUTO 策略
if (tableInfo.getIdType() == IdType.AUTO) {
// 使用 mybatis 原生的 Jdbc3KeyGenerator
keyGenerator = Jdbc3KeyGenerator.INSTANCE;
// 对应实体的主键属性
keyProperty = tableInfo.getKeyProperty();
keyColumn = SqlInjectionUtils.removeEscapeCharacter(tableInfo.getKeyColumn());
} else if (null != tableInfo.getKeySequence()) {
//...
}
}
// ...
// 将数据库主键生成的各种配置存放在 mappedStatement 里
return this.addInsertMappedStatement(mapperClass, modelClass, methodName, sqlSource, keyGenerator, keyProperty, keyColumn);
}
mybatis 原始逻辑
下面我们回顾一下 mybatis 是如何让数据库生成主键和回填主键值的。
主要分为两步：
- 执行 statement 时指定数据库生成主键的列名。
- statement 执行结束后，将返回值的主键值回填在 mapper 方法的参数上。（一般是回填在插入实体的主键属性上）
第一步：指定主键列名
在创建 preparedStatement 时会调用 org.apache.ibatis.executor.statement.PreparedStatementHandler#instantiateStatement
@Override
protected Statement instantiateStatement(Connection connection) throws SQLException {
String sql = boundSql.getSql();
if (mappedStatement.getKeyGenerator() instanceof Jdbc3KeyGenerator) {
String[] keyColumnNames = mappedStatement.getKeyColumns();
if (keyColumnNames == null) {
return connection.prepareStatement(sql, Statement.RETURN_GENERATED_KEYS);
} else {
// 一般会走到这一步，通过 jdbc 的原生方法指定主键列名
return connection.prepareStatement(sql, keyColumnNames);
}
}
// ...
}
可以注意到 mappedStatement.getKeyGenerator()和mappedStatement.getKeyColumns() 都是在 mp 的 sql 注入阶段赋值的。
第二步：回填主键值
入口在 org.apache.ibatis.executor.statement.PreparedStatementHandler#update
@Override
public int update(Statement statement) throws SQLException {
PreparedStatement ps = (PreparedStatement) statement;
ps.execute();
int rows = ps.getUpdateCount();
Object parameterObject = boundSql.getParameterObject();
KeyGenerator keyGenerator = mappedStatement.getKeyGenerator();
// statement 执行后回填主键值
keyGenerator.processAfter(executor, mappedStatement, ps, parameterObject);
return rows;
}
keyGenerator 是一个接口，实际执行的是 org.apache.ibatis.executor.keygen.Jdbc3KeyGenerator#processAfter。这个方法的调用链有点深，
最后调用的是 org.apache.ibatis.executor.keygen.Jdbc3KeyGenerator.KeyAssigner#assign 完成主键值回填。
protected void assign(ResultSet rs, Object param) {
// ...
MetaObject metaParam = configuration.newMetaObject(param);
try {
// ...
if (typeHandler == null) {
// Error?
} else {
// 关键代码
Object value = typeHandler.getResult(rs, columnPosition);
metaParam.setValue(propertyName, value);
}
} catch (SQLException e) {
throw new ExecutorException("Error getting generated key or setting result to parameter object. Cause: " + e,
e);
}
}
在 11 行完成了对原始参数主键属性的赋值。
雪花算法策略 ASSIGN_ID
mp 支持通过雪花算法生成主键。
为了实现这个功能，需要修改 mybatis 原本 DefaultParameterHandler 的一些逻辑，为此，mp 实现了一个 MybatisParameterHandler，继承原本的 DefaultParameterHandler 类。
在 MybatisParameterHandler 的构造方法，会调用 process 方法，完成主键的生成并赋值到实体参数的主键属性上。
private void process(Object parameter) {
if (parameter != null) {
// ...
if (tableInfo != null) {
MetaObject metaObject = this.configuration.newMetaObject(entity);
if (SqlCommandType.INSERT == this.sqlCommandType) {
// 核心逻辑：生成主键并赋值
populateKeys(tableInfo, metaObject, entity);
insertFill(metaObject, tableInfo);
} else {
updateFill(metaObject, tableInfo);
}
}
}
}
下面我们看看核心逻辑的代码：com.baomidou.mybatisplus.core.MybatisParameterHandler#populateKeys
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) {
// 默认获取到的是 DefaultIdentifierGenerator，里面持有 Sequence 来实现雪花算法
final IdentifierGenerator identifierGenerator = GlobalConfigUtils.getGlobalConfig(this.configuration).getIdentifierGenerator();
Object idValue = metaObject.getValue(keyProperty);
if (identifierGenerator.assignId(idValue)) {
if (idType.getKey() == IdType.ASSIGN_ID.getKey()) {
// 生成雪花算法的 id
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());
}
}
}
}
}
ASSIGN_ID 策略下防止 id 重复的方案
手动指定机器id
我们都知道，雪花算法想要真正实现分布式环境下的唯一性，需要让每台机器拥有唯一的 id，即唯一的 datacenterId 和 workerId 的组合。（最多支持 1024 个）
为了确保机器 id 的唯一性，我们可以在每台机器上设置 WORKER_ID 和 DATACENTER_ID 这样的环境变量，确保这两个变量值的组合在所有的机器中是唯一的，然后通过 yml 配置注入进 mp 的全局配置里，这样就不可能出现重复的 id 了。
mybatis-plus:
global-config:
sequence:
# 从环境变量读取下面的两个值
worker-id: ${WORKER_ID:0}
datacenter-id: ${DATACENTER_ID:0}
自定义 id 生成器
ASSIGN_ID 策略虽然默认使用雪花算法实现，但 mp 允许我们自定义 id 生成器，只需要将一个实现 IdentifierGenerator 接口的对象注册为一个 Bean，mp 自动配置类在启动时会将我们自定义的 id 生成器注入到全局配置中。
自动配置的代码见com.baomidou.mybatisplus.autoconfigure.MybatisPlusAutoConfiguration#sqlSessionFactory
@Bean
@ConditionalOnMissingBean
public SqlSessionFactory sqlSessionFactory(DataSource dataSource) throws Exception {
// ...
GlobalConfig globalConfig = this.properties.getGlobalConfig();
// ...
// 将实现 IdentifierGenerator 的 Bean 对象注入到全局配置
this.getBeanThen(IdentifierGenerator.class, globalConfig::setIdentifierGenerator);
factory.setGlobalConfig(globalConfig);
return factory.getObject();
}
例如，在公司，通常会有一个专门的分布式 id 生成系统，我们可以通过请求这个系统来获取实体的主键。 示例代码如下：
@Configuration
public class MybatisPlusConfiguration {
/**
* 自定义 id 生成器
* @return 实现 IdentifierGenerator 接口的对象
*/
@Bean
public IdentifierGenerator customIdentifierGenerator() {
return entity -> {
String tableName = entity.getClass().getSimpleName();
return getIdFromExternal(tableName);
};
}
/**
* 模拟从外部系统获取分布式 id
* @param namespace 命名空间（不同的表或实体类可以有不同的命名空间）
* @return 外部系统生成的 id
*/
private Long getIdFromExternal(String namespace) {
// 实际上需要根据 namespace 从外部系统获取 id
return 100L;
}
}
结语
mp 的主键自动生成是一个简单且实用的功能，但如果不了解其原理，很容易产生例如重复 id 这样的问题。希望通过本文的讲解，读者能够少踩一些坑。
如果这篇文章对你理解「主键自动生成原理」有帮助，麻烦点赞 + 关注支持一下～ 你的认可就是我持续输出源码解析系列的最大动力！❤️❤️❤️🤩🤩🤩🚀🚀🚀
