admin/mybatis-flex
1
0
Fork 0
mirror of https://gitee.com/mybatis-flex/mybatis-flex.git synced 2025-12-07 00:58:24 +08:00
Code Issues Packages Projects Releases Wiki Activity

feat: add DataSourceKey.use to support nested usage

Browse Source
This commit is contained in:
Michael Yang 2024-06-02 11:15:32 +08:00
parent c206540bf7
commit e7292f53b4
2 changed files with 16 additions and 7 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

14
mybatis-flex-core/src/main/java/com/mybatisflex/core/datasource/DataSourceKey.java
View File

@ -45,22 +45,32 @@ public class DataSourceKey {
} }
public static <T> T use(String dataSourceKey, Supplier<T> supplier) { public static <T> T use(String dataSourceKey, Supplier<T> supplier) {
String prevKey = manualKeyThreadLocal.get();
try { try {
use(dataSourceKey); manualKeyThreadLocal.set(dataSourceKey);
return supplier.get(); return supplier.get();
} finally { } finally {
if (prevKey != null) {
manualKeyThreadLocal.set(prevKey);
} else {
clear(); clear();
} }
} }
}
public static void use(String dataSourceKey, Runnable runnable) { public static void use(String dataSourceKey, Runnable runnable) {
String prevKey = manualKeyThreadLocal.get();
try { try {
use(dataSourceKey); manualKeyThreadLocal.set(dataSourceKey);
runnable.run(); runnable.run();
} finally { } finally {
if (prevKey != null) {
manualKeyThreadLocal.set(prevKey);
} else {
clear(); clear();
} }
} }
}
public static void clear() { public static void clear() {
annotationKeyThreadLocal.remove(); annotationKeyThreadLocal.remove();

5
mybatis-flex-test/mybatis-flex-native-test/src/main/java/com/mybatisflex/test/MultiDataSourceTester.java
View File

@ -29,7 +29,6 @@ import org.springframework.jdbc.datasource.embedded.EmbeddedDatabaseType;
import javax.sql.DataSource; import javax.sql.DataSource;
import java.util.List; import java.util.List;
import java.util.function.Supplier;
public class MultiDataSourceTester { public class MultiDataSourceTester {
@ -60,9 +59,9 @@ public class MultiDataSourceTester {
AuditManager.setMessageCollector(collector); AuditManager.setMessageCollector(collector);
Db.tx(() -> { Db.tx(() -> {
Db.selectAll(null, "tb_account"); Db.selectAll("tb_account");
DataSourceKey.use("ds2"); DataSourceKey.use("ds2");
Db.selectAll(null, "tb_account"); Db.selectAll("tb_account");
return true; return true;
}); });