chore: Introduce orgId param for non-reactive feature flagged methods to fetch correct org flags

app/server/appsmith-server/src/main/java/com/appsmith/server/actioncollections/base/ActionCollectionServiceCEImpl.java

@@ -11,10 +11,12 @@
 import com.appsmith.server.domains.ActionCollection;
 import com.appsmith.server.domains.NewAction;
 import com.appsmith.server.domains.NewPage;
+import com.appsmith.server.domains.User;
 import com.appsmith.server.dtos.ActionCollectionDTO;
 import com.appsmith.server.dtos.ActionCollectionViewDTO;
 import com.appsmith.server.exceptions.AppsmithError;
 import com.appsmith.server.exceptions.AppsmithException;
+import com.appsmith.server.helpers.ReactiveContextUtils;
 import com.appsmith.server.newactions.base.NewActionService;
 import com.appsmith.server.repositories.ActionCollectionRepository;
 import com.appsmith.server.services.AnalyticsService;
@@ -280,13 +282,18 @@ public Mono<ActionCollectionDTO> update(String id, ActionCollectionDTO actionCol
 
     @Override
     public Mono<ActionCollectionDTO> deleteWithoutPermissionUnpublishedActionCollection(String id) {
-        return deleteUnpublishedActionCollection(id, null, actionPermission.getDeletePermission());
+        return ReactiveContextUtils.getCurrentUser()
+                .flatMap(user -> deleteUnpublishedActionCollection(
+                        id, null, actionPermission.getDeletePermission(user.getOrganizationId())));
     }
 
     @Override
     public Mono<ActionCollectionDTO> deleteUnpublishedActionCollection(String id) {
-        return deleteUnpublishedActionCollection(
-                id, actionPermission.getDeletePermission(), actionPermission.getDeletePermission());
+        return ReactiveContextUtils.getCurrentUser()
+                .flatMap(user -> deleteUnpublishedActionCollection(
+                        id,
+                        actionPermission.getDeletePermission(user.getOrganizationId()),
+                        actionPermission.getDeletePermission(user.getOrganizationId())));
     }
 
     @Override
@@ -418,10 +425,17 @@ public Mono<ActionCollection> archiveById(String id) {
     }
 
     protected Mono<ActionCollection> archiveGivenActionCollection(ActionCollection actionCollection) {
-        Flux<NewAction> unpublishedJsActionsFlux = newActionService.findByCollectionIdAndViewMode(
-                actionCollection.getId(), false, actionPermission.getDeletePermission());
-        Flux<NewAction> publishedJsActionsFlux = newActionService.findByCollectionIdAndViewMode(
-                actionCollection.getId(), true, actionPermission.getDeletePermission());
+        Mono<User> currentUserMono = ReactiveContextUtils.getCurrentUser().cache();
+        Flux<NewAction> unpublishedJsActionsFlux =
+                currentUserMono.flatMapMany(user -> newActionService.findByCollectionIdAndViewMode(
+                        actionCollection.getId(),
+                        false,
+                        actionPermission.getDeletePermission(user.getOrganizationId())));
+        Flux<NewAction> publishedJsActionsFlux =
+                currentUserMono.flatMapMany(user -> newActionService.findByCollectionIdAndViewMode(
+                        actionCollection.getId(),
+                        true,
+                        actionPermission.getDeletePermission(user.getOrganizationId())));
         return unpublishedJsActionsFlux
                 .mergeWith(publishedJsActionsFlux)
                 .flatMap(toArchive -> newActionService

app/server/appsmith-server/src/main/java/com/appsmith/server/actioncollections/importable/ActionCollectionImportableServiceCEImpl.java

@@ -12,6 +12,7 @@
 import com.appsmith.server.dtos.ImportActionCollectionResultDTO;
 import com.appsmith.server.dtos.ImportingMetaDTO;
 import com.appsmith.server.dtos.MappedImportableResourcesDTO;
+import com.appsmith.server.helpers.ReactiveContextUtils;
 import com.appsmith.server.imports.importable.ImportableServiceCE;
 import com.appsmith.server.imports.importable.artifactbased.ArtifactBasedImportableService;
 import com.appsmith.server.repositories.ActionCollectionRepository;
@@ -162,10 +163,14 @@ private Mono<ImportActionCollectionResultDTO> importActionCollections(
                                 importedActionCollectionList, mappedImportableResourcesDTO);
                     }
 
-                    return Mono.zip(actionCollectionsInCurrentArtifactMono, actionCollectionsInBranchesMono)
+                    return Mono.zip(
+                                    actionCollectionsInCurrentArtifactMono,
+                                    actionCollectionsInBranchesMono,
+                                    ReactiveContextUtils.getCurrentUser())
                             .flatMap(objects -> {
                                 Map<String, ActionCollection> actionsCollectionsInCurrentArtifact = objects.getT1();
                                 Map<String, ActionCollection> actionsCollectionsInBranches = objects.getT2();
+                                String currentUserOrgId = objects.getT3().getOrganizationId();
 
                                 // set the existing action collections in the result DTO,
                                 // this will be required in next phases when we'll delete the outdated action
@@ -228,7 +233,7 @@ private Mono<ImportActionCollectionResultDTO> importActionCollections(
                                                 .put(idFromJsonFile, existingActionCollection);
                                     } else {
                                         artifactBasedImportableService.createNewResource(
-                                                importingMetaDTO, actionCollection, baseContext);
+                                                importingMetaDTO, actionCollection, baseContext, currentUserOrgId);
 
                                         populateDomainMappedReferences(mappedImportableResourcesDTO, actionCollection);
 

app/server/appsmith-server/src/main/java/com/appsmith/server/actioncollections/importable/applications/ActionCollectionApplicationImportableServiceCEImpl.java

@@ -89,8 +89,11 @@ public Context updateContextInResource(
 
     @Override
     public void createNewResource(
-            ImportingMetaDTO importingMetaDTO, ActionCollection actionCollection, Context baseContext) {
-        if (!importingMetaDTO.getPermissionProvider().canCreateAction((NewPage) baseContext)) {
+            ImportingMetaDTO importingMetaDTO,
+            ActionCollection actionCollection,
+            Context baseContext,
+            String organizationId) {
+        if (!importingMetaDTO.getPermissionProvider().canCreateAction((NewPage) baseContext, organizationId)) {
             throw new AppsmithException(
                     AppsmithError.ACL_NO_RESOURCE_FOUND, FieldName.PAGE, ((NewPage) baseContext).getId());
         }

app/server/appsmith-server/src/main/java/com/appsmith/server/applications/base/ApplicationServiceCEImpl.java

@@ -501,12 +501,13 @@ private Mono<? extends Application> generateAndSetPoliciesForPublicView(
                 .filter(applicationId -> !application.getId().equals(applicationId))
                 .cache();
 
-        List<Mono<Void>> updateInheritedDomainsList =
+        Mono<List<Mono<Void>>> updateInheritedDomainsListMono =
                 updatePoliciesForInheritingDomains(application, applicationPolicyMap, addViewAccess);
         List<Mono<Void>> updateIndependentDomainsList = updatePoliciesForIndependentDomains(
                 application, permissionGroupId, addViewAccess, otherApplicationsForThisRoleFlux);
 
-        return Flux.fromIterable(updateInheritedDomainsList)
+        return updateInheritedDomainsListMono
+                .flatMapMany(Flux::fromIterable)
                 .flatMap(voidMono -> voidMono)
                 .thenMany(Flux.fromIterable(updateIndependentDomainsList))
                 .flatMap(voidMono -> voidMono)
@@ -604,7 +605,7 @@ protected List<Mono<Void>> updatePoliciesForIndependentDomains(
         return list;
     }
 
-    protected List<Mono<Void>> updatePoliciesForInheritingDomains(
+    protected Mono<List<Mono<Void>>> updatePoliciesForInheritingDomains(
             Application application, Map<String, Policy> applicationPolicyMap, Boolean addViewAccess) {
 
         List<Mono<Void>> list = new ArrayList<>();
@@ -635,7 +636,7 @@ protected List<Mono<Void>> updatePoliciesForInheritingDomains(
                 .then();
         list.add(updatedThemesMono);
 
-        return list;
+        return Mono.just(list);
     }
 
     public Mono<Application> setTransientFields(Application application) {

app/server/appsmith-server/src/main/java/com/appsmith/server/applications/imports/ApplicationImportServiceCEImpl.java

@@ -24,6 +24,7 @@
 import com.appsmith.server.exceptions.AppsmithException;
 import com.appsmith.server.helpers.ImportArtifactPermissionProvider;
 import com.appsmith.server.helpers.ImportExportUtils;
+import com.appsmith.server.helpers.ReactiveContextUtils;
 import com.appsmith.server.imports.importable.ImportableService;
 import com.appsmith.server.imports.internal.artifactbased.ArtifactBasedImportServiceCE;
 import com.appsmith.server.layouts.UpdateLayoutService;
@@ -90,19 +91,23 @@ public class ApplicationImportServiceCEImpl
             Map.of(FieldName.ARTIFACT_CONTEXT, FieldName.APPLICATION, FieldName.ID, FieldName.APPLICATION_ID);
 
     @Override
-    public ImportArtifactPermissionProvider getImportArtifactPermissionProviderForImportingArtifact(
+    public Mono<ImportArtifactPermissionProvider> getImportArtifactPermissionProviderForImportingArtifact(
             Set<String> userPermissionGroups) {
-        return ImportArtifactPermissionProvider.builder(
-                        applicationPermission,
-                        pagePermission,
-                        actionPermission,
-                        datasourcePermission,
-                        workspacePermission)
-                .requiredPermissionOnTargetWorkspace(workspacePermission.getApplicationCreatePermission())
-                .permissionRequiredToCreateDatasource(true)
-                .permissionRequiredToEditDatasource(true)
-                .currentUserPermissionGroups(userPermissionGroups)
-                .build();
+
+        return ReactiveContextUtils.getCurrentUser().map(currentUser -> {
+            String userOrgId = currentUser.getOrganizationId();
+            return ImportArtifactPermissionProvider.builder(
+                            applicationPermission,
+                            pagePermission,
+                            actionPermission,
+                            datasourcePermission,
+                            workspacePermission)
+                    .requiredPermissionOnTargetWorkspace(workspacePermission.getApplicationCreatePermission(userOrgId))
+                    .permissionRequiredToCreateDatasource(true)
+                    .permissionRequiredToEditDatasource(true)
+                    .currentUserPermissionGroups(userPermissionGroups)
+                    .build();
+        });
     }
 
     @Override

app/server/appsmith-server/src/main/java/com/appsmith/server/artifacts/permissions/ArtifactPermissionCE.java

@@ -6,7 +6,7 @@ public interface ArtifactPermissionCE {
 
     AclPermission getEditPermission();
 
-    AclPermission getDeletePermission();
+    AclPermission getDeletePermission(String organizationId);
 
     AclPermission getGitConnectPermission();
 

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
@@ -42,7 +44,7 @@ public class FeatureFlaggedMethodInvokerAspect {
      * @see FeatureFlagged
      */
     @Around("execution(public * *(..)) && @annotation(com.appsmith.external.annotations.FeatureFlagged)")
-    public Object invokeMethodAtMethodLevelAnnotation(ProceedingJoinPoint joinPoint) throws IllegalAccessException {
+    public Object invokeMethodAtMethodLevelAnnotation(ProceedingJoinPoint joinPoint) {
         Method method = ((MethodSignature) joinPoint.getSignature()).getMethod();
         FeatureFlagged annotation = method.getAnnotation(FeatureFlagged.class);
         if (annotation.featureFlagName() == null) {
@@ -64,9 +66,19 @@ private Object invokeMethod(ProceedingJoinPoint joinPoint, FeatureFlagged annota
         } else if (Flux.class.isAssignableFrom(returnType)) {
             return featureFlagMono.flatMapMany(isSupported -> (Flux<?>) invokeMethod(isSupported, joinPoint, method));
         }
+        // Extract organizationId from method arguments
+        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 +118,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/datasources/base/DatasourceServiceCEImpl.java

@@ -23,6 +23,7 @@
 import com.appsmith.server.exceptions.AppsmithError;
 import com.appsmith.server.exceptions.AppsmithException;
 import com.appsmith.server.helpers.PluginExecutorHelper;
+import com.appsmith.server.helpers.ReactiveContextUtils;
 import com.appsmith.server.plugins.base.PluginService;
 import com.appsmith.server.ratelimiting.RateLimitService;
 import com.appsmith.server.repositories.DatasourceRepository;
@@ -152,7 +153,12 @@ public DatasourceServiceCEImpl(
 
     @Override
     public Mono<Datasource> create(Datasource datasource) {
-        return createEx(datasource, workspacePermission.getDatasourceCreatePermission(), false, null);
+        return ReactiveContextUtils.getCurrentUser()
+                .flatMap(user -> createEx(
+                        datasource,
+                        workspacePermission.getDatasourceCreatePermission(user.getOrganizationId()),
+                        false,
+                        null));
     }
 
     // TODO: Check usage
@@ -880,8 +886,9 @@ public Flux<Datasource> getAllByWorkspaceIdWithStorages(String workspaceId, AclP
 
     @Override
     public Mono<Datasource> archiveById(String id) {
-        return repository
-                .findById(id, datasourcePermission.getDeletePermission())
+        return ReactiveContextUtils.getCurrentUser()
+                .flatMap(user ->
+                        repository.findById(id, datasourcePermission.getDeletePermission(user.getOrganizationId())))
                 .switchIfEmpty(
                         Mono.error(new AppsmithException(AppsmithError.NO_RESOURCE_FOUND, FieldName.DATASOURCE, id)))
                 .zipWhen(datasource -> newActionRepository.countByDatasourceId(datasource.getId()))