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;
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.row.RowMapper;
025import com.mybatisflex.core.table.TableInfo;
026import com.mybatisflex.core.table.TableInfoFactory;
027import com.mybatisflex.core.util.StringUtil;
028
029import javax.sql.DataSource;
030import java.lang.reflect.InvocationHandler;
031import java.lang.reflect.InvocationTargetException;
032import java.lang.reflect.Method;
033
034/**
035 * @author michael
036 * @author norkts
037 */
038public class MapperInvocationHandler implements InvocationHandler {
039
040    private final Object mapper;
041    private final FlexDataSource dataSource;
042
043    public MapperInvocationHandler(Object mapper, DataSource dataSource) {
044        this.mapper = mapper;
045        if (dataSource instanceof FlexDataSource) {
046            this.dataSource = (FlexDataSource) dataSource;
047        } else {
048            this.dataSource = null;
049        }
050    }
051
052
053    @Override
054    public Object invoke(Object proxy, Method method, Object[] args) throws Throwable {
055        boolean needClearDsKey = false;
056        boolean needClearDbType = false;
057        try {
058            //获取用户动态指定，由用户指定数据源，则应该有用户清除
059            String dataSourceKey = DataSourceKey.get();
060            if (StringUtil.isBlank(dataSourceKey)) {
061                //通过 @UseDataSource 或者 @Table(dataSource) 去获取
062                String configDataSourceKey = getConfigDataSourceKey(method, proxy);
063                if (StringUtil.isNotBlank(configDataSourceKey)) {
064                    dataSourceKey = configDataSourceKey;
065                    DataSourceKey.use(dataSourceKey);
066                    needClearDsKey = true;
067                }
068            }
069
070            //最终通过数据源 自定义分片 策略去获取
071            String shardingDataSourceKey = DataSourceKey.getByShardingStrategy(dataSourceKey, proxy, method, args);
072            if (shardingDataSourceKey != null && !shardingDataSourceKey.equals(dataSourceKey)) {
073                DataSourceKey.use(shardingDataSourceKey);
074                needClearDsKey = true;
075            }
076
077            //优先获取用户自己配置的 dbType
078            DbType dbType = DialectFactory.getHintDbType();
079            if (dbType == null) {
080                if (dataSourceKey != null && dataSource != null) {
081                    dbType = dataSource.getDbType(dataSourceKey);
082                }
083                if (dbType == null) {
084                    dbType = FlexGlobalConfig.getDefaultConfig().getDbType();
085                }
086                DialectFactory.setHintDbType(dbType);
087                needClearDbType = true;
088            }
089            return method.invoke(mapper, args);
090        } catch (InvocationTargetException e) {
091            throw e.getCause();
092        } finally {
093            if (needClearDbType) {
094                DialectFactory.clearHintDbType();
095            }
096            if (needClearDsKey) {
097                DataSourceKey.clear();
098            }
099        }
100    }
101
102
103    private static String getConfigDataSourceKey(Method method, Object proxy) {
104        UseDataSource useDataSource = method.getAnnotation(UseDataSource.class);
105        if (useDataSource != null && StringUtil.isNotBlank(useDataSource.value())) {
106            return useDataSource.value();
107        }
108
109        Class<?>[] interfaces = proxy.getClass().getInterfaces();
110        for (Class<?> anInterface : interfaces) {
111            UseDataSource annotation = anInterface.getAnnotation(UseDataSource.class);
112            if (annotation != null) {
113                return annotation.value();
114            }
115        }
116
117        if (interfaces[0] != RowMapper.class) {
118            TableInfo tableInfo = TableInfoFactory.ofMapperClass(interfaces[0]);
119            if (tableInfo != null) {
120                String dataSourceKey = tableInfo.getDataSource();
121                if (StringUtil.isNotBlank(dataSourceKey)) {
122                    return dataSourceKey;
123                }
124            }
125        }
126        return null;
127    }
128
129
130}