001/*
002 *  Copyright (c) 2022-2023, Mybatis-Flex (fuhai999@gmail.com).
003 *  <p>
004 *  Licensed under the Apache License, Version 2.0 (the "License");
005 *  you may not use this file except in compliance with the License.
006 *  You may obtain a copy of the License at
007 *  <p>
008 *  http://www.apache.org/licenses/LICENSE-2.0
009 *  <p>
010 *  Unless required by applicable law or agreed to in writing, software
011 *  distributed under the License is distributed on an "AS IS" BASIS,
012 *  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
013 *  See the License for the specific language governing permissions and
014 *  limitations under the License.
015 */
016package com.mybatisflex.core.keygen;
017
018import com.mybatisflex.annotation.KeyType;
019import com.mybatisflex.core.FlexConsts;
020import com.mybatisflex.core.FlexGlobalConfig;
021import com.mybatisflex.core.exception.FlexExceptions;
022import com.mybatisflex.core.table.IdInfo;
023import com.mybatisflex.core.table.TableInfo;
024import com.mybatisflex.core.util.StringUtil;
025import org.apache.ibatis.executor.keygen.Jdbc3KeyGenerator;
026import org.apache.ibatis.executor.keygen.KeyGenerator;
027import org.apache.ibatis.executor.keygen.NoKeyGenerator;
028import org.apache.ibatis.executor.keygen.SelectKeyGenerator;
029import org.apache.ibatis.mapping.*;
030import org.apache.ibatis.session.Configuration;
031
032import java.util.ArrayList;
033import java.util.Arrays;
034import java.util.List;
035
036public class MybatisKeyGeneratorUtil {
037
038    private MybatisKeyGeneratorUtil() {
039    }
040
041    public static KeyGenerator createTableKeyGenerator(TableInfo tableInfo, MappedStatement ms) {
042        List<IdInfo> primaryKeyList = tableInfo.getPrimaryKeyList();
043
044        //无主键
045        if (primaryKeyList == null || primaryKeyList.isEmpty()) {
046            return NoKeyGenerator.INSTANCE;
047        }
048
049        //多主键的
050        if (primaryKeyList.size() > 1) {
051            return new MultiPrimaryKeyGenerator(ms, tableInfo, primaryKeyList);
052        }
053
054        return createIdKeyGenerator(tableInfo, ms, primaryKeyList.get(0));
055    }
056
057
058    public static KeyGenerator createIdKeyGenerator(TableInfo tableInfo, MappedStatement ms, IdInfo idInfo) {
059        FlexGlobalConfig.KeyConfig globalKeyConfig = FlexGlobalConfig.getConfig(ms.getConfiguration()).getKeyConfig();
060        KeyType keyType = getKeyType(idInfo, globalKeyConfig);
061
062        if (keyType == null || keyType == KeyType.None) {
063            return NoKeyGenerator.INSTANCE;
064        }
065
066        //自增主键
067        if (keyType == KeyType.Auto) {
068            return Jdbc3KeyGenerator.INSTANCE;
069        }
070
071        //通过 java 生成的主键
072        if (keyType == KeyType.Generator) {
073            return new CustomKeyGenerator(ms.getConfiguration(), tableInfo, idInfo);
074        }
075
076        //通过序列生成的注解
077        String sequence = getKeyValue(idInfo, globalKeyConfig);
078        if (StringUtil.isBlank(sequence)) {
079            throw FlexExceptions.wrap("Please config sequence by @Id(value=\"...\") for field: %s in class: %s"
080                , idInfo.getProperty()
081                , tableInfo.getEntityClass().getName());
082        }
083
084
085        String selectId = ms.getId() + SelectKeyGenerator.SELECT_KEY_SUFFIX;
086        SqlSource sqlSource = ms.getLang().createSqlSource(ms.getConfiguration(), sequence.trim(), idInfo.getPropertyType());
087        MappedStatement.Builder msBuilder = new MappedStatement.Builder(ms.getConfiguration(), selectId, sqlSource, SqlCommandType.SELECT)
088            .resource(ms.getResource())
089            .fetchSize(null)
090            .timeout(null)
091            .statementType(StatementType.PREPARED)
092            .keyGenerator(NoKeyGenerator.INSTANCE)
093            .keyProperty(FlexConsts.ENTITY + "." + idInfo.getProperty())
094            .keyColumn(idInfo.getColumn())
095            .databaseId(ms.getDatabaseId())
096            .lang(ms.getLang())
097            .resultOrdered(false)
098            .resultSets(null)
099            .resultMaps(createIdResultMaps(ms.getConfiguration(), selectId + "-Inline", idInfo.getPropertyType(), new ArrayList<>()))
100            .resultSetType(null)
101            .flushCacheRequired(false)
102            .useCache(false)
103            .cache(ms.getCache());
104
105        MappedStatement keyMappedStatement = msBuilder.build();
106        ms.getConfiguration().addMappedStatement(keyMappedStatement);
107
108        //看到有的框架把 keyGenerator 添加到 mybatis 的当前配置里去，其实是完全没必要的
109        //因为只有在 xml 解析的时候，才可能存在多一个 MappedStatement 拥有同一个 keyGenerator 的情况
110        //当前每个方法都拥有一个自己的 keyGenerator 了，没必要添加
111        //addKeyGenerator(selectId, keyGenerator)
112        return new SelectKeyGenerator(keyMappedStatement, isKeyBefore(idInfo, globalKeyConfig));
113    }
114
115
116    private static List<ResultMap> createIdResultMaps(Configuration configuration,
117                                                      String statementId, Class<?> resultType, List<ResultMapping> resultMappings) {
118        ResultMap resultMap = new ResultMap.Builder(configuration, statementId, resultType, resultMappings, null)
119            .build();
120        return Arrays.asList(resultMap);
121    }
122
123
124    /**
125     * 获取主键的 keyType，优先通过 @id 获取，获取不到通过全局配置获取
126     */
127    public static KeyType getKeyType(IdInfo idInfo, FlexGlobalConfig.KeyConfig globalKeyConfig) {
128        KeyType keyType = idInfo.getKeyType();
129        if (keyType != KeyType.None) {
130            return keyType;
131        }
132
133        if (globalKeyConfig != null) {
134            return globalKeyConfig.getKeyType();
135        }
136
137        return keyType;
138    }
139
140
141    public static String getKeyValue(IdInfo idInfo, FlexGlobalConfig.KeyConfig globalKeyConfig) {
142        String value = idInfo.getValue();
143        if (StringUtil.isBlank(value) && globalKeyConfig != null) {
144            value = globalKeyConfig.getValue();
145        }
146        return value;
147    }
148
149
150    public static boolean isKeyBefore(IdInfo idInfo, FlexGlobalConfig.KeyConfig globalKeyConfig) {
151        Boolean before = idInfo.getBefore();
152        if (before == null && globalKeyConfig != null) {
153            return globalKeyConfig.isBefore();
154        } else {
155            return before == null || before;
156        }
157    }
158
159}