chore: Segregate organization flags based on orgnization id

app/server/appsmith-server/src/main/java/com/appsmith/server/aspect/FeatureFlaggedMethodInvokerAspect.java

@@ -15,11 +15,13 @@
 import org.aspectj.lang.reflect.MethodSignature;
 import org.springframework.context.ApplicationContext;
 import org.springframework.stereotype.Component;
+import org.springframework.util.StringUtils;
 import reactor.core.publisher.Flux;
 import reactor.core.publisher.Mono;
 
 import java.beans.Introspector;
 import java.lang.reflect.Method;
+import java.lang.reflect.Parameter;
 
 @RequiredArgsConstructor
 @Aspect
@@ -64,9 +66,23 @@ private Object invokeMethod(ProceedingJoinPoint joinPoint, FeatureFlagged annota
         } else if (Flux.class.isAssignableFrom(returnType)) {
             return featureFlagMono.flatMapMany(isSupported -> (Flux<?>) invokeMethod(isSupported, joinPoint, method));
         }
+
+        // For supporting non-reactive methods with feature flagging annotation we need to extract organizationId from
+        // method arguments and then check if the feature flag is enabled for that organization. This hampers the code
+        // readability so avoid non-reactive methods for @FeatureFlagged unless the method is getting called from server
+        // internal flows where the user context is not provided.
+        String organizationId = extractOrganizationId(joinPoint.getArgs(), method.getParameters());
+
+        if (!StringUtils.hasLength(organizationId)) {
+            String errorMessage =
+                    "Add missing organizationId parameter and enforce non-null value for orgnization-specific feature flags retrieval in non-reactive methods";
+            AppsmithException exception = getInvalidAnnotationUsageException(method, errorMessage);
+            log.error(exception.getMessage());
+            throw exception;
+        }
         // For non-reactive methods with feature flagging annotation we will be using the in memory feature flag cache
         // which is getting updated whenever the organization feature flags are updated.
-        return invokeMethod(isFeatureFlagEnabled(flagName), joinPoint, method);
+        return invokeMethod(isFeatureFlagEnabled(flagName, organizationId), joinPoint, method);
     }
 
     private Object invokeMethod(Boolean isFeatureSupported, ProceedingJoinPoint joinPoint, Method method) {
@@ -106,10 +122,26 @@ AppsmithException getInvalidAnnotationUsageException(Method method, String error
                 error);
     }
 
-    boolean isFeatureFlagEnabled(FeatureFlagEnum flagName) {
-        CachedFeatures cachedFeatures = featureFlagService.getCachedOrganizationFeatureFlags();
+    boolean isFeatureFlagEnabled(FeatureFlagEnum flagName, String organizationId) {
+        CachedFeatures cachedFeatures = featureFlagService.getCachedOrganizationFeatureFlags(organizationId);
         return cachedFeatures != null
                 && !CollectionUtils.isNullOrEmpty(cachedFeatures.getFeatures())
                 && Boolean.TRUE.equals(cachedFeatures.getFeatures().get(flagName.name()));
     }
+
+    private String extractOrganizationId(Object[] args, Parameter[] parameters) {
+        if (args == null || parameters == null || args.length != parameters.length) {
+            return null;
+        }
+
+        // First try to find parameter named exactly "organizationId" or "orgId"
+        for (int i = 0; i < parameters.length; i++) {
+            String paramName = parameters[i].getName();
+            if ((paramName.equals("organizationId") || paramName.equals("orgId")) && args[i] instanceof String) {
+                return (String) args[i];
+            }
+        }
+
+        return null;
+    }
 }

app/server/appsmith-server/src/main/java/com/appsmith/server/plugins/solutions/PluginTriggerSolutionCEImpl.java

@@ -78,11 +78,6 @@ public Mono<TriggerResultDTO> trigger(
         Mono<PluginExecutor> pluginExecutorMono =
                 pluginMono.flatMap(plugin -> pluginExecutorHelper.getPluginExecutor(Mono.just(plugin)));
 
-        // TODO: Flags are needed here for google sheets integration to support shared drive behind a flag
-        // Once thoroughly tested, this flag can be removed
-        Map<String, Boolean> featureFlagMap =
-                featureFlagService.getCachedOrganizationFeatureFlags().getFeatures();
-
         /*
          * Since there is no datasource provided, we are passing the Datasource Context connection and datasourceConfiguration as null.
          * We will leave the execution to respective plugin executor.
@@ -91,9 +86,15 @@ public Mono<TriggerResultDTO> trigger(
             Plugin plugin = pair.getT1();
             PluginExecutor pluginExecutor = pair.getT2();
             setHeadersToTriggerRequest(plugin, httpHeaders, triggerRequestDTO);
-            return setOrganizationAndInstanceId(triggerRequestDTO)
-                    .flatMap(updatedTriggerRequestDTO -> ((PluginExecutor<Object>) pluginExecutor)
-                            .triggerWithFlags(null, null, updatedTriggerRequestDTO, featureFlagMap));
+            return setOrganizationAndInstanceId(triggerRequestDTO).flatMap(updatedTriggerRequestDTO -> {
+                // TODO: Flags are needed here for google sheets integration to support shared drive behind a
+                //  flag once thoroughly tested, this flag can be removed
+                Map<String, Boolean> featureFlags = featureFlagService
+                        .getCachedOrganizationFeatureFlags(updatedTriggerRequestDTO.getOrganizationId())
+                        .getFeatures();
+                return ((PluginExecutor<Object>) pluginExecutor)
+                        .triggerWithFlags(null, null, updatedTriggerRequestDTO, featureFlags);
+            });
         });
     }
 

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

@@ -42,5 +42,5 @@ Mono<Organization> getAllRemoteFeaturesForOrganizationAndUpdateFeatureFlagsWithP
 
     Mono<Organization> checkAndExecuteMigrationsForOrganizationFeatureFlags(Organization organization);
 
-    CachedFeatures getCachedOrganizationFeatureFlags();
+    CachedFeatures getCachedOrganizationFeatureFlags(String organizationId);
 }

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

@@ -13,7 +13,6 @@
 import com.appsmith.server.services.OrganizationService;
 import com.appsmith.server.services.SessionUserService;
 import com.appsmith.server.services.UserIdentifierService;
-import lombok.Getter;
 import lombok.RequiredArgsConstructor;
 import lombok.extern.slf4j.Slf4j;
 import reactor.core.publisher.Mono;
@@ -39,9 +38,7 @@ public class FeatureFlagServiceCEImpl implements FeatureFlagServiceCE {
     private final FeatureFlagMigrationHelper featureFlagMigrationHelper;
     private static final long FEATURE_FLAG_CACHE_TIME_MIN = 120;
 
-    // TODO @CloudBilling: Remove once all the helper methods consuming @FeatureFlagged are converted to reactive
-    @Getter
-    private CachedFeatures cachedOrganizationFeatureFlags;
+    private Map<String, CachedFeatures> cachedOrganizationFeatureFlags = new HashMap<>();
 
     /**
      * This function checks if the feature is enabled for the current user. In case the user object is not present,
@@ -167,7 +164,7 @@ public Mono<Map<String, Boolean>> getOrganizationFeatures(String organizationId)
         return cacheableFeatureFlagHelper
                 .fetchCachedOrganizationFeatures(organizationId)
                 .map(cachedFeatures -> {
-                    cachedOrganizationFeatureFlags = cachedFeatures;
+                    cachedOrganizationFeatureFlags.put(organizationId, cachedFeatures);
                     return cachedFeatures.getFeatures();
                 })
                 .switchIfEmpty(Mono.just(new HashMap<>()));
@@ -182,4 +179,8 @@ public Mono<Map<String, Boolean>> getOrganizationFeatures(String organizationId)
     public Mono<Organization> checkAndExecuteMigrationsForOrganizationFeatureFlags(Organization organization) {
         return organizationService.checkAndExecuteMigrationsForOrganizationFeatureFlags(organization);
     }
+
+    public CachedFeatures getCachedOrganizationFeatureFlags(String organizationId) {
+        return this.cachedOrganizationFeatureFlags.getOrDefault(organizationId, new CachedFeatures());
+    }
 }

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

@@ -254,7 +254,6 @@ private Mono<String> getTrueEnvironmentId(
     private Mono<ExecuteActionDTO> populateExecuteActionDTO(ExecuteActionDTO executeActionDTO, NewAction newAction) {
         Mono<String> instanceIdMono = configService.getInstanceId();
         Mono<String> organizationIdMono = organizationService.getCurrentUserOrganizationId();
-
         Mono<ExecuteActionDTO> systemInfoPopulatedExecuteActionDTOMono =
                 actionExecutionSolutionHelper.populateExecuteActionDTOWithSystemInfo(executeActionDTO);
 
@@ -755,17 +754,20 @@ protected Mono<ActionExecutionResult> verifyDatasourceAndMakeRequest(
                         .tag("plugin", plugin.getPackageName())
                         .name(ACTION_EXECUTION_DATASOURCE_CONTEXT)
                         .tap(Micrometer.observation(observationRegistry)))
-                .flatMap(tuple2 -> {
-                    DatasourceStorage datasourceStorage1 = tuple2.getT1();
-                    DatasourceContext<?> resourceContext = tuple2.getT2();
+                .zipWith(organizationService.getCurrentUserOrganizationId())
+                .flatMap(objects -> {
+                    DatasourceStorage datasourceStorage1 = objects.getT1().getT1();
+                    DatasourceContext<?> resourceContext = objects.getT1().getT2();
+                    String organizationId = objects.getT2();
                     // Now that we have the context (connection details), execute the action.
 
                     Instant requestedAt = Instant.now();
-                    Map<String, Boolean> features = featureFlagService.getCachedOrganizationFeatureFlags() != null
-                            ? featureFlagService
-                                    .getCachedOrganizationFeatureFlags()
-                                    .getFeatures()
-                            : Collections.emptyMap();
+                    Map<String, Boolean> features =
+                            featureFlagService.getCachedOrganizationFeatureFlags(organizationId) != null
+                                    ? featureFlagService
+                                            .getCachedOrganizationFeatureFlags(organizationId)
+                                            .getFeatures()
+                                    : Collections.emptyMap();
 
                     // TODO: Flags are needed here for google sheets integration to support shared drive behind a flag
                     // Once thoroughly tested, this flag can be removed

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

@@ -108,26 +108,32 @@ public Mono<TriggerResultDTO> trigger(
                     final PluginExecutor pluginExecutor = tuple.getT3();
                     final Datasource datasource = tuple.getT4();
 
-                    // TODO: Flags are needed here for google sheets integration to support shared drive behind a flag
-                    // Once thoroughly tested, this flag can be removed
-                    Map<String, Boolean> featureFlagMap = featureFlagService.getCachedOrganizationFeatureFlags() != null
-                            ? featureFlagService
-                                    .getCachedOrganizationFeatureFlags()
-                                    .getFeatures()
-                            : Collections.emptyMap();
-
                     return datasourceContextService
                             .getDatasourceContext(datasourceStorage, plugin)
                             // Now that we have the context (connection details), execute the action.
                             // datasource remains unevaluated for datasource of DBAuth Type Authentication,
                             // However the context comes from evaluated datasource.
                             .flatMap(resourceContext -> populateTriggerRequestDto(triggerRequestDTO, datasource)
-                                    .flatMap(updatedTriggerRequestDTO -> ((PluginExecutor<Object>) pluginExecutor)
-                                            .triggerWithFlags(
-                                                    resourceContext.getConnection(),
-                                                    datasourceStorage.getDatasourceConfiguration(),
-                                                    updatedTriggerRequestDTO,
-                                                    featureFlagMap)));
+                                    .flatMap(updatedTriggerRequestDTO -> {
+                                        String organizationId = updatedTriggerRequestDTO.getOrganizationId();
+                                        // TODO: Flags are needed here for google sheets integration to support shared
+                                        // drive behind a flag
+                                        // Once thoroughly tested, this flag can be removed
+                                        Map<String, Boolean> featureFlagMap =
+                                                featureFlagService.getCachedOrganizationFeatureFlags(organizationId)
+                                                                != null
+                                                        ? featureFlagService
+                                                                .getCachedOrganizationFeatureFlags(organizationId)
+                                                                .getFeatures()
+                                                        : Collections.emptyMap();
+
+                                        return ((PluginExecutor<Object>) pluginExecutor)
+                                                .triggerWithFlags(
+                                                        resourceContext.getConnection(),
+                                                        datasourceStorage.getDatasourceConfiguration(),
+                                                        updatedTriggerRequestDTO,
+                                                        featureFlagMap);
+                                    }));
                 });
 
         // If the plugin hasn't implemented the trigger function, go for the default implementation

app/server/appsmith-server/src/test/java/com/appsmith/server/aspect/FeatureFlaggedMethodInvokerAspectTest.java

@@ -1,5 +1,6 @@
 package com.appsmith.server.aspect;
 
+import com.appsmith.external.annotations.FeatureFlagged;
 import com.appsmith.external.enums.FeatureFlagEnum;
 import com.appsmith.server.aspect.component.TestComponent;
 import com.appsmith.server.exceptions.AppsmithError;
@@ -19,6 +20,7 @@
 import java.util.Map;
 
 import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.mockito.ArgumentMatchers.any;
 import static org.mockito.ArgumentMatchers.eq;
 
 @SpringBootTest
@@ -41,7 +43,8 @@ void setUp() {
 
         CachedFeatures cachedFeatures = new CachedFeatures();
         cachedFeatures.setFeatures(Map.of(FeatureFlagEnum.ORGANIZATION_TEST_FEATURE.name(), Boolean.FALSE));
-        Mockito.when(featureFlagService.getCachedOrganizationFeatureFlags()).thenReturn(cachedFeatures);
+        Mockito.when(featureFlagService.getCachedOrganizationFeatureFlags(any()))
+                .thenReturn(cachedFeatures);
     }
 
     @Test
@@ -107,18 +110,18 @@ void ceEeDiffMethodReturnsFlux_ceImplTest() {
 
     @Test
     void ceEeSyncMethod_eeImplTest() {
-        Mockito.when(featureFlagService.check(eq(FeatureFlagEnum.ORGANIZATION_TEST_FEATURE)))
-                .thenReturn(Mono.just(true));
-        StepVerifier.create(testComponent.ceEeSyncMethod("arg_"))
-                .assertNext(result -> assertEquals("arg_ee_impl_method", result))
-                .verifyComplete();
+        CachedFeatures cachedFeatures = new CachedFeatures();
+        cachedFeatures.setFeatures(Map.of(FeatureFlagEnum.ORGANIZATION_TEST_FEATURE.name(), Boolean.TRUE));
+        Mockito.when(featureFlagService.getCachedOrganizationFeatureFlags("organizationId"))
+                .thenReturn(cachedFeatures);
+        String result = testComponent.ceEeSyncMethod("arg_", "organizationId");
+        assertEquals("arg_ee_impl_method", result);
     }
 
     @Test
     void ceEeSyncMethod_ceImplTest() {
-        StepVerifier.create(testComponent.ceEeSyncMethod("arg_"))
-                .assertNext(result -> assertEquals("arg_ce_impl_method", result))
-                .verifyComplete();
+        String result = testComponent.ceEeSyncMethod("arg_", "organizationId");
+        assertEquals("arg_ce_impl_method", result);
     }
 
     @Test
@@ -142,4 +145,23 @@ void ceEeThrowNonAppsmithException_eeImplTest_throwExceptionFromAspect() {
                         && throwable.getMessage().equals("This is a test exception"))
                 .verify();
     }
+
+    @Test
+    void ceEeSyncWithoutOrganizationMethod_eeImplTest() {
+        CachedFeatures cachedFeatures = new CachedFeatures();
+        cachedFeatures.setFeatures(Map.of(FeatureFlagEnum.ORGANIZATION_TEST_FEATURE.name(), Boolean.TRUE));
+        Mockito.when(featureFlagService.getCachedOrganizationFeatureFlags("organizationId"))
+                .thenReturn(cachedFeatures);
+        try {
+            testComponent.ceEeSyncMethodWithoutOrganization("arg_");
+        } catch (AppsmithException e) {
+            assertEquals(
+                    AppsmithError.INVALID_METHOD_LEVEL_ANNOTATION_USAGE.getMessage(
+                            FeatureFlagged.class.getSimpleName(),
+                            "TestComponentImpl",
+                            "ceEeSyncMethodWithoutOrganization",
+                            "Add missing organizationId parameter and enforce non-null value for orgnization-specific feature flags retrieval in non-reactive methods"),
+                    e.getMessage());
+        }
+    }
 }

app/server/appsmith-server/src/test/java/com/appsmith/server/aspect/component/TestComponentImpl.java

@@ -40,8 +40,8 @@ public Flux<String> ceEeDiffMethodReturnsFlux() {
 
     @Override
     @FeatureFlagged(featureFlagName = FeatureFlagEnum.ORGANIZATION_TEST_FEATURE)
-    public Mono<String> ceEeSyncMethod(String arg) {
-        return Mono.just(arg + "ee_impl_method");
+    public String ceEeSyncMethod(String arg, String organizationId) {
+        return arg + "ee_impl_method";
     }
 
     @Override
@@ -55,4 +55,10 @@ public Mono<Void> ceEeThrowAppsmithException(String arg) {
     public Mono<Void> ceEeThrowNonAppsmithException(String arg) {
         return Mono.error(new RuntimeException("This is a test exception"));
     }
+
+    @Override
+    @FeatureFlagged(featureFlagName = FeatureFlagEnum.ORGANIZATION_TEST_FEATURE)
+    public String ceEeSyncMethodWithoutOrganization(String arg) {
+        return arg + "ee_impl_method";
+    }
 }

app/server/appsmith-server/src/test/java/com/appsmith/server/aspect/component/ce/TestComponentCE.java

@@ -13,9 +13,11 @@ public interface TestComponentCE {
 
     Flux<String> ceEeDiffMethodReturnsFlux();
 
-    Mono<String> ceEeSyncMethod(String arg);
+    String ceEeSyncMethod(String arg, String organizationId);
 
     Mono<Void> ceEeThrowAppsmithException(String arg);
 
     Mono<Void> ceEeThrowNonAppsmithException(String arg);
+
+    String ceEeSyncMethodWithoutOrganization(String arg);
 }

app/server/appsmith-server/src/test/java/com/appsmith/server/aspect/component/ce/TestComponentCEImpl.java

@@ -27,8 +27,8 @@ public Flux<String> ceEeDiffMethodReturnsFlux() {
     }
 
     @Override
-    public Mono<String> ceEeSyncMethod(String arg) {
-        return Mono.just(arg + "ce_impl_method");
+    public String ceEeSyncMethod(String arg, String organizationId) {
+        return arg + "ce_impl_method";
     }
 
     @Override
@@ -40,4 +40,9 @@ public Mono<Void> ceEeThrowAppsmithException(String arg) {
     public Mono<Void> ceEeThrowNonAppsmithException(String arg) {
         return Mono.empty();
     }
+
+    @Override
+    public String ceEeSyncMethodWithoutOrganization(String arg) {
+        return arg + "ce_impl_method";
+    }
 }