feat: Adding Tenant Information to Redis Cache for quick fetch

app/server/appsmith-server/src/main/java/com/appsmith/server/domains/Tenant.java

@@ -10,13 +10,15 @@
 import org.springframework.data.annotation.Transient;
 import org.springframework.data.mongodb.core.mapping.Document;
 
+import java.io.Serializable;
+
 @Getter
 @Setter
 @ToString
 @NoArgsConstructor
 @Document
 @FieldNameConstants
-public class Tenant extends BaseDomain {
+public class Tenant extends BaseDomain implements Serializable {
 
     @Unique String slug;
 

app/server/appsmith-server/src/main/java/com/appsmith/server/domains/ce/TenantConfigurationCE.java

@@ -10,12 +10,13 @@
 import lombok.Data;
 import org.apache.commons.lang3.ObjectUtils;
 
+import java.io.Serializable;
 import java.util.ArrayList;
 import java.util.List;
 import java.util.Map;
 
 @Data
-public class TenantConfigurationCE {
+public class TenantConfigurationCE implements Serializable {
 
     private String googleMapsKey;
 

app/server/appsmith-server/src/main/java/com/appsmith/server/repositories/ce/CacheableRepositoryHelperCE.java

@@ -1,5 +1,6 @@
 package com.appsmith.server.repositories.ce;
 
+import com.appsmith.server.domains.Tenant;
 import com.appsmith.server.domains.User;
 import reactor.core.publisher.Mono;
 
@@ -18,4 +19,8 @@ public interface CacheableRepositoryHelperCE {
     Mono<String> getDefaultTenantId();
 
     Mono<String> getInstanceAdminPermissionGroupId();
+
+    Mono<Tenant> fetchDefaultTenant(String tenantId);
+
+    Mono<Void> evictCachedTenant(String tenantId);
 }

app/server/appsmith-server/src/main/java/com/appsmith/server/repositories/ce/CacheableRepositoryHelperCEImpl.java

@@ -6,6 +6,7 @@
 import com.appsmith.server.domains.Config;
 import com.appsmith.server.domains.PermissionGroup;
 import com.appsmith.server.domains.Tenant;
+import com.appsmith.server.domains.TenantConfiguration;
 import com.appsmith.server.domains.User;
 import com.appsmith.server.domains.Workspace;
 import com.appsmith.server.exceptions.AppsmithError;
@@ -166,4 +167,33 @@ public Mono<String> getInstanceAdminPermissionGroupId() {
                 .doOnSuccess(permissionGroupId ->
                         inMemoryCacheableRepositoryHelper.setInstanceAdminPermissionGroupId(permissionGroupId));
     }
+
+    /**
+     * Returns the default tenant from the cache if present.
+     * If not present in cache, then it fetches the default tenant from the database and adds to redis.
+     * @param tenantId
+     * @return
+     */
+    @Cache(cacheName = "tenant", key = "{#tenantId}")
+    @Override
+    public Mono<Tenant> fetchDefaultTenant(String tenantId) {
+        BridgeQuery<Tenant> defaultTenantCriteria = Bridge.equal(Tenant.Fields.slug, FieldName.DEFAULT);
+        BridgeQuery<Tenant> notDeletedCriteria = notDeleted();
+        BridgeQuery<Tenant> andCriteria = Bridge.and(defaultTenantCriteria, notDeletedCriteria);
+        Query query = new Query();
+        query.addCriteria(andCriteria);
+
+        return mongoOperations.findOne(query, Tenant.class).map(tenant -> {
+            if (tenant.getTenantConfiguration() == null) {
+                tenant.setTenantConfiguration(new TenantConfiguration());
+            }
+            return tenant;
+        });
+    }
+
+    @CacheEvict(cacheName = "tenant", key = "{#tenantId}")
+    @Override
+    public Mono<Void> evictCachedTenant(String tenantId) {
+        return Mono.empty().then();
+    }
 }

app/server/appsmith-server/src/main/java/com/appsmith/server/repositories/ce/TenantRepositoryCE.java

@@ -5,6 +5,7 @@
 import reactor.core.publisher.Mono;
 
 public interface TenantRepositoryCE extends BaseRepository<Tenant, String>, CustomTenantRepositoryCE {
-
+    // Use tenantService.getDefaultTenant() instead of this method as it is cached to redis.
+    @Deprecated(forRemoval = true)
     Mono<Tenant> findBySlug(String slug);
 }

app/server/appsmith-server/src/main/java/com/appsmith/server/services/TenantServiceImpl.java

@@ -1,6 +1,7 @@
 package com.appsmith.server.services;
 
 import com.appsmith.server.helpers.FeatureFlagMigrationHelper;
+import com.appsmith.server.repositories.CacheableRepositoryHelper;
 import com.appsmith.server.repositories.TenantRepository;
 import com.appsmith.server.services.ce.TenantServiceCEImpl;
 import com.appsmith.server.solutions.EnvManager;
@@ -19,7 +20,15 @@ public TenantServiceImpl(
             AnalyticsService analyticsService,
             ConfigService configService,
             @Lazy EnvManager envManager,
-            FeatureFlagMigrationHelper featureFlagMigrationHelper) {
-        super(validator, repository, analyticsService, configService, envManager, featureFlagMigrationHelper);
+            FeatureFlagMigrationHelper featureFlagMigrationHelper,
+            CacheableRepositoryHelper cacheableRepositoryHelper) {
+        super(
+                validator,
+                repository,
+                analyticsService,
+                configService,
+                envManager,
+                featureFlagMigrationHelper,
+                cacheableRepositoryHelper);
     }
 }

app/server/appsmith-server/src/main/java/com/appsmith/server/services/ce/CacheableFeatureFlagHelperCEImpl.java

@@ -45,7 +45,6 @@
 @Slf4j
 @RequiredArgsConstructor
 public class CacheableFeatureFlagHelperCEImpl implements CacheableFeatureFlagHelperCE {
-
     private final TenantRepository tenantRepository;
     private final ConfigService configService;
     private final CloudServicesConfig cloudServicesConfig;

app/server/appsmith-server/src/main/java/com/appsmith/server/services/ce/TenantServiceCEImpl.java

@@ -12,6 +12,7 @@
 import com.appsmith.server.exceptions.AppsmithException;
 import com.appsmith.server.helpers.CollectionUtils;
 import com.appsmith.server.helpers.FeatureFlagMigrationHelper;
+import com.appsmith.server.repositories.CacheableRepositoryHelper;
 import com.appsmith.server.repositories.TenantRepository;
 import com.appsmith.server.services.AnalyticsService;
 import com.appsmith.server.services.BaseService;
@@ -39,17 +40,21 @@ public class TenantServiceCEImpl extends BaseService<TenantRepository, Tenant, S
 
     private final FeatureFlagMigrationHelper featureFlagMigrationHelper;
 
+    private final CacheableRepositoryHelper cacheableRepositoryHelper;
+
     public TenantServiceCEImpl(
             Validator validator,
             TenantRepository repository,
             AnalyticsService analyticsService,
             ConfigService configService,
             @Lazy EnvManager envManager,
-            FeatureFlagMigrationHelper featureFlagMigrationHelper) {
+            FeatureFlagMigrationHelper featureFlagMigrationHelper,
+            CacheableRepositoryHelper cacheableRepositoryHelper) {
         super(validator, repository, analyticsService);
         this.configService = configService;
         this.envManager = envManager;
         this.featureFlagMigrationHelper = featureFlagMigrationHelper;
+        this.cacheableRepositoryHelper = cacheableRepositoryHelper;
     }
 
     @Override
@@ -59,7 +64,6 @@ public Mono<String> getDefaultTenantId() {
         if (StringUtils.hasLength(tenantId)) {
             return Mono.just(tenantId);
         }
-
         return repository.findBySlug(FieldName.DEFAULT).map(Tenant::getId).map(tenantId -> {
             // Set the cache value before returning.
             this.tenantId = tenantId;
@@ -69,6 +73,7 @@ public Mono<String> getDefaultTenantId() {
 
     @Override
     public Mono<Tenant> updateTenantConfiguration(String tenantId, TenantConfiguration tenantConfiguration) {
+        Mono<Void> evictTenantCache = cacheableRepositoryHelper.evictCachedTenant(tenantId);
         return repository
                 .findById(tenantId, MANAGE_TENANT)
                 .switchIfEmpty(Mono.error(
@@ -89,14 +94,23 @@ public Mono<Tenant> updateTenantConfiguration(String tenantId, TenantConfigurati
                             return Mono.empty();
                         });
                     }
+
                     return envMono.then(Mono.zip(Mono.just(oldtenantConfiguration), Mono.just(tenant)));
                 })
                 .flatMap(tuple2 -> {
                     Tenant tenant = tuple2.getT2();
                     TenantConfiguration oldConfig = tuple2.getT1();
                     AppsmithBeanUtils.copyNestedNonNullProperties(tenantConfiguration, oldConfig);
                     tenant.setTenantConfiguration(oldConfig);
-                    return repository.updateById(tenantId, tenant, MANAGE_TENANT);
+                    Mono<Tenant> updatedTenantMono = repository
+                            .updateById(tenantId, tenant, MANAGE_TENANT)
+                            .cache();
+                    // Firstly updating the Tenant object in the database and then evicting the cache.
+                    // returning the updatedTenant, notice the updatedTenantMono is cached using .cache()
+                    // hence it will not be evaluated again
+                    return updatedTenantMono
+                            .then(Mono.defer(() -> evictTenantCache))
+                            .then(updatedTenantMono);
                 });
     }
 
@@ -151,17 +165,29 @@ public Mono<Tenant> getTenantConfiguration() {
 
     @Override
     public Mono<Tenant> getDefaultTenant() {
-        // Get the default tenant object from the DB and then populate the relevant user permissions in that
-        // We are doing this differently because `findBySlug` is a Mongo JPA query and not a custom Appsmith query
-        return repository
-                .findBySlug(FieldName.DEFAULT)
-                .map(tenant -> {
-                    if (tenant.getTenantConfiguration() == null) {
-                        tenant.setTenantConfiguration(new TenantConfiguration());
-                    }
-                    return tenant;
-                })
-                .flatMap(tenant -> repository.setUserPermissionsInObject(tenant).switchIfEmpty(Mono.just(tenant)));
+        // Fetching Tenant from redis cache
+        return getDefaultTenantId()
+                .flatMap(tenantId -> cacheableRepositoryHelper.fetchDefaultTenant(tenantId))
+                .flatMap(tenant -> repository.setUserPermissionsInObject(tenant).switchIfEmpty(Mono.just(tenant)))
+                .onErrorResume(e -> {
+                    log.error("Error fetching default tenant from redis!", e);
+                    // If there is an error fetching the tenant from the cache, then evict the cache and fetching from
+                    // the db. This handles the case for deserialization errors. This prevents the entire instance to
+                    // go down if tenant cache is corrupted.
+                    // More info - https://github.com/appsmithorg/appsmith/issues/33504
+                    log.info("Evicting the default tenant from cache and fetching from the database!");
+                    return cacheableRepositoryHelper
+                            .evictCachedTenant(tenantId)
+                            .then(repository.findBySlug(FieldName.DEFAULT).map(tenant -> {
+                                if (tenant.getTenantConfiguration() == null) {
+                                    tenant.setTenantConfiguration(new TenantConfiguration());
+                                }
+                                return tenant;
+                            }))
+                            .flatMap(tenant -> repository
+                                    .setUserPermissionsInObject(tenant)
+                                    .switchIfEmpty(Mono.just(tenant)));
+                });
     }
 
     @Override
@@ -192,9 +218,12 @@ protected Mono<Tenant> getClientPertinentTenant(Tenant dbTenant, Tenant clientTe
         return Mono.just(clientTenant);
     }
 
+    // This function is used to save the tenant object in the database and evict the cache
     @Override
     public Mono<Tenant> save(Tenant tenant) {
-        return repository.save(tenant);
+        Mono<Void> evictTenantCache = cacheableRepositoryHelper.evictCachedTenant(tenantId);
+        Mono<Tenant> savedTenantMono = repository.save(tenant).cache();
+        return savedTenantMono.then(Mono.defer(() -> evictTenantCache)).then(savedTenantMono);
     }
 
     /**
@@ -242,6 +271,19 @@ public Mono<Tenant> retrieveById(String id) {
         return repository.findById(id);
     }
 
+    /**
+     * This function updates the tenant object in the database and evicts the cache
+     * @param tenantId
+     * @param tenant
+     * @return
+     */
+    @Override
+    public Mono<Tenant> update(String tenantId, Tenant tenant) {
+        Mono<Void> evictTenantCache = cacheableRepositoryHelper.evictCachedTenant(tenantId);
+        Mono<Tenant> updatedTenantMono = super.update(tenantId, tenant).cache();
+        return updatedTenantMono.then(Mono.defer(() -> evictTenantCache)).then(updatedTenantMono);
+    }
+
     /**
      * This function checks if the tenant needs to be restarted and restarts after the feature flag migrations are
      * executed.

app/server/appsmith-server/src/test/java/com/appsmith/server/services/ce/TenantServiceCETest.java

@@ -1,5 +1,6 @@
 package com.appsmith.server.services.ce;
 
+import com.appsmith.caching.components.CacheManager;
 import com.appsmith.external.enums.FeatureFlagEnum;
 import com.appsmith.server.constants.FeatureMigrationType;
 import com.appsmith.server.constants.LicensePlan;
@@ -9,18 +10,22 @@
 import com.appsmith.server.helpers.FeatureFlagMigrationHelper;
 import com.appsmith.server.helpers.UserUtils;
 import com.appsmith.server.helpers.ce.bridge.Bridge;
+import com.appsmith.server.repositories.CacheableRepositoryHelper;
 import com.appsmith.server.repositories.TenantRepository;
 import com.appsmith.server.repositories.UserRepository;
 import com.appsmith.server.services.TenantService;
 import com.appsmith.server.solutions.EnvManager;
 import org.junit.jupiter.api.AfterEach;
+import org.junit.jupiter.api.Assertions;
 import org.junit.jupiter.api.BeforeEach;
 import org.junit.jupiter.api.Test;
 import org.junit.jupiter.api.extension.ExtendWith;
 import org.mockito.Mockito;
 import org.springframework.beans.factory.annotation.Autowired;
 import org.springframework.boot.test.context.SpringBootTest;
 import org.springframework.boot.test.mock.mockito.MockBean;
+import org.springframework.boot.test.mock.mockito.SpyBean;
+import org.springframework.data.redis.core.ReactiveRedisTemplate;
 import org.springframework.security.test.context.support.WithUserDetails;
 import org.springframework.test.context.junit.jupiter.SpringExtension;
 import reactor.core.publisher.Mono;
@@ -42,6 +47,8 @@
 import static java.lang.Boolean.TRUE;
 import static org.assertj.core.api.Assertions.assertThat;
 import static org.mockito.ArgumentMatchers.any;
+import static org.mockito.ArgumentMatchers.anyString;
+import static org.mockito.Mockito.doReturn;
 
 @SpringBootTest
 @ExtendWith(SpringExtension.class)
@@ -62,6 +69,15 @@ class TenantServiceCETest {
     @Autowired
     UserUtils userUtils;
 
+    @Autowired
+    ReactiveRedisTemplate<String, Object> reactiveRedisTemplate;
+
+    @SpyBean
+    CacheManager cacheManager;
+
+    @SpyBean
+    CacheableRepositoryHelper cacheableRepositoryHelper;
+
     @MockBean
     FeatureFlagMigrationHelper featureFlagMigrationHelper;
 
@@ -84,16 +100,19 @@ public void setup() throws IOException {
                 .findByEmail("api_user")
                 .flatMap(user -> userUtils.makeSuperUser(List.of(user)))
                 .block();
+        doReturn(Mono.empty()).when(cacheManager).get(anyString(), anyString());
     }
 
     @AfterEach
     public void cleanup() {
+        Tenant updatedTenant = new Tenant();
+        updatedTenant.setTenantConfiguration(originalTenantConfiguration);
         tenantService
-                .getDefaultTenant()
-                .flatMap(tenant -> tenantRepository.updateAndReturn(
-                        tenant.getId(),
-                        Bridge.update().set(Tenant.Fields.tenantConfiguration, originalTenantConfiguration),
-                        Optional.empty()))
+                .getDefaultTenantId()
+                .flatMap(tenantId -> tenantService.update(tenantId, updatedTenant))
+                .doOnError(error -> {
+                    System.err.println("Error during cleanup: " + error.getMessage());
+                })
                 .block();
     }
 
@@ -362,4 +381,21 @@ void updateTenantConfiguration_updateStrongPasswordPolicy_success() {
                 })
                 .verifyComplete();
     }
+
+    /**
+     * This test checks that the tenant cache is created and data is fetched without any deserialization errors
+     * This will ensure if any new nested user-defined classes are created in the tenant object in the future, and
+     * implements serializable is missed for that class, the deserialization will fail leads this test to fail.
+     */
+    @Test
+    @WithUserDetails(value = "api_user")
+    public void testDeserializationErrors() {
+        String tenantId = tenantService.getDefaultTenantId().block();
+        Mono<Void> evictTenantCache = cacheableRepositoryHelper.evictCachedTenant(tenantId);
+        Mono<Boolean> hasKeyMono = reactiveRedisTemplate.hasKey("tenant:" + tenantId);
+        Mono<Tenant> cachedTenant = tenantService.getDefaultTenant();
+        StepVerifier.create(evictTenantCache.then(cachedTenant).then(hasKeyMono))
+                .assertNext(Assertions::assertTrue)
+                .verifyComplete();
+    }
 }