From e5fecf7ed8acf541d1b96c4ceda6af0b7202d57c Mon Sep 17 00:00:00 2001
From: life <13122192336@163.com>
Date: Sat, 19 Aug 2023 10:09:23 +0800
Subject: [PATCH] =?UTF-8?q?=E4=BF=AE=E6=94=B9=E8=AE=BE=E7=BD=AE=E9=97=AE?=
 =?UTF-8?q?=E9=A2=98?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 mybatis-flex-core/pom.xml                            |  5 -----
 .../mybatisflex/core/datasource/DataSourceKey.java   | 12 +++++++++---
 mybatis-flex-test/mybatis-flex-native-test/pom.xml   |  6 ++++++
 .../java/com/mybatisflex/test/MultiThreadsTest.java  |  3 +++
 4 files changed, 18 insertions(+), 8 deletions(-)

diff --git a/mybatis-flex-core/pom.xml b/mybatis-flex-core/pom.xml
index 1d59b37a..cd315fbf 100644
--- a/mybatis-flex-core/pom.xml
+++ b/mybatis-flex-core/pom.xml
@@ -71,11 +71,6 @@
             <optional>true</optional>
         </dependency>
 
-        <dependency>
-            <groupId>com.alibaba</groupId>
-            <artifactId>transmittable-thread-local</artifactId>
-            <version>2.14.2</version>
-        </dependency>
         <!--optional end-->
 
         <dependency>
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 98f8122d..144a05fb 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
@@ -15,7 +15,6 @@
  */
 package com.mybatisflex.core.datasource;
 
-import com.alibaba.ttl.TransmittableThreadLocal;
 import java.util.function.Supplier;
 
 /**
@@ -26,12 +25,12 @@ public class DataSourceKey {
     /**
      * 通过注解设置的 key
      */
-    private static final ThreadLocal<String> annotationKeyThreadLocal = new TransmittableThreadLocal<>();
+    private static ThreadLocal<String> annotationKeyThreadLocal = new ThreadLocal<>();
 
     /**
      * 通过手动编码指定的 key
      */
-    private static final ThreadLocal<String> manualKeyThreadLocal = new TransmittableThreadLocal<>();
+    private static ThreadLocal<String> manualKeyThreadLocal = new ThreadLocal<>();
 
     public static String manualKey;
 
@@ -82,4 +81,11 @@ public class DataSourceKey {
         return key != null ? key : annotationKeyThreadLocal.get();
     }
 
+    public static void setAnnotationKeyThreadLocal(ThreadLocal<String> annotationKeyThreadLocal) {
+        DataSourceKey.annotationKeyThreadLocal = annotationKeyThreadLocal;
+    }
+
+    public static void setManualKeyThreadLocal(ThreadLocal<String> manualKeyThreadLocal) {
+        DataSourceKey.manualKeyThreadLocal = manualKeyThreadLocal;
+    }
 }
diff --git a/mybatis-flex-test/mybatis-flex-native-test/pom.xml b/mybatis-flex-test/mybatis-flex-native-test/pom.xml
index 48d41039..61da5f2b 100644
--- a/mybatis-flex-test/mybatis-flex-native-test/pom.xml
+++ b/mybatis-flex-test/mybatis-flex-native-test/pom.xml
@@ -108,6 +108,12 @@
 <!--            <scope>test</scope>-->
 <!--        </dependency>-->
 
+        <dependency>
+            <groupId>com.alibaba</groupId>
+            <artifactId>transmittable-thread-local</artifactId>
+            <version>2.14.2</version>
+        </dependency>
+
 
         <dependency>
             <groupId>org.assertj</groupId>
diff --git a/mybatis-flex-test/mybatis-flex-native-test/src/main/java/com/mybatisflex/test/MultiThreadsTest.java b/mybatis-flex-test/mybatis-flex-native-test/src/main/java/com/mybatisflex/test/MultiThreadsTest.java
index be742e01..ec61feee 100644
--- a/mybatis-flex-test/mybatis-flex-native-test/src/main/java/com/mybatisflex/test/MultiThreadsTest.java
+++ b/mybatis-flex-test/mybatis-flex-native-test/src/main/java/com/mybatisflex/test/MultiThreadsTest.java
@@ -1,5 +1,6 @@
 package com.mybatisflex.test;
 
+import com.alibaba.ttl.TransmittableThreadLocal;
 import com.mybatisflex.core.MybatisFlexBootstrap;
 import com.mybatisflex.core.audit.AuditManager;
 import com.mybatisflex.core.audit.ConsoleMessageCollector;
@@ -20,6 +21,8 @@ import org.springframework.jdbc.datasource.embedded.EmbeddedDatabaseType;
 public class MultiThreadsTest {
 
     public static void main(String[] args) {
+        DataSourceKey.setAnnotationKeyThreadLocal(new TransmittableThreadLocal<>());
+        DataSourceKey.setManualKeyThreadLocal(new TransmittableThreadLocal<>());
         DataSource dataSource = new EmbeddedDatabaseBuilder()
             .setType(EmbeddedDatabaseType.H2)
             .setName("db1")