Optimize tenant (DataLinkDC#2422)

* Optimize the process * Optimize system process tenant logic * formate code
aiwenmo · Oct 22, 2023 · f431fa5 · f431fa5
1 parent aba7120
commit f431fa5
Show file tree

Hide file tree

Showing 4 changed files with 17 additions and 10 deletions.
diff --git a/dinky-admin/src/main/java/org/dinky/configure/MybatisPlusConfig.java b/dinky-admin/src/main/java/org/dinky/configure/MybatisPlusConfig.java
@@ -115,6 +115,9 @@ public Expression getTenantId() {
 
             @Override
             public boolean ignoreTable(String tableName) {
+                if (TenantContextHolder.isIgnoreTenant()) {
+                    return true;
+                }
                 return !IGNORE_TABLE_NAMES.contains(tableName);
             }
         }));

diff --git a/dinky-admin/src/main/java/org/dinky/context/TenantContextHolder.java b/dinky-admin/src/main/java/org/dinky/context/TenantContextHolder.java
@@ -19,10 +19,21 @@
 
 package org.dinky.context;
 
+import java.util.Optional;
+
 /** TenantContextHolder */
 public class TenantContextHolder {
 
     private static final ThreadLocal<Object> TENANT_CONTEXT = new InheritableThreadLocal<>();
+    private static final ThreadLocal<Boolean> IGNORE_TENANT = new InheritableThreadLocal<>();
+
+    public static void ignoreTenant() {
+        IGNORE_TENANT.set(true);
+    }
+
+    public static boolean isIgnoreTenant() {
+        return Optional.ofNullable(IGNORE_TENANT.get()).orElse(false);
+    }
 
     public static void set(Object value) {
         TENANT_CONTEXT.set(value);

diff --git a/dinky-admin/src/main/java/org/dinky/init/SystemInit.java b/dinky-admin/src/main/java/org/dinky/init/SystemInit.java
@@ -57,8 +57,6 @@
 import java.util.List;
 import java.util.stream.Collectors;
 
-import org.slf4j.Logger;
-import org.slf4j.LoggerFactory;
 import org.springframework.boot.ApplicationArguments;
 import org.springframework.boot.ApplicationRunner;
 import org.springframework.context.annotation.Profile;
@@ -73,6 +71,7 @@
 import cn.hutool.core.lang.Singleton;
 import cn.hutool.core.util.StrUtil;
 import lombok.RequiredArgsConstructor;
+import lombok.extern.slf4j.Slf4j;
 
 /**
  * SystemInit
@@ -83,10 +82,10 @@
 @Order(value = 1)
 @RequiredArgsConstructor
 @Profile("!test")
+@Slf4j
 public class SystemInit implements ApplicationRunner {
     private final SystemConfiguration systemConfiguration = SystemConfiguration.getInstances();
 
-    private static final Logger log = LoggerFactory.getLogger(SystemInit.class);
     private final ProjectClient projectClient;
     private final SysConfigService sysConfigService;
     private final JobInstanceService jobInstanceService;
@@ -97,6 +96,7 @@ public class SystemInit implements ApplicationRunner {
 
     @Override
     public void run(ApplicationArguments args) {
+        TenantContextHolder.ignoreTenant();
         initResources();
 
         List<Tenant> tenants = tenantService.list();
@@ -230,15 +230,11 @@ public static Project getProject() {
     }
 
     public void registerUDF() {
-        // 设置admin用户 ，获取全部的udf代码，此地方没有租户隔离
-        TenantContextHolder.set(1);
         List<Task> allUDF = taskService.getAllUDF();
         if (CollUtil.isNotEmpty(allUDF)) {
             UdfCodePool.registerPool(allUDF.stream().map(UDFUtils::taskToUDF).collect(Collectors.toList()));
         }
         UdfCodePool.updateGitPool(gitProjectService.getGitPool());
-
-        TenantContextHolder.set(null);
     }
 
     public void updateGitBuildState() {

diff --git a/dinky-admin/src/main/java/org/dinky/job/FlinkJobTask.java b/dinky-admin/src/main/java/org/dinky/job/FlinkJobTask.java
@@ -20,7 +20,6 @@
 package org.dinky.job;
 
 import org.dinky.context.SpringContextUtils;
-import org.dinky.context.TenantContextHolder;
 import org.dinky.daemon.constant.FlinkTaskConstant;
 import org.dinky.daemon.task.DaemonTask;
 import org.dinky.daemon.task.DaemonTaskConfig;
@@ -55,7 +54,6 @@ public class FlinkJobTask implements DaemonTask {
     @Override
     public DaemonTask setConfig(DaemonTaskConfig config) {
         this.config = config;
-        TenantContextHolder.set(1);
         this.jobInfoDetail = jobInstanceService.getJobInfoDetail(config.getId());
         return this;
     }
@@ -74,7 +72,6 @@ public DaemonTask setConfig(DaemonTaskConfig config) {
     @Override
     public boolean dealTask() {
         volatilityBalance();
-        TenantContextHolder.set(1);
 
         boolean isDone = JobRefreshHandler.refreshJob(jobInfoDetail, isNeedSave());
         JobAlertHandler.getInstance().check(jobInfoDetail);