From cadae671d330331f985ce91a7ade6f669e0b553d Mon Sep 17 00:00:00 2001
From: Manish Kumar <107841575+sondermanish@users.noreply.github.com>
Date: Fri, 4 Aug 2023 20:39:52 +0530
Subject: [PATCH] chore: adding support for multiple-environment error messages
 (#26005)

## Description
> Support PR for EE error handling for multiple-environments
- To support the error handling in EE, this PR addresses refactoring of
methods

Fixes https://github.com/appsmithorg/appsmith/issues/25964
Ref: https://github.com/appsmithorg/appsmith-ee/pull/1980
---
 .../services/ce/DatasourceStorageServiceCEImpl.java      | 9 +++++++--
 .../server/services/DatasourceStorageServiceTest.java    | 3 ---
 2 files changed, 7 insertions(+), 5 deletions(-)

diff --git a/app/server/appsmith-server/src/main/java/com/appsmith/server/services/ce/DatasourceStorageServiceCEImpl.java b/app/server/appsmith-server/src/main/java/com/appsmith/server/services/ce/DatasourceStorageServiceCEImpl.java
index b94103d8d9..74eb89f40e 100644
--- a/app/server/appsmith-server/src/main/java/com/appsmith/server/services/ce/DatasourceStorageServiceCEImpl.java
+++ b/app/server/appsmith-server/src/main/java/com/appsmith/server/services/ce/DatasourceStorageServiceCEImpl.java
@@ -92,8 +92,7 @@ public class DatasourceStorageServiceCEImpl implements DatasourceStorageServiceC
 
                     return Mono.just(datasourceStorage);
                 })
-                .switchIfEmpty(Mono.error(new AppsmithException(
-                        AppsmithError.NO_RESOURCE_FOUND, FieldName.DATASOURCE, datasource.getName())));
+                .switchIfEmpty(Mono.defer(() -> errorMonoWhenDatasourceStorageNotFound(datasource, environmentId)));
     }
 
     @Override
@@ -341,4 +340,10 @@ public class DatasourceStorageServiceCEImpl implements DatasourceStorageServiceC
     public Mono<String> getEnvironmentNameFromEnvironmentIdForAnalytics(String environmentId) {
         return Mono.just(FieldName.UNUSED_ENVIRONMENT_ID);
     }
+
+    protected Mono<DatasourceStorage> errorMonoWhenDatasourceStorageNotFound(
+            Datasource datasource, String environmentId) {
+        return Mono.error(
+                new AppsmithException(AppsmithError.NO_RESOURCE_FOUND, FieldName.DATASOURCE, datasource.getName()));
+    }
 }
diff --git a/app/server/appsmith-server/src/test/java/com/appsmith/server/services/DatasourceStorageServiceTest.java b/app/server/appsmith-server/src/test/java/com/appsmith/server/services/DatasourceStorageServiceTest.java
index 88f53618cd..7e2b39530f 100644
--- a/app/server/appsmith-server/src/test/java/com/appsmith/server/services/DatasourceStorageServiceTest.java
+++ b/app/server/appsmith-server/src/test/java/com/appsmith/server/services/DatasourceStorageServiceTest.java
@@ -211,9 +211,6 @@ public class DatasourceStorageServiceTest {
         String datasourceId = "datasourceForUnsavedStorage";
         String environmentIdOne = FieldName.UNUSED_ENVIRONMENT_ID;
 
-        Mockito.when(pluginExecutorHelper.getPluginExecutor(Mockito.any()))
-                .thenReturn(Mono.just(new MockPluginExecutor()));
-
         Datasource datasource = new Datasource();
         datasource.setId(datasourceId);