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;
028import org.apache.ibatis.session.Configuration;
029
030import java.lang.reflect.InvocationHandler;
031import java.lang.reflect.InvocationTargetException;
032import java.lang.reflect.Method;
033
034public class MapperInvocationHandler implements InvocationHandler {
035
036    private final Object mapper;
037    private final FlexDataSource dataSource;
038
039    public MapperInvocationHandler(Object mapper, Configuration configuration) {
040        this.mapper = mapper;
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        boolean clearDsKey = false;
048        boolean clearDbType = false;
049        try {
050            //获取用户动态指定，由用户指定数据源，则应该有用户清除
051            String dataSourceKey = DataSourceKey.get();
052
053            if (StringUtil.isBlank(dataSourceKey)) {
054                //通过 @UseDataSource 或者 @Table(dataSource) 去获取
055                String configDataSourceKey = getConfigDataSourceKey(method, proxy);
056                if (StringUtil.isNotBlank(configDataSourceKey)) {
057                    dataSourceKey = configDataSourceKey;
058                    DataSourceKey.use(dataSourceKey);
059                    clearDsKey = true;
060                }
061            }
062
063            //优先获取用户自己配置的 dbType
064            DbType dbType = DialectFactory.getHintDbType();
065            if (dbType == null) {
066                if (dataSourceKey != null) {
067                    dbType = dataSource.getDbType(dataSourceKey);
068                }
069                if (dbType == null) {
070                    dbType = FlexGlobalConfig.getDefaultConfig().getDbType();
071                }
072                DialectFactory.setHintDbType(dbType);
073                clearDbType = true;
074            }
075            return method.invoke(mapper, args);
076        } catch (InvocationTargetException e1) {
077            throw e1.getCause();
078        } finally {
079            if (clearDbType) {
080                DialectFactory.clearHintDbType();
081            }
082            if (clearDsKey) {
083                DataSourceKey.clear();
084            }
085        }
086    }
087
088
089    private static String getConfigDataSourceKey(Method method, Object proxy) {
090        UseDataSource useDataSource = method.getAnnotation(UseDataSource.class);
091        if (useDataSource != null && StringUtil.isNotBlank(useDataSource.value())) {
092            return useDataSource.value();
093        }
094
095        Class<?>[] interfaces = proxy.getClass().getInterfaces();
096        for (Class<?> anInterface : interfaces) {
097            UseDataSource annotation = anInterface.getAnnotation(UseDataSource.class);
098            if (annotation != null) {
099                return annotation.value();
100            }
101        }
102
103        if (interfaces[0] != RowMapper.class) {
104            TableInfo tableInfo = TableInfoFactory.ofMapperClass(interfaces[0]);
105            if (tableInfo != null) {
106                String dataSourceKey = tableInfo.getDataSource();
107                if (StringUtil.isNotBlank(dataSourceKey)) {
108                    return dataSourceKey;
109                }
110            }
111        }
112        return null;
113    }
114
115
116}