From e7292f53b4adc7cd503172b2264eca6081d1384b Mon Sep 17 00:00:00 2001
From: Michael Yang <fuhai999@gmail.com>
Date: Sun, 2 Jun 2024 11:15:32 +0800
Subject: [PATCH] feat: add DataSourceKey.use to support nested usage

---
 .../core/datasource/DataSourceKey.java         | 18 ++++++++++++++----
 .../test/MultiDataSourceTester.java            |  5 ++---
 2 files changed, 16 insertions(+), 7 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 a8cd92e2..2b32d7d8 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
@@ -45,20 +45,30 @@ public class DataSourceKey {
     }
 
     public static <T> T use(String dataSourceKey, Supplier<T> supplier) {
+        String prevKey = manualKeyThreadLocal.get();
         try {
-            use(dataSourceKey);
+            manualKeyThreadLocal.set(dataSourceKey);
             return supplier.get();
         } finally {
-            clear();
+            if (prevKey != null) {
+                manualKeyThreadLocal.set(prevKey);
+            } else {
+                clear();
+            }
         }
     }
 
     public static void use(String dataSourceKey, Runnable runnable) {
+        String prevKey = manualKeyThreadLocal.get();
         try {
-            use(dataSourceKey);
+            manualKeyThreadLocal.set(dataSourceKey);
             runnable.run();
         } finally {
-            clear();
+            if (prevKey != null) {
+                manualKeyThreadLocal.set(prevKey);
+            } else {
+                clear();
+            }
         }
     }
 
diff --git a/mybatis-flex-test/mybatis-flex-native-test/src/main/java/com/mybatisflex/test/MultiDataSourceTester.java b/mybatis-flex-test/mybatis-flex-native-test/src/main/java/com/mybatisflex/test/MultiDataSourceTester.java
index 0e21d7ae..c9807a88 100644
--- a/mybatis-flex-test/mybatis-flex-native-test/src/main/java/com/mybatisflex/test/MultiDataSourceTester.java
+++ b/mybatis-flex-test/mybatis-flex-native-test/src/main/java/com/mybatisflex/test/MultiDataSourceTester.java
@@ -29,7 +29,6 @@ import org.springframework.jdbc.datasource.embedded.EmbeddedDatabaseType;
 
 import javax.sql.DataSource;
 import java.util.List;
-import java.util.function.Supplier;
 
 public class MultiDataSourceTester {
 
@@ -60,9 +59,9 @@ public class MultiDataSourceTester {
         AuditManager.setMessageCollector(collector);
 
         Db.tx(() -> {
-            Db.selectAll(null, "tb_account");
+            Db.selectAll("tb_account");
             DataSourceKey.use("ds2");
-            Db.selectAll(null, "tb_account");
+            Db.selectAll("tb_account");
             return true;
         });