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            //优先获取用户自己配置的 dbType
071            DbType dbType = DialectFactory.getHintDbType();
072            if (dbType == null) {
073                if (dataSourceKey != null && dataSource != null) {
074                    dbType = dataSource.getDbType(dataSourceKey);
075                }
076                if (dbType == null) {
077                    dbType = FlexGlobalConfig.getDefaultConfig().getDbType();
078                }
079                DialectFactory.setHintDbType(dbType);
080                needClearDbType = true;
081            }
082            return method.invoke(mapper, args);
083        } catch (InvocationTargetException e) {
084            throw e.getCause();
085        } finally {
086            if (needClearDbType) {
087                DialectFactory.clearHintDbType();
088            }
089            if (needClearDsKey) {
090                DataSourceKey.clear();
091            }
092        }
093    }
094
095
096    private static String getConfigDataSourceKey(Method method, Object proxy) {
097        UseDataSource useDataSource = method.getAnnotation(UseDataSource.class);
098        if (useDataSource != null && StringUtil.isNotBlank(useDataSource.value())) {
099            return useDataSource.value();
100        }
101
102        Class<?>[] interfaces = proxy.getClass().getInterfaces();
103        for (Class<?> anInterface : interfaces) {
104            UseDataSource annotation = anInterface.getAnnotation(UseDataSource.class);
105            if (annotation != null) {
106                return annotation.value();
107            }
108        }
109
110        if (interfaces[0] != RowMapper.class) {
111            TableInfo tableInfo = TableInfoFactory.ofMapperClass(interfaces[0]);
112            if (tableInfo != null) {
113                String dataSourceKey = tableInfo.getDataSource();
114                if (StringUtil.isNotBlank(dataSourceKey)) {
115                    return dataSourceKey;
116                }
117            }
118        }
119        return null;
120    }
121
122
123}