Use a single thread local with a custom object for the tenant

dbmalkovsky · Oct 9, 2023 · bcfdd1c · bcfdd1c
1 parent 476ad3b
commit bcfdd1c
Show file tree

Hide file tree

Showing 2 changed files with 14 additions and 18 deletions.
diff --git a/...common/src/main/java/org/flowable/common/engine/impl/tenant/ThreadLocalTenantContext.java b/...common/src/main/java/org/flowable/common/engine/impl/tenant/ThreadLocalTenantContext.java
@@ -12,38 +12,37 @@
  */
 package org.flowable.common.engine.impl.tenant;
 
-import org.flowable.common.engine.api.FlowableException;
 import org.flowable.common.engine.api.tenant.TenantContext;
 
 /**
  * @author Filip Hrisafov
  */
 public class ThreadLocalTenantContext implements TenantContext {
 
-    protected final ThreadLocal<String> tenantId = ThreadLocal.withInitial(() -> {
-        throw new FlowableException("Tenant value has not been set");
-    });
-    protected final ThreadLocal<Boolean> tenantIdSet = ThreadLocal.withInitial(() -> Boolean.FALSE);
+    protected final ThreadLocal<Tenant> tenantId = new ThreadLocal<>();
 
     @Override
     public String getTenantId() {
-        return tenantId.get();
+        Tenant tenant = tenantId.get();
+        return tenant != null ? tenant.tenantId() : null;
     }
 
     @Override
     public void setTenantId(String tenantId) {
-        this.tenantId.set(tenantId);
-        this.tenantIdSet.set(true);
+        this.tenantId.set(new Tenant(tenantId));
     }
 
     @Override
     public void clearTenantId() {
         tenantId.remove();
-        tenantIdSet.remove();
     }
 
     @Override
     public boolean isTenantIdSet() {
-        return tenantIdSet.get();
+        return tenantId.get() != null;
+    }
+
+    protected record Tenant(String tenantId) {
+
     }
 }
diff --git a/...on/src/test/java/org/flowable/common/engine/impl/tenant/ThreadLocalTenantContextTest.java b/...on/src/test/java/org/flowable/common/engine/impl/tenant/ThreadLocalTenantContextTest.java
@@ -13,9 +13,7 @@
 package org.flowable.common.engine.impl.tenant;
 
 import static org.assertj.core.api.Assertions.assertThat;
-import static org.assertj.core.api.Assertions.assertThatThrownBy;
 
-import org.flowable.common.engine.api.FlowableException;
 import org.flowable.common.engine.api.tenant.TenantContext;
 import org.junit.jupiter.api.Test;
 
@@ -28,9 +26,7 @@ class ThreadLocalTenantContextTest {
 
     @Test
     void getTenantId() {
-        assertThatThrownBy(() -> underTest.getTenantId())
-                .isExactlyInstanceOf(FlowableException.class)
-                .hasMessage("Tenant value has not been set");
+        assertThat(underTest.getTenantId()).isNull();
 
         underTest.setTenantId("acme");
         assertThat(underTest.getTenantId()).isEqualTo("acme");
@@ -39,9 +35,7 @@ void getTenantId() {
         assertThat(underTest.getTenantId()).isEqualTo("muppets");
 
         underTest.clearTenantId();
-        assertThatThrownBy(() -> underTest.getTenantId())
-                .isExactlyInstanceOf(FlowableException.class)
-                .hasMessage("Tenant value has not been set");
+        assertThat(underTest.getTenantId()).isNull();
     }
 
     @Test
@@ -53,5 +47,8 @@ void isTenantIdSet() {
 
         underTest.clearTenantId();
         assertThat(underTest.isTenantIdSet()).isFalse();
+
+        underTest.setTenantId(null);
+        assertThat(underTest.isTenantIdSet()).isTrue();
     }
 }