From bd17ecb7df171d80dcc79710ef8f34f003af6eec Mon Sep 17 00:00:00 2001
From: barql <272378774@qq.com>
Date: Sat, 29 Jul 2023 10:06:20 +0800
Subject: [PATCH] =?UTF-8?q?fix:=E5=A4=9A=E6=95=B0=E6=8D=AE=E6=BA=90?=
 =?UTF-8?q?=E6=B3=A8=E8=A7=A3=E5=B5=8C=E5=A5=97=E4=BD=BF=E7=94=A8=E9=97=AE?=
 =?UTF-8?q?=E9=A2=98=E5=A4=84=E7=90=86?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 .../core/datasource/DataSourceKey.java           | 16 ++++++++++++++--
 .../spring/datasource/DataSourceInterceptor.java |  4 ++--
 2 files changed, 16 insertions(+), 4 deletions(-)

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<String> keyThreadLocal = new ThreadLocal<>();
+
+    public static String manualKey;
+
     private DataSourceKey() {
     }
 
-    private static final ThreadLocal<String> 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> T use(String dataSourceKey, Supplier<T> 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 {