AnnatationDataSourceKeyResolver.java

package me.test.db.router;

import java.lang.annotation.Annotation;
import java.lang.reflect.Method;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;

import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.beans.BeansException;
import org.springframework.context.ApplicationContext;
import org.springframework.context.ApplicationContextAware;
import org.springframework.context.expression.BeanFactoryAccessor;
import org.springframework.context.expression.BeanFactoryResolver;
import org.springframework.context.expression.EnvironmentAccessor;
import org.springframework.context.expression.MapAccessor;
import org.springframework.core.LocalVariableTableParameterNameDiscoverer;
import org.springframework.core.ParameterNameDiscoverer;
import org.springframework.expression.Expression;
import org.springframework.expression.ExpressionParser;
import org.springframework.expression.spel.standard.SpelExpressionParser;
import org.springframework.expression.spel.support.StandardEvaluationContext;
import org.springframework.expression.spel.support.StandardTypeLocator;
import org.springframework.util.Assert;

/**
 * 根据方法上，或参数上的 ${link DataSourceKey @DataSourceKey} 注解的设定读取数据源路由Key。
 */
public class AnnatationDataSourceKeyResolver implements DataSourceKeyResolver, ApplicationContextAware {

    private Logger logger = LoggerFactory.getLogger(AnnatationDataSourceKeyResolver.class);
    private ApplicationContext appCtx;
    private ExpressionParser expressionParser = new SpelExpressionParser();
    private Map<Method, String[]> paramNameCache = new ConcurrentHashMap<Method, String[]>();
    private ParameterNameDiscoverer parameterNameDiscoverer = new LocalVariableTableParameterNameDiscoverer();


    @Override
    public String resolveKey(Object thisObj, Method targetMethod, Object[] args) {

        // 通过被@DataSourceKey注解的方法参数计算数据源Key
        String key = null;
        int usingCountOnMethodParam = 0;
        Annotation[][] annotations = targetMethod.getParameterAnnotations();
        for (int i = 0; i < args.length; i++) {
            for (int j = 0; j < annotations[i].length; j++) {
                if (annotations[i][j] instanceof DataSourceKey) {
                    usingCountOnMethodParam++;
                    Assert.isTrue(usingCountOnMethodParam <= 1,
                            "@DataSouceKey is found more than one on parameters of method : " + targetMethod);

                    DataSourceKey paramDsKey = (DataSourceKey) annotations[i][j];
                    Assert.isTrue(paramDsKey.value() == null || paramDsKey.value().length() == 0,
                            "@DataSouceKey can not set 'value' when using on method parameter. method : " + targetMethod);

                    if (paramDsKey.prefix().length() == 0) {
                        key = args[i] == null ? null : args[i].toString();
                    } else {
                        key = paramDsKey.prefix() + args[i].toString();
                    }
                }
            }
        }

        DataSourceKey methodDsKey = targetMethod.getAnnotation(DataSourceKey.class);
        Assert.isTrue(!(methodDsKey != null && usingCountOnMethodParam == 1),
                "@DataSouceKey can not using both on method and paramets. method : " + targetMethod);

        if (key != null && key.length() > 0) {
            logger.debug("using '{}' as datasource key from method @DataSouceKey ", key);
            return key;
        }

        // 通过被@DataSourceKey注解的方法计算数据源Key
        if (methodDsKey == null) {
            return null;
        }
        String exprStr = methodDsKey.value();
        Assert.isTrue(exprStr != null && exprStr.length() > 0, "@DataSouceKey's value is not specified. method : " + targetMethod);
        Assert.isTrue(methodDsKey.prefix().length() == 0, "@DataSouceKey can not set 'prefix' when using on method. method : " + targetMethod);

        // 查询是否被注解，如果被注解，则查询到注解的值
        Expression exp = expressionParser.parseExpression(exprStr);
        StandardEvaluationContext sec = new StandardEvaluationContext();
        sec.setRootObject(thisObj);
        sec.addPropertyAccessor(new BeanFactoryAccessor());
        sec.addPropertyAccessor(new MapAccessor());
        sec.addPropertyAccessor(new EnvironmentAccessor());
        sec.setBeanResolver(new BeanFactoryResolver(appCtx));
        sec.setTypeLocator(new StandardTypeLocator(this.getClass().getClassLoader()));

        // 注册参数(参考：org.springframework.cache.interceptor.LazyParamAwareEvaluationContext#loadArgsAsVariables())
        for (int i = 0; i < args.length; i++) {
            sec.setVariable("a" + i, args[i]);
            sec.setVariable("p" + i, args[i]);
        }

        String[] parameterNames = paramNameCache.get(targetMethod);
        if (parameterNames == null) {
            parameterNames = parameterNameDiscoverer.getParameterNames(targetMethod);
            paramNameCache.put(targetMethod, parameterNames);
        }
        if (parameterNames != null) {
            for (int i = 0; i < parameterNames.length; i++) {
                sec.setVariable(parameterNames[i], args[i]);
            }
        }

        Object value = exp.getValue(sec);
        if (value == null) {
            return null;
        }
        String strValue = value.toString();
        logger.debug("using '{}' as datasource key from parameter @DataSouceKey ", strValue);
        return strValue;
    }

    @Override
    public void setApplicationContext(ApplicationContext applicationContext) throws BeansException {

        this.appCtx = applicationContext;
    }
}