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.core.table.IdInfo;
019import com.mybatisflex.core.table.TableInfo;
020import org.apache.ibatis.executor.Executor;
021import org.apache.ibatis.executor.keygen.Jdbc3KeyGenerator;
022import org.apache.ibatis.executor.keygen.KeyGenerator;
023import org.apache.ibatis.mapping.MappedStatement;
024
025import java.sql.Statement;
026import java.util.ArrayList;
027import java.util.List;
028
029/**
030 * 多主键、复合主键 id 生成器
031 */
032public class MultiPrimaryKeyGenerator implements KeyGenerator, IMultiKeyGenerator {
033
034    private final List<KeyGenerator> keyGenerators;
035
036    //所有自增字段的名称
037    private final String[] autoGenKeyColumnNames;
038
039    public MultiPrimaryKeyGenerator(MappedStatement mappedStatement, TableInfo tableInfo, List<IdInfo> primaryKeyList) {
040        this.keyGenerators = new ArrayList<>();
041
042        List<String> autoGenKeyColumnNameList = new ArrayList<>();
043        for (IdInfo idInfo : primaryKeyList) {
044            KeyGenerator idKeyGenerator = MybatisKeyGeneratorUtil.createIdKeyGenerator(tableInfo, mappedStatement, idInfo);
045            keyGenerators.add(idKeyGenerator);
046            if (idKeyGenerator == Jdbc3KeyGenerator.INSTANCE) {
047                autoGenKeyColumnNameList.add(idInfo.getColumn());
048            }
049        }
050
051        this.autoGenKeyColumnNames = autoGenKeyColumnNameList.toArray(new String[0]);
052    }
053
054    @Override
055    public void processBefore(Executor executor, MappedStatement ms, Statement stmt, Object parameter) {
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     * 是否需要数据库 自动生成主键
071     *
072     * @return true: need generated keys
073     */
074    @Override
075    public boolean hasGeneratedKeys() {
076        for (KeyGenerator keyGenerator : keyGenerators) {
077            if (keyGenerator == Jdbc3KeyGenerator.INSTANCE) {
078                return true;
079            }
080        }
081        return false;
082    }
083
084    /**
085     * 自动生成主键的 columns 字段
086     *
087     * @return keyColumnNames
088     */
089    @Override
090    public String[] getKeyColumnNames() {
091        return autoGenKeyColumnNames;
092    }
093
094
095}