chore: Fixing CI post ME merge (#23668)

Fixing the way we consume some methods to avoid txn conflicts on EE

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

@@ -82,8 +82,8 @@ public interface DatasourceServiceCE {
     Mono<DatasourceDTO> convertToDatasourceDTO(Datasource datasource);
 
     // TODO: Remove the following snippet after client side API changes
-    Datasource convertToDatasource(DatasourceDTO datasourceDTO, String environmentId);
+    Mono<Datasource> convertToDatasource(DatasourceDTO datasourceDTO, String environmentId);
 
     // TODO: Remove the following snippet after client side API changes
-    String getTrueEnvironmentId(String environmentId);
+    Mono<String> getTrueEnvironmentId(String workspaceId, String environmentId);
 }

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

@@ -116,8 +116,8 @@ public class DatasourceServiceCEImpl implements DatasourceServiceCE {
     // TODO: Remove the following snippet after client side API changes
     @Override
     public Mono<DatasourceDTO> create(DatasourceDTO datasourceDTO, String environmentId) {
-        Datasource datasource = convertToDatasource(datasourceDTO, environmentId);
+        return convertToDatasource(datasourceDTO, environmentId)
-        return this.create(datasource)
+                .flatMap(datasource -> this.create(datasource))
                 .flatMap(this::convertToDatasourceDTO);
     }
 
@@ -187,12 +187,14 @@ public class DatasourceServiceCEImpl implements DatasourceServiceCE {
                 .flatMap(datasource1 -> {
                     // In case this was a newly created datasource, we need to update datasource reference first
                     return Flux.fromIterable(datasource.getDatasourceStorages().values())
-                            .map(datasourceStorageDTO -> {
+                            .flatMap(datasourceStorageDTO -> {
                                 DatasourceStorage datasourceStorage = new DatasourceStorage(datasourceStorageDTO);
                                 datasourceStorage.prepareTransientFields(datasource1);
-                                String trueEnvironmentId = getTrueEnvironmentId(datasourceStorageDTO.getEnvironmentId());
+                                return getTrueEnvironmentId(workspaceId, datasourceStorageDTO.getEnvironmentId())
-                                datasourceStorage.setEnvironmentId(getTrueEnvironmentId(trueEnvironmentId));
+                                        .map(trueEnvironmentId -> {
-                                return datasourceStorage;
+                                            datasourceStorage.setEnvironmentId(trueEnvironmentId);
+                                            return datasourceStorage;
+                                        });
                             })
                             .flatMap(datasourceStorage -> {
                                 // Make sure that we are creating entries only if the id is not already populated
@@ -237,9 +239,9 @@ public class DatasourceServiceCEImpl implements DatasourceServiceCE {
                                       DatasourceDTO datasourceDTO,
                                       String environmentId,
                                       Boolean isUserRefreshedUpdate) {
-        Datasource datasource = convertToDatasource(datasourceDTO, environmentId);
+        return convertToDatasource(datasourceDTO, environmentId)
-        return this.updateByEnvironmentId(id, datasource, environmentId, isUserRefreshedUpdate)
+                .flatMap(datasource -> updateByEnvironmentId(id, datasource, environmentId, isUserRefreshedUpdate))
-                .flatMap(datasource1 -> convertToDatasourceDTO(datasource1));
+                .flatMap(datasource -> convertToDatasourceDTO(datasource));
     }
 
     @Override
@@ -270,14 +272,15 @@ public class DatasourceServiceCEImpl implements DatasourceServiceCE {
         if (!datasource.getDatasourceStorages().isEmpty()) {
             // This is meant to be an update for storage
             return datasourceMono
-                    .flatMap(dbDatasource -> datasourceStorageService
+                    .flatMap(dbDatasource -> getTrueEnvironmentId(dbDatasource.getWorkspaceId(), environmentId)
-                            .updateByDatasourceAndEnvironmentId(datasource, getTrueEnvironmentId(environmentId), isUserRefreshedUpdate)
+                            .flatMap(trueEnvironmentId -> datasourceStorageService
-                            .map(datasourceStorage -> {
+                                    .updateByDatasourceAndEnvironmentId(datasource, trueEnvironmentId, isUserRefreshedUpdate)
-                                datasource.getDatasourceStorages()
+                                    .map(datasourceStorage -> {
-                                        .put(getTrueEnvironmentId(environmentId), new DatasourceStorageDTO(datasourceStorage));
+                                        datasource.getDatasourceStorages()
-                                copyNestedNonNullProperties(datasource, dbDatasource);
+                                                .put(trueEnvironmentId, new DatasourceStorageDTO(datasourceStorage));
-                                return dbDatasource;
+                                        copyNestedNonNullProperties(datasource, dbDatasource);
-                            }))
+                                        return dbDatasource;
+                                    })))
                     .flatMap(savedDatasource -> {
                         Map<String, Object> analyticsProperties = getAnalyticsProperties(savedDatasource);
                         if (isUserRefreshedUpdate.equals(Boolean.TRUE)) {
@@ -381,28 +384,40 @@ public class DatasourceServiceCEImpl implements DatasourceServiceCE {
      */
     @Override
     public Mono<DatasourceTestResult> testDatasource(DatasourceDTO datasourceDTO, String environmentId) {
-        Datasource datasource = convertToDatasource(datasourceDTO, environmentId);
         // the datasource has been created with datasourceStorageKey as output of getTrueEnvironmentId
-        String trueEnvironmentId = this.getTrueEnvironmentId(environmentId);
+        Mono<DatasourceStorage> datasourceStorageMono = this
-        DatasourceStorage datasourceStorage = datasourceStorageService
+                .getTrueEnvironmentId(datasourceDTO.getWorkspaceId(), environmentId)
-                .getDatasourceStorageFromDatasource(datasource, trueEnvironmentId);
+                .zipWhen(trueEnvironmentId -> convertToDatasource(datasourceDTO, trueEnvironmentId))
+                .flatMap(tuple2 -> {
+                    String trueEnvironmentId = tuple2.getT1();
+                    Datasource datasource = tuple2.getT2();
 
-        Mono<DatasourceStorage> datasourceStorageMono = Mono.just(datasourceStorage)
+                    DatasourceStorage datasourceStorage = datasourceStorageService
-                .flatMap(datasourceStorageService::checkEnvironment);
+                            .getDatasourceStorageFromDatasource(datasource, trueEnvironmentId);
-        // Fetch any fields that maybe encrypted from the db if the datasource being tested does not have those fields set.
+
-        // This scenario would happen whenever an existing datasource is being tested and no changes are present in the
+                    if (datasource.getId() == null) {
-        // encrypted field (because encrypted fields are not sent over the network after encryption back to the client
+                        return Mono.just(datasourceStorage);
-        if (datasourceStorage.getDatasourceId() != null && datasourceStorage.getDatasourceConfiguration() != null &&
+                    }
-                datasourceStorage.getDatasourceConfiguration().getAuthentication() != null) {
+                    // Check if we have execute access on this datasource
-            datasourceStorageMono =
+                    return this.findById(datasource.getId(), datasourcePermission.getExecutePermission())
-                    this.findById(datasource.getId(), datasourcePermission.getExecutePermission())
+                            .flatMap(dbDatasource -> {
-                            .flatMap(datasource1 -> datasourceStorageService.findByDatasourceAndEnvironmentId(datasource1, trueEnvironmentId))
+                                // Fetch any fields that maybe encrypted from the db if the datasource being tested does not have those fields set.
-                            .map(datasourceStorage1 -> {
+                                // This scenario would happen whenever an existing datasource is being tested and no changes are present in the
-                                copyNestedNonNullProperties(datasourceStorage, datasourceStorage1);
+                                // encrypted field (because encrypted fields are not sent over the network after encryption back to the client
-                                return datasourceStorage1;
+                                if (datasourceStorage.getDatasourceId() != null
+                                        && datasourceStorage.getDatasourceConfiguration() != null
+                                        && datasourceStorage.getDatasourceConfiguration().getAuthentication() != null) {
+                                    return datasourceStorageService.findByDatasourceAndEnvironmentId(dbDatasource, trueEnvironmentId)
+                                            .map(datasourceStorage1 -> {
+                                                copyNestedNonNullProperties(datasourceStorage, datasourceStorage1);
+                                                return datasourceStorage1;
+                                            });
+                                }
+                                return Mono.just(datasourceStorage);
                             })
-                            .switchIfEmpty(datasourceStorageMono);
+                            .switchIfEmpty(Mono.error(new AppsmithException(AppsmithError.UNAUTHORIZED_ACCESS)));
-        }
+                })
+                .flatMap(datasourceStorageService::checkEnvironment);
 
         return datasourceStorageMono
                 .flatMap(this::verifyDatasourceAndTest);
@@ -671,7 +686,7 @@ public class DatasourceServiceCEImpl implements DatasourceServiceCE {
 
     // TODO: Remove the following snippet after client side API changes
     @Override
-    public Datasource convertToDatasource(DatasourceDTO datasourceDTO, String environmentId) {
+    public Mono<Datasource> convertToDatasource(DatasourceDTO datasourceDTO, String environmentId) {
         Datasource datasource = new Datasource();
         datasource.setId(datasourceDTO.getId());
         datasource.setUserPermissions(datasourceDTO.getUserPermissions());
@@ -688,17 +703,22 @@ public class DatasourceServiceCEImpl implements DatasourceServiceCE {
 
         HashMap<String, DatasourceStorageDTO> storages = new HashMap<>();
         datasource.setDatasourceStorages(storages);
-        if (datasourceDTO.getDatasourceConfiguration() != null) {
-            storages.put(getTrueEnvironmentId(environmentId), new DatasourceStorageDTO(datasourceDTO, environmentId));
-        }
 
-        return datasource;
+
+        return getTrueEnvironmentId(datasource.getWorkspaceId(), environmentId)
+                .map(trueEnvironmentId -> {
+                    if (datasourceDTO.getDatasourceConfiguration() != null) {
+                        storages.put(trueEnvironmentId, new DatasourceStorageDTO(datasourceDTO, environmentId));
+                    }
+
+                    return datasource;
+                });
     }
 
     // TODO: Remove the following snippet after client side API changes
     @Override
-    public String getTrueEnvironmentId(String environmentId) {
+    public Mono<String> getTrueEnvironmentId(String workspaceId, String environmentId) {
-        return FieldName.UNUSED_ENVIRONMENT_ID;
+        return Mono.just(FieldName.UNUSED_ENVIRONMENT_ID);
     }
 
 }

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

@@ -126,14 +126,17 @@ public class MockDataServiceCEImpl implements MockDataServiceCE {
             datasource.setIsConfigured(true);
             datasource.setDatasourceConfiguration(datasourceConfiguration);
             HashMap<String, DatasourceStorageDTO> storages = new HashMap<>();
-            String trueEnvironmentId = datasourceService.getTrueEnvironmentId(environmentId);
-            DatasourceStorage datasourceStorage = new DatasourceStorage(datasource, trueEnvironmentId);
-            storages.put(trueEnvironmentId, new DatasourceStorageDTO(datasourceStorage));
-            datasource.setDatasourceStorages(storages);
 
-            return addAnalyticsForMockDataCreation(name, mockDataSource.getWorkspaceId())
+            return datasourceService.getTrueEnvironmentId(mockDataSource.getWorkspaceId(), environmentId)
-                    .then(createSuffixedDatasource(datasource, trueEnvironmentId))
+                    .flatMap(trueEnvironmentId -> {
-                    .flatMap(datasource1 -> datasourceService.convertToDatasourceDTO(datasource));
+                        DatasourceStorage datasourceStorage = new DatasourceStorage(datasource, trueEnvironmentId);
+                        storages.put(trueEnvironmentId, new DatasourceStorageDTO(datasourceStorage));
+                        datasource.setDatasourceStorages(storages);
+
+                        return addAnalyticsForMockDataCreation(name, mockDataSource.getWorkspaceId())
+                                .then(createSuffixedDatasource(datasource, trueEnvironmentId))
+                                .flatMap(datasource1 -> datasourceService.convertToDatasourceDTO(datasource));
+                    });
         });
 
     }

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

@@ -174,11 +174,14 @@ public class ActionExecutionSolutionCEImpl implements ActionExecutionSolutionCE
                                 branchName,
                                 executeActionDTO.getActionId(),
                                 actionPermission.getExecutePermission())
-                        .map(branchedAction -> {
+                        .flatMap(branchedAction -> {
                             executeActionDTO.setActionId(branchedAction.getId());
-                            return executeActionDTO;
+                            return Mono.just(executeActionDTO)
+                                    .zipWith(datasourceService.getTrueEnvironmentId(
+                                            branchedAction.getWorkspaceId(),
+                                            environmentId));
                         }))
-                .flatMap(executeActionDTO -> this.executeAction(executeActionDTO, datasourceService.getTrueEnvironmentId(environmentId))) // getTrue is temporary call
+                .flatMap(tuple2 -> this.executeAction(tuple2.getT1(), tuple2.getT2())) // getTrue is temporary call
                 .name(ACTION_EXECUTION_SERVER_EXECUTION)
                 .tap(Micrometer.observation(observationRegistry));
     }

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

@@ -45,11 +45,12 @@ public class DatasourceStructureSolutionCEImpl implements DatasourceStructureSol
 
     @Override
     public Mono<DatasourceStructure> getStructure(String datasourceId, boolean ignoreCache, String environmentId) {
-        String trueEnvironmentId = datasourceService.getTrueEnvironmentId(environmentId);
         return datasourceService.findById(datasourceId, datasourcePermission.getExecutePermission())
-                .flatMap(datasource1 -> datasourceStorageService.findByDatasourceAndEnvironmentId(
+                .zipWhen(datasource -> datasourceService
-                        datasource1,
+                        .getTrueEnvironmentId(datasource.getWorkspaceId(), environmentId))
-                        trueEnvironmentId))
+                .flatMap(tuple2 -> datasourceStorageService.findByDatasourceAndEnvironmentId(
+                        tuple2.getT1(),
+                        tuple2.getT2()))
                 .flatMap(datasourceStorage -> getStructure(datasourceStorage, ignoreCache))
                 .onErrorMap(
                         IllegalArgumentException.class,
@@ -130,7 +131,7 @@ public class DatasourceStructureSolutionCEImpl implements DatasourceStructureSol
 
                 })
                 .onErrorResume(error -> analyticsService.sendObjectEvent(AnalyticsEvents.DS_SCHEMA_FETCH_EVENT_FAILED, datasourceStorage,
-                        getAnalyticsPropertiesForTestEventStatus(datasourceStorage,false, error)).then(Mono.error(error)))
+                        getAnalyticsPropertiesForTestEventStatus(datasourceStorage, false, error)).then(Mono.error(error)))
                 .flatMap(structure -> analyticsService.sendObjectEvent(AnalyticsEvents.DS_SCHEMA_FETCH_EVENT_SUCCESS,
                                 datasourceStorage, getAnalyticsPropertiesForTestEventStatus(datasourceStorage, true, null))
                         .then(datasourceStorage.getDatasourceId() == null

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

@@ -856,7 +856,7 @@ public class ImportExportApplicationServiceCEImpl implements ImportExportApplica
                     }
                     return Mono.just(new ArrayList<Datasource>());
                 })
-                .zipWith(workspaceService.getDefaultEnvironmentId(workspaceId))
+                .zipWhen(datasourceList -> workspaceService.getDefaultEnvironmentId(workspaceId))
                 .flatMapMany(tuple2 -> {
                     List<Datasource> existingDatasources = tuple2.getT1();
                     String environmentId = tuple2.getT2();

app/server/appsmith-server/src/test/java/com/appsmith/server/solutions/ce/ActionExecutionSolutionCEImplTest.java

@@ -85,7 +85,7 @@ class ActionExecutionSolutionCEImplTest {
     ObjectMapper objectMapper;
     @MockBean
     NewActionRepository repository;
-    @MockBean
+    @SpyBean
     DatasourceService datasourceService;
     @MockBean
     PluginService pluginService;

app/server/appsmith-server/src/test/java/com/appsmith/server/solutions/ce/ActionExecutionSolutionCETest.java

@@ -1830,7 +1830,8 @@ public class ActionExecutionSolutionCETest {
         action.setActionConfiguration(actionConfiguration);
         action.setPageId(testPage.getId());
         action.setName("testActionExecuteDbQuery");
-        action.setDatasource(datasourceService.convertToDatasource(mockDatasource, defaultEnvironmentId));
+        Datasource datasource1 = datasourceService.convertToDatasource(mockDatasource, defaultEnvironmentId).block();
+        action.setDatasource(datasource1);
         ActionDTO createdAction = layoutActionService.createSingleAction(action, Boolean.FALSE).block();
 
         ExecuteActionDTO executeActionDTO = new ExecuteActionDTO();