001/*
002 *  Copyright (c) 2022-2025, 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.*;
036
037/**
038 * 为 row 的主键生成器
039 */
040public class RowKeyGenerator implements KeyGenerator, IMultiKeyGenerator {
041
042    private static final KeyGenerator[] NO_KEY_GENERATORS = new KeyGenerator[0];
043
044    private final MappedStatement ms;
045    private KeyGenerator[] keyGenerators;
046    private Set<String> autoKeyGeneratorNames;
047
048    public RowKeyGenerator(MappedStatement methodMappedStatement) {
049        this.ms = methodMappedStatement;
050    }
051
052    @Override
053    public void processBefore(Executor executor, MappedStatement ms, Statement stmt, Object parameter) {
054        Row row = (Row) ((Map<?, ?>) parameter).get(FlexConsts.ROW);
055        keyGenerators = buildRowKeyGenerators(RowCPI.obtainsPrimaryKeys(row));
056        for (KeyGenerator keyGenerator : keyGenerators) {
057            keyGenerator.processBefore(executor, ms, stmt, parameter);
058        }
059    }
060
061
062    @Override
063    public void processAfter(Executor executor, MappedStatement ms, Statement stmt, Object parameter) {
064        for (KeyGenerator keyGenerator : keyGenerators) {
065            keyGenerator.processAfter(executor, ms, stmt, parameter);
066        }
067    }
068
069
070    private KeyGenerator[] buildRowKeyGenerators(RowKey[] rowKeys) {
071        if (ArrayUtil.isEmpty(rowKeys)) {
072            return NO_KEY_GENERATORS;
073        }
074
075        KeyGenerator[] keyGenerators = new KeyGenerator[rowKeys.length];
076        for (int i = 0; i < rowKeys.length; i++) {
077            KeyGenerator keyGenerator = createByRowKey(rowKeys[i]);
078            keyGenerators[i] = keyGenerator;
079        }
080        return keyGenerators;
081    }
082
083
084    private KeyGenerator createByRowKey(RowKey rowKey) {
085        if (rowKey == null || rowKey.getKeyType() == KeyType.None) {
086            return NoKeyGenerator.INSTANCE;
087        }
088
089        String keyColumn = rowKey.getKeyColumn();
090        if (rowKey.getKeyType() == KeyType.Auto) {
091            if (autoKeyGeneratorNames == null) {
092                autoKeyGeneratorNames = new HashSet<>();
093            }
094            autoKeyGeneratorNames.add(keyColumn);
095            return new RowJdbc3KeyGenerator(keyColumn);
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 + "." + keyColumn)
113                .keyColumn(keyColumn)
114                .databaseId(ms.getDatabaseId())
115                .lang(ms.getLang())
116                .resultOrdered(false)
117                .resultSets(null)
118                .resultMaps(new ArrayList<>())
119                .resultSetType(null)
120                .flushCacheRequired(false)
121                .useCache(false)
122                .cache(ms.getCache());
123
124            MappedStatement keyMappedStatement = msBuilder.build();
125            ms.getConfiguration().addMappedStatement(keyMappedStatement);
126
127            //看到有的框架把 keyGenerator 添加到 mybatis 的当前配置里去，其实是完全没必要的
128            //因为只有在 xml 解析的时候，才可能存在多一个 MappedStatement 拥有同一个 keyGenerator 的情况
129            //当前每个方法都拥有一个自己的 keyGenerator 了，没必要添加
130            //addKeyGenerator(selectId, keyGenerator)
131            return new SelectKeyGenerator(keyMappedStatement, rowKey.isBefore());
132        }
133
134    }
135
136    /**
137     * 是否需要数据库生成主键
138     *
139     * @return true 需要生成
140     */
141    @Override
142    public boolean hasGeneratedKeys() {
143        return CollectionUtil.isNotEmpty(autoKeyGeneratorNames);
144    }
145
146    /**
147     * 数据库主键定义的 key
148     *
149     * @return key 数组
150     */
151    @Override
152    public String[] getKeyColumnNames() {
153        return autoKeyGeneratorNames == null ? new String[0] : autoKeyGeneratorNames.toArray(new String[0]);
154    }
155
156}