chore: refactor datasource storage to use custom repo method

app/server/appsmith-server/src/main/java/com/appsmith/server/repositories/ce/CustomDatasourceStorageRepositoryCE.java

@@ -2,5 +2,11 @@
 
 import com.appsmith.external.models.DatasourceStorage;
 import com.appsmith.server.repositories.AppsmithRepository;
+import reactor.core.publisher.Flux;
+import reactor.core.publisher.Mono;
 
-public interface CustomDatasourceStorageRepositoryCE extends AppsmithRepository<DatasourceStorage> {}
+public interface CustomDatasourceStorageRepositoryCE extends AppsmithRepository<DatasourceStorage> {
+    Mono<DatasourceStorage> findByDatasourceIdAndEnvironmentId(String datasourceId, String environmentId);
+
+    Flux<DatasourceStorage> findByDatasourceId(String datasourceId);
+}

app/server/appsmith-server/src/main/java/com/appsmith/server/repositories/ce/CustomDatasourceStorageRepositoryCEImpl.java

@@ -1,7 +1,26 @@
 package com.appsmith.server.repositories.ce;
 
 import com.appsmith.external.models.DatasourceStorage;
+import com.appsmith.server.helpers.ce.bridge.Bridge;
+import com.appsmith.server.helpers.ce.bridge.BridgeQuery;
 import com.appsmith.server.repositories.BaseAppsmithRepositoryImpl;
+import reactor.core.publisher.Flux;
+import reactor.core.publisher.Mono;
 
 public class CustomDatasourceStorageRepositoryCEImpl extends BaseAppsmithRepositoryImpl<DatasourceStorage>
-        implements CustomDatasourceStorageRepositoryCE {}
+        implements CustomDatasourceStorageRepositoryCE {
+    @Override
+    public Mono<DatasourceStorage> findByDatasourceIdAndEnvironmentId(String datasourceId, String environmentId) {
+        final BridgeQuery<DatasourceStorage> q = Bridge.<DatasourceStorage>equal(
+                        DatasourceStorage.Fields.datasourceId, datasourceId)
+                .equal(DatasourceStorage.Fields.environmentId, environmentId);
+        return queryBuilder().criteria(q).one();
+    }
+
+    @Override
+    public Flux<DatasourceStorage> findByDatasourceId(String datasourceId) {
+        final BridgeQuery<DatasourceStorage> q =
+                Bridge.<DatasourceStorage>equal(DatasourceStorage.Fields.datasourceId, datasourceId);
+        return queryBuilder().criteria(q).all();
+    }
+}

app/server/appsmith-server/src/main/java/com/appsmith/server/repositories/ce/DatasourceStorageRepositoryCE.java

@@ -2,11 +2,9 @@
 
 import com.appsmith.external.models.DatasourceStorage;
 import com.appsmith.server.repositories.BaseRepository;
-import reactor.core.publisher.Flux;
-import reactor.core.publisher.Mono;
+import com.appsmith.server.repositories.CustomDatasourceStorageRepository;
+import org.springframework.stereotype.Repository;
 
-public interface DatasourceStorageRepositoryCE extends BaseRepository<DatasourceStorage, String> {
-    Flux<DatasourceStorage> findByDatasourceId(String datasourceId);
-
-    Mono<DatasourceStorage> findByDatasourceIdAndEnvironmentId(String datasourceId, String environmentId);
-}
+@Repository
+public interface DatasourceStorageRepositoryCE
+        extends BaseRepository<DatasourceStorage, String>, CustomDatasourceStorageRepository {}