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 */
016
017package com.mybatisflex.spring.datasource;
018
019
020import com.mybatisflex.annotation.UseDataSource;
021import com.mybatisflex.core.datasource.DataSourceKey;
022import com.mybatisflex.core.util.StringUtil;
023import org.aopalliance.intercept.MethodInterceptor;
024import org.aopalliance.intercept.MethodInvocation;
025import org.springframework.core.MethodClassKey;
026
027import java.lang.reflect.Method;
028import java.util.Map;
029import java.util.concurrent.ConcurrentHashMap;
030
031/**
032 * 多数据源切换拦截器。
033 *
034 * @author 王帅
035 * @author barql
036 * @author michael
037 * 
038 * @since 2023-06-25
039 */
040public class DataSourceInterceptor implements MethodInterceptor {
041
042    /**
043     * 缓存方法对应的数据源。
044     */
045    private final Map<Object, String> dsCache = new ConcurrentHashMap<>();
046
047    @Override
048    public Object invoke(MethodInvocation invocation) throws Throwable {
049        String dsKey = DataSourceKey.getByManual();
050        if (StringUtil.isNotBlank(dsKey)) {
051            return invocation.proceed();
052        }
053
054        dsKey = findDataSourceKey(invocation.getMethod(), invocation.getThis().getClass());
055        if (StringUtil.isBlank(dsKey)) {
056            return invocation.proceed();
057        }
058
059        //方法嵌套时，挂起的 key
060        String suspendKey = DataSourceKey.getByAnnotation();
061
062        try {
063            DataSourceKey.useWithAnnotation(dsKey);
064            return invocation.proceed();
065        } finally {
066            //恢复挂起的 key
067            if (suspendKey != null) {
068                DataSourceKey.useWithAnnotation(suspendKey);
069            } else {
070                DataSourceKey.clear();
071            }
072        }
073    }
074
075    private String findDataSourceKey(Method method, Class<?> targetClass) {
076        Object cacheKey = new MethodClassKey(method, targetClass);
077        String dsKey = this.dsCache.get(cacheKey);
078        if (dsKey == null) {
079            dsKey = determineDataSourceKey(method, targetClass);
080            this.dsCache.put(cacheKey, dsKey);
081        }
082        return dsKey;
083    }
084
085
086    private String determineDataSourceKey(Method method, Class<?> targetClass) {
087
088        // 方法上定义有 UseDataSource 注解
089        UseDataSource annotation = method.getAnnotation(UseDataSource.class);
090        if (annotation != null) {
091            return annotation.value();
092        }
093
094        // 类上定义有 UseDataSource 注解
095        annotation = targetClass.getAnnotation(UseDataSource.class);
096        if (annotation != null) {
097            return annotation.value();
098        }
099
100        // 接口上定义有 UseDataSource 注解
101        Class<?>[] interfaces = targetClass.getInterfaces();
102        for (Class<?> anInterface : interfaces) {
103            annotation = anInterface.getAnnotation(UseDataSource.class);
104            if (annotation != null) {
105                return annotation.value();
106            }
107        }
108
109        return "";
110    }
111
112}