mybatis之KeyGenerator生成主键

讲mybatis主键生成之前，先说一下原生的主键回填吧。

原生写法

原生的需要在创建PreparedStatement时指定主键回填，Statement.RETURN_GENERATED_KEYS

PreparedStatement statement = connection.prepareStatement("insert into adx_material (adx_template_id,type) values (?,?)",
        Statement.RETURN_GENERATED_KEYS)

这样在返回的结果集中就会将主键返回

ResultSet rs = statement.getGeneratedKeys();
long id = rs.getLong(1);

mybatis主键回填

默认情况下，mybatis的insert语句不会返回自动生成的主键，而是返回插入的条数，如果需要获取到产生的自增主键，可以使用KeyGenerator接口

public interface KeyGenerator {
    // 执行insert之前执行，设置属性order="BEFORE"
    void processBefore(Executor executor, MappedStatement ms, Statement stmt, Object parameter);
    // 执行insert之前执行，设置属性order="AFTER"
    void processAfter(Executor executor, MappedStatement ms, Statement stmt, Object parameter);
}

该接口有三个实现类，NoKeyGenerator,Jdbc3KeyGenerator,SelectKeyGenerator

NoKeyGenerator

NoKeyGenerator该类中processBefore和processAfter方法都是空实现

Jdbc3KeyGenerator

Jdbc3KeyGenerator用于取回数据库生成的自增id，对应于mybatis-config.xml配置文件的useGenerateKeys全局配置以及映射配置文件中SQL节点中的useGenerateKeys属性。

Jdbc3KeyGenerator只实现了processAfter，适用于由数据库层面自动生成id的情况
需要配置useGeneratedKeys属性
mapper.xml示例

<insert id="insertUser" useGeneratedKeys="true" keyProperty="id">
    insert into user (user_name,sex,age) values 
    <foreach item="item" collection="list" separator=",">
        (#{item.userName},#{item.sex},#{item.age})
    </foreach>
</insert>

public class Jdbc3KeyGenerator implements KeyGenerator {
    private static final String SECOND_GENERIC_PARAM_NAME = "param2";
    public static final Jdbc3KeyGenerator INSTANCE = new Jdbc3KeyGenerator();
    private static final String MSG_TOO_MANY_KEYS = "Too many keys are generated. There are only %d target objects. You either specified a wrong 'keyProperty' or encountered a driver bug like #1523.";

    public Jdbc3KeyGenerator() {
    }

    public void processBefore(Executor executor, MappedStatement ms, Statement stmt, Object parameter) {
    }

    public void processAfter(Executor executor, MappedStatement ms, Statement stmt, Object parameter) {
        this.processBatch(ms, stmt, parameter);
    }

    public void processBatch(MappedStatement ms, Statement stmt, Object parameter) {
          // 获得keyProperties属性指定的属性名称，表示主键对应的属性名称
        String[] keyProperties = ms.getKeyProperties();
        if (keyProperties != null && keyProperties.length != 0) {
            try {
                  // 获取数据库自动生成的主键，如果没有生成主键，返回结果集为空
                ResultSet rs = stmt.getGeneratedKeys();
                Throwable var6 = null;

                try {
                      // 获取ResultSet的元数据
                    ResultSetMetaData rsmd = rs.getMetaData();
                    Configuration configuration = ms.getConfiguration();
                      // 检测数据库生成的主键列数和keyProperties的列数是否匹配
                    if (rsmd.getColumnCount() >= keyProperties.length) {
                        this.assignKeys(configuration, rs, rsmd, keyProperties, parameter);
                    }
                } catch (Throwable var17) {
                    var6 = var17;
                    throw var17;
                } finally {
                    if (rs != null) {
                        if (var6 != null) {
                            try {
                                rs.close();
                            } catch (Throwable var16) {
                                var6.addSuppressed(var16);
                            }
                        } else {
                            rs.close();
                        }
                    }

                }

            } catch (Exception var19) {
                throw new ExecutorException("Error getting generated key or setting result to parameter object. Cause: " + var19, var19);
            }
        }
    }

    private void assignKeys(Configuration configuration, ResultSet rs, ResultSetMetaData rsmd, String[] keyProperties, Object parameter) throws SQLException {
        if (!(parameter instanceof ParamMap) && !(parameter instanceof StrictMap)) {
            if (parameter instanceof ArrayList && !((ArrayList)parameter).isEmpty() && ((ArrayList)parameter).get(0) instanceof ParamMap) {
                this.assignKeysToParamMapList(configuration, rs, rsmd, keyProperties, (ArrayList)parameter);
            } else {
                this.assignKeysToParam(configuration, rs, rsmd, keyProperties, parameter);
            }
        } else {
            this.assignKeysToParamMap(configuration, rs, rsmd, keyProperties, (Map)parameter);
        }

    }

    private void assignKeysToParam(Configuration configuration, ResultSet rs, ResultSetMetaData rsmd, String[] keyProperties, Object parameter) throws SQLException {
        Collection<?> params = collectionize(parameter);
        if (!params.isEmpty()) {
            List<Jdbc3KeyGenerator.KeyAssigner> assignerList = new ArrayList();

            for(int i = 0; i < keyProperties.length; ++i) {
                assignerList.add(new Jdbc3KeyGenerator.KeyAssigner(configuration, rsmd, i + 1, (String)null, keyProperties[i]));
            }

            Iterator iterator = params.iterator();

            while(rs.next()) {
                if (!iterator.hasNext()) {
                    throw new ExecutorException(String.format("Too many keys are generated. There are only %d target objects. You either specified a wrong 'keyProperty' or encountered a driver bug like #1523.", params.size()));
                }

                Object param = iterator.next();
                assignerList.forEach((x) -> {
                    x.assign(rs, param);
                });
            }

        }
    }

    private void assignKeysToParamMapList(Configuration configuration, ResultSet rs, ResultSetMetaData rsmd, String[] keyProperties, ArrayList<ParamMap<?>> paramMapList) throws SQLException {
        Iterator<ParamMap<?>> iterator = paramMapList.iterator();
        List<Jdbc3KeyGenerator.KeyAssigner> assignerList = new ArrayList();

        for(long counter = 0L; rs.next(); ++counter) {
            if (!iterator.hasNext()) {
                throw new ExecutorException(String.format("Too many keys are generated. There are only %d target objects. You either specified a wrong 'keyProperty' or encountered a driver bug like #1523.", counter));
            }

            ParamMap<?> paramMap = (ParamMap)iterator.next();
            if (assignerList.isEmpty()) {
                for(int i = 0; i < keyProperties.length; ++i) {
                    assignerList.add(this.getAssignerForParamMap(configuration, rsmd, i + 1, paramMap, keyProperties[i], keyProperties, false).getValue());
                }
            }

            assignerList.forEach((x) -> {
                x.assign(rs, paramMap);
            });
        }

    }

    private void assignKeysToParamMap(Configuration configuration, ResultSet rs, ResultSetMetaData rsmd, String[] keyProperties, Map<String, ?> paramMap) throws SQLException {
        if (!paramMap.isEmpty()) {
            Map<String, Entry<Iterator<?>, List<Jdbc3KeyGenerator.KeyAssigner>>> assignerMap = new HashMap();

            for(int i = 0; i < keyProperties.length; ++i) {
                Entry<String, Jdbc3KeyGenerator.KeyAssigner> entry = this.getAssignerForParamMap(configuration, rsmd, i + 1, paramMap, keyProperties[i], keyProperties, true);
                Entry<Iterator<?>, List<Jdbc3KeyGenerator.KeyAssigner>> iteratorPair = (Entry)assignerMap.computeIfAbsent(entry.getKey(), (k) -> {
                    return entry(collectionize(paramMap.get(k)).iterator(), new ArrayList());
                });
                ((List)iteratorPair.getValue()).add(entry.getValue());
            }

            for(long counter = 0L; rs.next(); ++counter) {
                Iterator var13 = assignerMap.values().iterator();

                while(var13.hasNext()) {
                    Entry<Iterator<?>, List<Jdbc3KeyGenerator.KeyAssigner>> pair = (Entry)var13.next();
                    if (!((Iterator)pair.getKey()).hasNext()) {
                        throw new ExecutorException(String.format("Too many keys are generated. There are only %d target objects. You either specified a wrong 'keyProperty' or encountered a driver bug like #1523.", counter));
                    }

                    Object param = ((Iterator)pair.getKey()).next();
                    ((List)pair.getValue()).forEach((x) -> {
                        x.assign(rs, param);
                    });
                }
            }

        }
    }

    private Entry<String, Jdbc3KeyGenerator.KeyAssigner> getAssignerForParamMap(Configuration config, ResultSetMetaData rsmd, int columnPosition, Map<String, ?> paramMap, String keyProperty, String[] keyProperties, boolean omitParamName) {
        Set<String> keySet = paramMap.keySet();
        boolean singleParam = !keySet.contains("param2");
        int firstDot = keyProperty.indexOf(46);
        if (firstDot == -1) {
            if (singleParam) {
                return this.getAssignerForSingleParam(config, rsmd, columnPosition, paramMap, keyProperty, omitParamName);
            } else {
                throw new ExecutorException("Could not determine which parameter to assign generated keys to. Note that when there are multiple parameters, 'keyProperty' must include the parameter name (e.g. 'param.id'). Specified key properties are " + ArrayUtil.toString(keyProperties) + " and available parameters are " + keySet);
            }
        } else {
            String paramName = keyProperty.substring(0, firstDot);
            if (keySet.contains(paramName)) {
                String argParamName = omitParamName ? null : paramName;
                String argKeyProperty = keyProperty.substring(firstDot + 1);
                return entry(paramName, new Jdbc3KeyGenerator.KeyAssigner(config, rsmd, columnPosition, argParamName, argKeyProperty));
            } else if (singleParam) {
                return this.getAssignerForSingleParam(config, rsmd, columnPosition, paramMap, keyProperty, omitParamName);
            } else {
                throw new ExecutorException("Could not find parameter '" + paramName + "'. Note that when there are multiple parameters, 'keyProperty' must include the parameter name (e.g. 'param.id'). Specified key properties are " + ArrayUtil.toString(keyProperties) + " and available parameters are " + keySet);
            }
        }
    }

    private Entry<String, Jdbc3KeyGenerator.KeyAssigner> getAssignerForSingleParam(Configuration config, ResultSetMetaData rsmd, int columnPosition, Map<String, ?> paramMap, String keyProperty, boolean omitParamName) {
        String singleParamName = nameOfSingleParam(paramMap);
        String argParamName = omitParamName ? null : singleParamName;
        return entry(singleParamName, new Jdbc3KeyGenerator.KeyAssigner(config, rsmd, columnPosition, argParamName, keyProperty));
    }

    private static String nameOfSingleParam(Map<String, ?> paramMap) {
        return (String)paramMap.keySet().iterator().next();
    }

    private static Collection<?> collectionize(Object param) {
        if (param instanceof Collection) {
            return (Collection)param;
        } else {
            return param instanceof Object[] ? Arrays.asList((Object[])((Object[])param)) : Arrays.asList(param);
        }
    }

    private static <K, V> Entry<K, V> entry(K key, V value) {
        return new SimpleImmutableEntry(key, value);
    }

    private class KeyAssigner {
        private final Configuration configuration;
        private final ResultSetMetaData rsmd;
        private final TypeHandlerRegistry typeHandlerRegistry;
        private final int columnPosition;
        private final String paramName;
        private final String propertyName;
        private TypeHandler<?> typeHandler;

        protected KeyAssigner(Configuration configuration, ResultSetMetaData rsmd, int columnPosition, String paramName, String propertyName) {
            this.configuration = configuration;
            this.rsmd = rsmd;
            this.typeHandlerRegistry = configuration.getTypeHandlerRegistry();
            this.columnPosition = columnPosition;
            this.paramName = paramName;
            this.propertyName = propertyName;
        }

        protected void assign(ResultSet rs, Object param) {
            if (this.paramName != null) {
                param = ((ParamMap)param).get(this.paramName);
            }
                        // 为用户传入的实参创建相应的MetaObject对象
            MetaObject metaParam = this.configuration.newMetaObject(param);

            try {
                if (this.typeHandler == null) {
                    if (!metaParam.hasSetter(this.propertyName)) {
                        throw new ExecutorException("No setter found for the keyProperty '" + this.propertyName + "' in '" + metaParam.getOriginalObject().getClass().getName() + "'.");
                    }

                    Class<?> propertyType = metaParam.getSetterType(this.propertyName);
                    this.typeHandler = this.typeHandlerRegistry.getTypeHandler(propertyType, JdbcType.forCode(this.rsmd.getColumnType(this.columnPosition)));
                }

                if (this.typeHandler != null) {
                    Object value = this.typeHandler.getResult(rs, this.columnPosition);
                    metaParam.setValue(this.propertyName, value);
                }

            } catch (SQLException var5) {
                throw new ExecutorException("Error getting generated key or setting result to parameter object. Cause: " + var5, var5);
            }
        }
    }
}

SelectKeyGenerator

有些数据库不支持自动生成自增主键，可以使用SelectKeyGenerator来生成主键
执行selectKey节点的sql语句，获取insert语句所要主键

<insert id="insertUser">
      <selectKey keyProperty="id" resultType="int" order="BEFORE">
          select 10000
      </selectKey>
      insert into user (id,user_name) values (#{id},#{userName})
</insert>

SelectKeyGenerator会执行映射配置文件中定义的<selectKey>节点的SQL语句，该语句会获取insert语句所需要的主键。

public class SelectKeyGenerator implements KeyGenerator {
    public static final String SELECT_KEY_SUFFIX = "!selectKey";
      //标识是在insert语句之前执行还是之后执行
    private final boolean executeBefore;
      // <selectKey>节点中定义的SQL语句所对应的MappedStatement对象，解析<selectKey>节点时创建的，用于获取insert语句中使用的主键
    private final MappedStatement keyStatement;

    public SelectKeyGenerator(MappedStatement keyStatement, boolean executeBefore) {
        this.executeBefore = executeBefore;
        this.keyStatement = keyStatement;
    }

    public void processBefore(Executor executor, MappedStatement ms, Statement stmt, Object parameter) {
        if (this.executeBefore) {
            this.processGeneratedKeys(executor, ms, parameter);
        }

    }

    public void processAfter(Executor executor, MappedStatement ms, Statement stmt, Object parameter) {
        if (!this.executeBefore) {
            this.processGeneratedKeys(executor, ms, parameter);
        }

    }

    private void processGeneratedKeys(Executor executor, MappedStatement ms, Object parameter) {
        try {
            if (parameter != null && this.keyStatement != null && this.keyStatement.getKeyProperties() != null) {
                  // 获取selectKey中配置的属性名称，表示主键对应的属性
                String[] keyProperties = this.keyStatement.getKeyProperties();
                Configuration configuration = ms.getConfiguration();
                MetaObject metaParam = configuration.newMetaObject(parameter);
                  // 创建Executor对象，用于执行keyStatement中的sql,并得到主键对象
                Executor keyExecutor = configuration.newExecutor(executor.getTransaction(), ExecutorType.SIMPLE);
                List<Object> values = keyExecutor.query(this.keyStatement, parameter, RowBounds.DEFAULT, Executor.NO_RESULT_HANDLER);
                if (values.size() == 0) {
                    throw new ExecutorException("SelectKey returned no data.");
                }

                if (values.size() > 1) {
                    throw new ExecutorException("SelectKey returned more than one value.");
                }

                MetaObject metaResult = configuration.newMetaObject(values.get(0));
                if (keyProperties.length == 1) {
                    if (metaResult.hasGetter(keyProperties[0])) {
                        this.setValue(metaParam, keyProperties[0], metaResult.getValue(keyProperties[0]));
                    } else {
                        this.setValue(metaParam, keyProperties[0], values.get(0));
                    }
                } else {
                      // 联合主键的情况
                    this.handleMultipleProperties(keyProperties, metaParam, metaResult);
                }
            }

        } catch (ExecutorException var10) {
            throw var10;
        } catch (Exception var11) {
            throw new ExecutorException("Error selecting key or setting result to parameter object. Cause: " + var11, var11);
        }
    }

    private void handleMultipleProperties(String[] keyProperties, MetaObject metaParam, MetaObject metaResult) {
        String[] keyColumns = this.keyStatement.getKeyColumns();
        if (keyColumns != null && keyColumns.length != 0) {
            if (keyColumns.length != keyProperties.length) {
                throw new ExecutorException("If SelectKey has key columns, the number must match the number of key properties.");
            }

            for(int i = 0; i < keyProperties.length; ++i) {
                this.setValue(metaParam, keyProperties[i], metaResult.getValue(keyColumns[i]));
            }
        } else {
            String[] var5 = keyProperties;
            int var6 = keyProperties.length;

            for(int var7 = 0; var7 < var6; ++var7) {
                String keyProperty = var5[var7];
                this.setValue(metaParam, keyProperty, metaResult.getValue(keyProperty));
            }
        }

    }

    private void setValue(MetaObject metaParam, String property, Object value) {
        if (metaParam.hasSetter(property)) {
            metaParam.setValue(property, value);
        } else {
            throw new ExecutorException("No setter found for the keyProperty '" + property + "' in " + metaParam.getOriginalObject().getClass().getName() + ".");
        }
    }
}