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.isBlank(finalDsKey)) {
061                finalDsKey = getMethodDsKey(method, proxy);
062            }
063
064            //通过自定义分配策略去获取最终的数据源
065            finalDsKey = DataSourceKey.getShardingDsKey(finalDsKey, proxy, method, args);
066
067            if (StringUtil.isNotBlank(finalDsKey) && !finalDsKey.equals(userDsKey)) {
068                needClearDsKey = true;
069                DataSourceKey.use(finalDsKey);
070            }
071
072            DbType hintDbType = DialectFactory.getHintDbType();
073            if (hintDbType == null) {
074                if (finalDsKey != null && dataSource != null) {
075                    hintDbType = dataSource.getDbType(finalDsKey);
076                }
077
078                if (hintDbType == null) {
079                    hintDbType = FlexGlobalConfig.getDefaultConfig().getDbType();
080                }
081
082                needClearDbType = true;
083                DialectFactory.setHintDbType(hintDbType);
084            }
085            return cachedInvoker(method).invoke(proxy, method, args, sqlSession);
086        } catch (Throwable e) {
087            throw ExceptionUtil.unwrapThrowable(e);
088        } finally {
089            if (needClearDbType) {
090                DialectFactory.clearHintDbType();
091            }
092            if (needClearDsKey) {
093                if (userDsKey != null) {
094                    //恢复用户设置的数据源，并由用户主动去清除
095                    DataSourceKey.use(userDsKey);
096                } else {
097                    DataSourceKey.clear();
098                }
099            }
100        }
101    }
102
103
104    private static String getMethodDsKey(Method method, Object proxy) {
105        UseDataSource methodAnno = method.getAnnotation(UseDataSource.class);
106        if (methodAnno != null && StringUtil.isNotBlank(methodAnno.value())) {
107            return methodAnno.value();
108        }
109
110        Class<?>[] interfaces = proxy.getClass().getInterfaces();
111        for (Class<?> anInterface : interfaces) {
112            UseDataSource classAnno = anInterface.getAnnotation(UseDataSource.class);
113            if (classAnno != null && StringUtil.isNotBlank(classAnno.value())) {
114                return classAnno.value();
115            }
116        }
117
118        if (interfaces[0] != RowMapper.class) {
119            TableInfo tableInfo = TableInfoFactory.ofMapperClass(interfaces[0]);
120            if (tableInfo != null) {
121                String tableDsKey = tableInfo.getDataSource();
122                if (StringUtil.isNotBlank(tableDsKey)) {
123                    return tableDsKey;
124                }
125            }
126        }
127        return null;
128    }
129
130}