com.github.dennisit.vplus.data.dsource.shard.ShardDataSourceAspect Maven / Gradle / Ivy

Show more of this group Show more artifacts with this name
Show all versions of spring-saber-boxes Show documentation
spring-saber-boxes
There is a newer version: 2.0.8
package com.github.dennisit.vplus.data.dsource.shard;

import com.github.dennisit.vplus.data.dsource.ShardInfo;
import com.github.dennisit.vplus.data.dsource.ShardOn;
import com.github.dennisit.vplus.data.dsource.multi.MultipleDataSource;
import com.github.dennisit.vplus.data.dsource.multi.MultipleDataSourceAspect;
import lombok.extern.slf4j.Slf4j;
import org.aspectj.lang.ProceedingJoinPoint;
import org.aspectj.lang.Signature;
import org.aspectj.lang.reflect.MethodSignature;
import org.springframework.beans.factory.annotation.Required;
import org.springframework.core.annotation.AnnotationUtils;
import org.springframework.util.StringUtils;

import java.lang.annotation.Annotation;
import java.lang.reflect.Method;

@Slf4j
public class ShardDataSourceAspect extends MultipleDataSourceAspect {

    private ShardContext shardContext;

    public ShardDataSourceAspect() {
    }

    public Object doAround(ProceedingJoinPoint jp) throws Throwable {
        if (log.isDebugEnabled()) {
            log.debug("MultipleDataSourceAspectAdvice invoked!");
        }

        Signature signature = jp.getSignature();
        Lookup lookup = this.getLookup(signature, jp.getArgs());
        this.shardContext.setUpDataSourceTableLookup(lookup);
        String dataSourceKey = ShardContext.getDataSourceKey();
        if (StringUtils.hasText(dataSourceKey)) {
            MultipleDataSource.setDataSourceKey(dataSourceKey);
        }

        Object jpVal = jp.proceed();
        return jpVal;
    }

    private Lookup getLookup(Signature signature, Object[] args) {
        if (signature == null) {
            return null;
        } else {
            String dataSource = null;
            String tableName = null;
            Integer tableShardNum = null;
            Integer dbShardNum = null;
            String shardType = null;
            if (!(signature instanceof MethodSignature)) {
                return null;
            } else {
                MethodSignature methodSignature = (MethodSignature) signature;
                Method method = methodSignature.getMethod();
                Class declaringClazz = method.getDeclaringClass();
                dataSource = this.getDataSourceKey(signature);
                Annotation[][] annotations = method.getParameterAnnotations();
                int index = -1;

                for (int i = 0; i < annotations.length; ++i) {
                    Annotation[] paramAnnotations = annotations[i];

                    for (int j = 0; j < paramAnnotations.length; ++j) {
                        Annotation paramAnno = paramAnnotations[j];
                        if (paramAnno instanceof ShardOn) {
                            index = i;
                            ShardInfo shardInfo = (ShardInfo) AnnotationUtils.findAnnotation(declaringClazz, ShardInfo.class);
                            tableName = (String) AnnotationUtils.getValue(shardInfo, "tableName");
                            tableShardNum = (Integer) AnnotationUtils.getValue(shardInfo, "tableShardNum");
                            dbShardNum = (Integer) AnnotationUtils.getValue(shardInfo, "dbShardNum");
                            shardType = (String) AnnotationUtils.getValue(shardInfo, "shardType");
                            return index == -1 ? new ShardLookup(dataSource) : new ShardLookup(tableName, dataSource, tableShardNum, dbShardNum, shardType, args[index]);
                        }
                    }
                }

                return index == -1 ? new ShardLookup(dataSource) : new ShardLookup(tableName, dataSource, tableShardNum, dbShardNum, shardType, args[index]);
            }
        }
    }

    @Required
    public void setShardContext(ShardContext shardContext) {
        this.shardContext = shardContext;
    }

}