Mybatis Plus 3.x 增加 updateAllColumnById 扩展实现

问题背景

项目中有一个导入 Excel 文档的功能，读取 Excel 中的数据并合并到数据库中。如果是旧的数据就进行更新操作。但是导入的数据中，有一些字段是 null 值，而项目中使用的 ORM 框架是 MyBatis Plus，MP 的默认更新策略会跳过 null 值的更新。项目中导入 Excel 的功能要求是遇到需要更新的数据时，对数据库中的旧数据进行全量覆盖更新，那么就需要对这个功能配置单独的更新策略。

现在想来可以有三种方式进行实现：

1. 通过自定义表字段和值的 Map 映射编写 XML
2. 使用 @TableField 注解定义每个字段的更新策略
3. 定义自己的通用方法 updateAllColumnById ，继承实现类 DefaultSqlInjector ，然后继承 BaseMapper 或者 ServiceImpl 添加自定义方法，全局配置 sqlInjector 注入

我用了第 3 种方式来实现。

官方文档：SQL注入器

解决方案

实现 AbstractMethod：

import com.baomidou.mybatisplus.core.enums.SqlMethod;
import com.baomidou.mybatisplus.core.injector.AbstractMethod;
import com.baomidou.mybatisplus.core.metadata.TableFieldInfo;
import com.baomidou.mybatisplus.core.metadata.TableInfo;
import com.baomidou.mybatisplus.core.toolkit.sql.SqlScriptUtils;
import org.apache.ibatis.mapping.MappedStatement;
import org.apache.ibatis.mapping.SqlSource;

import java.util.function.Function;
import java.util.function.Predicate;

public class UpdateAllColumnById extends AbstractMethod {

    @Override
    public MappedStatement injectMappedStatement(Class<?> mapperClass, Class<?> modelClass, TableInfo tableInfo) {
        // sqlMethod 枚举中预设了一些 sql 语句模板，使用 UPDATE_BY_ID 预设的 sql 模板
        SqlMethod sqlMethod = SqlMethod.UPDATE_BY_ID;
        String additional = this.optlockVersion() + tableInfo.getLogicDeleteSql(true, false);
        String sqlSet = this.filterTableFieldInfo(tableInfo.getFieldList(), this.getPredicate(), getFunction(), "\n");
        sqlSet = SqlScriptUtils.convertSet(sqlSet);
        sqlSet = SqlScriptUtils.convertIf(sqlSet, String.format("%s != null", "et"), true);
        String sql = String.format(sqlMethod.getSql(), tableInfo.getTableName(), sqlSet, tableInfo.getKeyColumn(), "et." + tableInfo.getKeyProperty(), additional);
        SqlSource sqlSource = this.languageDriver.createSqlSource(this.configuration, sql, modelClass);
        // 定义自定义语句的 id，这里使用 UPDATE_ALL_COLUMN_BY_ID 作为注册的 id
        sqlMethod = SqlMethod.UPDATE_ALL_COLUMN_BY_ID;
        return this.addUpdateMappedStatement(mapperClass, modelClass, sqlMethod.getMethod(), sqlSource);
    }

    private Predicate<TableFieldInfo> getPredicate() {
        return (t) -> !t.isLogicDelete();
    }

    private Function<TableFieldInfo, String> getFunction() {
        return (i) -> i.getSqlSet(true, "et.");
    }
}

须要注意的是，代码所使用的 SqlMethod.UPDATE_ALL_COLUMN_BY_ID 是 MP 遗留的枚举（MP 从 2.x 迭代到 3.x 之后，废弃了 updateAllColumnById 的实现，但是 SqlMethod 中还遗留着 UPDATE_ALL_COLUMN_BY_ID 的枚举。如果所用的 MP 版本中不存在这个枚举，自定义一个就行。）

public enum MySqlMethod {
    UPDATE_ALL_COLUMN_BY_ID("updateAllColumnById", "根据ID 选择修改数据", "<script>\nUPDATE %s %s WHERE %s=#{%s} %s\n</script>");

    private final String method;
    private final String desc;
    private final String sql;
    
    private SqlMethod(String method, String desc, String sql) {
        this.method = method;
        this.desc = desc;
        this.sql = sql;
    }

    public String getMethod() {
        return this.method;
    }

    public String getDesc() {
        return this.desc;
    }

    public String getSql() {
        return this.sql;
    }
}

拓展 DefaultSqlInjector：

import com.baomidou.mybatisplus.core.injector.AbstractMethod;
import com.baomidou.mybatisplus.core.injector.DefaultSqlInjector;

import java.util.List;

public class MySqlInjector extends DefaultSqlInjector {

    /**
     * 如果只需增加方法，保留 MP 自带方法
     * 可以 super.getMethodList() 再 add
     *
     * @return
     */
    @Override
    public List<AbstractMethod> getMethodList(Class<?> mapperClass) {
        List<AbstractMethod> methodList = super.getMethodList(mapperClass);
        methodList.add(new UpdateAllColumnById());
        return methodList;
    }
}

注入 ISqlInjector Bean 到容器中:

import com.baomidou.mybatisplus.core.injector.ISqlInjector;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;

@Configuration
public class MybatisPlusConfig {

    @Bean
    public ISqlInjector iSqlInjector() {
        return new MySqlInjector();
    }
}

拓展 ServiceImpl：

import com.baomidou.mybatisplus.core.enums.SqlMethod;
import com.baomidou.mybatisplus.core.mapper.BaseMapper;
import com.baomidou.mybatisplus.core.toolkit.Assert;
import com.baomidou.mybatisplus.extension.service.impl.ServiceImpl;
import org.apache.ibatis.binding.MapperMethod;
import org.apache.ibatis.session.SqlSession;

import java.util.Collection;
import java.util.Iterator;

public class MyServiceImpl<M extends BaseMapper<T>, T> extends ServiceImpl<M, T> {

    public boolean updateAllColumnBatchById(Collection<T> entityList) {
        return this.updateAllColumnBatchById(entityList, 1000);
    }

    public boolean updateAllColumnBatchById(Collection<T> entityList, int batchSize) {
        Assert.notEmpty(entityList, "error: entityList must not be empty");
        // 这里传入的枚举参数需要和实现 `AbstractMethod` 中传入的自定义语句的 id 所对应枚举一致 
        String sqlStatement = this.sqlStatement(SqlMethod.UPDATE_ALL_COLUMN_BY_ID);
        SqlSession batchSqlSession = this.sqlSessionBatch();
        Throwable var5 = null;

        try {
            int i = 0;
            for (Iterator<T> var7 = entityList.iterator(); var7.hasNext(); ++i) {
                T anEntityList = var7.next();
                MapperMethod.ParamMap<T> param = new MapperMethod.ParamMap<>();
                param.put("et", anEntityList);
                batchSqlSession.update(sqlStatement, param);
                if (i >= 1 && i % batchSize == 0) {
                    batchSqlSession.flushStatements();
                }
            }
            batchSqlSession.flushStatements();
            return true;
        } catch (Throwable var17) {
            var5 = var17;
            throw var17;
        } finally {
            if (batchSqlSession != null) {
                if (var5 != null) {
                    try {
                        batchSqlSession.close();
                    } catch (Throwable var16) {
                        var5.addSuppressed(var16);
                    }
                } else {
                    batchSqlSession.close();
                }
            }
        }
    }
}

然后就可以通过继承 MyServiceImpl ，调用自定义的 UpdateAllColumnById 方法，实现全字段更新。

public class EntityServiceImpl extends MyServiceImpl<EntityMapper, Entity> implements IEntityService {
...
    private void updateAllColumnBatchBussiness(List<Entity> entities) {
        this.updateAllColumnBatchById(entities);
    }
...
}