chore: refactor customJs lib import flow for pg transaction support

app/server/appsmith-server/src/main/java/com/appsmith/server/dtos/ce/MappedImportableResourcesCE_DTO.java

@@ -3,6 +3,7 @@
 import com.appsmith.external.models.Datasource;
 import com.appsmith.external.models.DatasourceStorage;
 import com.appsmith.server.domains.Context;
+import com.appsmith.server.domains.CustomJSLib;
 import com.appsmith.server.dtos.CustomJSLibContextDTO;
 import com.appsmith.server.dtos.ImportActionCollectionResultDTO;
 import com.appsmith.server.dtos.ImportActionResultDTO;
@@ -50,4 +51,6 @@ public class MappedImportableResourcesCE_DTO {
     Map<String, List<Datasource>> datasourceDryRunQueries = new HashMap<>();
 
     Map<String, List<DatasourceStorage>> datasourceStorageDryRunQueries = new HashMap<>();
+
+    Map<String, List<CustomJSLib>> customJSLibsDryOps = new HashMap<>();
 }

app/server/appsmith-server/src/main/java/com/appsmith/server/jslibs/base/CustomJSLibServiceCE.java

@@ -9,6 +9,7 @@
 import reactor.core.publisher.Mono;
 
 import java.util.List;
+import java.util.Map;
 import java.util.Set;
 
 public interface CustomJSLibServiceCE extends CrudService<CustomJSLib, String> {
@@ -32,6 +33,12 @@ Mono<List<CustomJSLib>> getAllJSLibsInContext(
     Mono<CustomJSLibContextDTO> persistCustomJSLibMetaDataIfDoesNotExistAndGetDTO(
             CustomJSLib jsLib, Boolean isForceInstall);
 
+    Mono<CustomJSLibContextDTO> persistCustomJSLibMetaDataIfDoesNotExistAndGetDTO(
+            CustomJSLib jsLib,
+            Boolean isForceInstall,
+            Map<String, List<CustomJSLib>> customJSLibsDryOps,
+            boolean isDryOps);
+
     Flux<CustomJSLib> getAllVisibleJSLibsInContext(
             @NotNull String contextId, CreatorContextType contextType, String branchName, Boolean isViewMode);
 }

app/server/appsmith-server/src/main/java/com/appsmith/server/jslibs/base/CustomJSLibServiceCEImpl.java

@@ -4,6 +4,7 @@
 import com.appsmith.server.domains.Application;
 import com.appsmith.server.domains.CustomJSLib;
 import com.appsmith.server.dtos.CustomJSLibContextDTO;
+import com.appsmith.server.dtos.DBOpsType;
 import com.appsmith.server.jslibs.context.ContextBasedJsLibService;
 import com.appsmith.server.repositories.CustomJSLibRepository;
 import com.appsmith.server.services.AnalyticsService;
@@ -14,8 +15,10 @@
 import reactor.core.publisher.Flux;
 import reactor.core.publisher.Mono;
 
+import java.util.ArrayList;
 import java.util.Comparator;
 import java.util.List;
+import java.util.Map;
 import java.util.Set;
 import java.util.stream.Collectors;
 
@@ -73,11 +76,27 @@ public Mono<Boolean> addJSLibsToContext(
     @Override
     public Mono<CustomJSLibContextDTO> persistCustomJSLibMetaDataIfDoesNotExistAndGetDTO(
             CustomJSLib jsLib, Boolean isForceInstall) {
+        return persistCustomJSLibMetaDataIfDoesNotExistAndGetDTO(jsLib, isForceInstall, null, false);
+    }
+
+    @Override
+    public Mono<CustomJSLibContextDTO> persistCustomJSLibMetaDataIfDoesNotExistAndGetDTO(
+            CustomJSLib jsLib,
+            Boolean isForceInstall,
+            Map<String, List<CustomJSLib>> customJSLibsDryOps,
+            boolean isDryOps) {
         return repository
                 .findUniqueCustomJsLib(jsLib)
                 // Read more why Mono.defer is used here.
                 // https://stackoverflow.com/questions/54373920/mono-switchifempty-is-always-called
-                .switchIfEmpty(Mono.defer(() -> repository.save(jsLib)))
+                .switchIfEmpty(Mono.defer(() -> {
+                    if (isDryOps) {
+                        jsLib.updateForBulkWriteOperation();
+                        addDryOpsForEntity(DBOpsType.SAVE.name(), customJSLibsDryOps, jsLib);
+                        return Mono.just(jsLib);
+                    }
+                    return repository.save(jsLib);
+                }))
                 .flatMap(foundJSLib -> {
                     /*
                        The first check is to make sure that we are able to detect any previously truncated data and overwrite it the next time we receive valid data.
@@ -86,6 +105,10 @@ public Mono<CustomJSLibContextDTO> persistCustomJSLibMetaDataIfDoesNotExistAndGe
                     */
                     if ((jsLib.getDefs().length() > foundJSLib.getDefs().length()) || isForceInstall) {
                         jsLib.setId(foundJSLib.getId());
+                        if (isDryOps) {
+                            addDryOpsForEntity(DBOpsType.SAVE.name(), customJSLibsDryOps, jsLib);
+                            return Mono.just(jsLib);
+                        }
                         return repository.save(jsLib);
                     }
 
@@ -141,4 +164,15 @@ public Flux<CustomJSLib> getAllVisibleJSLibsInContext(
                 .getAllVisibleJSLibContextDTOFromContext(contextId, branchName, isViewMode)
                 .flatMapMany(repository::findCustomJsLibsInContext);
     }
+
+    private void addDryOpsForEntity(
+            String queryType, Map<String, List<CustomJSLib>> dryRunOpsMap, CustomJSLib createdCustomJsLib) {
+        if (dryRunOpsMap.containsKey(queryType)) {
+            dryRunOpsMap.get(queryType).add(createdCustomJsLib);
+        } else {
+            List<CustomJSLib> customJsLibList = new ArrayList<>();
+            customJsLibList.add(createdCustomJsLib);
+            dryRunOpsMap.put(queryType, customJsLibList);
+        }
+    }
 }

app/server/appsmith-server/src/main/java/com/appsmith/server/repositories/DryOperationRepository.java

@@ -2,6 +2,7 @@
 
 import com.appsmith.external.models.Datasource;
 import com.appsmith.external.models.DatasourceStorage;
+import com.appsmith.server.domains.CustomJSLib;
 import com.appsmith.server.dtos.MappedImportableResourcesDTO;
 import jakarta.annotation.PostConstruct;
 import lombok.RequiredArgsConstructor;
@@ -22,13 +23,16 @@ public class DryOperationRepository {
 
     private final DatasourceStorageRepository datasourceStorageRepository;
 
+    private final CustomJSLibRepository customJSLibRepository;
+
     private Map<Class<?>, AppsmithRepository<?>> repoByEntityClass;
 
     @PostConstruct
     public void init() {
         final Map<Class<?>, AppsmithRepository<?>> map = new HashMap<>();
         map.put(Datasource.class, datasourceRepository);
         map.put(DatasourceStorage.class, datasourceStorageRepository);
+        map.put(CustomJSLib.class, customJSLibRepository);
         repoByEntityClass = Collections.unmodifiableMap(map);
     }
 
@@ -44,6 +48,10 @@ public Flux<DatasourceStorage> saveDatasourceStorageToDb(List<DatasourceStorage>
         return datasourceStorageRepository.saveAll(datasourceStorage);
     }
 
+    private Flux<CustomJSLib> saveCustomJSLibToDb(List<CustomJSLib> customJSLibs) {
+        return customJSLibRepository.saveAll(customJSLibs);
+    }
+
     public Mono<Void> executeAllDbOps(MappedImportableResourcesDTO mappedImportableResourcesDTO) {
 
         Flux<List<Datasource>> datasourceFLux = Flux.fromIterable(mappedImportableResourcesDTO
@@ -65,6 +73,16 @@ public Mono<Void> executeAllDbOps(MappedImportableResourcesDTO mappedImportableR
                             .get(key);
                     return saveDatasourceStorageToDb(datasourceStorageList).collectList();
                 });
-        return Flux.merge(datasourceFLux, datasourceStorageFLux).then();
+
+        Flux<List<CustomJSLib>> customJSLibFLux = Flux.fromIterable(
+                        mappedImportableResourcesDTO.getCustomJSLibsDryOps().keySet())
+                .flatMap(key -> {
+                    List<CustomJSLib> customJSLibList =
+                            mappedImportableResourcesDTO.getCustomJSLibsDryOps().get(key);
+                    return saveCustomJSLibToDb(customJSLibList).collectList();
+                });
+
+        return Flux.merge(datasourceFLux, datasourceStorageFLux, customJSLibFLux)
+                .then();
     }
 }