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.mybatis.binding;
017
018import com.mybatisflex.annotation.UseDataSource;
019import com.mybatisflex.core.FlexGlobalConfig;
020import com.mybatisflex.core.datasource.DataSourceKey;
021import com.mybatisflex.core.datasource.FlexDataSource;
022import com.mybatisflex.core.dialect.DbType;
023import com.mybatisflex.core.dialect.DialectFactory;
024import com.mybatisflex.core.mybatis.FlexConfiguration;
025import com.mybatisflex.core.row.RowMapper;
026import com.mybatisflex.core.table.TableInfo;
027import com.mybatisflex.core.table.TableInfoFactory;
028import com.mybatisflex.core.util.StringUtil;
029import org.apache.ibatis.reflection.ExceptionUtil;
030import org.apache.ibatis.session.SqlSession;
031
032import java.lang.reflect.Method;
033import java.util.Map;
034
035public class FlexMapperProxy<T> extends MybatisMapperProxy<T> {
036    private final FlexDataSource dataSource;
037
038    public FlexMapperProxy(SqlSession sqlSession, Class<T> mapperInterface, Map<Method, MapperMethodInvoker> methodCache,
039                           FlexConfiguration configuration) {
040        super(sqlSession, mapperInterface, methodCache);
041        this.dataSource = (FlexDataSource) configuration.getEnvironment().getDataSource();
042    }
043
044
045    @Override
046    public Object invoke(Object proxy, Method method, Object[] args) throws Throwable {
047        if (Object.class.equals(method.getDeclaringClass())) {
048            return method.invoke(this, args);
049        }
050
051        boolean needClearDsKey = false;
052        boolean needClearDbType = false;
053
054        //由用户指定的数据
055        String userDsKey = DataSourceKey.get();
056        //最终使用的数据源
057        String finalDsKey = userDsKey;
058
059        try {
060            if (StringUtil.noText(finalDsKey)) {
061                // Mapper 方法上获取 UseDataSource的value值
062                finalDsKey = getMethodDsKey(method, proxy);
063                // 对数据源取值进行动态取值处理
064                if (StringUtil.hasText(finalDsKey)) {
065                    finalDsKey = DataSourceKey.processDataSourceKey(finalDsKey, proxy, method, args);
066                }
067            }
068
069            // 通过自定义分配策略去获取最终的数据源
070            finalDsKey = DataSourceKey.getShardingDsKey(finalDsKey, proxy, method, args);
071
072            if (StringUtil.hasText(finalDsKey) && !finalDsKey.equals(userDsKey)) {
073                needClearDsKey = true;
074                DataSourceKey.use(finalDsKey);
075            }
076
077            DbType hintDbType = DialectFactory.getHintDbType();
078            if (hintDbType == null) {
079                if (finalDsKey != null && dataSource != null) {
080                    hintDbType = dataSource.getDbType(finalDsKey);
081                }
082
083                if (hintDbType == null) {
084                    hintDbType = FlexGlobalConfig.getDefaultConfig().getDbType();
085                }
086
087                needClearDbType = true;
088                DialectFactory.setHintDbType(hintDbType);
089            }
090            return cachedInvoker(method).invoke(proxy, method, args, sqlSession);
091        } catch (Throwable e) {
092            throw ExceptionUtil.unwrapThrowable(e);
093        } finally {
094            if (needClearDbType) {
095                DialectFactory.clearHintDbType();
096            }
097            if (needClearDsKey) {
098                if (userDsKey != null) {
099                    //恢复用户设置的数据源，并由用户主动去清除
100                    DataSourceKey.use(userDsKey);
101                } else {
102                    DataSourceKey.clear();
103                }
104            }
105        }
106    }
107
108
109    private static String getMethodDsKey(Method method, Object proxy) {
110        UseDataSource methodAnno = method.getAnnotation(UseDataSource.class);
111        if (methodAnno != null && StringUtil.hasText(methodAnno.value())) {
112            return methodAnno.value();
113        }
114
115        Class<?>[] interfaces = proxy.getClass().getInterfaces();
116        for (Class<?> anInterface : interfaces) {
117            UseDataSource classAnno = anInterface.getAnnotation(UseDataSource.class);
118            if (classAnno != null && StringUtil.hasText(classAnno.value())) {
119                return classAnno.value();
120            }
121        }
122
123        if (interfaces[0] != RowMapper.class) {
124            TableInfo tableInfo = TableInfoFactory.ofMapperClass(interfaces[0]);
125            if (tableInfo != null) {
126                String tableDsKey = tableInfo.getDataSource();
127                if (StringUtil.hasText(tableDsKey)) {
128                    return tableDsKey;
129                }
130            }
131        }
132        return null;
133    }
134
135}