Bulk Implementation

eng/code-quality-reports/src/main/resources/revapi/revapi.json

@@ -89,6 +89,26 @@
           "new": "method <S extends T> (reactor\\.core\\.publisher\\.Mono<)?S>? com\\.azure\\.spring\\.data\\.cosmos\\.repository\\.(Reactive)?CosmosRepository<T, ID extends java\\.io\\.Serializable>::save\\(ID, com\\.azure\\.cosmos\\.models\\.PartitionKey, java\\.lang\\.Class<S>, com\\.azure\\.cosmos\\.models\\.CosmosPatchOperations(, com\\.azure\\.cosmos\\.models\\.CosmosPatchItemRequestOptions)?\\)",
           "justification": "Spring interfaces are allowed to add methods."
         },
+        {
+          "code": "java.method.addedToInterface",
+          "new": "method <S extends T, T> void com.azure.spring.data.cosmos.core.CosmosOperations::deleteEntities(com.azure.spring.data.cosmos.repository.support.CosmosEntityInformation<T, ?>, java.lang.Iterable<S>)",
+          "justification": "Spring interfaces are allowed to add methods."
+        },
+        {
+          "code": "java.method.addedToInterface",
+          "new": "method <S extends T, T> java.lang.Iterable<S> com.azure.spring.data.cosmos.core.CosmosOperations::insertAll(com.azure.spring.data.cosmos.repository.support.CosmosEntityInformation<T, ?>, java.lang.Iterable<S>)",
+          "justification": "Spring interfaces are allowed to add methods."
+        },
+        {
+          "code": "java.method.addedToInterface",
+          "new": "method <S extends T, T> reactor.core.publisher.Mono<java.lang.Void> com.azure.spring.data.cosmos.core.ReactiveCosmosOperations::deleteEntities(com.azure.spring.data.cosmos.repository.support.CosmosEntityInformation<T, ?>, reactor.core.publisher.Flux<S>)",
+          "justification": "Spring interfaces are allowed to add methods."
+        },
+        {
+          "code": "java.method.addedToInterface",
+          "new": "method <S extends T, T> reactor.core.publisher.Flux<S> com.azure.spring.data.cosmos.core.ReactiveCosmosOperations::insertAll(com.azure.spring.data.cosmos.repository.support.CosmosEntityInformation<T, ?>, reactor.core.publisher.Flux<S>)",
+          "justification": "Spring interfaces are allowed to add methods."
+        },
         {
           "regex": true,
           "code": "java\\.annotation\\.(added|attributeValueChanged)",

sdk/spring/azure-spring-data-cosmos/CHANGELOG.md

@@ -3,6 +3,7 @@
 ### 3.39.0-beta.1 (Unreleased)
 
 #### Features Added
+* Updated Spring and Reactive Spring repository `saveAll` and `deleteAll` APIs to use bulk functionality implementation. NOTE: `azure-spring-data-cosmos` is currently unable to set throughput control limits at the request level, which will need to be achieved by creating multiple clients. - See [PR 36611](https://github.com/Azure/azure-sdk-for-java/pull/36611).
 
 #### Breaking Changes
 

sdk/spring/azure-spring-data-cosmos/src/main/java/com/azure/spring/data/cosmos/core/CosmosOperations.java

@@ -147,6 +147,7 @@ public interface CosmosOperations {
      * @return the patched item
      */
     <T> T patch(Object id, PartitionKey partitionKey, Class<T> domainType, CosmosPatchOperations patchOperations, CosmosPatchItemRequestOptions options);
+
     /**
      * Inserts item
      *
@@ -158,6 +159,17 @@ public interface CosmosOperations {
      */
     <T> T insert(String containerName, T objectToSave, PartitionKey partitionKey);
 
+    /**
+     * Insert items using bulk operations.
+     *
+     * @param information must not be {@literal null}
+     * @param entities must not be {@literal null}
+     * @param <T> type class of domain type
+     * @param <S> type class of domain type
+     * @return the inserted item
+     */
+    <S extends T, T> Iterable<S> insertAll(CosmosEntityInformation<T, ?> information, Iterable<S> entities);
+
     /**
      * Inserts item
      * @param containerName must not be {@literal null}
@@ -213,7 +225,17 @@ public interface CosmosOperations {
     <T> void deleteEntity(String containerName, T entity);
 
     /**
-     * Delete all items in a container
+     * Delete using a list of entities with bulk
+     *
+     * @param <T> type class of domain type
+     * @param <S> type class of domain type
+     * @param information must not be {@literal null}
+     * @param entities must not be {@literal null}
+     */
+    <S extends T, T> void deleteEntities(CosmosEntityInformation<T, ?> information, Iterable<S> entities);
+
+    /**
+     * Delete all items in a container. Uses bulk if possible.
      *
      * @param containerName the container name
      * @param domainType the domainType
@@ -228,7 +250,7 @@ public interface CosmosOperations {
     void deleteContainer(String containerName);
 
     /**
-     * Delete items matching query
+     * Delete items matching query. Uses bulk if possible.
      *
      * @param query the document query
      * @param domainType type class
@@ -386,5 +408,4 @@ public interface CosmosOperations {
      * @return the Page
      */
     <T> Page<T> runPaginationQuery(SqlQuerySpec querySpec, Pageable pageable, Class<?> domainType, Class<T> returnType);
-
 }

sdk/spring/azure-spring-data-cosmos/src/main/java/com/azure/spring/data/cosmos/core/CosmosTemplate.java

@@ -6,9 +6,13 @@
 import com.azure.cosmos.CosmosAsyncClient;
 import com.azure.cosmos.CosmosAsyncContainer;
 import com.azure.cosmos.CosmosAsyncDatabase;
+import com.azure.cosmos.models.CosmosBulkExecutionOptions;
+import com.azure.cosmos.models.CosmosBulkItemRequestOptions;
+import com.azure.cosmos.models.CosmosBulkOperations;
 import com.azure.cosmos.models.CosmosContainerProperties;
 import com.azure.cosmos.models.CosmosContainerResponse;
 import com.azure.cosmos.models.CosmosDatabaseResponse;
+import com.azure.cosmos.models.CosmosItemOperation;
 import com.azure.cosmos.models.CosmosItemRequestOptions;
 import com.azure.cosmos.models.CosmosItemResponse;
 import com.azure.cosmos.models.CosmosPatchItemRequestOptions;
@@ -239,6 +243,54 @@ public <T> T insert(String containerName, T objectToSave, PartitionKey partition
         return toDomainObject(domainType, response.getItem());
     }
 
+    /**
+     * Insert all items with bulk.
+     *
+     * @param information the CosmosEntityInformation
+     * @param entities the Iterable entities to be inserted
+     * @param <T> type class of domain type
+     * @param <S> type class of domain type
+     * @return Flux of result
+     */
+    @SuppressWarnings("unchecked")
+    public <S extends T, T> Iterable<S> insertAll(CosmosEntityInformation<T, ?> information, Iterable<S> entities) {
+        Assert.notNull(entities, "entities to be inserted should not be null");
+
+        String containerName = information.getContainerName();
+        Class<T> domainType = information.getJavaType();
+
+        List<CosmosItemOperation> cosmosItemOperations = new ArrayList<>();
+        entities.forEach(entity -> {
+            JsonNode originalItem = mappingCosmosConverter.writeJsonNode(entity);
+            PartitionKey partitionKey = new PartitionKey(information.getPartitionKeyFieldValue(entity));
+            final CosmosBulkItemRequestOptions options = new CosmosBulkItemRequestOptions();
+            applyBulkVersioning(domainType, originalItem, options);
+            cosmosItemOperations.add(CosmosBulkOperations.getUpsertItemOperation(originalItem,
+                partitionKey, options));
+        });
+
+        // Default micro batch size is 100 which will be too high for most Spring cases, this configuration
+        // allows it to start at 1 and increase until it finds the appropriate batch size.
+        CosmosBulkExecutionOptions cosmosBulkExecutionOptions = new CosmosBulkExecutionOptions();
+        cosmosBulkExecutionOptions.setInitialMicroBatchSize(1);
+
+        return (Iterable<S>) this.getCosmosAsyncClient()
+            .getDatabase(this.getDatabaseName())
+            .getContainer(containerName)
+            .executeBulkOperations(Flux.fromIterable(cosmosItemOperations), cosmosBulkExecutionOptions)
+            .publishOn(Schedulers.parallel())
+            .onErrorResume(throwable ->
+                CosmosExceptionUtils.exceptionHandler("Failed to insert item(s)", throwable,
+                    this.responseDiagnosticsProcessor))
+            .flatMap(r -> {
+                CosmosUtils.fillAndProcessResponseDiagnostics(this.responseDiagnosticsProcessor,
+                    r.getResponse().getCosmosDiagnostics(), null);
+                JsonNode responseItem = r.getResponse().getItem(JsonNode.class);
+                return responseItem != null ? Flux.just(toDomainObject(domainType, responseItem)) : Flux.empty();
+            })
+            .collectList().block();
+    }
+
     /**
      * Patches item
      *
@@ -681,6 +733,51 @@ public <T> void deleteEntity(String containerName, T entity) {
         deleteItem(originalItem, containerName, domainType);
     }
 
+    /**
+     * Deletes the entities using bulk
+     *
+     * @param information the CosmosEntityInformation
+     * @param entities the Iterable entities to be inserted
+     * @param <T> type class of domain type
+     * @param <S> type class of domain type
+     */
+    @Override
+    public <S extends T, T> void deleteEntities(CosmosEntityInformation<T, ?> information, Iterable<S> entities) {
+        Assert.notNull(entities, "entities to be deleted should not be null");
+
+        String containerName = information.getContainerName();
+        Class<T> domainType = information.getJavaType();
+
+        List<CosmosItemOperation> cosmosItemOperations = new ArrayList<>();
+        entities.forEach(entity -> {
+            JsonNode originalItem = mappingCosmosConverter.writeJsonNode(entity);
+            PartitionKey partitionKey = new PartitionKey(information.getPartitionKeyFieldValue(entity));
+            final CosmosBulkItemRequestOptions options = new CosmosBulkItemRequestOptions();
+            applyBulkVersioning(domainType, originalItem, options);
+            cosmosItemOperations.add(CosmosBulkOperations.getDeleteItemOperation(String.valueOf(information.getId(entity)),
+                new PartitionKey(information.getPartitionKeyFieldValue(entity)), options));
+        });
+
+        // Default micro batch size is 100 which will be too high for most Spring cases, this configuration
+        // allows it to start at 1 and increase until it finds the appropriate batch size.
+        CosmosBulkExecutionOptions cosmosBulkExecutionOptions = new CosmosBulkExecutionOptions();
+        cosmosBulkExecutionOptions.setInitialMicroBatchSize(1);
+
+        this.getCosmosAsyncClient()
+                .getDatabase(this.getDatabaseName())
+                .getContainer(containerName)
+                .executeBulkOperations(Flux.fromIterable(cosmosItemOperations), cosmosBulkExecutionOptions)
+                .publishOn(Schedulers.parallel())
+                .onErrorResume(throwable ->
+                    CosmosExceptionUtils.exceptionHandler("Failed to delete item(s)", throwable,
+                        this.responseDiagnosticsProcessor))
+                .flatMap(response -> {
+                    CosmosUtils.fillAndProcessResponseDiagnostics(this.responseDiagnosticsProcessor,
+                        response.getResponse().getCosmosDiagnostics(), null);
+                    return Flux.empty();
+                }).blockLast();
+    }
+
     private void deleteById(String containerName, Object id, PartitionKey partitionKey,
                             CosmosItemRequestOptions options) {
         Assert.hasText(containerName, "containerName should not be null, empty or only whitespaces");
@@ -755,7 +852,7 @@ public <T> Boolean exists(@NonNull CosmosQuery query, @NonNull Class<T> domainTy
 
     /**
      * Delete the DocumentQuery, need to query the domains at first, then delete the item from the result. The cosmos db
-     * Sql API do _NOT_ support DELETE query, we cannot add one DeleteQueryGenerator.
+     * Sql API do _NOT_ support DELETE query, we cannot add one DeleteQueryGenerator. Uses bulk if possible.
      *
      * @param query The representation for query method.
      * @param domainType Class of domain
@@ -771,11 +868,46 @@ public <T> Iterable<T> delete(@NonNull CosmosQuery query, @NonNull Class<T> doma
         Assert.hasText(containerName, "container should not be null, empty or only whitespaces");
         String finalContainerName = getContainerNameOverride(containerName);
 
+        @SuppressWarnings("unchecked")
+        CosmosEntityInformation<T, Object> entityInfo = (CosmosEntityInformation<T, Object>) CosmosEntityInformation.getInstance(domainType);
+
         final List<JsonNode> results = findItemsAsFlux(query, finalContainerName, domainType).collectList().block();
         assert results != null;
-        return results.stream()
-            .map(item -> deleteItem(item, finalContainerName, domainType))
-            .collect(Collectors.toList());
+
+        if (entityInfo.getPartitionKeyFieldName() != null) {
+            Flux<CosmosItemOperation> cosmosItemOperationFlux = Flux.fromIterable(results).map(item -> {
+                T object = toDomainObject(domainType, item);
+                Object id = entityInfo.getId(object);
+                String idString = id != null ? id.toString() : "";
+                final CosmosBulkItemRequestOptions options = new CosmosBulkItemRequestOptions();
+                applyBulkVersioning(domainType, item, options);
+                return CosmosBulkOperations.getDeleteItemOperation(idString,
+                    new PartitionKey(entityInfo.getPartitionKeyFieldValue(object)), options);
+            });
+
+            // Default micro batch size is 100 which will be too high for most Spring cases, this configuration
+            // allows it to start at 1 and increase until it finds the appropriate batch size.
+            CosmosBulkExecutionOptions cosmosBulkExecutionOptions = new CosmosBulkExecutionOptions();
+            cosmosBulkExecutionOptions.setInitialMicroBatchSize(1);
+
+            this.getCosmosAsyncClient()
+                .getDatabase(this.getDatabaseName())
+                .getContainer(containerName)
+                .executeBulkOperations(cosmosItemOperationFlux, cosmosBulkExecutionOptions)
+                .publishOn(Schedulers.parallel())
+                .onErrorResume(throwable ->
+                    CosmosExceptionUtils.exceptionHandler("Failed to delete item(s)",
+                        throwable, this.responseDiagnosticsProcessor))
+                .collectList().block();
+
+            return results.stream()
+                .map(jsonNode -> toDomainObject(domainType, jsonNode))
+                .collect(Collectors.toList());
+        } else {
+            return results.stream()
+                .map(item -> deleteItem(item, finalContainerName, domainType))
+                .collect(Collectors.toList());
+        }
     }
 
     @Override
@@ -1129,6 +1261,15 @@ private void applyVersioning(Class<?> domainType,
         }
     }
 
+    private void applyBulkVersioning(Class<?> domainType,
+                                 JsonNode jsonNode,
+                                 CosmosBulkItemRequestOptions options) {
+        CosmosEntityInformation<?, ?> entityInformation = CosmosEntityInformation.getInstance(domainType);
+        if (entityInformation.isVersioned()) {
+            options.setIfMatchETag(jsonNode.get(Constants.ETAG_PROPERTY_DEFAULT_NAME).asText());
+        }
+    }
+
     private void maybeEmitEvent(CosmosMappingEvent<?> event) {
         if (canPublishEvent()) {
             this.applicationContext.publishEvent(event);

sdk/spring/azure-spring-data-cosmos/src/main/java/com/azure/spring/data/cosmos/core/ReactiveCosmosOperations.java

@@ -147,6 +147,17 @@ Mono<CosmosContainerProperties> replaceContainerProperties(String containerName,
      */
     <T> Mono<T> insert(String containerName, T objectToSave);
 
+    /**
+     * Insert all items with bulk.
+     *
+     * @param entityInformation must not be {@literal null}
+     * @param entities must not be {@literal null}
+     * @param <T> type class of domain type
+     * @param <S> type class of domain type
+     * @return Flux of result
+     */
+    <S extends T, T> Flux<S> insertAll(CosmosEntityInformation<T, ?> entityInformation, Flux<S> entities);
+
     /**
      * patches item
      * @param id must not be {@literal null}
@@ -210,7 +221,18 @@ Mono<CosmosContainerProperties> replaceContainerProperties(String containerName,
     <T> Mono<Void> deleteEntity(String containerName, T entity);
 
     /**
-     * Delete all items in a container
+     * Delete all items with bulk.
+     *
+     * @param entityInformation must not be {@literal null}
+     * @param entities must not be {@literal null}
+     * @param <T> type class of domain type
+     * @param <S> type class of domain type
+     * @return void Mono
+     */
+    <S extends T, T> Mono<Void> deleteEntities(CosmosEntityInformation<T, ?> entityInformation, Flux<S> entities);
+
+    /**
+     * Delete all items in a container. Uses bulk if possible.
      *
      * @param containerName the container name
      * @param domainType the domainType