optimize DataSourceInterceptor

This commit is contained in:
开源海哥 2023-06-25 14:37:46 +08:00
parent 7e014a4d14
commit 5e94ec87b3

View File

@ -19,6 +19,7 @@ package com.mybatisflex.spring.datasource;
import com.mybatisflex.annotation.UseDataSource; import com.mybatisflex.annotation.UseDataSource;
import com.mybatisflex.core.datasource.DataSourceKey; import com.mybatisflex.core.datasource.DataSourceKey;
import com.mybatisflex.core.util.StringUtil;
import org.aopalliance.intercept.MethodInterceptor; import org.aopalliance.intercept.MethodInterceptor;
import org.aopalliance.intercept.MethodInvocation; import org.aopalliance.intercept.MethodInvocation;
@ -32,7 +33,17 @@ public class DataSourceInterceptor implements MethodInterceptor {
@Override @Override
public Object invoke(MethodInvocation invocation) throws Throwable { public Object invoke(MethodInvocation invocation) throws Throwable {
String dsKey = determineDataSourceKey(invocation);
String dsKey = DataSourceKey.get();
if (StringUtil.isNotBlank(dsKey)) {
return invocation.proceed();
}
dsKey = determineDataSourceKey(invocation);
if (StringUtil.isBlank(dsKey)) {
return invocation.proceed();
}
DataSourceKey.use(dsKey); DataSourceKey.use(dsKey);
try { try {
return invocation.proceed(); return invocation.proceed();
@ -42,27 +53,35 @@ public class DataSourceInterceptor implements MethodInterceptor {
} }
private String determineDataSourceKey(MethodInvocation invocation) { private String determineDataSourceKey(MethodInvocation invocation) {
UseDataSource annotation;
Object aThis = invocation.getThis();
if (aThis != null) {
// 类上定义有 UseDataSource 注解
Class<?> aClass = aThis.getClass();
annotation = aClass.getAnnotation(UseDataSource.class);
if (annotation != null) {
return annotation.value();
}
}
// 方法上定义有 UseDataSource 注解 // 方法上定义有 UseDataSource 注解
annotation = invocation.getMethod().getAnnotation(UseDataSource.class); UseDataSource annotation = invocation.getMethod().getAnnotation(UseDataSource.class);
if (annotation != null) { if (annotation != null) {
return annotation.value(); return annotation.value();
} }
// 没有的话使用当前数据源 Object target = invocation.getThis();
return DataSourceKey.get();
if (target != null) {
// 类上定义有 UseDataSource 注解
Class<?> targetClass = target.getClass();
annotation = targetClass.getAnnotation(UseDataSource.class);
if (annotation != null) {
return annotation.value();
}
// 接口上定义有 UseDataSource 注解
Class<?>[] interfaces = targetClass.getInterfaces();
for (Class<?> anInterface : interfaces) {
annotation = anInterface.getAnnotation(UseDataSource.class);
if (annotation != null) {
return annotation.value();
}
}
}
return null;
} }
} }