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.row.Row;
021import com.mybatisflex.core.row.RowCPI;
022import com.mybatisflex.core.row.RowKey;
023import com.mybatisflex.core.util.ArrayUtil;
024import com.mybatisflex.core.util.CollectionUtil;
025import org.apache.ibatis.executor.Executor;
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.MappedStatement;
030import org.apache.ibatis.mapping.SqlCommandType;
031import org.apache.ibatis.mapping.SqlSource;
032import org.apache.ibatis.mapping.StatementType;
033
034import java.sql.Statement;
035import java.util.ArrayList;
036import java.util.List;
037import java.util.Map;
038
039/**
040 * 为 row 的主键生成器
041 */
042public class RowKeyGenerator implements KeyGenerator, IMultiKeyGenerator {
043    private static KeyGenerator[] NO_KEY_GENERATORS = new KeyGenerator[0];
044
045    private MappedStatement ms;
046    private KeyGenerator[] keyGenerators;
047    private List<String> autoKeyGeneratorNames;
048
049    public RowKeyGenerator(MappedStatement methodMappedStatement) {
050        this.ms = methodMappedStatement;
051    }
052
053    @Override
054    public void processBefore(Executor executor, MappedStatement ms, Statement stmt, Object parameter) {
055        Row row = (Row) ((Map) parameter).get(FlexConsts.ROW);
056        keyGenerators = buildRowKeyGenerators(RowCPI.obtainsPrimaryKeys(row));
057        for (KeyGenerator keyGenerator : keyGenerators) {
058            keyGenerator.processBefore(executor, ms, stmt, parameter);
059        }
060    }
061
062
063    @Override
064    public void processAfter(Executor executor, MappedStatement ms, Statement stmt, Object parameter) {
065        for (KeyGenerator keyGenerator : keyGenerators) {
066            keyGenerator.processAfter(executor, ms, stmt, parameter);
067        }
068    }
069
070
071    private KeyGenerator[] buildRowKeyGenerators(RowKey[] rowKeys) {
072        if (ArrayUtil.isEmpty(rowKeys)) {
073            return NO_KEY_GENERATORS;
074        }
075
076        KeyGenerator[] keyGenerators = new KeyGenerator[rowKeys.length];
077        for (int i = 0; i < rowKeys.length; i++) {
078            KeyGenerator keyGenerator = createByRowKey(rowKeys[i]);
079            keyGenerators[i] = keyGenerator;
080        }
081        return keyGenerators;
082    }
083
084
085    private KeyGenerator createByRowKey(RowKey rowKey) {
086        if (rowKey == null || rowKey.getKeyType() == KeyType.None) {
087            return NoKeyGenerator.INSTANCE;
088        }
089
090        if (rowKey.getKeyType() == KeyType.Auto) {
091            if (autoKeyGeneratorNames == null) {
092                autoKeyGeneratorNames = new ArrayList<>();
093            }
094            autoKeyGeneratorNames.add(rowKey.getKeyColumn());
095            return new RowJdbc3KeyGenerator(rowKey.getKeyColumn());
096        }
097
098        if (rowKey.getKeyType() == KeyType.Generator) {
099            return new RowCustomKeyGenerator(rowKey);
100        }
101        //通过数据库的 sequence 生成主键
102        else {
103            String selectId = "row." + SelectKeyGenerator.SELECT_KEY_SUFFIX;
104            String sequence = rowKey.getValue();
105            SqlSource sqlSource = ms.getLang().createSqlSource(ms.getConfiguration(), sequence.trim(), Object.class);
106            MappedStatement.Builder msBuilder = new MappedStatement.Builder(ms.getConfiguration(), selectId, sqlSource, SqlCommandType.SELECT)
107                    .resource(ms.getResource())
108                    .fetchSize(null)
109                    .timeout(null)
110                    .statementType(StatementType.PREPARED)
111                    .keyGenerator(NoKeyGenerator.INSTANCE)
112                    .keyProperty(FlexConsts.ROW + "." + rowKey.getKeyColumn())
113//                    .keyColumn(FlexConsts.ROW + "." + rowKey.getKeyColumn())
114                    .keyColumn(rowKey.getKeyColumn())
115                    .databaseId(ms.getDatabaseId())
116                    .lang(ms.getLang())
117                    .resultOrdered(false)
118                    .resultSets(null)
119                    .resultMaps(new ArrayList<>())
120                    .resultSetType(null)
121                    .flushCacheRequired(false)
122                    .useCache(false)
123                    .cache(ms.getCache());
124
125            MappedStatement keyMappedStatement = msBuilder.build();
126            ms.getConfiguration().addMappedStatement(keyMappedStatement);
127
128            //看到有的框架把 keyGenerator 添加到 mybatis 的当前配置里去，其实是完全没必要的
129            //因为只有在 xml 解析的时候，才可能存在多一个 MappedStatement 拥有同一个 keyGenerator 的情况
130            //当前每个方法都拥有一个自己的 keyGenerator 了，没必要添加
131            //this.addKeyGenerator(selectId, keyGenerator);
132            return new SelectKeyGenerator(keyMappedStatement, rowKey.isBefore());
133        }
134
135    }
136
137    /**
138     * 是否需要数据库生成主键
139     *
140     * @return true 需要生成
141     */
142    @Override
143    public boolean hasGeneratedKeys() {
144        return CollectionUtil.isNotEmpty(autoKeyGeneratorNames);
145    }
146
147    /**
148     * 数据库主键定义的 key
149     *
150     * @return key 数组
151     */
152    @Override
153    public String[] getKeyColumnNames() {
154        return autoKeyGeneratorNames.toArray(new String[0]);
155    }
156}