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.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.MapUtil;
029import com.mybatisflex.core.util.StringUtil;
030import org.apache.ibatis.reflection.ExceptionUtil;
031import org.apache.ibatis.session.SqlSession;
032
033import java.lang.reflect.Method;
034import java.util.Map;
035import java.util.concurrent.ConcurrentHashMap;
036
037public class FlexMapperProxy<T> extends MybatisMapperProxy<T> {
038    private static final String NULL_KEY = "@NK";
039    private static final Map<Method, String> methodDsKeyCache = new ConcurrentHashMap<>();
040
041    private final FlexDataSource dataSource;
042
043    public FlexMapperProxy(SqlSession sqlSession, Class<T> mapperInterface, Map<Method, MapperMethodInvoker> methodCache,
044                           FlexConfiguration configuration) {
045        super(sqlSession, mapperInterface, methodCache);
046        this.dataSource = (FlexDataSource) configuration.getEnvironment().getDataSource();
047    }
048
049
050    @Override
051    public Object invoke(Object proxy, Method method, Object[] args) throws Throwable {
052        if (Object.class.equals(method.getDeclaringClass())) {
053            return method.invoke(this, args);
054        }
055
056        boolean needClearDsKey = false;
057        boolean needClearDbType = false;
058        try {
059            //获取用户动态指定，由用户指定数据源，则应该有用户清除
060            String dataSourceKey = DataSourceKey.get();
061            if (StringUtil.isBlank(dataSourceKey)) {
062                //通过 @UseDataSource 或者 @Table(dataSource) 去获取
063                String configDataSourceKey = getConfigDataSourceKey(method, proxy);
064                if (StringUtil.isNotBlank(configDataSourceKey)) {
065                    dataSourceKey = configDataSourceKey;
066                    DataSourceKey.use(dataSourceKey);
067                    needClearDsKey = true;
068                }
069            }
070
071            //最终通过数据源 自定义分片 策略去获取
072            String shardingDataSourceKey = DataSourceKey.getByShardingStrategy(dataSourceKey, proxy, method, args);
073            if (shardingDataSourceKey != null && !shardingDataSourceKey.equals(dataSourceKey)) {
074                DataSourceKey.use(shardingDataSourceKey);
075                needClearDsKey = true;
076            }
077
078            //优先获取用户自己配置的 dbType
079            DbType dbType = DialectFactory.getHintDbType();
080
081            if (dbType == null) {
082                if (shardingDataSourceKey != null && dataSource != null) {
083                    //使用最终分片获取数据源类型
084                    dbType = dataSource.getDbType(shardingDataSourceKey);
085                }
086
087                if (dbType == null && dataSourceKey != null && dataSource != null) {
088                    dbType = dataSource.getDbType(dataSourceKey);
089                }
090
091                //设置了dbTypeGlobal，那么就使用全局的dbTypeGlobal
092                if (dbType == null) {
093                    dbType = DialectFactory.getGlobalDbType();
094                }
095
096                if (dbType == null) {
097                    dbType = FlexGlobalConfig.getDefaultConfig().getDbType();
098                }
099
100                DialectFactory.setHintDbType(dbType);
101                needClearDbType = true;
102            }
103//            return method.invoke(mapper, args);
104            return cachedInvoker(method).invoke(proxy, method, args, sqlSession);
105        } catch (Throwable e) {
106            throw ExceptionUtil.unwrapThrowable(e);
107        } finally {
108            if (needClearDbType) {
109                DialectFactory.clearHintDbType();
110            }
111            if (needClearDsKey) {
112                DataSourceKey.clear();
113            }
114        }
115    }
116
117
118    private static String getConfigDataSourceKey(Method method, Object proxy) {
119        String result = MapUtil.computeIfAbsent(methodDsKeyCache, method, method1 -> {
120            UseDataSource useDataSource = method1.getAnnotation(UseDataSource.class);
121            if (useDataSource != null && StringUtil.isNotBlank(useDataSource.value())) {
122                return useDataSource.value();
123            }
124
125            Class<?>[] interfaces = proxy.getClass().getInterfaces();
126            for (Class<?> anInterface : interfaces) {
127                UseDataSource annotation = anInterface.getAnnotation(UseDataSource.class);
128                if (annotation != null) {
129                    return annotation.value();
130                }
131            }
132
133            if (interfaces[0] != RowMapper.class) {
134                TableInfo tableInfo = TableInfoFactory.ofMapperClass(interfaces[0]);
135                if (tableInfo != null) {
136                    String dataSourceKey = tableInfo.getDataSource();
137                    if (StringUtil.isNotBlank(dataSourceKey)) {
138                        return dataSourceKey;
139                    }
140                }
141            }
142            return NULL_KEY;
143        });
144
145        return NULL_KEY.equals(result) ? null : result;
146    }
147
148}