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.FlexGlobalConfig;
021import com.mybatisflex.core.exception.FlexExceptions;
022import com.mybatisflex.core.table.IdInfo;
023import com.mybatisflex.core.table.TableInfo;
024import com.mybatisflex.core.util.ConvertUtil;
025import com.mybatisflex.core.util.StringUtil;
026import org.apache.ibatis.executor.Executor;
027import org.apache.ibatis.executor.ExecutorException;
028import org.apache.ibatis.executor.keygen.KeyGenerator;
029import org.apache.ibatis.mapping.MappedStatement;
030import org.apache.ibatis.reflection.MetaObject;
031import org.apache.ibatis.session.Configuration;
032
033import java.sql.Statement;
034import java.util.Map;
035
036/**
037 * 通过 java 编码的方式生成主键
038 * 当主键类型配置为 KeyType#Generator 时，使用此生成器生成
039 * {@link KeyType#Generator}
040 */
041public class CustomKeyGenerator implements KeyGenerator {
042
043    protected Configuration configuration;
044    protected IKeyGenerator keyGenerator;
045    protected TableInfo tableInfo;
046    protected IdInfo idInfo;
047
048
049    public CustomKeyGenerator(Configuration configuration, TableInfo tableInfo, IdInfo idInfo) {
050        this.configuration = configuration;
051        FlexGlobalConfig.KeyConfig globalKeyConfig = FlexGlobalConfig.getConfig(configuration).getKeyConfig();
052        String keyValue = MybatisKeyGeneratorUtil.getKeyValue(idInfo, globalKeyConfig);
053        this.keyGenerator = KeyGeneratorFactory.getKeyGenerator(keyValue);
054        this.tableInfo = tableInfo;
055        this.idInfo = idInfo;
056
057        ensuresKeyGeneratorNotNull();
058    }
059
060    private void ensuresKeyGeneratorNotNull() {
061        if (keyGenerator == null) {
062            throw FlexExceptions.wrap("The name of \"%s\" key generator not exist.\n" +
063                    "please check annotation @Id(value=\"%s\") at field: %s#%s"
064                , idInfo.getValue(), idInfo.getValue(), tableInfo.getEntityClass().getName(), idInfo.getProperty());
065        }
066    }
067
068
069    @Override
070    public void processBefore(Executor executor, MappedStatement ms, Statement stmt, Object parameter) {
071        Object entity = ((Map<?, ?>) parameter).get(FlexConsts.ENTITY);
072        try {
073            Object existId = tableInfo.getValue(entity, idInfo.getProperty());
074            // 若用户主动设置了主键，则使用用户自己设置的主键，不再生成主键
075            // 只有主键为 null 或者 空字符串时，对主键进行设置
076            if (existId == null || (existId instanceof String && StringUtil.isBlank((String) existId))) {
077                Configuration msConfiguration = ms.getConfiguration();
078                MetaObject metaParam = msConfiguration.newMetaObject(parameter);
079                Object generateId = keyGenerator.generate(entity, idInfo.getColumn());
080                MetaObject metaObjectForProperty = metaParam.metaObjectForProperty(FlexConsts.ENTITY);
081                Class<?> setterType = tableInfo.getReflector().getSetterType(idInfo.getProperty());
082                Object id = ConvertUtil.convert(generateId, setterType);
083                this.setValue(metaObjectForProperty, this.idInfo.getProperty(), id);
084            }
085        } catch (Exception e) {
086            throw FlexExceptions.wrap(e);
087        }
088    }
089
090
091    @Override
092    public void processAfter(Executor executor, MappedStatement ms, Statement stmt, Object parameter) {
093        // do nothing
094    }
095
096    private void setValue(MetaObject metaParam, String property, Object value) {
097        if (!metaParam.hasSetter(property)) {
098            throw new ExecutorException("No setter found for the keyProperty '" + property + "' in " + metaParam.getOriginalObject().getClass().getName() + ".");
099        } else {
100            metaParam.setValue(property, value);
101        }
102    }
103
104}