diff --git a/mybatis-flex-core/src/main/java/com/mybatisflex/core/datasource/DataSourceKey.java b/mybatis-flex-core/src/main/java/com/mybatisflex/core/datasource/DataSourceKey.java index 97da09f7..306fbcff 100644 --- a/mybatis-flex-core/src/main/java/com/mybatisflex/core/datasource/DataSourceKey.java +++ b/mybatis-flex-core/src/main/java/com/mybatisflex/core/datasource/DataSourceKey.java @@ -19,13 +19,24 @@ import java.util.function.Supplier; public class DataSourceKey { + private static final ThreadLocal keyThreadLocal = new ThreadLocal<>(); + + public static String manualKey; + private DataSourceKey() { } - private static final ThreadLocal keyThreadLocal = new ThreadLocal<>(); - public static void use(String dataSourceKey) { keyThreadLocal.set(dataSourceKey.trim()); + manualKey = dataSourceKey; + } + + public static void use(String dataSourceKey, boolean isManual) { + if (isManual) + use(dataSourceKey); + else { + keyThreadLocal.set(dataSourceKey.trim()); + } } public static T use(String dataSourceKey, Supplier supplier) { @@ -48,6 +59,7 @@ public class DataSourceKey { public static void clear() { keyThreadLocal.remove(); + manualKey = null; } public static String get() { diff --git a/mybatis-flex-spring/src/main/java/com/mybatisflex/spring/datasource/DataSourceInterceptor.java b/mybatis-flex-spring/src/main/java/com/mybatisflex/spring/datasource/DataSourceInterceptor.java index 813cf430..29e14368 100644 --- a/mybatis-flex-spring/src/main/java/com/mybatisflex/spring/datasource/DataSourceInterceptor.java +++ b/mybatis-flex-spring/src/main/java/com/mybatisflex/spring/datasource/DataSourceInterceptor.java @@ -43,7 +43,7 @@ public class DataSourceInterceptor implements MethodInterceptor { @Override public Object invoke(MethodInvocation invocation) throws Throwable { - String dsKey = DataSourceKey.get(); + String dsKey = DataSourceKey.manualKey; if (StringUtil.isNotBlank(dsKey)) { return invocation.proceed(); } @@ -53,7 +53,7 @@ public class DataSourceInterceptor implements MethodInterceptor { return invocation.proceed(); } - DataSourceKey.use(dsKey); + DataSourceKey.use(dsKey, false); try { return invocation.proceed(); } finally {