diff --git a/app/server/appsmith-server/src/main/java/com/appsmith/server/repositories/ce/CustomDatasourceStorageStructureRepositoryCE.java b/app/server/appsmith-server/src/main/java/com/appsmith/server/repositories/ce/CustomDatasourceStorageStructureRepositoryCE.java
index d9fb719744..d8ef2313e2 100644
--- a/app/server/appsmith-server/src/main/java/com/appsmith/server/repositories/ce/CustomDatasourceStorageStructureRepositoryCE.java
+++ b/app/server/appsmith-server/src/main/java/com/appsmith/server/repositories/ce/CustomDatasourceStorageStructureRepositoryCE.java
@@ -5,5 +5,5 @@ import reactor.core.publisher.Mono;
 
 public interface CustomDatasourceStorageStructureRepositoryCE {
 
-    Mono<Void> updateStructure(String datasourceId, String environmentId, DatasourceStructure structure);
+    Mono<Integer> updateStructure(String datasourceId, String environmentId, DatasourceStructure structure);
 }
diff --git a/app/server/appsmith-server/src/main/java/com/appsmith/server/repositories/ce/CustomDatasourceStorageStructureRepositoryCEImpl.java b/app/server/appsmith-server/src/main/java/com/appsmith/server/repositories/ce/CustomDatasourceStorageStructureRepositoryCEImpl.java
index 88e96e74cc..30d388e673 100644
--- a/app/server/appsmith-server/src/main/java/com/appsmith/server/repositories/ce/CustomDatasourceStorageStructureRepositoryCEImpl.java
+++ b/app/server/appsmith-server/src/main/java/com/appsmith/server/repositories/ce/CustomDatasourceStorageStructureRepositoryCEImpl.java
@@ -2,30 +2,22 @@ package com.appsmith.server.repositories.ce;
 
 import com.appsmith.external.models.DatasourceStorageStructure;
 import com.appsmith.external.models.DatasourceStructure;
-import com.appsmith.server.helpers.ce.bridge.Bridge;
 import com.appsmith.server.repositories.BaseAppsmithRepositoryImpl;
 import com.appsmith.server.repositories.CacheableRepositoryHelper;
-import com.appsmith.server.repositories.DatasourceStorageStructureRepository;
-import org.springframework.beans.factory.annotation.Autowired;
-import org.springframework.context.annotation.Lazy;
 import org.springframework.data.mongodb.core.ReactiveMongoOperations;
 import org.springframework.data.mongodb.core.convert.MongoConverter;
+import org.springframework.data.mongodb.core.query.Criteria;
+import org.springframework.data.mongodb.core.query.Query;
 import org.springframework.data.mongodb.core.query.Update;
 import org.springframework.stereotype.Component;
 import reactor.core.publisher.Mono;
 
-import static com.appsmith.server.helpers.ce.bridge.Bridge.bridge;
+import static org.springframework.data.mongodb.core.query.Criteria.where;
 
 @Component
 public class CustomDatasourceStorageStructureRepositoryCEImpl
         extends BaseAppsmithRepositoryImpl<DatasourceStorageStructure>
         implements CustomDatasourceStorageStructureRepositoryCE {
-
-    @SuppressWarnings("SpringJavaAutowiredFieldsWarningInspection") // Can't do lazy wiring in constructor params.
-    @Autowired
-    @Lazy
-    DatasourceStorageStructureRepository repository;
-
     public CustomDatasourceStorageStructureRepositoryCEImpl(
             ReactiveMongoOperations mongoOperations,
             MongoConverter mongoConverter,
@@ -33,22 +25,20 @@ public class CustomDatasourceStorageStructureRepositoryCEImpl
         super(mongoOperations, mongoConverter, cacheableRepositoryHelper);
     }
 
+    public static Criteria getDatasourceIdAndEnvironmentIdCriteria(String datasourceId, String environmentId) {
+        return new Criteria()
+                .andOperator(
+                        where(DatasourceStorageStructure.Fields.datasourceId).is(datasourceId),
+                        where(DatasourceStorageStructure.Fields.environmentId).is(environmentId));
+    }
+
     @Override
-    public Mono<Void> updateStructure(String datasourceId, String environmentId, DatasourceStructure structure) {
-        final Bridge criteria = bridge().equal(DatasourceStorageStructure.Fields.datasourceId, datasourceId)
-                .equal(DatasourceStorageStructure.Fields.environmentId, environmentId);
-
-        final Update update = Update.update(DatasourceStorageStructure.Fields.structure, structure);
-
-        return queryBuilder().criteria(criteria).updateFirst(update).flatMap(count -> {
-            if (count == 0) {
-                DatasourceStorageStructure dss = new DatasourceStorageStructure();
-                dss.setDatasourceId(datasourceId);
-                dss.setEnvironmentId(environmentId);
-                dss.setStructure(structure);
-                return repository.save(dss).then();
-            }
-            return Mono.empty();
-        });
+    public Mono<Integer> updateStructure(String datasourceId, String environmentId, DatasourceStructure structure) {
+        return mongoOperations
+                .upsert(
+                        new Query().addCriteria(getDatasourceIdAndEnvironmentIdCriteria(datasourceId, environmentId)),
+                        Update.update(DatasourceStorageStructure.Fields.structure, structure),
+                        DatasourceStorageStructure.class)
+                .map(updateResult -> Math.toIntExact(updateResult.getModifiedCount()));
     }
 }
diff --git a/app/server/appsmith-server/src/main/java/com/appsmith/server/services/ce/DatasourceStructureServiceCE.java b/app/server/appsmith-server/src/main/java/com/appsmith/server/services/ce/DatasourceStructureServiceCE.java
index d54ccc0a4e..62def45a09 100644
--- a/app/server/appsmith-server/src/main/java/com/appsmith/server/services/ce/DatasourceStructureServiceCE.java
+++ b/app/server/appsmith-server/src/main/java/com/appsmith/server/services/ce/DatasourceStructureServiceCE.java
@@ -10,5 +10,5 @@ public interface DatasourceStructureServiceCE {
 
     Mono<DatasourceStorageStructure> save(DatasourceStorageStructure datasourceStorageStructure);
 
-    Mono<Void> saveStructure(String datasourceId, String environmentId, DatasourceStructure structure);
+    Mono<Integer> saveStructure(String datasourceId, String environmentId, DatasourceStructure structure);
 }
diff --git a/app/server/appsmith-server/src/main/java/com/appsmith/server/services/ce/DatasourceStructureServiceCEImpl.java b/app/server/appsmith-server/src/main/java/com/appsmith/server/services/ce/DatasourceStructureServiceCEImpl.java
index 8e64b86f63..7862fc3c35 100644
--- a/app/server/appsmith-server/src/main/java/com/appsmith/server/services/ce/DatasourceStructureServiceCEImpl.java
+++ b/app/server/appsmith-server/src/main/java/com/appsmith/server/services/ce/DatasourceStructureServiceCEImpl.java
@@ -25,7 +25,7 @@ public class DatasourceStructureServiceCEImpl implements DatasourceStructureServ
     }
 
     @Override
-    public Mono<Void> saveStructure(String datasourceId, String environmentId, DatasourceStructure structure) {
+    public Mono<Integer> saveStructure(String datasourceId, String environmentId, DatasourceStructure structure) {
         return repository.updateStructure(datasourceId, environmentId, structure);
     }
 }