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