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