chore: Code split service classes for enabling feature based migrations for SSO (#27404)

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

@@ -47,6 +47,10 @@ public class TenantConfigurationCE {
     // 2. Because of grandfathering via cron where tenant level feature flags are fetched
     Map<FeatureFlagEnum, FeatureMigrationType> featuresWithPendingMigration;
 
+    // This variable is used to indicate if the server needs to be restarted after the migration based on feature flags
+    // is complete.
+    Boolean isRestartRequired;
+
     public void addThirdPartyAuth(String auth) {
         if (thirdPartyAuths == null) {
             thirdPartyAuths = new ArrayList<>();

app/server/appsmith-server/src/main/java/com/appsmith/server/helpers/FeatureFlagMigrationHelper.java

@@ -1,254 +1,5 @@
 package com.appsmith.server.helpers;
 
-import com.appsmith.server.constants.FeatureMigrationType;
+import com.appsmith.server.helpers.ce.FeatureFlagMigrationHelperCE;
-import com.appsmith.server.domains.Tenant;
-import com.appsmith.server.domains.TenantConfiguration;
-import com.appsmith.server.featureflags.CachedFeatures;
-import com.appsmith.server.featureflags.FeatureFlagEnum;
-import com.appsmith.server.services.CacheableFeatureFlagHelper;
-import com.appsmith.server.solutions.ce.ScheduledTaskCEImpl;
-import lombok.RequiredArgsConstructor;
-import lombok.extern.slf4j.Slf4j;
-import org.springframework.stereotype.Component;
-import reactor.core.publisher.Mono;
 
-import java.time.Instant;
+public interface FeatureFlagMigrationHelper extends FeatureFlagMigrationHelperCE {}
-import java.time.temporal.ChronoUnit;
-import java.util.HashMap;
-import java.util.Map;
-
-import static com.appsmith.server.constants.FeatureMigrationType.DISABLE;
-import static com.appsmith.server.constants.FeatureMigrationType.ENABLE;
-import static java.lang.Boolean.FALSE;
-import static java.lang.Boolean.TRUE;
-
-@Component
-@RequiredArgsConstructor
-@Slf4j
-public class FeatureFlagMigrationHelper {
-
-    private final CacheableFeatureFlagHelper cacheableFeatureFlagHelper;
-
-    /**
-     * To avoid race condition keep the refresh rate lower than cron execution interval {@link ScheduledTaskCEImpl}
-     * to update the tenant level feature flags
-     */
-    private static final long TENANT_FEATURES_CACHE_TIME_MIN = 115;
-
-    public Mono<Map<FeatureFlagEnum, FeatureMigrationType>> getUpdatedFlagsWithPendingMigration(Tenant defaultTenant) {
-        return getUpdatedFlagsWithPendingMigration(defaultTenant, FALSE);
-    }
-
-    /**
-     * Method to get the updated feature flags with pending migrations. This method finds and registers the flags for
-     * migration by comparing the diffs between the feature flags stored in cache and the latest one pulled from CS
-     * @param tenant        Tenant for which the feature flags need to be updated
-     * @param forceUpdate   Flag to force update the tenant level feature flags
-     * @return              Map of feature flags with pending migrations
-     */
-    public Mono<Map<FeatureFlagEnum, FeatureMigrationType>> getUpdatedFlagsWithPendingMigration(
-            Tenant tenant, boolean forceUpdate) {
-
-        /*
-         *  1. Fetch current/saved feature flags from cache
-         *  2. Force update the tenant flags keeping existing flags as fallback in case the API call to fetch the flags fails for some reason
-         *  3. Get the diff and update the flags with pending migrations to be used to run migrations selectively
-         */
-        return cacheableFeatureFlagHelper
-                .fetchCachedTenantFeatures(tenant.getId())
-                .zipWhen(existingCachedFlags -> {
-                    if (existingCachedFlags.getRefreshedAt().until(Instant.now(), ChronoUnit.MINUTES)
-                                    < TENANT_FEATURES_CACHE_TIME_MIN
-                            && !forceUpdate) {
-                        return Mono.just(existingCachedFlags);
-                    }
-                    return this.refreshTenantFeatures(tenant, existingCachedFlags);
-                })
-                .map(tuple2 -> {
-                    CachedFeatures existingCachedFlags = tuple2.getT1();
-                    CachedFeatures latestFlags = tuple2.getT2();
-                    return this.getUpdatedFlagsWithPendingMigration(tenant, latestFlags, existingCachedFlags);
-                });
-    }
-
-    /**
-     * Method to force update the tenant level feature flags. This will be utilised in scenarios where we don't want
-     * to wait for the flags to get updated for cron scheduled time
-     *
-     * @param tenant    tenant for which the feature flags need to be updated
-     * @return          Cached features
-     */
-    private Mono<CachedFeatures> refreshTenantFeatures(Tenant tenant, CachedFeatures existingCachedFeatures) {
-        /*
-        1. Force update the flag
-            a. Evict the cache
-            b. Fetch and save latest flags from CS
-        2. In case the tenant is unable to fetch the latest flags save the existing flags from step 1 to cache (fallback)
-         */
-        String tenantId = tenant.getId();
-        return cacheableFeatureFlagHelper
-                .evictCachedTenantFeatures(tenantId)
-                .then(cacheableFeatureFlagHelper.fetchCachedTenantFeatures(tenantId))
-                .flatMap(features -> {
-                    if (CollectionUtils.isNullOrEmpty(features.getFeatures())) {
-                        // In case the retrieval of the latest flags from CS encounters an error, the previous flags
-                        // will serve as a fallback value.
-                        return cacheableFeatureFlagHelper.updateCachedTenantFeatures(tenantId, existingCachedFeatures);
-                    }
-                    return Mono.just(features);
-                });
-    }
-
-    /**
-     * Method to check the diffs between the existing feature flags and the latest flags pulled from CS. If there are
-     * any diffs save the flags with required migration types:
-     * Flag transitions:
-     * 1. false -> true : Migration to enable the feature flag
-     * 2. true -> false : Migration to disable the feature flag
-     * 3. There is a scenario when the migrations will be blocked on user input and may end up in a case where we just
-     * have to remove the entry as migration it's no longer needed:
-     *      Step 1: Feature gets enabled by adding a valid licence and enable migration gets registered
-     *      Step 2: License expires which results in feature getting disabled so migration entry gets registered with
-     *              disable type (This will happen via cron to check the license status)
-     *      Step 3: As the migration will be blocked by the user input for downgrade migration, DB state will be
-     *              maintained
-     *      Step 4: User adds the valid key or renews the subscription again which results in enabling the feature and
-     *              ends up in nullifying the effect for step 2
-     *
-     * @param tenant                    Tenant for which the feature flag migrations stats needs to be stored
-     * @param latestFlags               Latest flags pulled in from CS
-     * @param existingCachedFlags       Flags which are already stored in cache
-     * @return                          updated tenant with the required flags with pending migrations
-     */
-    private Map<FeatureFlagEnum, FeatureMigrationType> getUpdatedFlagsWithPendingMigration(
-            Tenant tenant, CachedFeatures latestFlags, CachedFeatures existingCachedFlags) {
-
-        // 1. Check if there are any diffs for the feature flags
-        // 2. Update the flags for pending migration within provided tenant object
-        Map<FeatureFlagEnum, FeatureMigrationType> featureDiffsWithMigrationType = new HashMap<>();
-        Map<String, Boolean> existingFeatureMap = existingCachedFlags.getFeatures();
-        latestFlags.getFeatures().forEach((key, value) -> {
-            if (value != null && !value.equals(existingFeatureMap.get(key))) {
-                try {
-                    featureDiffsWithMigrationType.put(
-                            FeatureFlagEnum.valueOf(key), Boolean.TRUE.equals(value) ? ENABLE : DISABLE);
-                } catch (Exception e) {
-                    // Ignore IllegalArgumentException as all the feature flags are not added on
-                    // server side
-                }
-            }
-        });
-        return getUpdatedFlagsWithPendingMigration(featureDiffsWithMigrationType, tenant);
-    }
-
-    private Map<FeatureFlagEnum, FeatureMigrationType> getUpdatedFlagsWithPendingMigration(
-            Map<FeatureFlagEnum, FeatureMigrationType> latestFeatureDiffsWithMigrationType, Tenant dbTenant) {
-
-        Map<FeatureFlagEnum, FeatureMigrationType> featuresWithPendingMigrationDB =
-                dbTenant.getTenantConfiguration().getFeaturesWithPendingMigration() == null
-                        ? new HashMap<>()
-                        : dbTenant.getTenantConfiguration().getFeaturesWithPendingMigration();
-
-        Map<FeatureFlagEnum, FeatureMigrationType> updatedFlagsForMigrations =
-                new HashMap<>(featuresWithPendingMigrationDB);
-
-        // We should expect the following state after the latest run:
-        // featuresWithPendingMigrationDB       => {feature1 : enable, feature2 : disable}
-        // latestFeatureDiffsWithMigrationType  => {feature1 : enable, feature2 : enable, feature3 : disable}
-        // updatedFlagsForMigrations            => {feature1 : enable, feature3 : disable}
-
-        updatedFlagsForMigrations.forEach((featureFlagEnum, featureMigrationType) -> {
-            if (latestFeatureDiffsWithMigrationType.containsKey(featureFlagEnum)
-                    && !featureMigrationType.equals(latestFeatureDiffsWithMigrationType.get(featureFlagEnum))) {
-                /*
-                Scenario when the migrations will be blocked on user input and may end up in a case where we just have
-                to remove the entry as migration it's no longer needed:
-                    Step 1: Feature gets enabled by adding a valid licence and enable migration gets registered
-                    Step 2: License expires which results in feature getting disabled so migration entry gets registered
-                            with disable type (This will happen via cron to check the license status)
-                    Step 3: As the migration will be blocked by the user input for downgrade migration, DB state will be
-                            maintained
-                    Step 4: User adds the valid key or renews the subscription again which results in enabling the
-                            feature and ends up in nullifying the effect for step 2
-                 */
-                updatedFlagsForMigrations.remove(featureFlagEnum);
-                latestFeatureDiffsWithMigrationType.remove(featureFlagEnum);
-            }
-        });
-        // Add the latest flags which were not part of earlier check.
-        updatedFlagsForMigrations.putAll(latestFeatureDiffsWithMigrationType);
-        return updatedFlagsForMigrations;
-    }
-
-    /**
-     * Method to check and execute if the migrations are required for the provided feature flag.
-     * @param tenant            Tenant for which the migrations need to be executed
-     * @param featureFlagEnum   Feature flag for which the migrations need to be executed
-     * @return                  Boolean indicating if the migrations is successfully executed or not
-     */
-    public Mono<Boolean> checkAndExecuteMigrationsForFeatureFlag(Tenant tenant, FeatureFlagEnum featureFlagEnum) {
-
-        TenantConfiguration tenantConfiguration = tenant.getTenantConfiguration();
-        if (featureFlagEnum == null
-                || tenantConfiguration == null
-                || CollectionUtils.isNullOrEmpty(tenantConfiguration.getFeaturesWithPendingMigration())) {
-            return Mono.just(TRUE);
-        }
-        return isMigrationRequired(tenant, featureFlagEnum).flatMap(isMigrationRequired -> {
-            if (FALSE.equals(isMigrationRequired)) {
-                return Mono.just(TRUE);
-            }
-            return this.executeMigrationsBasedOnFeatureFlag(tenantConfiguration, featureFlagEnum);
-        });
-    }
-
-    /**
-     * Method to check if the migrations are required for the provided feature flag.
-     * @param tenant            Tenant for which the migrations need to be executed
-     * @param featureFlagEnum   Feature flag for which the migrations need to be executed
-     * @return                  Boolean indicating if the migrations is required or not
-     */
-    private Mono<Boolean> isMigrationRequired(Tenant tenant, FeatureFlagEnum featureFlagEnum) {
-        Map<FeatureFlagEnum, FeatureMigrationType> featureMigrationTypeMap =
-                tenant.getTenantConfiguration().getFeaturesWithPendingMigration();
-        if (CollectionUtils.isNullOrEmpty(featureMigrationTypeMap)) {
-            return Mono.just(FALSE);
-        }
-        return cacheableFeatureFlagHelper
-                .fetchCachedTenantFeatures(tenant.getId())
-                .map(cachedFeatures -> {
-                    Map<String, Boolean> featureFlags = cachedFeatures.getFeatures();
-                    if (featureFlags.containsKey(featureFlagEnum.name())) {
-                        return (TRUE.equals(featureFlags.get(featureFlagEnum.name()))
-                                        && FeatureMigrationType.ENABLE.equals(
-                                                featureMigrationTypeMap.get(featureFlagEnum)))
-                                || (FALSE.equals(featureFlags.get(featureFlagEnum.name()))
-                                        && FeatureMigrationType.DISABLE.equals(
-                                                featureMigrationTypeMap.get(featureFlagEnum)));
-                    }
-                    return FALSE;
-                });
-    }
-
-    /**
-     * Method to execute the migrations for the given feature flag.
-     * @param tenantConfiguration   Tenant configuration for which the migrations need to be executed
-     * @param featureFlagEnum       Feature flag for which the migrations need to be executed
-     * @return                      Boolean indicating if the migrations is successfully executed or not
-     */
-    private Mono<Boolean> executeMigrationsBasedOnFeatureFlag(
-            TenantConfiguration tenantConfiguration, FeatureFlagEnum featureFlagEnum) {
-        // TODO implement migrations as per the supported features in license plan
-        Map<FeatureFlagEnum, FeatureMigrationType> featuresWithPendingMigration =
-                tenantConfiguration.getFeaturesWithPendingMigration();
-        if (CollectionUtils.isNullOrEmpty(featuresWithPendingMigration)
-                || !featuresWithPendingMigration.containsKey(featureFlagEnum)) {
-            return Mono.just(TRUE);
-        }
-        log.debug(
-                "Running the migration for flag {} with migration type {}",
-                featureFlagEnum.name(),
-                featuresWithPendingMigration.get(featureFlagEnum));
-        return Mono.just(TRUE);
-    }
-}

app/server/appsmith-server/src/main/java/com/appsmith/server/helpers/FeatureFlagMigrationHelperImpl.java

@@ -0,0 +1,13 @@
+package com.appsmith.server.helpers;
+
+import com.appsmith.server.helpers.ce.FeatureFlagMigrationHelperCEImpl;
+import com.appsmith.server.services.CacheableFeatureFlagHelper;
+import org.springframework.stereotype.Component;
+
+@Component
+public class FeatureFlagMigrationHelperImpl extends FeatureFlagMigrationHelperCEImpl
+        implements FeatureFlagMigrationHelper {
+    public FeatureFlagMigrationHelperImpl(CacheableFeatureFlagHelper cacheableFeatureFlagHelper) {
+        super(cacheableFeatureFlagHelper);
+    }
+}

app/server/appsmith-server/src/main/java/com/appsmith/server/helpers/ce/FeatureFlagMigrationHelperCE.java

@@ -0,0 +1,20 @@
+package com.appsmith.server.helpers.ce;
+
+import com.appsmith.server.constants.FeatureMigrationType;
+import com.appsmith.server.domains.Tenant;
+import com.appsmith.server.featureflags.FeatureFlagEnum;
+import reactor.core.publisher.Mono;
+
+import java.util.Map;
+
+public interface FeatureFlagMigrationHelperCE {
+
+    Mono<Map<FeatureFlagEnum, FeatureMigrationType>> getUpdatedFlagsWithPendingMigration(Tenant defaultTenant);
+
+    Mono<Map<FeatureFlagEnum, FeatureMigrationType>> getUpdatedFlagsWithPendingMigration(
+            Tenant tenant, boolean forceUpdate);
+
+    Mono<Boolean> checkAndExecuteMigrationsForFeatureFlag(Tenant tenant, FeatureFlagEnum featureFlagEnum);
+
+    Mono<Boolean> executeMigrationsBasedOnFeatureFlag(Tenant tenant, FeatureFlagEnum featureFlagEnum);
+}

app/server/appsmith-server/src/main/java/com/appsmith/server/helpers/ce/FeatureFlagMigrationHelperCEImpl.java

@@ -0,0 +1,255 @@
+package com.appsmith.server.helpers.ce;
+
+import com.appsmith.server.constants.FeatureMigrationType;
+import com.appsmith.server.domains.Tenant;
+import com.appsmith.server.domains.TenantConfiguration;
+import com.appsmith.server.featureflags.CachedFeatures;
+import com.appsmith.server.featureflags.FeatureFlagEnum;
+import com.appsmith.server.helpers.CollectionUtils;
+import com.appsmith.server.services.CacheableFeatureFlagHelper;
+import com.appsmith.server.solutions.ce.ScheduledTaskCEImpl;
+import lombok.RequiredArgsConstructor;
+import lombok.extern.slf4j.Slf4j;
+import reactor.core.publisher.Mono;
+
+import java.time.Instant;
+import java.time.temporal.ChronoUnit;
+import java.util.HashMap;
+import java.util.Map;
+
+import static com.appsmith.server.constants.FeatureMigrationType.DISABLE;
+import static com.appsmith.server.constants.FeatureMigrationType.ENABLE;
+import static java.lang.Boolean.FALSE;
+import static java.lang.Boolean.TRUE;
+
+@RequiredArgsConstructor
+@Slf4j
+public class FeatureFlagMigrationHelperCEImpl implements FeatureFlagMigrationHelperCE {
+
+    private final CacheableFeatureFlagHelper cacheableFeatureFlagHelper;
+
+    /**
+     * To avoid race condition keep the refresh rate lower than cron execution interval {@link ScheduledTaskCEImpl}
+     * to update the tenant level feature flags
+     */
+    private static final long TENANT_FEATURES_CACHE_TIME_MIN = 115;
+
+    @Override
+    public Mono<Map<FeatureFlagEnum, FeatureMigrationType>> getUpdatedFlagsWithPendingMigration(Tenant defaultTenant) {
+        return getUpdatedFlagsWithPendingMigration(defaultTenant, FALSE);
+    }
+
+    /**
+     * Method to get the updated feature flags with pending migrations. This method finds and registers the flags for
+     * migration by comparing the diffs between the feature flags stored in cache and the latest one pulled from CS
+     * @param tenant        Tenant for which the feature flags need to be updated
+     * @param forceUpdate   Flag to force update the tenant level feature flags
+     * @return              Map of feature flags with pending migrations
+     */
+    @Override
+    public Mono<Map<FeatureFlagEnum, FeatureMigrationType>> getUpdatedFlagsWithPendingMigration(
+            Tenant tenant, boolean forceUpdate) {
+
+        /*
+         *  1. Fetch current/saved feature flags from cache
+         *  2. Force update the tenant flags keeping existing flags as fallback in case the API call to fetch the flags fails for some reason
+         *  3. Get the diff and update the flags with pending migrations to be used to run migrations selectively
+         */
+        return cacheableFeatureFlagHelper
+                .fetchCachedTenantFeatures(tenant.getId())
+                .zipWhen(existingCachedFlags -> {
+                    if (existingCachedFlags.getRefreshedAt().until(Instant.now(), ChronoUnit.MINUTES)
+                                    < TENANT_FEATURES_CACHE_TIME_MIN
+                            && !forceUpdate) {
+                        return Mono.just(existingCachedFlags);
+                    }
+                    return this.refreshTenantFeatures(tenant, existingCachedFlags);
+                })
+                .map(tuple2 -> {
+                    CachedFeatures existingCachedFlags = tuple2.getT1();
+                    CachedFeatures latestFlags = tuple2.getT2();
+                    return this.getUpdatedFlagsWithPendingMigration(tenant, latestFlags, existingCachedFlags);
+                });
+    }
+
+    /**
+     * Method to force update the tenant level feature flags. This will be utilised in scenarios where we don't want
+     * to wait for the flags to get updated for cron scheduled time
+     *
+     * @param tenant    tenant for which the feature flags need to be updated
+     * @return          Cached features
+     */
+    private Mono<CachedFeatures> refreshTenantFeatures(Tenant tenant, CachedFeatures existingCachedFeatures) {
+        /*
+        1. Force update the flag
+            a. Evict the cache
+            b. Fetch and save latest flags from CS
+        2. In case the tenant is unable to fetch the latest flags save the existing flags from step 1 to cache (fallback)
+         */
+        String tenantId = tenant.getId();
+        return cacheableFeatureFlagHelper
+                .evictCachedTenantFeatures(tenantId)
+                .then(cacheableFeatureFlagHelper.fetchCachedTenantFeatures(tenantId))
+                .flatMap(features -> {
+                    if (CollectionUtils.isNullOrEmpty(features.getFeatures())) {
+                        // In case the retrieval of the latest flags from CS encounters an error, the previous flags
+                        // will serve as a fallback value.
+                        return cacheableFeatureFlagHelper.updateCachedTenantFeatures(tenantId, existingCachedFeatures);
+                    }
+                    return Mono.just(features);
+                });
+    }
+
+    /**
+     * Method to check the diffs between the existing feature flags and the latest flags pulled from CS. If there are
+     * any diffs save the flags with required migration types:
+     * Flag transitions:
+     * 1. false -> true : Migration to enable the feature flag
+     * 2. true -> false : Migration to disable the feature flag
+     * 3. There is a scenario when the migrations will be blocked on user input and may end up in a case where we just
+     * have to remove the entry as migration it's no longer needed:
+     *      Step 1: Feature gets enabled by adding a valid licence and enable migration gets registered
+     *      Step 2: License expires which results in feature getting disabled so migration entry gets registered with
+     *              disable type (This will happen via cron to check the license status)
+     *      Step 3: As the migration will be blocked by the user input for downgrade migration, DB state will be
+     *              maintained
+     *      Step 4: User adds the valid key or renews the subscription again which results in enabling the feature and
+     *              ends up in nullifying the effect for step 2
+     *
+     * @param tenant                    Tenant for which the feature flag migrations stats needs to be stored
+     * @param latestFlags               Latest flags pulled in from CS
+     * @param existingCachedFlags       Flags which are already stored in cache
+     * @return                          updated tenant with the required flags with pending migrations
+     */
+    private Map<FeatureFlagEnum, FeatureMigrationType> getUpdatedFlagsWithPendingMigration(
+            Tenant tenant, CachedFeatures latestFlags, CachedFeatures existingCachedFlags) {
+
+        // 1. Check if there are any diffs for the feature flags
+        // 2. Update the flags for pending migration within provided tenant object
+        Map<FeatureFlagEnum, FeatureMigrationType> featureDiffsWithMigrationType = new HashMap<>();
+        Map<String, Boolean> existingFeatureMap = existingCachedFlags.getFeatures();
+        latestFlags.getFeatures().forEach((key, value) -> {
+            if (value != null && !value.equals(existingFeatureMap.get(key))) {
+                try {
+                    featureDiffsWithMigrationType.put(
+                            FeatureFlagEnum.valueOf(key), Boolean.TRUE.equals(value) ? ENABLE : DISABLE);
+                } catch (Exception e) {
+                    // Ignore IllegalArgumentException as all the feature flags are not added on
+                    // server side
+                }
+            }
+        });
+        return getUpdatedFlagsWithPendingMigration(featureDiffsWithMigrationType, tenant);
+    }
+
+    private Map<FeatureFlagEnum, FeatureMigrationType> getUpdatedFlagsWithPendingMigration(
+            Map<FeatureFlagEnum, FeatureMigrationType> latestFeatureDiffsWithMigrationType, Tenant dbTenant) {
+
+        Map<FeatureFlagEnum, FeatureMigrationType> featuresWithPendingMigrationDB =
+                dbTenant.getTenantConfiguration().getFeaturesWithPendingMigration() == null
+                        ? new HashMap<>()
+                        : dbTenant.getTenantConfiguration().getFeaturesWithPendingMigration();
+
+        Map<FeatureFlagEnum, FeatureMigrationType> updatedFlagsForMigrations =
+                new HashMap<>(featuresWithPendingMigrationDB);
+
+        // We should expect the following state after the latest run:
+        // featuresWithPendingMigrationDB       => {feature1 : enable, feature2 : disable}
+        // latestFeatureDiffsWithMigrationType  => {feature1 : enable, feature2 : enable, feature3 : disable}
+        // updatedFlagsForMigrations            => {feature1 : enable, feature3 : disable}
+
+        updatedFlagsForMigrations.forEach((featureFlagEnum, featureMigrationType) -> {
+            if (latestFeatureDiffsWithMigrationType.containsKey(featureFlagEnum)
+                    && !featureMigrationType.equals(latestFeatureDiffsWithMigrationType.get(featureFlagEnum))) {
+                /*
+                Scenario when the migrations will be blocked on user input and may end up in a case where we just have
+                to remove the entry as migration it's no longer needed:
+                    Step 1: Feature gets enabled by adding a valid licence and enable migration gets registered
+                    Step 2: License expires which results in feature getting disabled so migration entry gets registered
+                            with disable type (This will happen via cron to check the license status)
+                    Step 3: As the migration will be blocked by the user input for downgrade migration, DB state will be
+                            maintained
+                    Step 4: User adds the valid key or renews the subscription again which results in enabling the
+                            feature and ends up in nullifying the effect for step 2
+                 */
+                updatedFlagsForMigrations.remove(featureFlagEnum);
+                latestFeatureDiffsWithMigrationType.remove(featureFlagEnum);
+            }
+        });
+        // Add the latest flags which were not part of earlier check.
+        updatedFlagsForMigrations.putAll(latestFeatureDiffsWithMigrationType);
+        return updatedFlagsForMigrations;
+    }
+
+    /**
+     * Method to check and execute if the migrations are required for the provided feature flag.
+     * @param tenant            Tenant for which the migrations need to be executed
+     * @param featureFlagEnum   Feature flag for which the migrations need to be executed
+     * @return                  Boolean indicating if the migrations is successfully executed or not
+     */
+    public Mono<Boolean> checkAndExecuteMigrationsForFeatureFlag(Tenant tenant, FeatureFlagEnum featureFlagEnum) {
+
+        TenantConfiguration tenantConfiguration = tenant.getTenantConfiguration();
+        if (featureFlagEnum == null
+                || tenantConfiguration == null
+                || CollectionUtils.isNullOrEmpty(tenantConfiguration.getFeaturesWithPendingMigration())) {
+            return Mono.just(TRUE);
+        }
+        return isMigrationRequired(tenant, featureFlagEnum).flatMap(isMigrationRequired -> {
+            if (FALSE.equals(isMigrationRequired)) {
+                return Mono.just(TRUE);
+            }
+
+            Map<FeatureFlagEnum, FeatureMigrationType> featuresWithPendingMigration =
+                    tenantConfiguration.getFeaturesWithPendingMigration();
+            if (CollectionUtils.isNullOrEmpty(featuresWithPendingMigration)
+                    || !featuresWithPendingMigration.containsKey(featureFlagEnum)) {
+                return Mono.just(TRUE);
+            }
+            log.debug(
+                    "Running the migration for flag {} with migration type {}",
+                    featureFlagEnum.name(),
+                    featuresWithPendingMigration.get(featureFlagEnum));
+            return this.executeMigrationsBasedOnFeatureFlag(tenant, featureFlagEnum);
+        });
+    }
+
+    /**
+     * Method to check if the migrations are required for the provided feature flag.
+     * @param tenant            Tenant for which the migrations need to be executed
+     * @param featureFlagEnum   Feature flag for which the migrations need to be executed
+     * @return                  Boolean indicating if the migrations is required or not
+     */
+    private Mono<Boolean> isMigrationRequired(Tenant tenant, FeatureFlagEnum featureFlagEnum) {
+        Map<FeatureFlagEnum, FeatureMigrationType> featureMigrationTypeMap =
+                tenant.getTenantConfiguration().getFeaturesWithPendingMigration();
+        if (CollectionUtils.isNullOrEmpty(featureMigrationTypeMap)) {
+            return Mono.just(FALSE);
+        }
+        return cacheableFeatureFlagHelper
+                .fetchCachedTenantFeatures(tenant.getId())
+                .map(cachedFeatures -> {
+                    Map<String, Boolean> featureFlags = cachedFeatures.getFeatures();
+                    if (featureFlags.containsKey(featureFlagEnum.name())) {
+                        return (TRUE.equals(featureFlags.get(featureFlagEnum.name()))
+                                        && FeatureMigrationType.ENABLE.equals(
+                                                featureMigrationTypeMap.get(featureFlagEnum)))
+                                || (FALSE.equals(featureFlags.get(featureFlagEnum.name()))
+                                        && FeatureMigrationType.DISABLE.equals(
+                                                featureMigrationTypeMap.get(featureFlagEnum)));
+                    }
+                    return FALSE;
+                });
+    }
+
+    /**
+     * Method to execute the migrations for the given feature flag.
+     * @param tenant                Tenant for which the migrations need to be executed
+     * @param featureFlagEnum       Feature flag for which the migrations need to be executed
+     * @return                      Boolean indicating if the migrations is successfully executed or not
+     */
+    @Override
+    public Mono<Boolean> executeMigrationsBasedOnFeatureFlag(Tenant tenant, FeatureFlagEnum featureFlagEnum) {
+        return Mono.just(TRUE);
+    }
+}

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

@@ -2,7 +2,9 @@ package com.appsmith.server.services.ce;
 
 import com.appsmith.server.domains.User;
 import org.springframework.web.server.ServerWebExchange;
+import reactor.core.publisher.Flux;
 import reactor.core.publisher.Mono;
+import reactor.util.function.Tuple2;
 
 import java.util.List;
 
@@ -17,4 +19,6 @@ public interface SessionUserServiceCE {
     Mono<List<String>> getSessionKeysByUserEmail(String email);
 
     Mono<Long> deleteSessionsByKeys(List<String> keys);
+
+    Flux<Tuple2<String, User>> getSessionKeysWithUserSessions();
 }

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

@@ -16,6 +16,7 @@ import org.springframework.security.core.context.SecurityContext;
 import org.springframework.security.oauth2.client.authentication.OAuth2AuthenticationToken;
 import org.springframework.web.server.ServerWebExchange;
 import org.springframework.web.server.WebSession;
+import reactor.core.publisher.Flux;
 import reactor.core.publisher.Mono;
 import reactor.util.function.Tuple2;
 
@@ -78,9 +79,28 @@ public class SessionUserServiceCEImpl implements SessionUserServiceCE {
                 .then();
     }
 
+    /**
+     * This method returns a list of session keys, for the given user email.
+     * @param email The email of the user whose sessions keys should be fetched.
+     * @return      A Mono of list of session keys.
+     */
     @Override
     public Mono<List<String>> getSessionKeysByUserEmail(String email) {
         // This pattern string comes from calling `ReactiveRedisSessionRepository.getSessionKey("*")` private method.
+        return getSessionKeysWithUserSessions()
+                // Now we have tuples of session keys, and the corresponding user objects.
+                // Filter the ones we need to clear out.
+                .filter(tuple ->
+                        StringUtils.equalsIgnoreCase(email, tuple.getT2().getEmail()))
+                .map(Tuple2::getT1)
+                .collectList();
+    }
+
+    /**
+     * This method returns a Flux of tuples, where the first element is the session key, and the second element is the
+     * corresponding User object.
+     */
+    public Flux<Tuple2<String, User>> getSessionKeysWithUserSessions() {
         return redisOperations
                 .keys(SPRING_SESSION_PATTERN)
                 .flatMap(key -> Mono.zip(
@@ -96,13 +116,7 @@ public class SessionUserServiceCEImpl implements SessionUserServiceCE {
                                 .map(e -> (User) ((SecurityContext) e.getValue())
                                         .getAuthentication()
                                         .getPrincipal())
-                                .next()))
+                                .next()));
-                // Now we have tuples of session keys, and the corresponding user objects.
-                // Filter the ones we need to clear out.
-                .filter(tuple ->
-                        StringUtils.equalsIgnoreCase(email, tuple.getT2().getEmail()))
-                .map(Tuple2::getT1)
-                .collectList();
     }
 
     @Override

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

@@ -23,4 +23,8 @@ public interface TenantServiceCE extends CrudService<Tenant, String> {
     Mono<Tenant> save(Tenant tenant);
 
     Mono<Tenant> checkAndExecuteMigrationsForTenantFeatureFlags(Tenant tenant);
+
+    Mono<Tenant> retrieveById(String id);
+
+    Mono<Void> restartTenant();
 }

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

@@ -18,6 +18,7 @@ import com.appsmith.server.services.BaseService;
 import com.appsmith.server.services.ConfigService;
 import com.appsmith.server.solutions.EnvManager;
 import jakarta.validation.Validator;
+import lombok.extern.slf4j.Slf4j;
 import org.springframework.context.annotation.Lazy;
 import org.springframework.data.mongodb.core.ReactiveMongoTemplate;
 import org.springframework.data.mongodb.core.convert.MongoConverter;
@@ -30,6 +31,7 @@ import java.util.Map;
 import static com.appsmith.server.acl.AclPermission.MANAGE_TENANT;
 import static java.lang.Boolean.TRUE;
 
+@Slf4j
 public class TenantServiceCEImpl extends BaseService<TenantRepository, Tenant, String> implements TenantServiceCE {
 
     private String tenantId = null;
@@ -233,6 +235,35 @@ public class TenantServiceCEImpl extends BaseService<TenantRepository, Tenant, S
                 });
     }
 
+    @Override
+    public Mono<Tenant> retrieveById(String id) {
+        if (!StringUtils.hasLength(id)) {
+            return Mono.error(new AppsmithException(AppsmithError.INVALID_PARAMETER, FieldName.ID));
+        }
+        return repository.retrieveById(id);
+    }
+
+    /**
+     * This function checks if the tenant needs to be restarted and restarts after the feature flag migrations are
+     * executed.
+     *
+     * @return
+     */
+    @Override
+    public Mono<Void> restartTenant() {
+        // Avoid dependency on user context as this method will be called internally by the server
+        Mono<Tenant> defaultTenantMono = this.getDefaultTenantId().flatMap(this::retrieveById);
+        return defaultTenantMono.flatMap(updatedTenant -> {
+            if (TRUE.equals(updatedTenant.getTenantConfiguration().getIsRestartRequired())) {
+                log.debug("Triggering tenant restart after the feature flag migrations are executed");
+                TenantConfiguration tenantConfiguration = updatedTenant.getTenantConfiguration();
+                tenantConfiguration.setIsRestartRequired(false);
+                return this.update(updatedTenant.getId(), updatedTenant).then(envManager.restartWithoutAclCheck());
+            }
+            return Mono.empty();
+        });
+    }
+
     private boolean isMigrationRequired(Tenant tenant) {
         return tenant.getTenantConfiguration() != null
                 && (!CollectionUtils.isNullOrEmpty(

app/server/appsmith-server/src/main/java/com/appsmith/server/solutions/ce/EnvManagerCE.java

@@ -16,6 +16,8 @@ public interface EnvManagerCE {
 
     Mono<Void> applyChanges(Map<String, String> changes, String originHeader);
 
+    Mono<Map<String, String>> applyChangesToEnvFileWithoutAclCheck(Map<String, String> changes);
+
     Mono<Void> applyChangesFromMultipartFormData(MultiValueMap<String, Part> formData, String originHeader);
 
     void setAnalyticsEventAction(
@@ -25,6 +27,8 @@ public interface EnvManagerCE {
 
     Map<String, String> parseToMap(String content);
 
+    Mono<Map<String, String>> getAllWithoutAclCheck();
+
     Mono<Map<String, String>> getAll();
 
     Mono<Map<String, String>> getAllNonEmpty();
@@ -33,6 +37,8 @@ public interface EnvManagerCE {
 
     Mono<Void> restart();
 
+    Mono<Void> restartWithoutAclCheck();
+
     Mono<Boolean> sendTestEmail(TestEmailConfigRequestDTO requestDTO);
 
     Mono<Void> download(ServerWebExchange exchange);

app/server/appsmith-server/src/main/java/com/appsmith/server/solutions/ce/EnvManagerCEImpl.java

@@ -13,6 +13,7 @@ import com.appsmith.server.dtos.TestEmailConfigRequestDTO;
 import com.appsmith.server.exceptions.AppsmithError;
 import com.appsmith.server.exceptions.AppsmithException;
 import com.appsmith.server.helpers.CollectionUtils;
+import com.appsmith.server.helpers.FeatureFlagMigrationHelper;
 import com.appsmith.server.helpers.FileUtils;
 import com.appsmith.server.helpers.TextUtils;
 import com.appsmith.server.helpers.UserUtils;
@@ -269,6 +270,8 @@ public class EnvManagerCEImpl implements EnvManagerCE {
         return valueBuilder.toString();
     }
 
+    // Expect user object to be null when this method is getting called to run the tenant specific migrations without
+    // user context
     private Mono<Void> validateChanges(User user, Map<String, String> changes) {
         if (changes.containsKey(APPSMITH_ADMIN_EMAILS.name())) {
             String emailCsv = StringUtils.trimAllWhitespace(changes.get(APPSMITH_ADMIN_EMAILS.name()));
@@ -278,7 +281,7 @@ public class EnvManagerCEImpl implements EnvManagerCE {
                 return Mono.error(new AppsmithException(AppsmithError.INVALID_PARAMETER, "Admin Emails"));
             } else { // make sure user is not removing own email
                 Set<String> adminEmails = TextUtils.csvToSet(emailCsv);
-                if (!adminEmails.contains(user.getEmail())) { // user can not remove own email address
+                if (user != null && !adminEmails.contains(user.getEmail())) { // user can not remove own email address
                     return Mono.error(new AppsmithException(
                             AppsmithError.GENERIC_BAD_REQUEST, "Removing own email from Admin Email is not allowed"));
                 }
@@ -347,43 +350,14 @@ public class EnvManagerCEImpl implements EnvManagerCE {
         // configuration
         return verifyCurrentUserIsSuper()
                 .flatMap(user -> validateChanges(user, changes).thenReturn(user))
-                .flatMap(user -> {
+                .flatMap(user -> applyChangesToEnvFileWithoutAclCheck(changes)
-                    // Write the changes to the env file.
+                        // For configuration variables, save the variables to the config collection instead of .env file
-                    final String originalContent;
+                        // We ideally want to migrate all variables from .env file to the config collection for better
-                    final Path envFilePath = Path.of(commonConfig.getEnvFilePath());
+                        // scalability
-
+                        // Write the changes to the tenant collection in configuration field
-                    try {
+                        .flatMap(originalVariables -> updateTenantConfiguration(user.getTenantId(), changes)
-                        originalContent = Files.readString(envFilePath);
+                                .then(sendAnalyticsEvent(user, originalVariables, changes))
-                    } catch (IOException e) {
+                                .thenReturn(originalVariables)))
-                        log.error("Unable to read env file " + envFilePath, e);
-                        return Mono.error(e);
-                    }
-                    Map<String, String> originalVariables = parseToMap(originalContent);
-
-                    final Map<String, String> envFileChanges = new HashMap<>(changes);
-                    final Set<String> tenantConfigurationKeys = allowedTenantConfiguration();
-                    for (final String key : changes.keySet()) {
-                        if (tenantConfigurationKeys.contains(key)) {
-                            envFileChanges.remove(key);
-                        }
-                    }
-                    final List<String> changedContent = transformEnvContent(originalContent, envFileChanges);
-
-                    try {
-                        Files.write(envFilePath, changedContent);
-                    } catch (IOException e) {
-                        log.error("Unable to write to env file " + envFilePath, e);
-                        return Mono.error(e);
-                    }
-
-                    // For configuration variables, save the variables to the config collection instead of .env file
-                    // We ideally want to migrate all variables from .env file to the config collection for better
-                    // scalability
-                    // Write the changes to the tenant collection in configuration field
-                    return updateTenantConfiguration(user.getTenantId(), changes)
-                            .then(sendAnalyticsEvent(user, originalVariables, changes))
-                            .thenReturn(originalVariables);
-                })
                 .flatMap(originalValues -> {
                     Mono<Void> dependentTasks = Mono.empty();
 
@@ -455,6 +429,45 @@ public class EnvManagerCEImpl implements EnvManagerCE {
                 });
     }
 
+    /**
+     * This method applies the changes to the env file and should be called internally within the server as the ACL
+     * checks are skipped. For client side calls please use {@link EnvManagerCEImpl#applyChanges(Map, String)}.
+     * Please refer {@link FeatureFlagMigrationHelper} for the use case where ACL checks
+     * should be skipped.
+     *
+     * @param changes       Map of changes to be applied to the env file
+     * @return              Map of original variables before the changes were applied
+     */
+    @Override
+    public Mono<Map<String, String>> applyChangesToEnvFileWithoutAclCheck(Map<String, String> changes) {
+        final Path envFilePath = Path.of(commonConfig.getEnvFilePath());
+        String originalContent;
+        try {
+            originalContent = Files.readString(envFilePath);
+        } catch (IOException e) {
+            log.error("Unable to read env file " + envFilePath, e);
+            return Mono.error(e);
+        }
+        Map<String, String> originalVariables = parseToMap(originalContent);
+
+        final Map<String, String> envFileChanges = new HashMap<>(changes);
+        final Set<String> tenantConfigurationKeys = allowedTenantConfiguration();
+        for (final String key : changes.keySet()) {
+            if (tenantConfigurationKeys.contains(key)) {
+                envFileChanges.remove(key);
+            }
+        }
+        final List<String> changedContent = transformEnvContent(originalContent, envFileChanges);
+
+        try {
+            Files.write(envFilePath, changedContent);
+        } catch (IOException e) {
+            log.error("Unable to write to env file " + envFilePath, e);
+            return Mono.error(e);
+        }
+        return Mono.just(originalVariables);
+    }
+
     @Override
     public Mono<Void> applyChangesFromMultipartFormData(MultiValueMap<String, Part> formData, String originHeader) {
         return Flux.fromIterable(formData.entrySet())
@@ -656,22 +669,28 @@ public class EnvManagerCEImpl implements EnvManagerCE {
 
     @Override
     public Mono<Map<String, String>> getAll() {
-        return verifyCurrentUserIsSuper().flatMap(user -> {
+        return verifyCurrentUserIsSuper().then(getAllWithoutAclCheck());
-            final String originalContent;
+    }
-            try {
-                originalContent = Files.readString(Path.of(commonConfig.getEnvFilePath()));
-            } catch (NoSuchFileException e) {
-                return Mono.error(new AppsmithException(AppsmithError.ENV_FILE_NOT_FOUND));
-            } catch (IOException e) {
-                log.error("Unable to read env file " + commonConfig.getEnvFilePath(), e);
-                return Mono.error(e);
-            }
 
-            // set the default values to response
+    /**
-            Map<String, String> envKeyValueMap = parseToMap(originalContent);
+     * This function is used to get all the env variables from the env file and should be called internally within the
-
+     * server as the ACL checks are skipped. For client side calls please use {@link EnvManagerCEImpl#getAll()}.
-            return Mono.justOrEmpty(envKeyValueMap);
+     *
-        });
+     * @return  Returns a map of all the env variables
+     */
+    @Override
+    public Mono<Map<String, String>> getAllWithoutAclCheck() {
+        String originalContent;
+        try {
+            originalContent = Files.readString(Path.of(commonConfig.getEnvFilePath()));
+        } catch (NoSuchFileException e) {
+            return Mono.error(new AppsmithException(AppsmithError.ENV_FILE_NOT_FOUND));
+        } catch (IOException e) {
+            log.error("Unable to read env file " + commonConfig.getEnvFilePath(), e);
+            return Mono.error(e);
+        }
+        // set the default values to response
+        return Mono.just(parseToMap(originalContent));
     }
 
     /**
@@ -704,18 +723,27 @@ public class EnvManagerCEImpl implements EnvManagerCE {
 
     @Override
     public Mono<Void> restart() {
-        return verifyCurrentUserIsSuper().flatMap(user -> {
+        return verifyCurrentUserIsSuper().then(restartWithoutAclCheck());
-            log.warn("Initiating restart via supervisor.");
+    }
-            try {
+
-                Runtime.getRuntime().exec(new String[] {
+    /**
-                    "supervisorctl", "restart", "backend", "editor", "rts",
+     * This function is used to restart the server using supervisorctl command and should be called internally within
-                });
+     * the server as the ACL checks are skipped. For client side calls we should use {@link EnvManagerCEImpl#restart()}
-            } catch (IOException e) {
+     *
-                log.error("Error invoking supervisorctl to restart.", e);
+     * @return  Returns a Mono<Void>
-                return Mono.error(new AppsmithException(AppsmithError.INTERNAL_SERVER_ERROR));
+     */
-            }
+    @Override
-            return Mono.empty();
+    public Mono<Void> restartWithoutAclCheck() {
-        });
+        log.warn("Initiating restart via supervisor.");
+        try {
+            Runtime.getRuntime().exec(new String[] {
+                "supervisorctl", "restart", "backend", "editor", "rts",
+            });
+        } catch (IOException e) {
+            log.error("Error invoking supervisorctl to restart.", e);
+            return Mono.error(new AppsmithException(AppsmithError.INTERNAL_SERVER_ERROR));
+        }
+        return Mono.empty();
     }
 
     @Override

app/server/appsmith-server/src/main/java/com/appsmith/server/solutions/ce/ScheduledTaskCEImpl.java

@@ -26,7 +26,8 @@ public class ScheduledTaskCEImpl implements ScheduledTaskCE {
                 .getAllRemoteFeaturesForTenantAndUpdateFeatureFlagsWithPendingMigrations()
                 .then(tenantService
                         .getDefaultTenant()
-                        .flatMap(featureFlagService::checkAndExecuteMigrationsForTenantFeatureFlags))
+                        .flatMap(featureFlagService::checkAndExecuteMigrationsForTenantFeatureFlags)
+                        .then(tenantService.restartTenant()))
                 .doOnError(error -> log.error("Error while fetching features from Cloud Services {0}", error))
                 .subscribeOn(scheduler)
                 .subscribe();

app/server/appsmith-server/src/test/java/com/appsmith/server/helpers/FeatureFlagMigrationHelperTest.java

@@ -25,6 +25,7 @@ import java.util.UUID;
 
 import static com.appsmith.server.constants.FeatureMigrationType.DISABLE;
 import static com.appsmith.server.constants.FeatureMigrationType.ENABLE;
+import static com.appsmith.server.constants.MigrationStatus.PENDING;
 import static com.appsmith.server.featureflags.FeatureFlagEnum.TENANT_TEST_FEATURE;
 import static org.assertj.core.api.AssertionsForInterfaceTypes.assertThat;
 import static org.mockito.ArgumentMatchers.any;
@@ -43,7 +44,7 @@ class FeatureFlagMigrationHelperTest {
     void setUp() {}
 
     @Test
-    void getUpdatedFlagsWithPendingMigration_diffForExistingAndLatestFlag_pendingMigrationReported() {
+    void getUpdatedFlagsWithPendingMigration_diffForExistingAndLatestFlag_pendingMigrationReportedWithDisableStatus() {
         Tenant defaultTenant = new Tenant();
         defaultTenant.setId(UUID.randomUUID().toString());
         defaultTenant.setTenantConfiguration(new TenantConfiguration());
@@ -80,6 +81,44 @@ class FeatureFlagMigrationHelperTest {
                 .verifyComplete();
     }
 
+    @Test
+    void getUpdatedFlagsWithPendingMigration_diffForExistingAndLatestFlag_pendingMigrationReportedWithEnableStatus() {
+        Tenant defaultTenant = new Tenant();
+        defaultTenant.setId(UUID.randomUUID().toString());
+        defaultTenant.setTenantConfiguration(new TenantConfiguration());
+
+        CachedFeatures existingCachedFeatures = new CachedFeatures();
+        Map<String, Boolean> featureMap = new HashMap<>();
+        featureMap.put(TENANT_TEST_FEATURE.name(), false);
+        existingCachedFeatures.setFeatures(featureMap);
+        existingCachedFeatures.setRefreshedAt(Instant.now().minus(1, ChronoUnit.DAYS));
+
+        CachedFeatures latestCachedFeatures = new CachedFeatures();
+        Map<String, Boolean> latestFeatureMap = new HashMap<>();
+        latestFeatureMap.put(TENANT_TEST_FEATURE.name(), true);
+        latestCachedFeatures.setFeatures(latestFeatureMap);
+        latestCachedFeatures.setRefreshedAt(Instant.now());
+
+        Mockito.when(cacheableFeatureFlagHelper.fetchCachedTenantFeatures(any()))
+                .thenReturn(Mono.just(existingCachedFeatures))
+                .thenReturn(Mono.just(latestCachedFeatures));
+
+        Mockito.when(cacheableFeatureFlagHelper.evictCachedTenantFeatures(any()))
+                .thenReturn(Mono.empty());
+
+        Mono<Map<FeatureFlagEnum, FeatureMigrationType>> getUpdatedFlagsWithPendingMigration =
+                featureFlagMigrationHelper.getUpdatedFlagsWithPendingMigration(defaultTenant);
+
+        StepVerifier.create(getUpdatedFlagsWithPendingMigration)
+                .assertNext(featureFlagEnumFeatureMigrationTypeMap -> {
+                    assertThat(featureFlagEnumFeatureMigrationTypeMap).isNotEmpty();
+                    assertThat(featureFlagEnumFeatureMigrationTypeMap.size()).isEqualTo(1);
+                    assertThat(featureFlagEnumFeatureMigrationTypeMap.get(TENANT_TEST_FEATURE))
+                            .isEqualTo(ENABLE);
+                })
+                .verifyComplete();
+    }
+
     @Test
     void getUpdatedFlagsWithPendingMigration_noDiffForExistingAndLatestFlag_noPendingMigrations() {
         Tenant defaultTenant = new Tenant();
@@ -163,6 +202,8 @@ class FeatureFlagMigrationHelperTest {
         Tenant defaultTenant = new Tenant();
         TenantConfiguration tenantConfiguration = new TenantConfiguration();
         tenantConfiguration.setFeaturesWithPendingMigration(Map.of(TENANT_TEST_FEATURE, ENABLE));
+        tenantConfiguration.setMigrationStatus(PENDING);
+        defaultTenant.setTenantConfiguration(tenantConfiguration);
 
         CachedFeatures existingCachedFeatures = new CachedFeatures();
         Map<String, Boolean> featureMap = new HashMap<>();
@@ -175,7 +216,14 @@ class FeatureFlagMigrationHelperTest {
         Mono<Boolean> resultMono =
                 featureFlagMigrationHelper.checkAndExecuteMigrationsForFeatureFlag(defaultTenant, TENANT_TEST_FEATURE);
         StepVerifier.create(resultMono)
-                .assertNext(result -> assertThat(result).isTrue())
+                .assertNext(result -> {
+                    assertThat(result).isTrue();
+                    assertThat(tenantConfiguration
+                                    .getFeaturesWithPendingMigration()
+                                    .size())
+                            .isEqualTo(1);
+                    assertThat(tenantConfiguration.getMigrationStatus()).isEqualTo(PENDING);
+                })
                 .verifyComplete();
     }
 }