From 1ce5af2a0fcdd35701399f9800c1dfcf14a26b5e Mon Sep 17 00:00:00 2001 From: Rakesh Kumar Date: Mon, 28 Sep 2020 19:26:56 +0530 Subject: [PATCH 01/22] Transactional batch seperate Signed-off-by: Rakesh Kumar --- .../azure/cosmos/CosmosAsyncContainer.java | 85 ++++ .../com/azure/cosmos/CosmosContainer.java | 83 ++++ .../com/azure/cosmos/TransactionalBatch.java | 306 ++++++++++++++ .../TransactionalBatchItemRequestOptions.java | 58 +++ .../TransactionalBatchOperationResult.java | 224 ++++++++++ .../TransactionalBatchRequestOptions.java | 133 ++++++ .../cosmos/TransactionalBatchResponse.java | 340 +++++++++++++++ .../implementation/AsyncDocumentClient.java | 22 + .../cosmos/implementation/HttpConstants.java | 5 + .../cosmos/implementation/OperationType.java | 4 +- .../implementation/RxDocumentClientImpl.java | 96 +++++ .../RxDocumentServiceResponse.java | 2 +- .../implementation/RxGatewayStoreModel.java | 1 + .../cosmos/implementation/TracerProvider.java | 10 + .../azure/cosmos/implementation/Utils.java | 2 +- .../implementation/batch/BatchExecUtils.java | 113 +++++ .../implementation/batch/BatchExecutor.java | 53 +++ .../batch/BatchRequestResponseConstant.java | 42 ++ .../batch/BatchResponseParser.java | 245 +++++++++++ .../batch/ItemBatchOperation.java | 199 +++++++++ .../batch/ServerBatchRequest.java | 137 ++++++ .../SinglePartitionKeyServerBatchRequest.java | 59 +++ .../HttpTransportClient.java | 8 + .../directconnectivity/WFConstants.java | 1 + .../rntbd/RntbdConstants.java | 10 +- .../rntbd/RntbdRequestFrame.java | 2 + .../rntbd/RntbdRequestHeaders.java | 16 + .../java/com/azure/cosmos/BatchTestBase.java | 216 ++++++++++ .../azure/cosmos/TransactionalBatchTest.java | 390 ++++++++++++++++++ .../batch/BatchOperationResultTests.java | 68 +++ .../batch/BatchResponsePayloadWriter.java | 46 +++ .../batch/PartitionKeyBatchResponseTests.java | 61 +++ 32 files changed, 3033 insertions(+), 4 deletions(-) create mode 100644 sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/TransactionalBatch.java create mode 100644 sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/TransactionalBatchItemRequestOptions.java create mode 100644 sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/TransactionalBatchOperationResult.java create mode 100644 sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/TransactionalBatchRequestOptions.java create mode 100644 sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/TransactionalBatchResponse.java create mode 100644 sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/batch/BatchExecUtils.java create mode 100644 sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/batch/BatchExecutor.java create mode 100644 sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/batch/BatchRequestResponseConstant.java create mode 100644 sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/batch/BatchResponseParser.java create mode 100644 sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/batch/ItemBatchOperation.java create mode 100644 sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/batch/ServerBatchRequest.java create mode 100644 sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/batch/SinglePartitionKeyServerBatchRequest.java create mode 100644 sdk/cosmos/azure-cosmos/src/test/java/com/azure/cosmos/BatchTestBase.java create mode 100644 sdk/cosmos/azure-cosmos/src/test/java/com/azure/cosmos/TransactionalBatchTest.java create mode 100644 sdk/cosmos/azure-cosmos/src/test/java/com/azure/cosmos/implementation/batch/BatchOperationResultTests.java create mode 100644 sdk/cosmos/azure-cosmos/src/test/java/com/azure/cosmos/implementation/batch/BatchResponsePayloadWriter.java create mode 100644 sdk/cosmos/azure-cosmos/src/test/java/com/azure/cosmos/implementation/batch/PartitionKeyBatchResponseTests.java diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncContainer.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncContainer.java index 19c30c48d536..0bcbffeebedb 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncContainer.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncContainer.java @@ -15,6 +15,7 @@ import com.azure.cosmos.implementation.TracerProvider; import com.azure.cosmos.implementation.Utils; import com.azure.cosmos.implementation.apachecommons.lang.StringUtils; +import com.azure.cosmos.implementation.batch.BatchExecutor; import com.azure.cosmos.implementation.query.QueryInfo; import com.azure.cosmos.models.CosmosConflictProperties; import com.azure.cosmos.models.CosmosContainerProperties; @@ -67,6 +68,7 @@ public class CosmosAsyncContainer { private final String queryItemsSpanName; private final String readAllConflictsSpanName; private final String queryConflictsSpanName; + private final String batchSpanName; private CosmosAsyncScripts scripts; CosmosAsyncContainer(String id, CosmosAsyncDatabase database) { @@ -87,6 +89,7 @@ public class CosmosAsyncContainer { this.queryItemsSpanName = "queryItems." + this.id; this.readAllConflictsSpanName = "readAllConflicts." + this.id; this.queryConflictsSpanName = "queryConflicts." + this.id; + this.batchSpanName = "transactionalBatch." + this.id; } /** @@ -493,6 +496,88 @@ private T transform(Object object, Class classType) { return Utils.getSimpleObjectMapper().convertValue(object, classType); } + /** + * Executes the transactional batch at the Azure Cosmos service as an asynchronous operation. + * + * @param transactionalBatch Batch having list of operation and partition key which will be executed by this container. + * + * @return A Mono response which contains details of execution of the transactional batch. + *

+ * If the transactional batch executes successfully, the value returned by {@link + * TransactionalBatchResponse#getResponseStatus} on the response returned will be set to 200}. + *

+ * If an operation within the transactional batch fails during execution, no changes from the batch will be + * committed and the status of the failing operation is made available by {@link + * TransactionalBatchResponse#getResponseStatus}. To obtain information about the operations that failed, the + * response can be enumerated. This returns {@link TransactionalBatchOperationResult} instances corresponding to + * each operation in the transactional batch in the order they were added to the transactional batch. For a result + * corresponding to an operation within the transactional batch, use + * {@link TransactionalBatchOperationResult#getResponseStatus} + * to access the status of the operation. If the operation was not executed or it was aborted due to the failure of + * another operation within the transactional batch, the value of this field will be 424; + * for the operation that caused the batch to abort, the value of this field + * will indicate the cause of failure. + *

+ * The value returned by {@link TransactionalBatchResponse#getResponseStatus} on the response returned may also have + * values such as 500 in case of server errors and 429. + *

+ * Use {@link TransactionalBatchResponse#isSuccessStatusCode} on the response returned to ensure that the + * transactional batch succeeded. + */ + @Beta(Beta.SinceVersion.V4_4_0) + public Mono executeTransactionalBatch(TransactionalBatch transactionalBatch) { + return executeTransactionalBatch(transactionalBatch, new TransactionalBatchRequestOptions()); + } + + /** + * Executes the transactional batch at the Azure Cosmos service as an asynchronous operation. + * + * @param transactionalBatch Batch having list of operation and partition key which will be executed by this container. + * @param requestOptions Options that apply specifically to batch request. + * + * @return A Mono response which contains details of execution of the transactional batch. + *

+ * If the transactional batch executes successfully, the value returned by {@link + * TransactionalBatchResponse#getResponseStatus} on the response returned will be set to 200}. + *

+ * If an operation within the transactional batch fails during execution, no changes from the batch will be + * committed and the status of the failing operation is made available by {@link + * TransactionalBatchResponse#getResponseStatus}. To obtain information about the operations that failed, the + * response can be enumerated. This returns {@link TransactionalBatchOperationResult} instances corresponding to + * each operation in the transactional batch in the order they were added to the transactional batch. For a result + * corresponding to an operation within the transactional batch, use + * {@link TransactionalBatchOperationResult#getResponseStatus} + * to access the status of the operation. If the operation was not executed or it was aborted due to the failure of + * another operation within the transactional batch, the value of this field will be 424; + * for the operation that caused the batch to abort, the value of this field + * will indicate the cause of failure. + *

+ * The value returned by {@link TransactionalBatchResponse#getResponseStatus} on the response returned may also have + * values such as 500 in case of server errors and 429. + *

+ * Use {@link TransactionalBatchResponse#isSuccessStatusCode} on the response returned to ensure that the + * transactional batch succeeded. + */ + @Beta(Beta.SinceVersion.V4_4_0) + public Mono executeTransactionalBatch( + TransactionalBatch transactionalBatch, + TransactionalBatchRequestOptions requestOptions) { + + if (requestOptions == null) { + requestOptions = new TransactionalBatchRequestOptions(); + } + + final BatchExecutor executor = new BatchExecutor(this, transactionalBatch, requestOptions); + final Mono responseMono = executor.executeAsync(); + + return withContext(context -> database.getClient().getTracerProvider(). + traceEnabledBatchResponsePublisher(responseMono, + context, + this.batchSpanName, + database.getId(), + database.getClient().getServiceEndpoint())); + } + /** * Reads an item. *

diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosContainer.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosContainer.java index a9f14e603b01..789d304cbf09 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosContainer.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosContainer.java @@ -270,6 +270,19 @@ private CosmosItemResponse blockDeleteItemResponse(Mono batchResponseMono) { + try { + return batchResponseMono.block(); + } catch (Exception ex) { + final Throwable throwable = Exceptions.unwrap(ex); + if (throwable instanceof CosmosException) { + throw (CosmosException) throwable; + } else { + throw ex; + } + } + } + /** * Read all items as {@link CosmosPagedIterable} in the current container. * @@ -453,6 +466,76 @@ public CosmosItemResponse deleteItem(T item, CosmosItemRequestOption return this.blockDeleteItemResponse(asyncContainer.deleteItem(item, options)); } + /** + * Executes the transactional batch at the Azure Cosmos service as an asynchronous operation. + * + * @param transactionalBatch Batch having list of operation and partition key which will be executed by this container. + * + * @return A Mono response which contains details of execution of the transactional batch. + *

+ * If the transactional batch executes successfully, the value returned by {@link + * TransactionalBatchResponse#getResponseStatus} on the response returned will be set to 200}. + *

+ * If an operation within the transactional batch fails during execution, no changes from the batch will be + * committed and the status of the failing operation is made available by {@link + * TransactionalBatchResponse#getResponseStatus}. To obtain information about the operations that failed, the + * response can be enumerated. This returns {@link TransactionalBatchOperationResult} instances corresponding to + * each operation in the transactional batch in the order they were added to the transactional batch. For a result + * corresponding to an operation within the transactional batch, use + * {@link TransactionalBatchOperationResult#getResponseStatus} + * to access the status of the operation. If the operation was not executed or it was aborted due to the failure of + * another operation within the transactional batch, the value of this field will be 424; + * for the operation that caused the batch to abort, the value of this field + * will indicate the cause of failure. + *

+ * The value returned by {@link TransactionalBatchResponse#getResponseStatus} on the response returned may also have + * values such as 500 in case of server errors and 429. + *

+ * Use {@link TransactionalBatchResponse#isSuccessStatusCode} on the response returned to ensure that the + * transactional batch succeeded. + */ + @Beta(Beta.SinceVersion.V4_4_0) + public TransactionalBatchResponse executeTransactionalBatch(TransactionalBatch transactionalBatch) { + return this.blockBatchResponse(asyncContainer.executeTransactionalBatch(transactionalBatch)); + } + + /** + * Executes the transactional batch at the Azure Cosmos service as an asynchronous operation. + * + * @param transactionalBatch Batch having list of operation and partition key which will be executed by this container. + * @param requestOptions Options that apply specifically to batch request. + * + * @return A Mono response which contains details of execution of the transactional batch. + *

+ * If the transactional batch executes successfully, the value returned by {@link + * TransactionalBatchResponse#getResponseStatus} on the response returned will be set to 200}. + *

+ * If an operation within the transactional batch fails during execution, no changes from the batch will be + * committed and the status of the failing operation is made available by {@link + * TransactionalBatchResponse#getResponseStatus}. To obtain information about the operations that failed, the + * response can be enumerated. This returns {@link TransactionalBatchOperationResult} instances corresponding to + * each operation in the transactional batch in the order they were added to the transactional batch. For a result + * corresponding to an operation within the transactional batch, use + * {@link TransactionalBatchOperationResult#getResponseStatus} + * to access the status of the operation. If the operation was not executed or it was aborted due to the failure of + * another operation within the transactional batch, the value of this field will be 424; + * for the operation that caused the batch to abort, the value of this field + * will indicate the cause of failure. + *

+ * The value returned by {@link TransactionalBatchResponse#getResponseStatus} on the response returned may also have + * values such as 500 in case of server errors and 429. + *

+ * Use {@link TransactionalBatchResponse#isSuccessStatusCode} on the response returned to ensure that the + * transactional batch succeeded. + */ + @Beta(Beta.SinceVersion.V4_4_0) + public TransactionalBatchResponse executeTransactionalBatch( + TransactionalBatch transactionalBatch, + TransactionalBatchRequestOptions requestOptions) { + + return this.blockBatchResponse(asyncContainer.executeTransactionalBatch(transactionalBatch, requestOptions)); + } + /** * Gets the Cosmos scripts using the current container as context. * diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/TransactionalBatch.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/TransactionalBatch.java new file mode 100644 index 000000000000..602e391a1905 --- /dev/null +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/TransactionalBatch.java @@ -0,0 +1,306 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +package com.azure.cosmos; + +import com.azure.cosmos.implementation.OperationType; +import com.azure.cosmos.implementation.batch.ItemBatchOperation; +import com.azure.cosmos.models.PartitionKey; + +import java.util.ArrayList; + +import static com.azure.cosmos.implementation.guava25.base.Preconditions.checkNotNull; + +/** + * Represents a batch of operations against items with the same {@link PartitionKey} in a container that will be performed + * in a transactional manner at the Azure Cosmos DB service. + *

+ * Use {@link TransactionalBatch#createTransactionalBatch(PartitionKey)} or new {@link #TransactionalBatch(PartitionKey)} + * to create an instance of TransactionalBatch + * Example + * This example atomically modifies a set of documents as a batch. + * 

{@code
+ * public class ToDoActivity {
+ *     public final String type;
+ *     public final String id;
+ *     public final String status;
+ *     public ToDoActivity(String type, String id, String status) {
+ *         this.type = type;
+ *         this.id = id;
+ *         this.status = status;
+ *     }
+ * }
+ *
+ * String activityType = "personal";
+ *
+ * ToDoActivity test1 = new ToDoActivity(activityType, "learning", "ToBeDone");
+ * ToDoActivity test2 = new ToDoActivity(activityType, "shopping", "Done");
+ * ToDoActivity test3 = new ToDoActivity(activityType, "swimming", "ToBeDone");
+ *
+ * TransactionalBatch batch = TransactionalBatch.createTransactionalBatch(new Cosmos.PartitionKey(activityType))
+ *     .createItem(test1)
+ *     .replaceItem(test2.id, test2)
+ *     .upsertItem(test3)
+ *     .deleteItem("reading");
+ *
+ * try (TransactionalBatchResponse response = container.executeTransactionalBatch(batch) {
+ *
+ *     if (!response.IsSuccessStatusCode) {
+ *        // Handle and log exception
+ *        return;
+ *     }
+ *
+ *     // Look up interested results - e.g., via typed access on operation results
+ *
+ *     TransactionalBatchOperationResult result = response.getOperationResultAtIndex(0, ToDoActivity.class);
+ *     ToDoActivity readActivity = result.getItem();
+ * }
+ * }
+ * + * Example + *

This example atomically reads a set of documents as a batch. + * 

{@code
+ * String activityType = "personal";
+ *
+ * TransactionalBatch batch = TransactionalBatch.createTransactionalBatch(new Cosmos.PartitionKey(activityType))
+ *     .readItem("playing")
+ *     .readItem("walking")
+ *     .readItem("jogging")
+ *     .readItem("running")
+ *
+ * try (TransactionalBatchResponse response = container.executeTransactionalBatch(new Cosmos.PartitionKey(activityType)) {
+ *
+ *     // Look up interested results - eg. via direct access to operation result stream
+ *
+ *     List resultItems = new ArrayList();
+ *
+ *     for (TransactionalBatchOperationResult result : response) {
+ *         resultItems.add(result.getResourceObject().toString())
+ *     }
+ * }
+ * }
+ *

+ * See: + * Limits on TransactionalBatch requests. + */ +public final class TransactionalBatch { + + private final ArrayList> operations; + private final PartitionKey partitionKey; + + public TransactionalBatch(PartitionKey partitionKey) { + this.operations = new ArrayList<>(); + this.partitionKey = partitionKey; + } + + /** + * Initializes a new instance of {@link TransactionalBatch} + * that will contain operations to be performed across multiple items in the container with the provided partition + * key in a transactional manner + * + * @param partitionKey the partition key for all items in the batch. + * + * @return A new instance of {@link TransactionalBatch}. + */ + public static TransactionalBatch createTransactionalBatch(PartitionKey partitionKey) { + return new TransactionalBatch(partitionKey); + } + + /** + * Adds an operation to create an item into the batch. + * + * @param item A JSON serializable object that must contain an id property. + * @param The type of item to be created. + * + * @return The transactional batch instance with the operation added. + */ + public TransactionalBatch createItem(T item) { + checkNotNull(item, "expected non-null item"); + return this.createItem(item, new TransactionalBatchItemRequestOptions()); + } + + /** + * Adds an operation to create an item into the batch. + * + * @param The type of item to be created. + * + * @param item A JSON serializable object that must contain an id property. + * @param requestOptions The options for the item request. + * @return The transactional batch instance with the operation added. + */ + public TransactionalBatch createItem(T item, TransactionalBatchItemRequestOptions requestOptions) { + + checkNotNull(item, "expected non-null item"); + if (requestOptions == null) { + requestOptions = new TransactionalBatchItemRequestOptions(); + } + + this.operations.add( + new ItemBatchOperation.Builder( + OperationType.Create, + this.operations.size()) + .requestOptions(requestOptions.toRequestOptions()) + .resource(item) + .build()); + + return this; + } + + /** + * Adds an operation to delete an item into the batch. + * + * @param id The unique id of the item. + * + * @return The transactional batch instance with the operation added. + */ + public TransactionalBatch deleteItem(String id) { + checkNotNull(id, "expected non-null id"); + return this.deleteItem(id, new TransactionalBatchItemRequestOptions()); + } + + /** + * Adds an operation to delete an item into the batch. + * + * @param id The unique id of the item. + * @param requestOptions The options for the item request. + * + * @return The transactional batch instance with the operation added. + */ + public TransactionalBatch deleteItem(String id, TransactionalBatchItemRequestOptions requestOptions) { + + checkNotNull(id, "expected non-null id"); + if (requestOptions == null) { + requestOptions = new TransactionalBatchItemRequestOptions(); + } + + this.operations.add(new ItemBatchOperation.Builder(OperationType.Delete, this.operations.size()) + .requestOptions(requestOptions.toRequestOptions()) + .id(id) + .build()); + + return this; + } + + /** + * Adds an operation to read an item into the batch. + * + * @param id The unique id of the item. + * + * @return The transactional batch instance with the operation added. + */ + public TransactionalBatch readItem(String id) { + checkNotNull(id, "expected non-null id"); + return this.readItem(id, new TransactionalBatchItemRequestOptions()); + } + + /** + * Adds an operation to read an item into the batch. + * + * @param id The unique id of the item. + * @param requestOptions The options for the item request. + * + * @return The transactional batch instance with the operation added. + */ + public TransactionalBatch readItem(String id, TransactionalBatchItemRequestOptions requestOptions) { + + checkNotNull(id, "expected non-null id"); + if (requestOptions == null) { + requestOptions = new TransactionalBatchItemRequestOptions(); + } + + this.operations.add(new ItemBatchOperation.Builder(OperationType.Read, this.operations.size()) + .requestOptions(requestOptions.toRequestOptions()) + .id(id) + .build()); + + return this; + } + + /** + * Adds an operation to replace an item into the batch. + * + * @param id The unique id of the item. + * @param item A JSON serializable object that must contain an id property. + * @param The type of item to be replaced. + * + * @return The transactional batch instance with the operation added. + */ + public TransactionalBatch replaceItem(String id, T item) { + checkNotNull(id, "expected non-null id"); + checkNotNull(item, "expected non-null item"); + return this.replaceItem(id, item, new TransactionalBatchItemRequestOptions()); + } + + /** + * Adds an operation to replace an item into the batch. + * + * @param The type of item to be replaced. + * + * @param id The unique id of the item. + * @param item A JSON serializable object that must contain an id property. + * @param requestOptions The options for the item request. + * @return The transactional batch instance with the operation added. + */ + public TransactionalBatch replaceItem( + String id, T item, TransactionalBatchItemRequestOptions requestOptions) { + + checkNotNull(id, "expected non-null id"); + checkNotNull(item, "expected non-null item"); + if (requestOptions == null) { + requestOptions = new TransactionalBatchItemRequestOptions(); + } + + this.operations.add(new ItemBatchOperation.Builder(OperationType.Replace, this.operations.size()) + .requestOptions(requestOptions.toRequestOptions()) + .resource(item) + .id(id) + .build()); + + return this; + } + + /** + * Adds an operation to upsert an item into the batch. + * + * @param item A JSON serializable object that must contain an id property. + * @param The type of item to be upserted. + * + * @return The transactional batch instance with the operation added. + */ + public TransactionalBatch upsertItem(T item) { + checkNotNull(item, "expected non-null item"); + return this.upsertItem(item, new TransactionalBatchItemRequestOptions()); + } + + /** + * Adds an operation to upsert an item into the batch. + * + * @param The type of item to be upserted. + * + * @param item A JSON serializable object that must contain an id property. + * @param requestOptions The options for the item request. + * @return The transactional batch instance with the operation added. + */ + public TransactionalBatch upsertItem(T item, TransactionalBatchItemRequestOptions requestOptions) { + + checkNotNull(item, "expected non-null item"); + if (requestOptions == null) { + requestOptions = new TransactionalBatchItemRequestOptions(); + } + + this.operations.add(new ItemBatchOperation.Builder(OperationType.Upsert, this.operations.size()) + .requestOptions(requestOptions.toRequestOptions()) + .resource(item) + .build()); + + return this; + } + + public ArrayList> getOperations() { + return operations; + } + + public PartitionKey getPartitionKey() { + return partitionKey; + } +} diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/TransactionalBatchItemRequestOptions.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/TransactionalBatchItemRequestOptions.java new file mode 100644 index 000000000000..79e706fab39b --- /dev/null +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/TransactionalBatchItemRequestOptions.java @@ -0,0 +1,58 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +package com.azure.cosmos; + +import com.azure.cosmos.implementation.RequestOptions; + +public final class TransactionalBatchItemRequestOptions { + private String ifMatchETag; + private String ifNoneMatchETag; + + /** + * Gets the If-Match (ETag) associated with the request in the Azure Cosmos DB service. + * + * @return ifMatchETag the ifMatchETag associated with the request. + */ + public String getIfMatchETag() { + return this.ifMatchETag; + } + + /** + * Sets the If-Match (ETag) associated with the request in the Azure Cosmos DB service. + * + * @param ifMatchETag the ifMatchETag associated with the request. + * @return the current request options + */ + public TransactionalBatchItemRequestOptions setIfMatchETag(final String ifMatchETag) { + this.ifMatchETag = ifMatchETag; + return this; + } + + /** + * Gets the If-None-Match (ETag) associated with the request in the Azure Cosmos DB service. + * + * @return the ifNoneMatchETag associated with the request. + */ + public String getIfNoneMatchETag() { + return this.ifNoneMatchETag; + } + + /** + * Sets the If-None-Match (ETag) associated with the request in the Azure Cosmos DB service. + * + * @param ifNoneMatchEtag the ifNoneMatchETag associated with the request. + * @return the current request options + */ + public TransactionalBatchItemRequestOptions setIfNoneMatchETag(final String ifNoneMatchEtag) { + this.ifNoneMatchETag = ifNoneMatchEtag; + return this; + } + + public RequestOptions toRequestOptions() { + final RequestOptions requestOptions = new RequestOptions(); + requestOptions.setIfMatchETag(getIfMatchETag()); + requestOptions.setIfNoneMatchETag(getIfNoneMatchETag()); + return requestOptions; + } +} diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/TransactionalBatchOperationResult.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/TransactionalBatchOperationResult.java new file mode 100644 index 000000000000..91f61df63695 --- /dev/null +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/TransactionalBatchOperationResult.java @@ -0,0 +1,224 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +package com.azure.cosmos; + +import com.fasterxml.jackson.databind.node.ObjectNode; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +import java.time.Duration; + +import static com.azure.cosmos.implementation.guava25.base.Preconditions.checkNotNull; + +/** + * Represents a result for a specific operation that was part of a {@link TransactionalBatch} request. + * + * @param the type parameter + */ +public final class TransactionalBatchOperationResult implements AutoCloseable { + + private final static Logger logger = LoggerFactory.getLogger(TransactionalBatchOperationResult.class); + + private String eTag; + private double requestCharge; + private TResource item; + private ObjectNode resourceObject; + private int responseStatus; + private Duration retryAfter; + private int subStatusCode; + + /** + * Instantiates a new Transactional batch operation result. + * + * @param responseStatus the response status + */ + public TransactionalBatchOperationResult(final int responseStatus) { + this.responseStatus = responseStatus; + } + + /** + * Instantiates a new Transactional batch operation result. + * + * @param other the other + */ + public TransactionalBatchOperationResult(final TransactionalBatchOperationResult other) { + + checkNotNull(other, "expected non-null other"); + + this.responseStatus = other.responseStatus; + this.subStatusCode = other.subStatusCode; + this.eTag = other.eTag; + this.requestCharge = other.requestCharge; + this.retryAfter = other.retryAfter; + this.resourceObject = other.resourceObject; + } + + /** + * Instantiates a new Transactional batch operation result. + * + * @param result the result + * @param item the item + */ + public TransactionalBatchOperationResult(TransactionalBatchOperationResult result, TResource item) { + this(result); + this.item = item; + } + + /** + * Initializes a new instance of the {@link TransactionalBatchOperationResult} class. + */ + public TransactionalBatchOperationResult() { + } + + /** + * Gets the entity tag associated with the current item. + *

+ * ETags are used for concurrency checking when updating resources. + * + * @return Entity tag associated with the current item. + */ + public String getETag() { + return this.eTag; + } + + /** + * Sets e tag. + * + * @param value the value + * + * @return the e tag + */ + public TransactionalBatchOperationResult setETag(final String value) { + this.eTag = value; + return this; + } + + /** + * Gets the request charge in request units for the current operation. + * + * @return Request charge in request units for the current operation. + */ + public double getRequestCharge() { + return requestCharge; + } + + /** + * Sets request charge. + * + * @param value the value + * + * @return the request charge + */ + public TransactionalBatchOperationResult setRequestCharge(final double value) { + this.requestCharge = value; + return this; + } + + /** + * Gets the item associated with the current result. + * + * @return Resource associated with the current result. + */ + public TResource getItem() { + return this.item; + } + + /** + * Sets item. + * + * @param value the value + * + * @return the item + */ + public TransactionalBatchOperationResult setItem(final TResource value) { + this.item = value; + return this; + } + + /** + * Gets retry after. + * + * @return the retry after + */ + public Duration getRetryAfter() { + return this.retryAfter; + } + + /** + * Sets retry after. + * + * @param value the value + * + * @return the retry after + */ + public TransactionalBatchOperationResult setRetryAfter(final Duration value) { + this.retryAfter = value; + return this; + } + + /** + * Gets sub status code. + * + * @return the sub status code + */ + public int getSubStatusCode() { + return this.subStatusCode; + } + + /** + * Sets sub status code. + * + * @param value the value + * + * @return the sub status code + */ + public TransactionalBatchOperationResult setSubStatusCode(final int value) { + this.subStatusCode = value; + return this; + } + + /** + * Gets a value indicating whether the current operation completed successfully. + * + * @return {@code true} if the current operation completed successfully; {@code false} otherwise. + */ + public boolean isSuccessStatusCode() { + return 200 <= this.responseStatus && this.responseStatus <= 299; + } + + /** + * Gets response status. + * + * @return the response status + */ + public int getResponseStatus() { + return this.responseStatus; + } + + public void setResponseStatus(int value) { + this.responseStatus = value; + } + + public ObjectNode getResourceObject() { + return resourceObject; + } + + public void setResourceObject(ObjectNode resourceObject) { + this.resourceObject = resourceObject; + } + + @Override + public void close() { + try { + if (this.item instanceof AutoCloseable) { + ((AutoCloseable) this.item).close(); // assumes an idempotent close implementation + } + } catch (Exception ex) { + logger.debug("Unexpected failure in closing item", ex); + } + + this.item = null; + this.resourceObject = null; + } +} diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/TransactionalBatchRequestOptions.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/TransactionalBatchRequestOptions.java new file mode 100644 index 000000000000..b92db378eea5 --- /dev/null +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/TransactionalBatchRequestOptions.java @@ -0,0 +1,133 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +package com.azure.cosmos; + +import com.azure.cosmos.implementation.RequestOptions; +import com.azure.cosmos.models.IndexingDirective; + +public final class TransactionalBatchRequestOptions { + private ConsistencyLevel consistencyLevel; + private IndexingDirective indexingDirective; + private String sessionToken; + private String ifMatchETag; + private String ifNoneMatchETag; + + /** + * Constructor + */ + public TransactionalBatchRequestOptions() { + super(); + } + + /** + * Gets the If-Match (ETag) associated with the request in the Azure Cosmos DB service. + * + * @return the ifMatchETag associated with the request. + */ + public String getIfMatchETag() { + return this.ifMatchETag; + } + + /** + * Sets the If-Match (ETag) associated with the request in the Azure Cosmos DB service. + * + * @param ifMatchETag the ifMatchETag associated with the request. + * @return the current request options + */ + public TransactionalBatchRequestOptions setIfMatchETag(String ifMatchETag) { + this.ifMatchETag = ifMatchETag; + return this; + } + + /** + * Gets the If-None-Match (ETag) associated with the request in the Azure Cosmos DB service. + * + * @return the ifNoneMatchETag associated with the request. + */ + public String getIfNoneMatchETag() { + return this.ifNoneMatchETag; + } + + /** + * Sets the If-None-Match (ETag) associated with the request in the Azure Cosmos DB service. + * + * @param ifNoneMatchETag the ifNoneMatchETag associated with the request. + * @return the current request options + */ + public TransactionalBatchRequestOptions setIfNoneMatchETag(String ifNoneMatchETag) { + this.ifNoneMatchETag = ifNoneMatchETag; + return this; + } + + /** + * Gets the consistency level required for the request. + * + * @return the consistency level. + */ + + public ConsistencyLevel getConsistencyLevel() { + return consistencyLevel; + } + + /** + * Sets the consistency level required for the request. + * + * @param consistencyLevel the consistency level. + * @return the CosmosItemRequestOptions. + */ + TransactionalBatchRequestOptions setConsistencyLevel(ConsistencyLevel consistencyLevel) { + this.consistencyLevel = consistencyLevel; + return this; + } + + /** + * Gets the indexing directive (index, do not index etc). + * + * @return the indexing directive. + */ + public IndexingDirective getIndexingDirective() { + return indexingDirective; + } + + /** + * Sets the indexing directive (index, do not index etc). + * + * @param indexingDirective the indexing directive. + * @return the CosmosItemRequestOptions. + */ + public TransactionalBatchRequestOptions setIndexingDirective(IndexingDirective indexingDirective) { + this.indexingDirective = indexingDirective; + return this; + } + + /** + * Gets the token for use with session consistency. + * + * @return the session token. + */ + public String getSessionToken() { + return sessionToken; + } + + /** + * Sets the token for use with session consistency. + * + * @param sessionToken the session token. + * @return the CosmosItemRequestOptions. + */ + public TransactionalBatchRequestOptions setSessionToken(String sessionToken) { + this.sessionToken = sessionToken; + return this; + } + + public RequestOptions toRequestOptions() { + final RequestOptions requestOptions = new RequestOptions(); + requestOptions.setIfMatchETag(getIfMatchETag()); + requestOptions.setIfNoneMatchETag(getIfNoneMatchETag()); + requestOptions.setConsistencyLevel(getConsistencyLevel()); + requestOptions.setIndexingDirective(indexingDirective); + requestOptions.setSessionToken(sessionToken); + return requestOptions; + } +} diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/TransactionalBatchResponse.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/TransactionalBatchResponse.java new file mode 100644 index 000000000000..f3e90b2c6e90 --- /dev/null +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/TransactionalBatchResponse.java @@ -0,0 +1,340 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +package com.azure.cosmos; + +import com.azure.cosmos.implementation.HttpConstants; +import com.azure.cosmos.implementation.JsonSerializable; +import com.azure.cosmos.implementation.apachecommons.collections.list.UnmodifiableList; +import com.azure.cosmos.implementation.batch.ItemBatchOperation; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +import java.io.IOException; +import java.time.Duration; +import java.util.ArrayList; +import java.util.Collection; +import java.util.Iterator; +import java.util.List; +import java.util.ListIterator; +import java.util.Map; +import java.util.stream.Stream; + +import static com.azure.cosmos.implementation.guava25.base.Preconditions.checkArgument; +import static com.azure.cosmos.implementation.guava25.base.Preconditions.checkNotNull; + +/** + * Response of a {@link TransactionalBatch} request. + */ +public class TransactionalBatchResponse implements AutoCloseable, List> { + + private final static Logger logger = LoggerFactory.getLogger(TransactionalBatchResponse.class); + + private Map responseHeaders; + private final int responseStatus; + private String errorMessage; + private List> results; + private int subStatusCode; + private List> operations; + private CosmosDiagnostics cosmosDiagnostics; + + /** + * Initializes a new instance of the {@link TransactionalBatchResponse} class. + * + * @param responseStatus the response status. + * @param subStatusCode the response sub-status code. + * @param errorMessage an error message or {@code null}. + * @param responseHeaders the response http headers + * @param cosmosDiagnostics the diagnostic + * @param operations a {@link List list} of {@link ItemBatchOperation batch operations}. + */ + public TransactionalBatchResponse( + final int responseStatus, + final int subStatusCode, + final String errorMessage, + final Map responseHeaders, + final CosmosDiagnostics cosmosDiagnostics, + final List> operations) { + + checkNotNull(responseStatus, "expected non-null responseStatus"); + checkNotNull(responseHeaders, "expected non-null responseHeaders"); + checkNotNull(operations, "expected non-null operations"); + + this.responseStatus = responseStatus; + this.subStatusCode = subStatusCode; + this.errorMessage = errorMessage; + this.responseHeaders = responseHeaders; + this.cosmosDiagnostics = cosmosDiagnostics; + this.operations = UnmodifiableList.unmodifiableList(operations); + this.results = new ArrayList<>(); + } + + public void createAndPopulateResults(final List> operations, final int retryAfterMilliseconds) { + for (int i = 0; i < operations.size(); i++) { + this.results.add( + new TransactionalBatchOperationResult<>(this.getResponseStatus()) + .setSubStatusCode(this.getSubStatusCode()) + .setRetryAfter(Duration.ofMillis(retryAfterMilliseconds))); + } + } + + /** + * Gets the result of the operation at the provided index in the current {@link TransactionalBatchResponse batch}. + *

+ * @param the type parameter. + * @param index 0-based index of the operation in the batch whose result needs to be returned. + * de-serialized, when present. + * @param type class type for which deserialization is needed. + * + * @return TransactionalBatchOperationResult containing the individual result of operation. + * @throws IOException if the body of the resource cannot be read. + */ + public TransactionalBatchOperationResult getOperationResultAtIndex( + final int index, + final Class type) throws IOException { + + checkArgument(index >= 0, "expected non-negative index"); + checkNotNull(type, "expected non-null type"); + + final TransactionalBatchOperationResult result = this.results.get(index); + T item = null; + + if (result.getResourceObject() != null) { + item = new JsonSerializable(result.getResourceObject()).toObject(type); + } + + return new TransactionalBatchOperationResult(result, item); + } + + public CosmosDiagnostics getCosmosDiagnostics() { + return cosmosDiagnostics; + } + + /** + * Gets the number of operation results. + */ + public int size() { + return this.results == null ? 0 : this.results.size(); + } + + /** + * Returns a value indicating whether the batch was successfully processed. + * + * @return a value indicating whether the batch was successfully processed. + */ + public boolean isSuccessStatusCode() { + return this.responseStatus >= 200 && this.responseStatus <= 299; + } + + /** + * Gets all the activity IDs associated with the response. + * + * @return an enumerable that contains the activity IDs. + */ + public Iterable getActivityIds() { + return Stream.of(this.getActivityId())::iterator; + } + + /** + * Gets the activity ID that identifies the server request made to execute the batch. + * + * @return the activity ID that identifies the server request made to execute the batch. + */ + public String getActivityId() { + return this.responseHeaders.get(HttpConstants.HttpHeaders.ACTIVITY_ID); + } + + public final List> getBatchOperations() { + return this.operations; + } + + /** + * Gets the reason for the failure of the batch request, if any, or {@code null}. + * + * @return the reason for the failure of the batch request, if any, or {@code null}. + */ + public String getErrorMessage() { + return this.errorMessage; + } + + public void setErrorMessage(String value) { + this.errorMessage = value; + } + + /** + * Gets the request charge for the batch request. + * + * @return the request charge measured in request units. + */ + public double getRequestCharge() { + final String value = this.responseHeaders.get(HttpConstants.HttpHeaders.REQUEST_CHARGE); + try { + return Double.valueOf(value); + } catch (NumberFormatException e) { + logger.warn("INVALID x-ms-request-charge value {}.", value); + return 0; + } + } + + /** + * Gets the response status code of the batch request. + * + * @return the response status code of the batch request. + */ + public int getResponseStatus() { + return this.responseStatus; + } + + /** + * Gets the response headers. + * + * @return the response header map. + */ + public Map getResponseHeaders() { + return responseHeaders; + } + + /** + * Gets the amount of time to wait before retrying this or any other request due to throttling. + * + * @return the amount of time to wait before retrying this or any other request due to throttling. + */ + public Duration getRetryAfter() { + if (this.responseHeaders.containsKey(HttpConstants.HttpHeaders.RETRY_AFTER)) { + return Duration.parse(this.responseHeaders.get(HttpConstants.HttpHeaders.RETRY_AFTER)); + } + + return null; + } + + public int getSubStatusCode() { + return this.subStatusCode; + } + + /** + * Gets the result of the operation at the provided index in the batch. + * + * @param index 0-based index of the operation in the batch whose result needs to be returned. + * + * @return Result of operation at the provided index in the batch. + */ + @Override + public TransactionalBatchOperationResult get(int index) { + return this.results.get(index); + } + + @Override + public int indexOf(Object o) { + return this.results.indexOf(o); + } + + @Override + public Iterator> iterator() { + return this.results.iterator(); + } + + @Override + public int lastIndexOf(Object o) { + return 0; + } + + @Override + public ListIterator> listIterator() { + return null; + } + + @Override + public ListIterator> listIterator(int index) { + return null; + } + + @Override + public boolean remove(Object result) { + return this.results.remove(result); + } + + @Override + public TransactionalBatchOperationResult remove(int index) { + return null; + } + + @Override + public boolean removeAll(Collection collection) { + return this.results.removeAll(collection); + } + + @Override + public boolean retainAll(Collection collection) { + return this.results.retainAll(collection); + } + + @Override + public TransactionalBatchOperationResult set(int index, TransactionalBatchOperationResult result) { + return this.results.set(index, result); + } + + @Override + public List> subList(int fromIndex, int toIndex) { + return this.results.subList(fromIndex, toIndex); + } + + @Override + public Object[] toArray() { + return this.results.toArray(); + } + + @Override + public T[] toArray(T[] a) { + return this.results.toArray(a); + } + + @Override + public boolean isEmpty() { + return this.results.isEmpty(); + } + + @Override + public boolean add(TransactionalBatchOperationResult result) { + return this.results.add(result); + } + + @Override + public void add(int index, TransactionalBatchOperationResult element) { + this.results.add(index, element); + } + + @Override + public boolean addAll(Collection> collection) { + return this.results.addAll(collection); + } + + @Override + public boolean addAll(int index, Collection> collection) { + return this.results.addAll(index, collection); + } + + @Override + public void clear() { + this.results.clear(); + } + + @Override + public boolean contains(Object result) { + return this.results.contains(result); + } + + @Override + public boolean containsAll(Collection c) { + return false; + } + + /** + * Closes the current {@link TransactionalBatchResponse}. + */ + public void close() { + this.operations = null; + this.responseHeaders = null; + this.results = null; + this.cosmosDiagnostics = null; + } +} diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/AsyncDocumentClient.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/AsyncDocumentClient.java index 8cc921f0aa02..a49e3c049b7d 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/AsyncDocumentClient.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/AsyncDocumentClient.java @@ -4,7 +4,11 @@ import com.azure.core.credential.AzureKeyCredential; import com.azure.cosmos.ConsistencyLevel; +import com.azure.cosmos.implementation.batch.ServerBatchRequest; +import com.azure.cosmos.TransactionalBatchResponse; import com.azure.cosmos.implementation.apachecommons.lang.StringUtils; +import com.azure.cosmos.implementation.caches.RxClientCollectionCache; +import com.azure.cosmos.implementation.caches.RxPartitionKeyRangeCache; import com.azure.cosmos.models.CosmosItemIdentity; import com.azure.cosmos.models.CosmosQueryRequestOptions; import com.azure.cosmos.models.FeedResponse; @@ -791,6 +795,24 @@ Flux> queryStoredProcedures(String collectionLink, Mono executeStoredProcedure(String storedProcedureLink, RequestOptions options, List procedureParams); + /** + * Executes a batch request + *

+ * After subscription the operation will be performed. + * The {@link Mono} upon successful completion will contain a batch response which will have individual responses. + * In case of failure the {@link Mono} will error. + * + * @param collectionLink the link to the parent document collection. + * @param serverBatchRequest the batch request with the content and flags. + * @param options the request options. + * @param disableAutomaticIdGeneration the flag for disabling automatic id generation. + * @return a {@link Mono} containing the single resource response with the created document or an error. + */ + Mono executeBatchRequest(String collectionLink, + ServerBatchRequest serverBatchRequest, + RequestOptions options, + boolean disableAutomaticIdGeneration); + /** * Creates a trigger. *

diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/HttpConstants.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/HttpConstants.java index 3f1cd14e93be..c5d4562f4a25 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/HttpConstants.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/HttpConstants.java @@ -249,6 +249,11 @@ public static class HttpHeaders { public static final String API_TYPE = "x-ms-cosmos-apitype"; public static final String QUERY_METRICS = "x-ms-documentdb-query-metrics"; + // Batch operations + public static final String IS_BATCH_ATOMIC = "x-ms-cosmos-batch-atomic"; + public static final String IS_BATCH_ORDERED = "x-ms-cosmos-batch-ordered"; + public static final String IS_BATCH_REQUEST = "x-ms-cosmos-is-batch-request"; + public static final String SHOULD_BATCH_CONTINUE_ON_ERROR = "x-ms-cosmos-batch-continue-on-error"; } public static class A_IMHeaderValues { diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/OperationType.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/OperationType.java index 01bc4e85ec03..d3fda30412e7 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/OperationType.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/OperationType.java @@ -10,6 +10,7 @@ public enum OperationType { AbortPartitionMigration, AbortSplit, AddComputeGatewayRequestCharges, + Batch, BatchApply, BatchReportThroughputUtilization, CompletePartitionMigration, @@ -49,6 +50,7 @@ public boolean isWriteOperation() { this == ExecuteJavaScript || this == Replace || this == Upsert || - this == Update; + this == Update || + this == Batch; } } diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/RxDocumentClientImpl.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/RxDocumentClientImpl.java index 26b6e8965ef5..7227b6b4aa38 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/RxDocumentClientImpl.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/RxDocumentClientImpl.java @@ -7,6 +7,10 @@ import com.azure.cosmos.ConnectionMode; import com.azure.cosmos.ConsistencyLevel; import com.azure.cosmos.DirectConnectionConfig; +import com.azure.cosmos.implementation.batch.BatchResponseParser; +import com.azure.cosmos.implementation.batch.ServerBatchRequest; +import com.azure.cosmos.implementation.batch.SinglePartitionKeyServerBatchRequest; +import com.azure.cosmos.TransactionalBatchResponse; import com.azure.cosmos.implementation.apachecommons.lang.StringUtils; import com.azure.cosmos.implementation.caches.RxClientCollectionCache; import com.azure.cosmos.implementation.caches.RxCollectionCache; @@ -51,6 +55,7 @@ import java.net.URI; import java.net.URLEncoder; import java.nio.ByteBuffer; +import java.nio.charset.StandardCharsets; import java.time.Instant; import java.util.ArrayList; import java.util.Collections; @@ -72,6 +77,8 @@ import static com.azure.cosmos.BridgeInternal.toFeedResponsePage; import static com.azure.cosmos.BridgeInternal.toResourceResponse; import static com.azure.cosmos.BridgeInternal.toStoredProcedureResponse; +import static com.azure.cosmos.implementation.guava25.base.Preconditions.checkArgument; +import static com.azure.cosmos.implementation.guava25.base.Preconditions.checkNotNull; import static com.azure.cosmos.models.ModelBridgeInternal.serializeJsonToByteBuffer; import static com.azure.cosmos.models.ModelBridgeInternal.toDatabaseAccount; @@ -1176,6 +1183,64 @@ private Mono getCreateDocumentRequest(DocumentClientRe return addPartitionKeyInformation(request, content, document, options, collectionObs); } + private RxDocumentServiceRequest getBatchDocumentRequest(DocumentClientRetryPolicy requestRetryPolicy, + String documentCollectionLink, + ServerBatchRequest serverBatchRequest, + RequestOptions options, + boolean disableAutomaticIdGeneration) { + + checkArgument(StringUtils.isNotEmpty(documentCollectionLink), "expected non empty documentCollectionLink"); + checkNotNull(serverBatchRequest, "expected non null serverBatchRequest"); + + Instant serializationStartTimeUTC = Instant.now(); + ByteBuffer content = ByteBuffer.wrap(serverBatchRequest.transferRequestBody().getBytes(StandardCharsets.UTF_8)); + Instant serializationEndTimeUTC = Instant.now(); + SerializationDiagnosticsContext.SerializationDiagnostics serializationDiagnostics = new SerializationDiagnosticsContext.SerializationDiagnostics( + serializationStartTimeUTC, + serializationEndTimeUTC, + SerializationDiagnosticsContext.SerializationType.ITEM_SERIALIZATION); + + String path = Utils.joinPath(documentCollectionLink, Paths.DOCUMENTS_PATH_SEGMENT); + Map requestHeaders = this.getRequestHeaders(options, ResourceType.Document, OperationType.Batch); + + RxDocumentServiceRequest request = RxDocumentServiceRequest.create( + OperationType.Batch, + ResourceType.Document, + path, + requestHeaders, + options, + content); + + if (requestRetryPolicy != null) { + requestRetryPolicy.onBeforeSendRequest(request); + } + + SerializationDiagnosticsContext serializationDiagnosticsContext = BridgeInternal.getSerializationDiagnosticsContext(request.requestContext.cosmosDiagnostics); + if (serializationDiagnosticsContext != null) { + serializationDiagnosticsContext.addSerializationDiagnostics(serializationDiagnostics); + } + + return addBatchHeaders(request, serverBatchRequest); + } + + private RxDocumentServiceRequest addBatchHeaders(RxDocumentServiceRequest request, + ServerBatchRequest serverBatchRequest) { + + if(serverBatchRequest instanceof SinglePartitionKeyServerBatchRequest) { + PartitionKeyInternal partitionKeyInternal = BridgeInternal.getPartitionKeyInternal(((SinglePartitionKeyServerBatchRequest) serverBatchRequest).getPartitionKey()); + request.setPartitionKeyInternal(partitionKeyInternal); + request.getHeaders().put(HttpConstants.HttpHeaders.PARTITION_KEY, Utils.escapeNonAscii(partitionKeyInternal.toJson())); + } else { + throw new UnsupportedOperationException("Unknown Server request."); + } + + request.getHeaders().put(HttpConstants.HttpHeaders.IS_BATCH_REQUEST, Boolean.TRUE.toString()); + request.getHeaders().put(HttpConstants.HttpHeaders.IS_BATCH_ATOMIC, String.valueOf(serverBatchRequest.isAtomicBatch())); + request.getHeaders().put(HttpConstants.HttpHeaders.SHOULD_BATCH_CONTINUE_ON_ERROR, String.valueOf(serverBatchRequest.isShouldContinueOnError())); + + return request; + } + private void populateHeaders(RxDocumentServiceRequest request, RequestVerb httpMethod) { request.getHeaders().put(HttpConstants.HttpHeaders.X_DATE, Utils.nowAsRFC1123()); if (this.masterKeyOrResourceToken != null || this.resourceTokensMap != null @@ -2216,6 +2281,15 @@ public Mono executeStoredProcedure(String storedProcedu return ObservableHelper.inlineIfPossibleAsObs(() -> executeStoredProcedureInternal(storedProcedureLink, options, procedureParams, documentClientRetryPolicy), documentClientRetryPolicy); } + @Override + public Mono executeBatchRequest(String collectionLink, + ServerBatchRequest serverBatchRequest, + RequestOptions options, + boolean disableAutomaticIdGeneration) { + DocumentClientRetryPolicy documentClientRetryPolicy = this.resetSessionTokenRetryPolicy.getRequestPolicy(); + return ObservableHelper.inlineIfPossibleAsObs(() -> executeBatchRequestInternal(collectionLink, serverBatchRequest, options, documentClientRetryPolicy, disableAutomaticIdGeneration), documentClientRetryPolicy); + } + private Mono executeStoredProcedureInternal(String storedProcedureLink, RequestOptions options, List procedureParams, DocumentClientRetryPolicy retryPolicy) { @@ -2248,6 +2322,28 @@ private Mono executeStoredProcedureInternal(String stor } } + private Mono executeBatchRequestInternal(String collectionLink, + ServerBatchRequest serverBatchRequest, + RequestOptions options, + DocumentClientRetryPolicy requestRetryPolicy, + boolean disableAutomaticIdGeneration) { + + try { + logger.debug("Executing a Batch request with number of operations {}", serverBatchRequest.getOperations().size()); + + RxDocumentServiceRequest documentServiceRequest = getBatchDocumentRequest(requestRetryPolicy, collectionLink, serverBatchRequest, options, disableAutomaticIdGeneration); + Mono responseObservable = create(documentServiceRequest, requestRetryPolicy); + + return responseObservable + .flatMap(serviceResponse -> BatchResponseParser.fromDocumentServiceResponseAsync(serviceResponse, serverBatchRequest, true)) + .onErrorResume(throwable -> BatchResponseParser.fromErrorResponseAsync(throwable, serverBatchRequest)); + + } catch (Exception ex) { + logger.debug("Failure in executing a batch due to [{}]", ex.getMessage(), ex); + return BatchResponseParser.fromErrorResponseAsync(ex, serverBatchRequest); + } + } + @Override public Mono> createTrigger(String collectionLink, Trigger trigger, RequestOptions options) { diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/RxDocumentServiceResponse.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/RxDocumentServiceResponse.java index 55c018e5db3f..5b5baab4a5eb 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/RxDocumentServiceResponse.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/RxDocumentServiceResponse.java @@ -180,7 +180,7 @@ private String getOwnerFullName() { return null; } - CosmosDiagnostics getCosmosDiagnostics() { + public CosmosDiagnostics getCosmosDiagnostics() { if (this.storeResponse == null) { return null; } diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/RxGatewayStoreModel.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/RxGatewayStoreModel.java index 43e529c7a2d6..74dde8562617 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/RxGatewayStoreModel.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/RxGatewayStoreModel.java @@ -339,6 +339,7 @@ private void validateOrThrow(RxDocumentServiceRequest request, private Mono invokeAsyncInternal(RxDocumentServiceRequest request) { switch (request.getOperationType()) { case Create: + case Batch: return this.create(request); case Upsert: return this.upsert(request); diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/TracerProvider.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/TracerProvider.java index 87c51d14412a..b15fcb29b76f 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/TracerProvider.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/TracerProvider.java @@ -5,6 +5,7 @@ import com.azure.core.util.Context; import com.azure.core.util.tracing.Tracer; import com.azure.cosmos.CosmosException; +import com.azure.cosmos.TransactionalBatchResponse; import com.azure.cosmos.models.CosmosItemResponse; import com.azure.cosmos.models.CosmosResponse; import reactor.core.publisher.Mono; @@ -108,6 +109,15 @@ public > Mono traceEnabledCosmosResponsePublisher (T response) -> response.getStatusCode()); } + public Mono traceEnabledBatchResponsePublisher(Mono resultPublisher, + Context context, + String spanName, + String databaseId, + String endpoint) { + return traceEnabledPublisher(resultPublisher, context, spanName, databaseId, endpoint, + TransactionalBatchResponse::getResponseStatus); + } + public Mono> traceEnabledCosmosItemResponsePublisher(Mono> resultPublisher, Context context, String spanName, diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/Utils.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/Utils.java index a4387415b5d5..348e5f28a09a 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/Utils.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/Utils.java @@ -340,7 +340,7 @@ public static boolean isCollectionChild(ResourceType type) { public static boolean isWriteOperation(OperationType operationType) { return operationType == OperationType.Create || operationType == OperationType.Upsert || operationType == OperationType.Delete || operationType == OperationType.Replace - || operationType == OperationType.ExecuteJavaScript; + || operationType == OperationType.ExecuteJavaScript || operationType == OperationType.Batch; } public static boolean isFeedRequest(OperationType requestOperationType) { diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/batch/BatchExecUtils.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/batch/BatchExecUtils.java new file mode 100644 index 000000000000..27e8cc1d2eab --- /dev/null +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/batch/BatchExecUtils.java @@ -0,0 +1,113 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +package com.azure.cosmos.implementation.batch; + +import com.azure.cosmos.implementation.HttpConstants; +import com.azure.cosmos.implementation.OperationType; +import com.azure.cosmos.implementation.RequestOptions; +import com.azure.cosmos.implementation.directconnectivity.WFConstants; +import io.netty.buffer.ByteBufUtil; + +import java.util.List; +import java.util.Map; + +import static com.azure.cosmos.implementation.batch.BatchRequestResponseConstant.OPERATION_CREATE; +import static com.azure.cosmos.implementation.batch.BatchRequestResponseConstant.OPERATION_DELETE; +import static com.azure.cosmos.implementation.batch.BatchRequestResponseConstant.OPERATION_READ; +import static com.azure.cosmos.implementation.batch.BatchRequestResponseConstant.OPERATION_REPLACE; +import static com.azure.cosmos.implementation.batch.BatchRequestResponseConstant.OPERATION_UPSERT; +import static com.azure.cosmos.implementation.guava25.base.Preconditions.checkArgument; +import static com.azure.cosmos.implementation.guava27.Strings.lenientFormat; + +/** + * Util methods for batch requests/response. + */ +public class BatchExecUtils { + + static void ensureValid( + final List> operations, + final RequestOptions options) { + + final String errorMessage = BatchExecUtils.isValid(operations, options); + checkArgument(errorMessage == null, errorMessage); + } + + public static String isValid(final List> operations, final RequestOptions batchOptions) { + + String errorMessage = null; + + if (operations == null) { + errorMessage = "expected non-null operations"; + } + + if (errorMessage == null && operations.size() == 0) { + errorMessage = "expected operations.size > 0"; + } + + if (errorMessage == null && batchOptions != null) { + if (batchOptions.getIfMatchETag() != null || batchOptions.getIfNoneMatchETag() != null) { + errorMessage = "one or more request options provided on the batch request are not supported"; + } + } + + if (errorMessage == null) { + for (ItemBatchOperation operation : operations) { + + final RequestOptions batchOperationOptions = operation.getRequestOptions(); + final Map batchOperationProperties = batchOperationOptions != null ? batchOperationOptions.getProperties() : null; + + if (batchOperationProperties != null + && (batchOperationProperties.containsKey(WFConstants.BackendHeaders.EFFECTIVE_PARTITION_KEY) + || batchOperationProperties.containsKey(WFConstants.BackendHeaders.EFFECTIVE_PARTITION_KEY_STRING) + || batchOperationProperties.containsKey(HttpConstants.HttpHeaders.PARTITION_KEY))) { + + final String epkString = (String) batchOperationProperties.computeIfPresent( + WFConstants.BackendHeaders.EFFECTIVE_PARTITION_KEY_STRING, + (k, v) -> v instanceof String ? v : null); + + final byte[] epk = (byte[]) batchOperationProperties.computeIfPresent( + WFConstants.BackendHeaders.EFFECTIVE_PARTITION_KEY, + (k, v) -> v instanceof byte[] ? v : null); + + final String pkString = (String) batchOperationProperties.computeIfPresent( + HttpConstants.HttpHeaders.PARTITION_KEY, + (k, v) -> v instanceof String ? v : null); + + if ((epk == null && pkString == null) || epkString == null) { + return lenientFormat( + "expected byte[] value for %s and string value for %s, not (%s, %s)", + WFConstants.BackendHeaders.EFFECTIVE_PARTITION_KEY, + WFConstants.BackendHeaders.EFFECTIVE_PARTITION_KEY_STRING, + epk == null + ? (pkString == null ? "null" : pkString) + : ByteBufUtil.hexDump(epk), epkString == null ? "null" : epkString); + } + + if (operation.getPartitionKey() != null ) { + errorMessage = "partition key and effective partition key may not both be set."; + } + } + } + } + + return errorMessage; + } + + static String getStringOperationType(OperationType operationType) { + switch (operationType) { + case Create: + return OPERATION_CREATE; + case Delete: + return OPERATION_DELETE; + case Read: + return OPERATION_READ; + case Replace: + return OPERATION_REPLACE; + case Upsert: + return OPERATION_UPSERT; + } + + return null; + } +} diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/batch/BatchExecutor.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/batch/BatchExecutor.java new file mode 100644 index 000000000000..ed0d07c294ca --- /dev/null +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/batch/BatchExecutor.java @@ -0,0 +1,53 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +package com.azure.cosmos.implementation.batch; + +import com.azure.cosmos.BridgeInternal; +import com.azure.cosmos.CosmosAsyncContainer; +import com.azure.cosmos.CosmosBridgeInternal; +import com.azure.cosmos.TransactionalBatch; +import com.azure.cosmos.TransactionalBatchRequestOptions; +import com.azure.cosmos.TransactionalBatchResponse; +import com.azure.cosmos.implementation.RequestOptions; +import reactor.core.publisher.Mono; + +import java.util.List; + +public final class BatchExecutor { + + private final CosmosAsyncContainer container; + private final RequestOptions options; + private final TransactionalBatch transactionalBatch; + + public BatchExecutor( + final CosmosAsyncContainer container, + final TransactionalBatch transactionalBatch, + final TransactionalBatchRequestOptions options) { + + this.container = container; + this.transactionalBatch = transactionalBatch; + this.options = options.toRequestOptions(); + } + + /** + * Create a batch request from list of operations and executes it. + * + * @return Response from the server. + */ + public final Mono executeAsync() { + + List> operations = this.transactionalBatch.getOperations(); + + BatchExecUtils.ensureValid(operations, this.options); + + final SinglePartitionKeyServerBatchRequest request = SinglePartitionKeyServerBatchRequest.createAsync( + this.transactionalBatch.getPartitionKey(), + operations); + request.setAtomicBatch(true); + request.setShouldContinueOnError(false); + + return CosmosBridgeInternal.getAsyncDocumentClient(container.getDatabase()) + .executeBatchRequest(BridgeInternal.getLink(container), request, options, false); + } +} diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/batch/BatchRequestResponseConstant.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/batch/BatchRequestResponseConstant.java new file mode 100644 index 000000000000..d5fcc9588ac1 --- /dev/null +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/batch/BatchRequestResponseConstant.java @@ -0,0 +1,42 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +package com.azure.cosmos.implementation.batch; + +/** + * This contains all the extra constants needed for batch/bulk. This will be usefull even if Hybrid row comes in. + * This contains all the constants we have in Backend. Any addition to backend should be added here. + */ +public class BatchRequestResponseConstant { + + // Size limits: + public static final int MAX_DIRECT_MODE_BATCH_REQUEST_BODY_SIZE_IN_BYTES = 220201; + public static final int MAX_OPERATIONS_IN_DIRECT_MODE_BATCH_REQUEST = 100; + + static final String FIELD_OPERATION_TYPE = "operationType"; + static final String FIELD_RESOURCE_TYPE = "resourceType"; + static final String FIELD_TIME_TO_LIVE_IN_SECONDS = "timeToLiveInSeconds"; + static final String FIELD_ID = "id"; + static final String FIELD_INDEXING_DIRECTIVE = "indexingDirective"; + static final String FIELD_IF_MATCH = "ifMatch"; + static final String FIELD_IF_NONE_MATCH = "ifNoneMatch"; + static final String FIELD_PARTITION_KEY = "partitionKey"; + static final String FIELD_RESOURCE_BODY = "resourceBody"; + static final String FIELD_BINARY_ID = "binaryId"; + static final String FIELD_EFFECTIVE_PARTITIONKEY = "effectivePartitionKey"; + static final String FIELD_STATUS_CODE = "statusCode"; + static final String FIELD_SUBSTATUS_CODE = "subStatusCode"; + static final String FIELD_REQUEST_CHARGE = "requestCharge"; + static final String FIELD_RETRY_AFTER_MILLISECONDS = "retryAfterMilliseconds"; + static final String FIELD_ETAG = "eTag"; + static final String FIELD_MINIMAL_RETURN_PREFERENCE = "minimalReturnPreference"; + static final String FIELD_IS_CLIENTENCRYPTED = "isClientEncrypted"; + + // Batch supported operation type for json + static final String OPERATION_CREATE = "Create"; + static final String OPERATION_PATCH = "Patch"; + static final String OPERATION_READ = "Read"; + static final String OPERATION_UPSERT = "Upsert"; + static final String OPERATION_DELETE = "Delete"; + static final String OPERATION_REPLACE = "Replace"; +} diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/batch/BatchResponseParser.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/batch/BatchResponseParser.java new file mode 100644 index 000000000000..fc39e6a0a21b --- /dev/null +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/batch/BatchResponseParser.java @@ -0,0 +1,245 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +package com.azure.cosmos.implementation.batch; + +import com.azure.cosmos.CosmosException; +import com.azure.cosmos.TransactionalBatchOperationResult; +import com.azure.cosmos.TransactionalBatchResponse; +import com.azure.cosmos.implementation.HttpConstants; +import com.azure.cosmos.implementation.JsonSerializable; +import com.azure.cosmos.implementation.RxDocumentServiceResponse; +import com.azure.cosmos.implementation.Utils; +import com.azure.cosmos.implementation.apachecommons.lang.StringUtils; +import com.fasterxml.jackson.databind.ObjectMapper; +import com.fasterxml.jackson.databind.node.ObjectNode; +import io.netty.handler.codec.http.HttpResponseStatus; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; +import reactor.core.publisher.Mono; + +import java.io.IOException; +import java.time.Duration; +import java.util.ArrayList; + +import static com.azure.cosmos.implementation.HttpConstants.HttpHeaders.RETRY_AFTER_IN_MILLISECONDS; +import static com.azure.cosmos.implementation.HttpConstants.HttpHeaders.SUB_STATUS; +import static com.azure.cosmos.implementation.guava25.base.Preconditions.checkState; + +public final class BatchResponseParser { + + private final static Logger logger = LoggerFactory.getLogger(BatchResponseParser.class); + private final static char HYBRID_V1 = 129; + + /** Creates a transactional batch response} from a exception + * + * @param throwable the {@link Throwable error}. + * @param request the {@link ServerBatchRequest batch request} that produced {@code message}. + * + * @return a Mono that provides the {@link TransactionalBatchResponse transactional batch response} created + * from {@link TransactionalBatchResponse message} when the asynchronous operation completes. + */ + public static Mono fromErrorResponseAsync( + final Throwable throwable, + final ServerBatchRequest request) { + + if (throwable instanceof CosmosException) { + final CosmosException cosmosException = (CosmosException) throwable; + final TransactionalBatchResponse response = new TransactionalBatchResponse( + cosmosException.getStatusCode(), + cosmosException.getSubStatusCode(), + cosmosException.getMessage(), + cosmosException.getResponseHeaders(), + cosmosException.getDiagnostics(), + request.getOperations()); + + response.createAndPopulateResults(request.getOperations(), 0); + return Mono.just(response); + } else { + return Mono.error(throwable); + } + } + + /** Creates a transactional batch response} from a response message + * + * @param documentServiceResponse the {@link RxDocumentServiceResponse response message}. + * @param request the {@link ServerBatchRequest batch request} that produced {@code message}. + * @param shouldPromoteOperationStatus indicates whether the operation status should be promoted. + * + * @return a Mono that provides the {@link TransactionalBatchResponse transactional batch response} created + * from {@link RxDocumentServiceResponse message} when the asynchronous operation completes. + */ + public static Mono fromDocumentServiceResponseAsync( + final RxDocumentServiceResponse documentServiceResponse, + final ServerBatchRequest request, + final boolean shouldPromoteOperationStatus) { + + TransactionalBatchResponse response = null; + final String responseContent = documentServiceResponse.getResponseBodyAsString(); + + if (StringUtils.isNotEmpty(responseContent)) { + response = BatchResponseParser.populateFromResponseContentAsync(documentServiceResponse, request, shouldPromoteOperationStatus); + + if (response == null) { + // Convert any payload read failures as InternalServerError + response = new TransactionalBatchResponse( + HttpResponseStatus.INTERNAL_SERVER_ERROR.code(), + 0, + "ServerResponseDeserializationFailure", + documentServiceResponse.getResponseHeaders(), + documentServiceResponse.getCosmosDiagnostics(), + request.getOperations()); + } + } + + int responseStatusCode = documentServiceResponse.getStatusCode(); + int responseSubStatusCode = Integer.parseInt( + documentServiceResponse.getResponseHeaders().getOrDefault(SUB_STATUS, String.valueOf(0))); + + if (response == null) { + response = new TransactionalBatchResponse( + responseStatusCode, + responseSubStatusCode, + null, + documentServiceResponse.getResponseHeaders(), + documentServiceResponse.getCosmosDiagnostics(), + request.getOperations()); + } + + if (response.size() != request.getOperations().size()) { + if (responseStatusCode >= 200 && responseStatusCode <= 299) { + // Server should be guaranteeing number of results equal to operations when + // batch request is successful - so fail as InternalServerError if this is not the case. + response = new TransactionalBatchResponse( + HttpResponseStatus.INTERNAL_SERVER_ERROR.code(), + 0, + "Invalid server response", + documentServiceResponse.getResponseHeaders(), + documentServiceResponse.getCosmosDiagnostics(), + request.getOperations()); + } + + // When the overall response status code is TooManyRequests, propagate the RetryAfter into the individual operations. + int retryAfterMilliseconds = 0; + + if (responseStatusCode == HttpResponseStatus.TOO_MANY_REQUESTS.code()) { + String retryResponseValue = documentServiceResponse.getResponseHeaders().getOrDefault(RETRY_AFTER_IN_MILLISECONDS, null); + if (StringUtils.isNotEmpty(retryResponseValue)) { + try { + retryAfterMilliseconds = Integer.parseInt(retryResponseValue); + } catch (NumberFormatException ex) { + // Do nothing. It's number format exception + } + } + } + + response.createAndPopulateResults(request.getOperations(), retryAfterMilliseconds); + } + + checkState(response.size() == request.getOperations().size(), + "Number of responses should be equal to number of operations in request."); + + return Mono.just(response); + } + + private static TransactionalBatchResponse populateFromResponseContentAsync( + final RxDocumentServiceResponse documentServiceResponse, + final ServerBatchRequest request, + final boolean shouldPromoteOperationStatus) { + + final ArrayList> results = new ArrayList<>(); + final String responseContent = documentServiceResponse.getResponseBodyAsString(); + + if (responseContent.charAt(0) != HYBRID_V1) { + // Read from a json response body. To enable hybrid row just complete the else part + final ObjectMapper mapper = Utils.getSimpleObjectMapper(); + + try{ + final ObjectNode[] objectNodes = mapper.readValue(documentServiceResponse.getResponseBodyAsString(), ObjectNode[].class); + for (ObjectNode objectInArray : objectNodes) { + final TransactionalBatchOperationResult batchOperationResult = BatchResponseParser.createBatchOperationResultFromJson(objectInArray); + results.add(batchOperationResult); + } + } catch (IOException ex) { + logger.error("Exception in parsing response", ex); + } + + } else { + // TODO(rakkuma): Implement hybrid row response parsing logic here. Parse the response hybrid row buffer + // into array list of TransactionalBatchOperationResult. Remaining part is taken care from the caller function. + logger.error("Hybrid row is not implemented right now"); + return null; + } + + int responseStatusCode = documentServiceResponse.getStatusCode(); + int responseSubStatusCode = Integer.parseInt( + documentServiceResponse.getResponseHeaders().getOrDefault(SUB_STATUS, String.valueOf(HttpConstants.SubStatusCodes.UNKNOWN))); + + // Status code of the exact operation which failed. + if (responseStatusCode == HttpResponseStatus.MULTI_STATUS.code() + && shouldPromoteOperationStatus) { + for (TransactionalBatchOperationResult result : results) { + if (result.getResponseStatus()!= HttpResponseStatus.FAILED_DEPENDENCY.code()) { + responseStatusCode = result.getResponseStatus(); + responseSubStatusCode = result.getSubStatusCode(); + break; + } + } + } + + final TransactionalBatchResponse response = new TransactionalBatchResponse( + responseStatusCode, + responseSubStatusCode, + null, + documentServiceResponse.getResponseHeaders(), + documentServiceResponse.getCosmosDiagnostics(), + request.getOperations()); + + response.addAll(results); + + return response; + } + + /** + * Read batch operation result result. + * + * TODO(rakkuma): Similarly hybrid row result needs to be parsed. + * + * @param objectNode having response for a single operation. + * + * @return the result + */ + private static TransactionalBatchOperationResult createBatchOperationResultFromJson(ObjectNode objectNode) { + final TransactionalBatchOperationResult transactionalBatchOperationResult = new TransactionalBatchOperationResult<>(); + + final JsonSerializable jsonSerializable = new JsonSerializable(objectNode); + transactionalBatchOperationResult.setResponseStatus(jsonSerializable.getInt(BatchRequestResponseConstant.FIELD_STATUS_CODE)); + + final Integer subStatusCode = jsonSerializable.getInt(BatchRequestResponseConstant.FIELD_SUBSTATUS_CODE); + if(subStatusCode != null) { + transactionalBatchOperationResult.setSubStatusCode(subStatusCode); + } + + final Double requestCharge = jsonSerializable.getDouble(BatchRequestResponseConstant.FIELD_REQUEST_CHARGE); + if(requestCharge != null) { + transactionalBatchOperationResult.setRequestCharge(requestCharge); + } + + final Integer retryAfterMilliseconds = jsonSerializable.getInt(BatchRequestResponseConstant.FIELD_RETRY_AFTER_MILLISECONDS); + if(retryAfterMilliseconds != null) { + transactionalBatchOperationResult.setRetryAfter(Duration.ofMillis(retryAfterMilliseconds)); + } + + final String etag = jsonSerializable.getString(BatchRequestResponseConstant.FIELD_ETAG); + if(etag != null) { + transactionalBatchOperationResult.setETag(etag); + } + + final ObjectNode resourceBody = jsonSerializable.getObject(BatchRequestResponseConstant.FIELD_RESOURCE_BODY); + if(resourceBody != null) { + transactionalBatchOperationResult.setResourceObject(resourceBody); + } + + return transactionalBatchOperationResult; + } +} diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/batch/ItemBatchOperation.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/batch/ItemBatchOperation.java new file mode 100644 index 000000000000..d5218c244b2f --- /dev/null +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/batch/ItemBatchOperation.java @@ -0,0 +1,199 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +package com.azure.cosmos.implementation.batch; + +import com.azure.cosmos.implementation.JsonSerializable; +import com.azure.cosmos.implementation.OperationType; +import com.azure.cosmos.implementation.RequestOptions; +import com.azure.cosmos.implementation.apachecommons.lang.StringUtils; +import com.azure.cosmos.models.PartitionKey; + +import static com.azure.cosmos.implementation.guava25.base.Preconditions.checkArgument; +import static com.azure.cosmos.implementation.guava25.base.Preconditions.checkNotNull; + +/** + * Represents an operation on an item which will be executed as part of a batch request on a container. + */ +public final class ItemBatchOperation implements AutoCloseable { + + private TResource resource; + private String materialisedResource; + + private String id; + private int operationIndex; + private PartitionKey partitionKey; + private String partitionKeyJson; + private final OperationType operationType; + private RequestOptions requestOptions; + + private ItemBatchOperation( + final OperationType operationType, + final int operationIndex, + final PartitionKey partitionKey, + final String id, + final TResource resource, + final RequestOptions requestOptions) { + + checkArgument(operationIndex >= 0, "expected operationIndex >= 0, not %s", operationIndex); + checkNotNull(operationType, "expected non-null operationType"); + + this.operationType = operationType; + this.operationIndex = operationIndex; + this.partitionKey = partitionKey; + this.id = id; + this.resource = resource; + this.requestOptions = requestOptions; + } + + // TODO(rakkuma): Similarly for hybrid row, operation needs to be written in Hybrid row. + static JsonSerializable writeOperation(final ItemBatchOperation operation) { + final JsonSerializable jsonSerializable = new JsonSerializable(); + + jsonSerializable.set(BatchRequestResponseConstant.FIELD_OPERATION_TYPE, BatchExecUtils.getStringOperationType(operation.getOperationType())); + + if (StringUtils.isNotEmpty(operation.getPartitionKeyJson())) { + // This is set in BatchAsyncContainerExecutor.resolvePartitionKeyRangeIdAsync. For transactional no need to + // pass partition key in operations as batch will have it. + jsonSerializable.set(BatchRequestResponseConstant.FIELD_PARTITION_KEY, operation.getPartitionKeyJson()); + } + + if (StringUtils.isNotEmpty(operation.getId())) { + jsonSerializable.set(BatchRequestResponseConstant.FIELD_ID, operation.getId()); + } + + if (operation.getResource() != null) { + jsonSerializable.set(BatchRequestResponseConstant.FIELD_RESOURCE_BODY, operation.getResource()); + } + + if (operation.getRequestOptions() != null) { + RequestOptions requestOptions = operation.getRequestOptions(); + + if (StringUtils.isNotEmpty(requestOptions.getIfMatchETag())) { + jsonSerializable.set(BatchRequestResponseConstant.FIELD_IF_MATCH, requestOptions.getIfMatchETag()); + } + + if (StringUtils.isNotEmpty(requestOptions.getIfNoneMatchETag())) { + jsonSerializable.set(BatchRequestResponseConstant.FIELD_IF_NONE_MATCH, requestOptions.getIfNoneMatchETag()); + } + } + + return jsonSerializable; + } + + public String getId() { + return this.id; + } + + int getOperationIndex() { + return this.operationIndex; + } + + ItemBatchOperation setOperationIndex(final int value) { + this.operationIndex = value; + return this; + } + + public OperationType getOperationType() { + return this.operationType; + } + + public PartitionKey getPartitionKey() { + return partitionKey; + } + + public ItemBatchOperation setPartitionKey(final PartitionKey value) { + partitionKey = value; + return this; + } + + private String getPartitionKeyJson() { + return partitionKeyJson; + } + + ItemBatchOperation setPartitionKeyJson(final String value) { + partitionKeyJson = value; + return this; + } + + public RequestOptions getRequestOptions() { + return requestOptions; + } + + public TResource getResource() { + return resource; + } + + public String getMaterialisedResource() { + return materialisedResource; + } + + public void setMaterialisedResource(String materialisedResource) { + this.materialisedResource = materialisedResource; + } + + /** + * Closes this {@link ItemBatchOperation}. + */ + public void close() { + try { + if (this.resource instanceof AutoCloseable) { + ((AutoCloseable) this.resource).close(); // assumes an idempotent close implementation + } + this.resource = null; + } catch (Exception ex) { + // + } + + this.materialisedResource = null; + } + + public static final class Builder { + + private final OperationType operationType; + private final int operationIndex; + private String id; + private PartitionKey partitionKey; + private RequestOptions requestOptions; + private TResource resource; + + public Builder(final OperationType type, final int index) { + + checkNotNull(type, "expected non-null type"); + checkArgument(index >= 0, "expected index >= 0, not %s", index); + + this.operationType = type; + this.operationIndex = index; + } + + public Builder id(String value) { + this.id = value; + return this; + } + + public Builder partitionKey(PartitionKey value) { + this.partitionKey = value; + return this; + } + + public Builder requestOptions(RequestOptions value) { + this.requestOptions = value; + return this; + } + + public Builder resource(TResource value) { + this.resource = value; + return this; + } + + public ItemBatchOperation build() { + return new ItemBatchOperation<>( + this.operationType, + this.operationIndex, + this.partitionKey, + this.id, + this.resource, + this.requestOptions); + } + } +} diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/batch/ServerBatchRequest.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/batch/ServerBatchRequest.java new file mode 100644 index 000000000000..133fa74ab660 --- /dev/null +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/batch/ServerBatchRequest.java @@ -0,0 +1,137 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +package com.azure.cosmos.implementation.batch; + +import com.azure.cosmos.implementation.JsonSerializable; +import com.azure.cosmos.implementation.Utils; +import com.azure.cosmos.implementation.apachecommons.collections.list.UnmodifiableList; +import com.fasterxml.jackson.databind.node.ArrayNode; + +import java.util.List; + +import static com.azure.cosmos.implementation.guava25.base.Preconditions.checkNotNull; +import static com.azure.cosmos.implementation.guava25.base.Preconditions.checkState; + +/** + * This class represents a server batch request. + */ +public abstract class ServerBatchRequest { + + private final int maxBodyLength; + private final int maxOperationCount; + + private String requestBody; + private List> operations; + private boolean isAtomicBatch = false; + private boolean shouldContinueOnError = false; + + /** + * Initializes a new {@link ServerBatchRequest request} instance. + * + * @param maxBodyLength Maximum length allowed for the request body. + * @param maxOperationCount Maximum number of operations allowed in the request. + */ + ServerBatchRequest(int maxBodyLength, int maxOperationCount) { + this.maxBodyLength = maxBodyLength; + this.maxOperationCount = maxOperationCount; + } + + /** + * Adds as many operations as possible from the given list of operations. + *

+ * Operations are added in order while ensuring the request stream never exceeds {@link #maxBodyLength}. + * + * @param operations Operations to be added; read-only. + * + * @return Any pending operations that were not included in the request. + */ + final List> createBodyStreamAsync( + final List> operations) { + return createBodyStreamAsync(operations, false); + } + + /** + * Adds as many operations as possible from the given list of operations. + * TODO(rakkuma): Similarly for hybrid row, request needs to be parsed to create a request body in any form. + * + *

+ * Operations are added in order while ensuring the request body never exceeds {@link #maxBodyLength}. + * + * @param operations operations to be added; read-only. + * @param ensureContinuousOperationIndexes specifies whether to stop adding operations to the request once there is + * non-continuity in the operation indexes. + * + * @return Any pending operations that were not included in the request. + */ + final List> createBodyStreamAsync( + final List> operations, + final boolean ensureContinuousOperationIndexes) { + + checkNotNull(operations, "expected non-null operations"); + + int totalSerializedLength = 0; + int totalOperationCount = 0; + + final ArrayNode arrayNode = Utils.getSimpleObjectMapper().createArrayNode(); + + for(ItemBatchOperation operation : operations) { + + final JsonSerializable operationJsonSerializable = ItemBatchOperation.writeOperation(operation); + final int operationSerializedLength = operationJsonSerializable.toString().length(); + + if (totalOperationCount != 0 && + (totalSerializedLength + operationSerializedLength > this.maxBodyLength || totalOperationCount + 1 > this.maxOperationCount)) { + // Apply the limit only if at least there is one operation in selected operations. + break; + } + + totalSerializedLength += operationSerializedLength; + totalOperationCount++; + + arrayNode.add(operationJsonSerializable.getPropertyBag()); + } + + this.requestBody = arrayNode.toString(); + this.operations = operations.subList(0, totalOperationCount); + List> pendingOperations = operations.subList(totalOperationCount, operations.size()); + return pendingOperations; + } + + public final String transferRequestBody() { + + checkState(this.requestBody != null, "expected non-null body"); + + final String requestBody = this.requestBody; + this.requestBody = null; + + return requestBody; + } + + /** + * Gets the list of {@link ItemBatchOperation operations} in this {@link ServerBatchRequest batch request}. + * + * The list returned by this method is unmodifiable. + * + * @return the list of {@link ItemBatchOperation operations} in this {@link ServerBatchRequest batch request}. + */ + public final List> getOperations() { + return UnmodifiableList.unmodifiableList(this.operations); + } + + public boolean isAtomicBatch() { + return this.isAtomicBatch; + } + + void setAtomicBatch(boolean atomicBatch) { + this.isAtomicBatch = atomicBatch; + } + + public boolean isShouldContinueOnError() { + return this.shouldContinueOnError; + } + + void setShouldContinueOnError(boolean shouldContinueOnError) { + this.shouldContinueOnError = shouldContinueOnError; + } +} diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/batch/SinglePartitionKeyServerBatchRequest.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/batch/SinglePartitionKeyServerBatchRequest.java new file mode 100644 index 000000000000..a499b089b4c9 --- /dev/null +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/batch/SinglePartitionKeyServerBatchRequest.java @@ -0,0 +1,59 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +package com.azure.cosmos.implementation.batch; + +import com.azure.cosmos.models.PartitionKey; + +import java.util.List; + +import static com.azure.cosmos.implementation.guava25.base.Preconditions.checkNotNull; + +public final class SinglePartitionKeyServerBatchRequest extends ServerBatchRequest { + + private final PartitionKey partitionKey; + + /** + * Initializes a new instance of the {@link SinglePartitionKeyServerBatchRequest} class. Single partition key server + * request. + * + * @param partitionKey Partition key that applies to all operations in this request. + */ + private SinglePartitionKeyServerBatchRequest(final PartitionKey partitionKey) { + super(Integer.MAX_VALUE, Integer.MAX_VALUE); + this.partitionKey = partitionKey; + } + + /** + * Creates an instance of {@link SinglePartitionKeyServerBatchRequest}. The body of the request is populated with + * operations till it reaches the provided maxBodyLength. + * + * @param partitionKey Partition key of the request. + * @param operations Operations to be added into this batch request. + * + * @return A newly created instance of {@link SinglePartitionKeyServerBatchRequest}. + */ + static SinglePartitionKeyServerBatchRequest createAsync( + final PartitionKey partitionKey, + final List> operations) { + + checkNotNull(partitionKey, "expected non-null partitionKey"); + checkNotNull(operations, "expected non-null operations"); + + final SinglePartitionKeyServerBatchRequest request = new SinglePartitionKeyServerBatchRequest(partitionKey); + request.createBodyStreamAsync(operations); + + return request; + } + + /** + * Returns the {@link PartitionKey partition key} that applies to all operations in this {@link + * SinglePartitionKeyServerBatchRequest batch request}. + * + * @return the {@link PartitionKey partition key} that applies to all operations in this {@link + * SinglePartitionKeyServerBatchRequest batch request}. + */ + public PartitionKey getPartitionKey() { + return this.partitionKey; + } +} diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/directconnectivity/HttpTransportClient.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/directconnectivity/HttpTransportClient.java index 3bc2af841710..b1a636c30e08 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/directconnectivity/HttpTransportClient.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/directconnectivity/HttpTransportClient.java @@ -277,6 +277,7 @@ private HttpRequest prepareHttpMessage( // HttpRequestMessage -> StreamContent -> MemoryStream all get disposed, the original stream will be left open. switch (resourceOperation.operationType) { case Create: + case Batch: requestUri = getResourceFeedUri(resourceOperation.resourceType, physicalAddress.getURIAsString(), request); method = HttpMethod.POST; assert request.getContentAsByteBufFlux() != null; @@ -472,6 +473,13 @@ private HttpRequest prepareHttpMessage( HttpTransportClient.addHeader(httpRequestHeaders, CustomHeaders.HttpHeaders.EXCLUDE_SYSTEM_PROPERTIES, request); + if (resourceOperation.operationType == OperationType.Batch) { + HttpTransportClient.addHeader(httpRequestHeaders, HttpConstants.HttpHeaders.IS_BATCH_REQUEST, request); + HttpTransportClient.addHeader(httpRequestHeaders, HttpConstants.HttpHeaders.SHOULD_BATCH_CONTINUE_ON_ERROR, request); + HttpTransportClient.addHeader(httpRequestHeaders, HttpConstants.HttpHeaders.IS_BATCH_ORDERED, request); + HttpTransportClient.addHeader(httpRequestHeaders, HttpConstants.HttpHeaders.IS_BATCH_ATOMIC, request); + } + return httpRequestMessage; } diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/directconnectivity/WFConstants.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/directconnectivity/WFConstants.java index f40dc6706405..26277ceeebc1 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/directconnectivity/WFConstants.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/directconnectivity/WFConstants.java @@ -69,6 +69,7 @@ public static class BackendHeaders { public static final String BINARY_ID = "x-ms-binary-id"; public static final String TIME_TO_LIVE_IN_SECONDS = "x-ms-time-to-live-in-seconds"; public static final String EFFECTIVE_PARTITION_KEY = "x-ms-effective-partition-key"; + public static final String EFFECTIVE_PARTITION_KEY_STRING = "x-ms-effective-partition-key-string"; public static final String BINARY_PASSTHROUGH_REQUEST = "x-ms-binary-passthrough-request"; public static final String FANOUT_OPERATION_STATE = "x-ms-fanout-operation-state"; public static final String CONTENT_SERIALIZATION_FORMAT = "x-ms-documentdb-content-serialization-format"; diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/directconnectivity/rntbd/RntbdConstants.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/directconnectivity/rntbd/RntbdConstants.java index e37ef30552b9..67c2908f8cff 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/directconnectivity/rntbd/RntbdConstants.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/directconnectivity/rntbd/RntbdConstants.java @@ -259,7 +259,8 @@ public enum RntbdOperationType { AbortPartitionMigration((short) 0x001F, OperationType.AbortPartitionMigration), PreReplaceValidation((short) 0x0020, OperationType.PreReplaceValidation), AddComputeGatewayRequestCharges((short) 0x0021, OperationType.AddComputeGatewayRequestCharges), - MigratePartition((short) 0x0022, OperationType.MigratePartition); + MigratePartition((short) 0x0022, OperationType.MigratePartition), + Batch((short) 0x0025, OperationType.Batch); private final short id; private final OperationType type; @@ -341,6 +342,8 @@ public static RntbdOperationType fromId(final short id) { return RntbdOperationType.AddComputeGatewayRequestCharges; case 0x0022: return RntbdOperationType.MigratePartition; + case 0x0025: + return RntbdOperationType.Batch; default: throw new DecoderException(lenientFormat("expected byte value matching %s value, not %s", RntbdOperationType.class.getSimpleName(), @@ -416,6 +419,8 @@ public static RntbdOperationType fromType(OperationType type) { return RntbdOperationType.MigratePartition; case AddComputeGatewayRequestCharges: return RntbdOperationType.AddComputeGatewayRequestCharges; + case Batch: + return RntbdOperationType.Batch; default: throw new IllegalArgumentException(lenientFormat("unrecognized operation type: %s", type)); } @@ -572,6 +577,9 @@ public enum RntbdRequestHeader implements RntbdHeader { AllowTentativeWrites((short) 0x0066, RntbdTokenType.Byte, false), IsUserRequest((short) 0x0067, RntbdTokenType.Byte, false), SharedOfferThroughput((short) 0x0068, RntbdTokenType.ULong, false), + IsBatchAtomic((short) 0x0073, RntbdTokenType.Byte, false), + ShouldBatchContinueOnError((short) 0x0074, RntbdTokenType.Byte, false), + IsBatchOrdered((short) 0x0075, RntbdTokenType.Byte, false), ReturnPreference((short) 0x0082, RntbdTokenType.Byte, false); public static final ImmutableMap map; diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/directconnectivity/rntbd/RntbdRequestFrame.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/directconnectivity/rntbd/RntbdRequestFrame.java index 87483ce79c4f..dd3d8379501b 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/directconnectivity/rntbd/RntbdRequestFrame.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/directconnectivity/rntbd/RntbdRequestFrame.java @@ -203,6 +203,8 @@ private static RntbdOperationType map(final OperationType operationType) { return RntbdOperationType.MigratePartition; case AddComputeGatewayRequestCharges: return RntbdOperationType.AddComputeGatewayRequestCharges; + case Batch: + return RntbdOperationType.Batch; default: final String reason = Strings.lenientFormat("Unrecognized operation type: %s", operationType); throw new UnsupportedOperationException(reason); diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/directconnectivity/rntbd/RntbdRequestHeaders.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/directconnectivity/rntbd/RntbdRequestHeaders.java index fc7e7f0c6906..9178a9fafa92 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/directconnectivity/rntbd/RntbdRequestHeaders.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/directconnectivity/rntbd/RntbdRequestHeaders.java @@ -151,6 +151,9 @@ final class RntbdRequestHeaders extends RntbdTokenStream { this.fillTokenFromHeader(headers, this::getTargetLsn, HttpHeaders.TARGET_LSN); this.fillTokenFromHeader(headers, this::getTimeToLiveInSeconds, BackendHeaders.TIME_TO_LIVE_IN_SECONDS); this.fillTokenFromHeader(headers, this::getTransportRequestID, HttpHeaders.TRANSPORT_REQUEST_ID); + this.fillTokenFromHeader(headers, this::isBatchAtomic, HttpHeaders.IS_BATCH_ATOMIC); + this.fillTokenFromHeader(headers, this::shouldBatchContinueOnError, HttpHeaders.SHOULD_BATCH_CONTINUE_ON_ERROR); + this.fillTokenFromHeader(headers, this::isBatchOrdered, HttpHeaders.IS_BATCH_ORDERED); // Will be null in case of direct, which is fine - BE will use the value slice the connection context this. // When this is used in Gateway, the header value will be populated with the proxied HTTP request's header, @@ -564,6 +567,19 @@ private RntbdToken getUserName() { return this.get(RntbdRequestHeader.UserName); } + // Batch + private RntbdToken isBatchAtomic() { + return this.get(RntbdRequestHeader.IsBatchAtomic); + } + + private RntbdToken shouldBatchContinueOnError() { + return this.get(RntbdRequestHeader.ShouldBatchContinueOnError); + } + + private RntbdToken isBatchOrdered() { + return this.get(RntbdRequestHeader.IsBatchOrdered); + } + private void addAimHeader(final Map headers) { final String value = headers.get(HttpHeaders.A_IM); diff --git a/sdk/cosmos/azure-cosmos/src/test/java/com/azure/cosmos/BatchTestBase.java b/sdk/cosmos/azure-cosmos/src/test/java/com/azure/cosmos/BatchTestBase.java new file mode 100644 index 000000000000..c2ee856a209a --- /dev/null +++ b/sdk/cosmos/azure-cosmos/src/test/java/com/azure/cosmos/BatchTestBase.java @@ -0,0 +1,216 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +package com.azure.cosmos; + +import com.azure.cosmos.implementation.ISessionToken; +import com.azure.cosmos.implementation.SessionTokenHelper; +import com.azure.cosmos.implementation.apachecommons.lang.StringUtils; +import com.azure.cosmos.implementation.directconnectivity.WFConstants; +import com.azure.cosmos.models.CosmosContainerResponse; +import com.azure.cosmos.models.CosmosDatabaseResponse; +import com.azure.cosmos.models.CosmosItemResponse; +import com.azure.cosmos.models.PartitionKey; +import com.azure.cosmos.models.ThroughputProperties; +import com.azure.cosmos.rx.TestSuiteBase; +import com.fasterxml.jackson.annotation.JsonProperty; +import io.netty.handler.codec.http.HttpResponseStatus; + +import java.util.Random; +import java.util.UUID; + +import static org.testng.Assert.assertEquals; +import static org.testng.Assert.assertTrue; + +public class BatchTestBase extends TestSuiteBase { + + private Random random = new Random(); + String partitionKey1 = "TBD1"; + + // Documents in partitionKey1 + TestDoc TestDocPk1ExistingA; + TestDoc TestDocPk1ExistingB ; + TestDoc TestDocPk1ExistingC; + TestDoc TestDocPk1ExistingD; + + public BatchTestBase(CosmosClientBuilder clientBuilder) { + super(clientBuilder); + } + + CosmosAsyncContainer createSharedThroughputContainer(CosmosAsyncClient client) { + CosmosAsyncContainer sharedThroughputContainer = null; + CosmosDatabaseResponse cosmosDatabaseResponse = client.createDatabaseIfNotExists( + "Shared_" + UUID.randomUUID().toString(), + ThroughputProperties.createManualThroughput(12000)).block(); + + CosmosAsyncDatabase db = client.getDatabase(cosmosDatabaseResponse.getProperties().getId()); + + for (int index = 0; index < 5; index++) { + + CosmosContainerResponse cosmosContainerResponse = db.createContainerIfNotExists(getCollectionDefinition()).block(); + assertTrue(Boolean.parseBoolean(cosmosContainerResponse.getResponseHeaders().get(WFConstants.BackendHeaders.SHARE_THROUGHPUT))); + + if (index == 2) { + sharedThroughputContainer = db.getContainer(cosmosContainerResponse.getProperties().getId()); + } + } + + return sharedThroughputContainer; + } + + void createJsonTestDocsAsync(CosmosAsyncContainer container) { + this.TestDocPk1ExistingA = this.createJsonTestDocAsync(container, this.partitionKey1); + this.TestDocPk1ExistingB = this.createJsonTestDocAsync(container, this.partitionKey1); + this.TestDocPk1ExistingC = this.createJsonTestDocAsync(container, this.partitionKey1); + this.TestDocPk1ExistingD = this.createJsonTestDocAsync(container, this.partitionKey1); + } + + TestDoc populateTestDoc(String partitionKey) { + return populateTestDoc(partitionKey, 20); + } + + TestDoc populateTestDoc(String partitionKey, int minDesiredSize) { + String description = StringUtils.repeat("x", minDesiredSize); + return new TestDoc(UUID.randomUUID().toString(), this.random.nextInt(), description, partitionKey); + } + + public TestDoc populateTestDoc(String id, String partitionKey) { + String description = StringUtils.repeat("x", 20); + return new TestDoc(id, this.random.nextInt(), description, partitionKey); + } + + TestDoc getTestDocCopy(TestDoc testDoc) { + return new TestDoc(testDoc.getId(), testDoc.getCost(), testDoc.getDescription(), testDoc.getStatus()); + } + + void verifyByReadAsync(CosmosAsyncContainer container, TestDoc doc) { + verifyByReadAsync(container, doc, null); + } + + void verifyByReadAsync(CosmosAsyncContainer container, TestDoc doc, String eTag) { + PartitionKey partitionKey = this.getPartitionKey(doc.getStatus()); + + CosmosItemResponse response = container.readItem(doc.getId(), partitionKey, TestDoc.class).block(); + + assertEquals(HttpResponseStatus.OK.code(), response.getStatusCode()); + assertEquals(doc, response.getItem()); + + if (eTag != null) { + assertEquals(eTag, response.getETag()); + } + } + + void verifyNotFoundAsync(CosmosAsyncContainer container, TestDoc doc) { + String id = doc.getId(); + PartitionKey partitionKey = this.getPartitionKey(doc.getStatus()); + + try { + CosmosItemResponse response = container.readItem(id, partitionKey, TestDoc.class).block(); + + // Gateway returns response instead of exception + assertEquals(HttpResponseStatus.NOT_FOUND.code(), response.getStatusCode()); + } catch (CosmosException ex) { + assertEquals(HttpResponseStatus.NOT_FOUND.code(), ex.getStatusCode()); + } + } + + PartitionKey getPartitionKey(String partitionKey) { + return new PartitionKey(partitionKey); + } + + private TestDoc createJsonTestDocAsync(CosmosAsyncContainer container, String partitionKey) { + return createJsonTestDocAsync(container, partitionKey, 20); + } + + private TestDoc createJsonTestDocAsync(CosmosAsyncContainer container, String partitionKey, int minDesiredSize) { + TestDoc doc = this.populateTestDoc(partitionKey, minDesiredSize); + CosmosItemResponse createResponse = container.createItem(doc, this.getPartitionKey(partitionKey), null).block(); + assertEquals(HttpResponseStatus.CREATED.code(), createResponse.getStatusCode()); + return doc; + } + + public Random getRandom() { + return random; + } + + ISessionToken getSessionToken(String sessionToken) { + String[] tokenParts = sessionToken.split(":"); + return SessionTokenHelper.parse(tokenParts[1]); + } + + public static class TestDoc { + public String id; + public int cost; + public String description; + + @JsonProperty("mypk") + public String status; + + public TestDoc() { + + } + + public TestDoc(String id, int cost, String description, String status) { + this.id = id; + this.cost = cost; + this.description = description; + this.status = status; + } + + public boolean equals(Object obj) { + if (this == obj) { + return true; + } + if (obj == null || getClass() != obj.getClass()) { + return false; + } + TestDoc testDoc2 = (TestDoc) obj; + return (this.getId().equals(testDoc2.getId()) && + this.getCost() == testDoc2.getCost()) && + this.getDescription().equals(testDoc2.getDescription()) && + this.getStatus().equals(testDoc2.getStatus()); + } + + @Override + public int hashCode() { + int hashCode = 1652434776; + hashCode = (hashCode * -1521134295) + this.id.hashCode(); + hashCode = (hashCode * -1521134295) + this.cost; + hashCode = (hashCode * -1521134295) + this.description.hashCode(); + hashCode = (hashCode * -1521134295) + this.status.hashCode(); + return hashCode; + } + + public String getId() { + return id; + } + + public void setId(String id) { + this.id = id; + } + + public int getCost() { + return cost; + } + + public void setCost(int cost) { + this.cost = cost; + } + + public String getDescription() { + return description; + } + + public void setDescription(String description) { + this.description = description; + } + + public String getStatus() { + return status; + } + + public void setStatus(String status) { + this.status = status; + } + } +} diff --git a/sdk/cosmos/azure-cosmos/src/test/java/com/azure/cosmos/TransactionalBatchTest.java b/sdk/cosmos/azure-cosmos/src/test/java/com/azure/cosmos/TransactionalBatchTest.java new file mode 100644 index 000000000000..dfef7a78b31e --- /dev/null +++ b/sdk/cosmos/azure-cosmos/src/test/java/com/azure/cosmos/TransactionalBatchTest.java @@ -0,0 +1,390 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +package com.azure.cosmos; + +import com.azure.cosmos.implementation.HttpConstants; +import com.azure.cosmos.implementation.ISessionToken; +import com.azure.cosmos.implementation.apachecommons.lang.StringUtils; +import com.azure.cosmos.implementation.guava25.base.Function; +import com.azure.cosmos.models.CosmosItemResponse; +import io.netty.handler.codec.http.HttpResponseStatus; +import org.testng.Assert; +import org.testng.annotations.AfterClass; +import org.testng.annotations.BeforeClass; +import org.testng.annotations.Factory; +import org.testng.annotations.Test; + +import java.util.UUID; + +import static com.azure.cosmos.implementation.batch.BatchRequestResponseConstant.MAX_OPERATIONS_IN_DIRECT_MODE_BATCH_REQUEST; +import static org.assertj.core.api.Assertions.assertThat; + +public class TransactionalBatchTest extends BatchTestBase { + + private CosmosAsyncClient batchClient; + private CosmosAsyncContainer batchContainer; + private CosmosAsyncContainer sharedThroughputContainer; + + @Factory(dataProvider = "simpleClientBuildersWithDirect") + public TransactionalBatchTest(CosmosClientBuilder clientBuilder) { + super(clientBuilder); + } + + @BeforeClass(groups = {"emulator"}, timeOut = SETUP_TIMEOUT) + public void before_TransactionalBatchTest() { + assertThat(this.batchClient).isNull(); + this.batchClient = getClientBuilder().buildAsyncClient(); + CosmosAsyncContainer asyncContainer = getSharedMultiPartitionCosmosContainer(this.batchClient); + batchContainer = batchClient.getDatabase(asyncContainer.getDatabase().getId()).getContainer(asyncContainer.getId()); + sharedThroughputContainer = super.createSharedThroughputContainer(this.batchClient); + } + + @AfterClass(groups = {"emulator"}, timeOut = SHUTDOWN_TIMEOUT, alwaysRun = true) + public void afterClass() { + assertThat(this.batchClient).isNotNull(); + + if (sharedThroughputContainer != null) { + // Delete the database created in this test + sharedThroughputContainer.getDatabase().delete().block(); + } + + this.batchClient.close(); + } + + @Test(groups = {"emulator"}, timeOut = TIMEOUT) + public void batchCrud() throws Exception { + this.runCrudAsync(batchContainer); + } + + @Test(groups = {"emulator"}, timeOut = TIMEOUT) + public void batchSharedThroughputCrud() throws Exception { + this.runCrudAsync(sharedThroughputContainer); + } + + @Test(groups = {"emulator"}, timeOut = TIMEOUT) + public void batchOrdered() { + CosmosAsyncContainer container = batchContainer; + this.createJsonTestDocsAsync(container); + + TestDoc firstDoc = this.populateTestDoc(this.partitionKey1); + + TestDoc replaceDoc = this.getTestDocCopy(firstDoc); + replaceDoc.setCost(replaceDoc.getCost() + 1); + + TransactionalBatchResponse batchResponse = container.executeTransactionalBatch( + TransactionalBatch.createTransactionalBatch(this.getPartitionKey(this.partitionKey1)) + .createItem(firstDoc) + .replaceItem(replaceDoc.getId(), replaceDoc)) + .block(); + + this.verifyBatchProcessed(batchResponse, 2); + + Assert.assertEquals(HttpResponseStatus.CREATED.code(), batchResponse.get(0).getResponseStatus()); + Assert.assertEquals(HttpResponseStatus.OK.code(), batchResponse.get(1).getResponseStatus()); + + // Ensure that the replace overwrote the doc from the first operation + this.verifyByReadAsync(container, replaceDoc); + } + + @Test(groups = {"emulator"}, timeOut = TIMEOUT) + public void batchItemETagAsync() { + CosmosAsyncContainer container = batchContainer; + this.createJsonTestDocsAsync(container); + + { + BatchTestBase.TestDoc testDocToCreate = this.populateTestDoc(this.partitionKey1); + + BatchTestBase.TestDoc testDocToReplace = this.getTestDocCopy(this.TestDocPk1ExistingA); + testDocToReplace.setCost(testDocToReplace.getCost() + 1); + + CosmosItemResponse response = container.readItem( + this.TestDocPk1ExistingA.getId(), + this.getPartitionKey(this.partitionKey1), + TestDoc.class).block(); + + Assert.assertEquals(HttpResponseStatus.OK.code(), response.getStatusCode()); + + TransactionalBatchItemRequestOptions firstReplaceOptions = new TransactionalBatchItemRequestOptions(); + firstReplaceOptions.setIfMatchETag(response.getETag()); + + TransactionalBatchResponse batchResponse = container.executeTransactionalBatch( + TransactionalBatch.createTransactionalBatch(this.getPartitionKey(this.partitionKey1)) + .createItem(testDocToCreate) + .replaceItem(testDocToReplace.getId(), testDocToReplace, firstReplaceOptions)) + .block(); + + this.verifyBatchProcessed(batchResponse, 2); + + Assert.assertEquals(HttpResponseStatus.CREATED.code(), batchResponse.get(0).getResponseStatus()); + Assert.assertEquals(HttpResponseStatus.OK.code(), batchResponse.get(1).getResponseStatus()); + + // Ensure that the replace overwrote the doc from the first operation + this.verifyByReadAsync(container, testDocToCreate, batchResponse.get(0).getETag()); + this.verifyByReadAsync(container, testDocToReplace, batchResponse.get(1).getETag()); + } + + { + TestDoc testDocToReplace = this.getTestDocCopy(this.TestDocPk1ExistingB); + testDocToReplace.setCost(testDocToReplace.getCost() + 1); + + TransactionalBatchItemRequestOptions replaceOptions = new TransactionalBatchItemRequestOptions(); + replaceOptions.setIfMatchETag(String.valueOf(this.getRandom().nextInt())); + + TransactionalBatchResponse batchResponse = container.executeTransactionalBatch( + TransactionalBatch.createTransactionalBatch(this.getPartitionKey(this.partitionKey1)) + .replaceItem(testDocToReplace.getId(), testDocToReplace, replaceOptions)) + .block(); + + this.verifyBatchProcessed(batchResponse, 1, HttpResponseStatus.PRECONDITION_FAILED); + + Assert.assertEquals(HttpResponseStatus.PRECONDITION_FAILED.code(), batchResponse.get(0).getResponseStatus()); + + // ensure the document was not updated + this.verifyByReadAsync(container, this.TestDocPk1ExistingB); + } + } + + @Test(groups = {"emulator"}, timeOut = TIMEOUT) + public void batchItemSessionTokenAsync() { + CosmosAsyncContainer container = batchContainer; + this.createJsonTestDocsAsync(container); + + TestDoc testDocToCreate = this.populateTestDoc(this.partitionKey1); + + BatchTestBase.TestDoc testDocToReplace = this.getTestDocCopy(this.TestDocPk1ExistingA); + testDocToReplace.setCost(testDocToReplace.getCost() + 1); + + CosmosItemResponse readResponse = container.readItem( + this.TestDocPk1ExistingA.getId(), + this.getPartitionKey(this.partitionKey1), + TestDoc.class).block(); + + Assert.assertEquals(HttpResponseStatus.OK.code(), readResponse.getStatusCode()); + + ISessionToken beforeRequestSessionToken = this.getSessionToken(readResponse.getResponseHeaders().get(HttpConstants.HttpHeaders.SESSION_TOKEN)); + + TransactionalBatchResponse batchResponse = container.executeTransactionalBatch( + TransactionalBatch.createTransactionalBatch(this.getPartitionKey(this.partitionKey1)) + .createItem(testDocToCreate) + .replaceItem(testDocToReplace.getId(), testDocToReplace)) + .block(); + + this.verifyBatchProcessed(batchResponse, 2); + + Assert.assertEquals(HttpResponseStatus.CREATED.code(), batchResponse.get(0).getResponseStatus()); + Assert.assertEquals(HttpResponseStatus.OK.code(), batchResponse.get(1).getResponseStatus()); + + ISessionToken afterRequestSessionToken = this.getSessionToken(batchResponse.getResponseHeaders().get(HttpConstants.HttpHeaders.SESSION_TOKEN)); + Assert.assertTrue(afterRequestSessionToken.getLSN() > beforeRequestSessionToken.getLSN(), "Response session token should be more than request session token"); + } + + @Test(groups = {"emulator"}, timeOut = TIMEOUT) + public void batchWithTooManyOperationsAsync() { + CosmosAsyncContainer container = batchContainer; + int operationCount = MAX_OPERATIONS_IN_DIRECT_MODE_BATCH_REQUEST + 1; + + // Increase the doc size by a bit so all docs won't fit in one server request. + TransactionalBatch batch = TransactionalBatch.createTransactionalBatch(this.getPartitionKey(this.partitionKey1)); + + for (int i = 0; i < operationCount; i++) { + batch.readItem("someId"); + } + + TransactionalBatchResponse batchResponse = container.executeTransactionalBatch(batch).block(); + Assert.assertEquals(HttpResponseStatus.BAD_REQUEST.code(), batchResponse.getResponseStatus()); + } + + @Test(groups = {"emulator"}, timeOut = TIMEOUT) + public void batchReadsOnlyAsync() throws Exception { + CosmosAsyncContainer container = batchContainer; + this.createJsonTestDocsAsync(container); + + TransactionalBatchResponse batchResponse = container.executeTransactionalBatch( + TransactionalBatch.createTransactionalBatch(this.getPartitionKey(this.partitionKey1)) + .readItem(this.TestDocPk1ExistingA.getId()) + .readItem(this.TestDocPk1ExistingB.getId()) + .readItem(this.TestDocPk1ExistingC.getId())) + .block(); + + this.verifyBatchProcessed(batchResponse, 3); + + Assert.assertEquals(HttpResponseStatus.OK.code(), batchResponse.get(0).getResponseStatus()); + Assert.assertEquals(HttpResponseStatus.OK.code(), batchResponse.get(1).getResponseStatus()); + Assert.assertEquals(HttpResponseStatus.OK.code(), batchResponse.get(2).getResponseStatus()); + + Assert.assertEquals(this.TestDocPk1ExistingA, batchResponse.getOperationResultAtIndex(0, TestDoc.class).getItem()); + Assert.assertEquals(this.TestDocPk1ExistingB, batchResponse.getOperationResultAtIndex(1, TestDoc.class).getItem()); + Assert.assertEquals(this.TestDocPk1ExistingC, batchResponse.getOperationResultAtIndex(2, TestDoc.class).getItem()); + } + + private TransactionalBatchResponse runCrudAsync(CosmosAsyncContainer container) throws Exception { + this.createJsonTestDocsAsync(container); + + BatchTestBase.TestDoc testDocToCreate = this.populateTestDoc(this.partitionKey1); + BatchTestBase.TestDoc testDocToUpsert = this.populateTestDoc(this.partitionKey1); + + BatchTestBase.TestDoc anotherTestDocToUpsert = this.getTestDocCopy(this.TestDocPk1ExistingA); + anotherTestDocToUpsert.setCost(anotherTestDocToUpsert.getCost() + 1); + + BatchTestBase.TestDoc testDocToReplace = this.getTestDocCopy(this.TestDocPk1ExistingB); + testDocToReplace.setCost(testDocToReplace.getCost() + 1); + + // We run CRUD operations where all are expected to return HTTP 2xx. + TransactionalBatchResponse batchResponse = container.executeTransactionalBatch( + TransactionalBatch.createTransactionalBatch(this.getPartitionKey(this.partitionKey1)) + .createItem(testDocToCreate) + .readItem(this.TestDocPk1ExistingC.getId()) + .replaceItem(testDocToReplace.getId(), testDocToReplace) + .upsertItem(testDocToUpsert) + .upsertItem(anotherTestDocToUpsert) + .deleteItem(this.TestDocPk1ExistingD.getId())) + .block(); + + this.verifyBatchProcessed(batchResponse, 6); + + Assert.assertEquals(HttpResponseStatus.CREATED.code(), batchResponse.get(0).getResponseStatus()); + Assert.assertEquals(HttpResponseStatus.OK.code(), batchResponse.get(1).getResponseStatus()); + Assert.assertEquals(HttpResponseStatus.OK.code(), batchResponse.get(2).getResponseStatus()); + Assert.assertEquals(HttpResponseStatus.CREATED.code(), batchResponse.get(3).getResponseStatus()); + Assert.assertEquals(HttpResponseStatus.OK.code(), batchResponse.get(4).getResponseStatus()); + Assert.assertEquals(HttpResponseStatus.NO_CONTENT.code(), batchResponse.get(5).getResponseStatus()); + + Assert.assertEquals(this.TestDocPk1ExistingC, batchResponse.getOperationResultAtIndex(1, TestDoc.class).getItem()); + + this.verifyByReadAsync(container, testDocToCreate); + this.verifyByReadAsync(container, testDocToReplace); + this.verifyByReadAsync(container, testDocToUpsert); + this.verifyByReadAsync(container, anotherTestDocToUpsert); + this.verifyNotFoundAsync(container, this.TestDocPk1ExistingD); + + return batchResponse; + } + + @Test(groups = {"emulator"}, timeOut = TIMEOUT) + public void batchWithCreateConflictAsync() { + this.runBatchWithCreateConflictAsync(batchContainer); + } + + @Test(groups = {"emulator"}, timeOut = TIMEOUT) + public void batchWithCreateConflictSharedThroughputAsync() { + this.runBatchWithCreateConflictAsync(this.sharedThroughputContainer); + } + + @Test(groups = {"emulator"}, timeOut = TIMEOUT) + public void batchWithInvalidCreateAsync() { + CosmosAsyncContainer container = batchContainer; + + // partition key mismatch between doc and and value passed in to the operation + this.runWithErrorAsync( + container, + batch -> batch.createItem(this.populateTestDoc(UUID.randomUUID().toString())), + HttpResponseStatus.BAD_REQUEST); + } + + @Test(groups = {"emulator"}, timeOut = TIMEOUT) + public void batchWithReadOfNonExistentEntityAsync() { + CosmosAsyncContainer container = batchContainer; + this.runWithErrorAsync( + container, + batch -> batch.readItem(UUID.randomUUID().toString()), + HttpResponseStatus.NOT_FOUND); + } + + @Test(groups = {"emulator"}, timeOut = TIMEOUT) + public void batchWithReplaceOfStaleEntityAsync() { + CosmosAsyncContainer container = batchContainer; + this.createJsonTestDocsAsync(container); + + TestDoc staleTestDocToReplace = this.getTestDocCopy(this.TestDocPk1ExistingA); + staleTestDocToReplace.setCost(staleTestDocToReplace.getCost() + 1); + + TransactionalBatchItemRequestOptions staleReplaceOptions = new TransactionalBatchItemRequestOptions(); + staleReplaceOptions.setIfMatchETag(UUID.randomUUID().toString()); + + this.runWithErrorAsync( + container, + batch -> batch.replaceItem(staleTestDocToReplace.getId(), staleTestDocToReplace, staleReplaceOptions), + HttpResponseStatus.PRECONDITION_FAILED); + + // make sure the stale doc hasn't changed + this.verifyByReadAsync(container, this.TestDocPk1ExistingA); + } + + @Test(groups = {"emulator"}, timeOut = TIMEOUT) + public void batchWithDeleteOfNonExistentEntityAsync() { + CosmosAsyncContainer container = batchContainer; + + this.runWithErrorAsync( + container, + batch -> batch.deleteItem(UUID.randomUUID().toString()), + HttpResponseStatus.NOT_FOUND); + } + + private void runBatchWithCreateConflictAsync(CosmosAsyncContainer container) { + this.createJsonTestDocsAsync(container); + + // try to create a doc with id that already exists (should return a Conflict) + TestDoc conflictingTestDocToCreate = this.getTestDocCopy(this.TestDocPk1ExistingA); + conflictingTestDocToCreate.setCost(conflictingTestDocToCreate.getCost()); + + this.runWithErrorAsync( + container, + batch -> batch.createItem(conflictingTestDocToCreate), + HttpResponseStatus.CONFLICT); + + // make sure the conflicted doc hasn't changed + this.verifyByReadAsync(container, this.TestDocPk1ExistingA); + } + + + private void runWithErrorAsync( + CosmosAsyncContainer container, + Function appendOperation, + HttpResponseStatus expectedFailedOperationStatusCode) { + + TestDoc testDocToCreate = this.populateTestDoc(this.partitionKey1); + TestDoc anotherTestDocToCreate = this.populateTestDoc(this.partitionKey1); + + TransactionalBatch batch = TransactionalBatch.createTransactionalBatch(this.getPartitionKey(this.partitionKey1)) + .createItem(testDocToCreate); + + appendOperation.apply(batch); + + TransactionalBatchResponse batchResponse = container.executeTransactionalBatch( + batch.createItem(anotherTestDocToCreate)) + .block(); + + this.verifyBatchProcessed(batchResponse, 3, expectedFailedOperationStatusCode); + + Assert.assertEquals(HttpResponseStatus.FAILED_DEPENDENCY.code(), batchResponse.get(0).getResponseStatus()); + Assert.assertEquals(expectedFailedOperationStatusCode.code(), batchResponse.get(1).getResponseStatus()); + Assert.assertEquals(HttpResponseStatus.FAILED_DEPENDENCY.code(), batchResponse.get(2).getResponseStatus()); + + this.verifyNotFoundAsync(container, testDocToCreate); + this.verifyNotFoundAsync(container, anotherTestDocToCreate); + } + + private void verifyBatchProcessed(TransactionalBatchResponse batchResponse, int numberOfOperations) { + this.verifyBatchProcessed(batchResponse, numberOfOperations, HttpResponseStatus.OK); + } + + private void verifyBatchProcessed(TransactionalBatchResponse batchResponse, int numberOfOperations, HttpResponseStatus expectedStatusCode) { + Assert.assertNotNull(batchResponse); + Assert.assertEquals( + batchResponse.getResponseStatus(), + expectedStatusCode.code(), + "Batch server response had StatusCode {0} instead of {1} expected and had ErrorMessage {2}"); + + Assert.assertEquals(numberOfOperations, batchResponse.size()); + Assert.assertTrue(batchResponse.getRequestCharge() > 0); + Assert.assertTrue(StringUtils.isNotEmpty(batchResponse.getCosmosDiagnostics().toString())); + + // Allow a delta since we round both the total charge and the individual operation + // charges to 2 decimal places. + Assert.assertEquals( + batchResponse.getRequestCharge(), + batchResponse.stream().mapToDouble(result -> result.getRequestCharge()).sum(), + 0.1); + } +} diff --git a/sdk/cosmos/azure-cosmos/src/test/java/com/azure/cosmos/implementation/batch/BatchOperationResultTests.java b/sdk/cosmos/azure-cosmos/src/test/java/com/azure/cosmos/implementation/batch/BatchOperationResultTests.java new file mode 100644 index 000000000000..f9ce0a38668e --- /dev/null +++ b/sdk/cosmos/azure-cosmos/src/test/java/com/azure/cosmos/implementation/batch/BatchOperationResultTests.java @@ -0,0 +1,68 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +package com.azure.cosmos.implementation.batch; + +import com.azure.cosmos.TransactionalBatchOperationResult; +import com.azure.cosmos.implementation.HttpConstants; +import com.azure.cosmos.implementation.Utils; +import io.netty.handler.codec.http.HttpResponseStatus; +import org.testng.annotations.Test; + +import java.time.Duration; + +import static org.testng.Assert.assertEquals; +import static org.testng.Assert.assertSame; + +public class BatchOperationResultTests { + + private static final int TIMEOUT = 40000; + + private TransactionalBatchOperationResult createTestResult() { + TransactionalBatchOperationResult result = new TransactionalBatchOperationResult(HttpResponseStatus.OK.code()); + result.setSubStatusCode(HttpConstants.SubStatusCodes.NAME_CACHE_IS_STALE); + result.setETag("TestETag"); + result.setRequestCharge(1.4); + result.setResourceObject(Utils.getSimpleObjectMapper().createObjectNode()); + result.setRetryAfter(Duration.ofMillis(1234)); + + return result; + } + + @Test(groups = {"unit"}, timeOut = TIMEOUT) + public void propertiesAreSetThroughCopyCtor() { + TransactionalBatchOperationResult other = createTestResult(); + TransactionalBatchOperationResult result = new TransactionalBatchOperationResult(other); + + assertEquals(other.getResponseStatus(), result.getResponseStatus()); + assertEquals(other.getSubStatusCode(), result.getSubStatusCode()); + assertEquals(other.getETag(), result.getETag()); + assertEquals(other.getRequestCharge(), result.getRequestCharge()); + assertEquals(other.getRetryAfter(), result.getRetryAfter()); + assertSame(other.getResourceObject(), result.getResourceObject()); + } + + @Test(groups = {"unit"}, timeOut = TIMEOUT) + public void propertiesAreSetThroughGenericCtor() { + TransactionalBatchOperationResult other = createTestResult(); + Object testObject = new Object(); + TransactionalBatchOperationResult result = new TransactionalBatchOperationResult(other, testObject); + + assertEquals(other.getResponseStatus(), result.getResponseStatus()); + assertEquals(other.getSubStatusCode(), result.getSubStatusCode()); + assertEquals(other.getETag(), result.getETag()); + assertEquals(other.getRequestCharge(), result.getRequestCharge()); + assertEquals(other.getRetryAfter(), result.getRetryAfter()); + assertSame(other.getResourceObject(), result.getResourceObject()); + assertSame(testObject, result.getItem()); + } + + @Test(groups = {"unit"}, timeOut = TIMEOUT) + public void isSuccessStatusCodeTrueFor200To299() { + for (int x = 100; x < 999; ++x) { + TransactionalBatchOperationResult result = new TransactionalBatchOperationResult(x); + boolean success = x >= 200 && x <= 299; + assertEquals(success, result.isSuccessStatusCode()); + } + } +} diff --git a/sdk/cosmos/azure-cosmos/src/test/java/com/azure/cosmos/implementation/batch/BatchResponsePayloadWriter.java b/sdk/cosmos/azure-cosmos/src/test/java/com/azure/cosmos/implementation/batch/BatchResponsePayloadWriter.java new file mode 100644 index 000000000000..cc89220ae281 --- /dev/null +++ b/sdk/cosmos/azure-cosmos/src/test/java/com/azure/cosmos/implementation/batch/BatchResponsePayloadWriter.java @@ -0,0 +1,46 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +package com.azure.cosmos.implementation.batch; + +import com.azure.cosmos.TransactionalBatchOperationResult; +import com.azure.cosmos.implementation.JsonSerializable; +import com.azure.cosmos.implementation.Utils; +import com.fasterxml.jackson.databind.node.ArrayNode; + +import java.util.List; + +public class BatchResponsePayloadWriter { + + private List> results; + + public BatchResponsePayloadWriter(List> results) { + this.results = results; + } + + public String generatePayload() { + return writeOperationResult().toString(); + } + + private ArrayNode writeOperationResult() { + ArrayNode arrayNode = Utils.getSimpleObjectMapper().createArrayNode(); + + for(TransactionalBatchOperationResult result : results) { + JsonSerializable operationJsonSerializable = writeResult(result); + + arrayNode.add(operationJsonSerializable.getPropertyBag()); + } + return arrayNode; + } + + private JsonSerializable writeResult(TransactionalBatchOperationResult result) { + + JsonSerializable jsonSerializable = new JsonSerializable(); + jsonSerializable.set(BatchRequestResponseConstant.FIELD_STATUS_CODE, result.getResponseStatus()); + jsonSerializable.set(BatchRequestResponseConstant.FIELD_SUBSTATUS_CODE, result.getSubStatusCode()); + jsonSerializable.set(BatchRequestResponseConstant.FIELD_ETAG, result.getETag()); + jsonSerializable.set(BatchRequestResponseConstant.FIELD_RESOURCE_BODY, result.getResourceObject()); + + return jsonSerializable; + } +} diff --git a/sdk/cosmos/azure-cosmos/src/test/java/com/azure/cosmos/implementation/batch/PartitionKeyBatchResponseTests.java b/sdk/cosmos/azure-cosmos/src/test/java/com/azure/cosmos/implementation/batch/PartitionKeyBatchResponseTests.java new file mode 100644 index 000000000000..c2b691399a45 --- /dev/null +++ b/sdk/cosmos/azure-cosmos/src/test/java/com/azure/cosmos/implementation/batch/PartitionKeyBatchResponseTests.java @@ -0,0 +1,61 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +package com.azure.cosmos.implementation.batch; + +import com.azure.cosmos.TransactionalBatchOperationResult; +import com.azure.cosmos.TransactionalBatchResponse; +import com.azure.cosmos.implementation.OperationType; +import com.azure.cosmos.implementation.RxDocumentServiceResponse; +import com.azure.cosmos.implementation.directconnectivity.StoreResponse; +import com.azure.cosmos.models.PartitionKey; +import io.netty.handler.codec.http.HttpResponseStatus; +import org.testng.annotations.Test; + +import java.nio.charset.StandardCharsets; +import java.util.ArrayList; +import java.util.Arrays; +import java.util.List; + +import static org.testng.AssertJUnit.assertEquals; + +public class PartitionKeyBatchResponseTests { + + private static final int TIMEOUT = 40000; + + @Test(groups = {"unit"}, timeOut = TIMEOUT) + public void statusCodesAreSetThroughResponseAsync() { + List> results = new ArrayList<>(); + ItemBatchOperation[] arrayOperations = new ItemBatchOperation[1]; + + ItemBatchOperation operation = new ItemBatchOperation.Builder(OperationType.Read,0) + .partitionKey(PartitionKey.NONE) + .id("0") + .build(); + + TransactionalBatchOperationResult transactionalBatchOperationResult = new TransactionalBatchOperationResult(HttpResponseStatus.OK.code()); + transactionalBatchOperationResult.setETag(operation.getId()); + + results.add(transactionalBatchOperationResult); + + arrayOperations[0] = operation; + + String responseContent = new BatchResponsePayloadWriter(results).generatePayload(); + + SinglePartitionKeyServerBatchRequest batchRequest = SinglePartitionKeyServerBatchRequest.createAsync( + PartitionKey.NONE, + Arrays.asList(arrayOperations)); + + StoreResponse storeResponse = new StoreResponse( + HttpResponseStatus.OK.code(), + new ArrayList<>(), + responseContent.getBytes(StandardCharsets.UTF_8)); + + TransactionalBatchResponse batchResponse = BatchResponseParser.fromDocumentServiceResponseAsync( + new RxDocumentServiceResponse(storeResponse), + batchRequest, + true).block(); + + assertEquals(HttpResponseStatus.OK.code(), batchResponse.getResponseStatus()); + } +} From d8865398ced1e9abeb107f1d5759a54cfeff76e6 Mon Sep 17 00:00:00 2001 From: Rakesh Kumar Date: Thu, 1 Oct 2020 00:25:08 +0530 Subject: [PATCH 02/22] Code review changes Signed-off-by: Rakesh Kumar --- .../java/com/azure/cosmos/BridgeInternal.java | 61 +++++++ .../com/azure/cosmos/TransactionalBatch.java | 22 ++- .../TransactionalBatchItemRequestOptions.java | 2 +- .../TransactionalBatchOperationResult.java | 118 +++---------- .../TransactionalBatchRequestOptions.java | 69 +------- .../cosmos/TransactionalBatchResponse.java | 165 +++--------------- .../implementation/batch/BatchExecUtils.java | 80 +-------- .../implementation/batch/BatchExecutor.java | 16 +- .../batch/BatchResponseParser.java | 90 +++++----- .../batch/ItemBatchOperation.java | 18 +- .../directconnectivity/WFConstants.java | 1 - .../cosmos/BatchOperationResultTests.java | 67 +++++++ .../java/com/azure/cosmos/BatchTestBase.java | 40 +---- .../azure/cosmos/TransactionalBatchTest.java | 99 +++++------ .../batch/BatchOperationResultTests.java | 68 -------- .../batch/BatchResponsePayloadWriter.java | 6 +- .../batch/PartitionKeyBatchResponseTests.java | 15 +- 17 files changed, 312 insertions(+), 625 deletions(-) create mode 100644 sdk/cosmos/azure-cosmos/src/test/java/com/azure/cosmos/BatchOperationResultTests.java delete mode 100644 sdk/cosmos/azure-cosmos/src/test/java/com/azure/cosmos/implementation/batch/BatchOperationResultTests.java diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/BridgeInternal.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/BridgeInternal.java index ad715703c411..465a757722c6 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/BridgeInternal.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/BridgeInternal.java @@ -14,6 +14,7 @@ import com.azure.cosmos.implementation.MetadataDiagnosticsContext; import com.azure.cosmos.implementation.QueryMetrics; import com.azure.cosmos.implementation.ReplicationPolicy; +import com.azure.cosmos.implementation.RequestOptions; import com.azure.cosmos.implementation.RequestTimeline; import com.azure.cosmos.implementation.Resource; import com.azure.cosmos.implementation.ResourceResponse; @@ -23,6 +24,7 @@ import com.azure.cosmos.implementation.ServiceUnavailableException; import com.azure.cosmos.implementation.StoredProcedureResponse; import com.azure.cosmos.implementation.Warning; +import com.azure.cosmos.implementation.batch.ItemBatchOperation; import com.azure.cosmos.implementation.directconnectivity.StoreResponse; import com.azure.cosmos.implementation.directconnectivity.StoreResult; import com.azure.cosmos.implementation.directconnectivity.Uri; @@ -595,4 +597,63 @@ public static Duration getRequestTimeoutFromDirectConnectionConfig(DirectConnect public static Duration getRequestTimeoutFromGatewayConnectionConfig(GatewayConnectionConfig gatewayConnectionConfig) { return gatewayConnectionConfig.getRequestTimeout(); } + + @Warning(value = INTERNAL_USE_ONLY_WARNING) + public static List> getOperationsFromTransactionalBatch(TransactionalBatch transactionalBatch) { + return transactionalBatch.getOperations(); + } + + @Warning(value = INTERNAL_USE_ONLY_WARNING) + public static PartitionKey getPartitionKeyFromTransactionalBatch(TransactionalBatch transactionalBatch) { + return transactionalBatch.getPartitionKey(); + } + + @Warning(value = INTERNAL_USE_ONLY_WARNING) + public static RequestOptions toRequestOptions(TransactionalBatchRequestOptions transactionalBatchRequestOptions) { + return transactionalBatchRequestOptions.toRequestOptions(); + } + + @Warning(value = INTERNAL_USE_ONLY_WARNING) + public static TransactionalBatchOperationResult createTransactionBatchResult( + String eTag, + Double requestCharge, + ObjectNode resourceObject, + int responseStatus, + Duration retryAfter, + Integer subStatusCode) { + + return new TransactionalBatchOperationResult( + eTag, + requestCharge, + resourceObject, + responseStatus, + retryAfter, + subStatusCode); + } + + @Warning(value = INTERNAL_USE_ONLY_WARNING) + public static TransactionalBatchResponse createTransactionBatchResponse( + int responseStatusCode, + Integer responseSubStatusCode, + String errorMessage, + Map responseHeaders, + CosmosDiagnostics cosmosDiagnostics, + List> operations) { + + return new TransactionalBatchResponse( + responseStatusCode, + responseSubStatusCode, + errorMessage, + responseHeaders, + cosmosDiagnostics, + operations); + } + + @Warning(value = INTERNAL_USE_ONLY_WARNING) + public static void addTransactionBatchResultInResponse( + TransactionalBatchResponse transactionalBatchResponse, + List> transactionalBatchOperationResult) { + + transactionalBatchResponse.addAll(transactionalBatchOperationResult); + } } diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/TransactionalBatch.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/TransactionalBatch.java index 602e391a1905..6bec92566fe2 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/TransactionalBatch.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/TransactionalBatch.java @@ -4,10 +4,12 @@ package com.azure.cosmos; import com.azure.cosmos.implementation.OperationType; +import com.azure.cosmos.implementation.apachecommons.collections.list.UnmodifiableList; import com.azure.cosmos.implementation.batch.ItemBatchOperation; import com.azure.cosmos.models.PartitionKey; import java.util.ArrayList; +import java.util.List; import static com.azure.cosmos.implementation.guava25.base.Preconditions.checkNotNull; @@ -68,7 +70,7 @@ * .readItem("jogging") * .readItem("running") * - * try (TransactionalBatchResponse response = container.executeTransactionalBatch(new Cosmos.PartitionKey(activityType)) { + * try (TransactionalBatchResponse response = container.executeTransactionalBatch(batch) { * * // Look up interested results - eg. via direct access to operation result stream * @@ -89,6 +91,8 @@ public final class TransactionalBatch { private final PartitionKey partitionKey; public TransactionalBatch(PartitionKey partitionKey) { + checkNotNull(partitionKey, "expected non-null partitionKey"); + this.operations = new ArrayList<>(); this.partitionKey = partitionKey; } @@ -296,11 +300,21 @@ public TransactionalBatch upsertItem(T item, TransactionalBatchItemRequestOp return this; } - public ArrayList> getOperations() { - return operations; + /** + * Return the list of operation in an unmodifiable instace so no one can change it in the down path. + * + * @return The list of operations which are to be executed. + */ + List> getOperations() { + return UnmodifiableList.unmodifiableList(operations); } - public PartitionKey getPartitionKey() { + /** + * Return the partition key for this batch. + * + * @return The partition key for this batch. + */ + PartitionKey getPartitionKey() { return partitionKey; } } diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/TransactionalBatchItemRequestOptions.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/TransactionalBatchItemRequestOptions.java index 79e706fab39b..bde6b127bfda 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/TransactionalBatchItemRequestOptions.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/TransactionalBatchItemRequestOptions.java @@ -49,7 +49,7 @@ public TransactionalBatchItemRequestOptions setIfNoneMatchETag(final String ifNo return this; } - public RequestOptions toRequestOptions() { + RequestOptions toRequestOptions() { final RequestOptions requestOptions = new RequestOptions(); requestOptions.setIfMatchETag(getIfMatchETag()); requestOptions.setIfNoneMatchETag(getIfNoneMatchETag()); diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/TransactionalBatchOperationResult.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/TransactionalBatchOperationResult.java index 91f61df63695..202ca5150753 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/TransactionalBatchOperationResult.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/TransactionalBatchOperationResult.java @@ -4,9 +4,6 @@ package com.azure.cosmos; import com.fasterxml.jackson.databind.node.ObjectNode; -import org.slf4j.Logger; -import org.slf4j.LoggerFactory; - import java.time.Duration; import static com.azure.cosmos.implementation.guava25.base.Preconditions.checkNotNull; @@ -16,24 +13,22 @@ * * @param the type parameter */ -public final class TransactionalBatchOperationResult implements AutoCloseable { - - private final static Logger logger = LoggerFactory.getLogger(TransactionalBatchOperationResult.class); +public final class TransactionalBatchOperationResult { private String eTag; - private double requestCharge; + private Double requestCharge; private TResource item; private ObjectNode resourceObject; private int responseStatus; private Duration retryAfter; - private int subStatusCode; + private Integer subStatusCode; /** * Instantiates a new Transactional batch operation result. * * @param responseStatus the response status */ - public TransactionalBatchOperationResult(final int responseStatus) { + TransactionalBatchOperationResult(final int responseStatus) { this.responseStatus = responseStatus; } @@ -42,7 +37,7 @@ public TransactionalBatchOperationResult(final int responseStatus) { * * @param other the other */ - public TransactionalBatchOperationResult(final TransactionalBatchOperationResult other) { + TransactionalBatchOperationResult(final TransactionalBatchOperationResult other) { checkNotNull(other, "expected non-null other"); @@ -60,7 +55,7 @@ public TransactionalBatchOperationResult(final TransactionalBatchOperationResult * @param result the result * @param item the item */ - public TransactionalBatchOperationResult(TransactionalBatchOperationResult result, TResource item) { + TransactionalBatchOperationResult(TransactionalBatchOperationResult result, TResource item) { this(result); this.item = item; } @@ -68,7 +63,18 @@ public TransactionalBatchOperationResult(TransactionalBatchOperationResult re /** * Initializes a new instance of the {@link TransactionalBatchOperationResult} class. */ - public TransactionalBatchOperationResult() { + TransactionalBatchOperationResult(String eTag, + Double requestCharge, + ObjectNode resourceObject, + int responseStatus, + Duration retryAfter, + Integer subStatusCode) { + this.eTag = eTag; + this.requestCharge = requestCharge; + this.resourceObject = resourceObject; + this.responseStatus = responseStatus; + this.retryAfter = retryAfter; + this.subStatusCode = subStatusCode; } /** @@ -82,37 +88,13 @@ public String getETag() { return this.eTag; } - /** - * Sets e tag. - * - * @param value the value - * - * @return the e tag - */ - public TransactionalBatchOperationResult setETag(final String value) { - this.eTag = value; - return this; - } - /** * Gets the request charge in request units for the current operation. * * @return Request charge in request units for the current operation. */ - public double getRequestCharge() { - return requestCharge; - } - - /** - * Sets request charge. - * - * @param value the value - * - * @return the request charge - */ - public TransactionalBatchOperationResult setRequestCharge(final double value) { - this.requestCharge = value; - return this; + public Double getRequestCharge() { + return this.requestCharge; } /** @@ -124,18 +106,6 @@ public TResource getItem() { return this.item; } - /** - * Sets item. - * - * @param value the value - * - * @return the item - */ - public TransactionalBatchOperationResult setItem(final TResource value) { - this.item = value; - return this; - } - /** * Gets retry after. * @@ -145,39 +115,15 @@ public Duration getRetryAfter() { return this.retryAfter; } - /** - * Sets retry after. - * - * @param value the value - * - * @return the retry after - */ - public TransactionalBatchOperationResult setRetryAfter(final Duration value) { - this.retryAfter = value; - return this; - } - /** * Gets sub status code. * * @return the sub status code */ - public int getSubStatusCode() { + public Integer getSubStatusCode() { return this.subStatusCode; } - /** - * Sets sub status code. - * - * @param value the value - * - * @return the sub status code - */ - public TransactionalBatchOperationResult setSubStatusCode(final int value) { - this.subStatusCode = value; - return this; - } - /** * Gets a value indicating whether the current operation completed successfully. * @@ -196,29 +142,7 @@ public int getResponseStatus() { return this.responseStatus; } - public void setResponseStatus(int value) { - this.responseStatus = value; - } - public ObjectNode getResourceObject() { return resourceObject; } - - public void setResourceObject(ObjectNode resourceObject) { - this.resourceObject = resourceObject; - } - - @Override - public void close() { - try { - if (this.item instanceof AutoCloseable) { - ((AutoCloseable) this.item).close(); // assumes an idempotent close implementation - } - } catch (Exception ex) { - logger.debug("Unexpected failure in closing item", ex); - } - - this.item = null; - this.resourceObject = null; - } } diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/TransactionalBatchRequestOptions.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/TransactionalBatchRequestOptions.java index b92db378eea5..c15bf7abe827 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/TransactionalBatchRequestOptions.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/TransactionalBatchRequestOptions.java @@ -4,14 +4,10 @@ package com.azure.cosmos; import com.azure.cosmos.implementation.RequestOptions; -import com.azure.cosmos.models.IndexingDirective; public final class TransactionalBatchRequestOptions { private ConsistencyLevel consistencyLevel; - private IndexingDirective indexingDirective; private String sessionToken; - private String ifMatchETag; - private String ifNoneMatchETag; /** * Constructor @@ -20,46 +16,6 @@ public TransactionalBatchRequestOptions() { super(); } - /** - * Gets the If-Match (ETag) associated with the request in the Azure Cosmos DB service. - * - * @return the ifMatchETag associated with the request. - */ - public String getIfMatchETag() { - return this.ifMatchETag; - } - - /** - * Sets the If-Match (ETag) associated with the request in the Azure Cosmos DB service. - * - * @param ifMatchETag the ifMatchETag associated with the request. - * @return the current request options - */ - public TransactionalBatchRequestOptions setIfMatchETag(String ifMatchETag) { - this.ifMatchETag = ifMatchETag; - return this; - } - - /** - * Gets the If-None-Match (ETag) associated with the request in the Azure Cosmos DB service. - * - * @return the ifNoneMatchETag associated with the request. - */ - public String getIfNoneMatchETag() { - return this.ifNoneMatchETag; - } - - /** - * Sets the If-None-Match (ETag) associated with the request in the Azure Cosmos DB service. - * - * @param ifNoneMatchETag the ifNoneMatchETag associated with the request. - * @return the current request options - */ - public TransactionalBatchRequestOptions setIfNoneMatchETag(String ifNoneMatchETag) { - this.ifNoneMatchETag = ifNoneMatchETag; - return this; - } - /** * Gets the consistency level required for the request. * @@ -81,26 +37,6 @@ TransactionalBatchRequestOptions setConsistencyLevel(ConsistencyLevel consistenc return this; } - /** - * Gets the indexing directive (index, do not index etc). - * - * @return the indexing directive. - */ - public IndexingDirective getIndexingDirective() { - return indexingDirective; - } - - /** - * Sets the indexing directive (index, do not index etc). - * - * @param indexingDirective the indexing directive. - * @return the CosmosItemRequestOptions. - */ - public TransactionalBatchRequestOptions setIndexingDirective(IndexingDirective indexingDirective) { - this.indexingDirective = indexingDirective; - return this; - } - /** * Gets the token for use with session consistency. * @@ -121,12 +57,9 @@ public TransactionalBatchRequestOptions setSessionToken(String sessionToken) { return this; } - public RequestOptions toRequestOptions() { + RequestOptions toRequestOptions() { final RequestOptions requestOptions = new RequestOptions(); - requestOptions.setIfMatchETag(getIfMatchETag()); - requestOptions.setIfNoneMatchETag(getIfNoneMatchETag()); requestOptions.setConsistencyLevel(getConsistencyLevel()); - requestOptions.setIndexingDirective(indexingDirective); requestOptions.setSessionToken(sessionToken); return requestOptions; } diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/TransactionalBatchResponse.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/TransactionalBatchResponse.java index f3e90b2c6e90..d4e03ca55221 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/TransactionalBatchResponse.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/TransactionalBatchResponse.java @@ -6,6 +6,7 @@ import com.azure.cosmos.implementation.HttpConstants; import com.azure.cosmos.implementation.JsonSerializable; import com.azure.cosmos.implementation.apachecommons.collections.list.UnmodifiableList; +import com.azure.cosmos.implementation.apachecommons.lang.StringUtils; import com.azure.cosmos.implementation.batch.ItemBatchOperation; import org.slf4j.Logger; import org.slf4j.LoggerFactory; @@ -14,11 +15,8 @@ import java.time.Duration; import java.util.ArrayList; import java.util.Collection; -import java.util.Iterator; import java.util.List; -import java.util.ListIterator; import java.util.Map; -import java.util.stream.Stream; import static com.azure.cosmos.implementation.guava25.base.Preconditions.checkArgument; import static com.azure.cosmos.implementation.guava25.base.Preconditions.checkNotNull; @@ -26,7 +24,7 @@ /** * Response of a {@link TransactionalBatch} request. */ -public class TransactionalBatchResponse implements AutoCloseable, List> { +public class TransactionalBatchResponse { private final static Logger logger = LoggerFactory.getLogger(TransactionalBatchResponse.class); @@ -34,7 +32,7 @@ public class TransactionalBatchResponse implements AutoCloseable, List> results; - private int subStatusCode; + private Integer subStatusCode; private List> operations; private CosmosDiagnostics cosmosDiagnostics; @@ -48,9 +46,9 @@ public class TransactionalBatchResponse implements AutoCloseable, List responseHeaders, final CosmosDiagnostics cosmosDiagnostics, @@ -69,15 +67,6 @@ public TransactionalBatchResponse( this.results = new ArrayList<>(); } - public void createAndPopulateResults(final List> operations, final int retryAfterMilliseconds) { - for (int i = 0; i < operations.size(); i++) { - this.results.add( - new TransactionalBatchOperationResult<>(this.getResponseStatus()) - .setSubStatusCode(this.getSubStatusCode()) - .setRetryAfter(Duration.ofMillis(retryAfterMilliseconds))); - } - } - /** * Gets the result of the operation at the provided index in the current {@link TransactionalBatchResponse batch}. *

@@ -126,15 +115,6 @@ public boolean isSuccessStatusCode() { return this.responseStatus >= 200 && this.responseStatus <= 299; } - /** - * Gets all the activity IDs associated with the response. - * - * @return an enumerable that contains the activity IDs. - */ - public Iterable getActivityIds() { - return Stream.of(this.getActivityId())::iterator; - } - /** * Gets the activity ID that identifies the server request made to execute the batch. * @@ -144,10 +124,6 @@ public String getActivityId() { return this.responseHeaders.get(HttpConstants.HttpHeaders.ACTIVITY_ID); } - public final List> getBatchOperations() { - return this.operations; - } - /** * Gets the reason for the failure of the batch request, if any, or {@code null}. * @@ -157,10 +133,6 @@ public String getErrorMessage() { return this.errorMessage; } - public void setErrorMessage(String value) { - this.errorMessage = value; - } - /** * Gets the request charge for the batch request. * @@ -168,6 +140,10 @@ public void setErrorMessage(String value) { */ public double getRequestCharge() { final String value = this.responseHeaders.get(HttpConstants.HttpHeaders.REQUEST_CHARGE); + if (StringUtils.isEmpty(value)) { + return 0; + } + try { return Double.valueOf(value); } catch (NumberFormatException e) { @@ -191,7 +167,7 @@ public int getResponseStatus() { * @return the response header map. */ public Map getResponseHeaders() { - return responseHeaders; + return this.responseHeaders; } /** @@ -207,10 +183,19 @@ public Duration getRetryAfter() { return null; } - public int getSubStatusCode() { + public Integer getSubStatusCode() { return this.subStatusCode; } + /** + * Get all the results of the operations in batch. + * + * @return Results of operation in batch. + */ + public List> getResults() { + return this.results; + } + /** * Gets the result of the operation at the provided index in the batch. * @@ -218,123 +203,15 @@ public int getSubStatusCode() { * * @return Result of operation at the provided index in the batch. */ - @Override public TransactionalBatchOperationResult get(int index) { return this.results.get(index); } - @Override - public int indexOf(Object o) { - return this.results.indexOf(o); - } - - @Override - public Iterator> iterator() { - return this.results.iterator(); - } - - @Override - public int lastIndexOf(Object o) { - return 0; - } - - @Override - public ListIterator> listIterator() { - return null; - } - - @Override - public ListIterator> listIterator(int index) { - return null; - } - - @Override - public boolean remove(Object result) { - return this.results.remove(result); - } - - @Override - public TransactionalBatchOperationResult remove(int index) { - return null; - } - - @Override - public boolean removeAll(Collection collection) { - return this.results.removeAll(collection); - } - - @Override - public boolean retainAll(Collection collection) { - return this.results.retainAll(collection); - } - - @Override - public TransactionalBatchOperationResult set(int index, TransactionalBatchOperationResult result) { - return this.results.set(index, result); - } - - @Override - public List> subList(int fromIndex, int toIndex) { - return this.results.subList(fromIndex, toIndex); - } - - @Override - public Object[] toArray() { - return this.results.toArray(); - } - - @Override - public T[] toArray(T[] a) { - return this.results.toArray(a); - } - - @Override public boolean isEmpty() { return this.results.isEmpty(); } - @Override - public boolean add(TransactionalBatchOperationResult result) { - return this.results.add(result); - } - - @Override - public void add(int index, TransactionalBatchOperationResult element) { - this.results.add(index, element); - } - - @Override - public boolean addAll(Collection> collection) { + boolean addAll(Collection> collection) { return this.results.addAll(collection); } - - @Override - public boolean addAll(int index, Collection> collection) { - return this.results.addAll(index, collection); - } - - @Override - public void clear() { - this.results.clear(); - } - - @Override - public boolean contains(Object result) { - return this.results.contains(result); - } - - @Override - public boolean containsAll(Collection c) { - return false; - } - - /** - * Closes the current {@link TransactionalBatchResponse}. - */ - public void close() { - this.operations = null; - this.responseHeaders = null; - this.results = null; - this.cosmosDiagnostics = null; - } } diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/batch/BatchExecUtils.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/batch/BatchExecUtils.java index 27e8cc1d2eab..18b2e8f1a5c7 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/batch/BatchExecUtils.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/batch/BatchExecUtils.java @@ -3,96 +3,18 @@ package com.azure.cosmos.implementation.batch; -import com.azure.cosmos.implementation.HttpConstants; import com.azure.cosmos.implementation.OperationType; -import com.azure.cosmos.implementation.RequestOptions; -import com.azure.cosmos.implementation.directconnectivity.WFConstants; -import io.netty.buffer.ByteBufUtil; - -import java.util.List; -import java.util.Map; import static com.azure.cosmos.implementation.batch.BatchRequestResponseConstant.OPERATION_CREATE; import static com.azure.cosmos.implementation.batch.BatchRequestResponseConstant.OPERATION_DELETE; import static com.azure.cosmos.implementation.batch.BatchRequestResponseConstant.OPERATION_READ; import static com.azure.cosmos.implementation.batch.BatchRequestResponseConstant.OPERATION_REPLACE; import static com.azure.cosmos.implementation.batch.BatchRequestResponseConstant.OPERATION_UPSERT; -import static com.azure.cosmos.implementation.guava25.base.Preconditions.checkArgument; -import static com.azure.cosmos.implementation.guava27.Strings.lenientFormat; /** * Util methods for batch requests/response. */ -public class BatchExecUtils { - - static void ensureValid( - final List> operations, - final RequestOptions options) { - - final String errorMessage = BatchExecUtils.isValid(operations, options); - checkArgument(errorMessage == null, errorMessage); - } - - public static String isValid(final List> operations, final RequestOptions batchOptions) { - - String errorMessage = null; - - if (operations == null) { - errorMessage = "expected non-null operations"; - } - - if (errorMessage == null && operations.size() == 0) { - errorMessage = "expected operations.size > 0"; - } - - if (errorMessage == null && batchOptions != null) { - if (batchOptions.getIfMatchETag() != null || batchOptions.getIfNoneMatchETag() != null) { - errorMessage = "one or more request options provided on the batch request are not supported"; - } - } - - if (errorMessage == null) { - for (ItemBatchOperation operation : operations) { - - final RequestOptions batchOperationOptions = operation.getRequestOptions(); - final Map batchOperationProperties = batchOperationOptions != null ? batchOperationOptions.getProperties() : null; - - if (batchOperationProperties != null - && (batchOperationProperties.containsKey(WFConstants.BackendHeaders.EFFECTIVE_PARTITION_KEY) - || batchOperationProperties.containsKey(WFConstants.BackendHeaders.EFFECTIVE_PARTITION_KEY_STRING) - || batchOperationProperties.containsKey(HttpConstants.HttpHeaders.PARTITION_KEY))) { - - final String epkString = (String) batchOperationProperties.computeIfPresent( - WFConstants.BackendHeaders.EFFECTIVE_PARTITION_KEY_STRING, - (k, v) -> v instanceof String ? v : null); - - final byte[] epk = (byte[]) batchOperationProperties.computeIfPresent( - WFConstants.BackendHeaders.EFFECTIVE_PARTITION_KEY, - (k, v) -> v instanceof byte[] ? v : null); - - final String pkString = (String) batchOperationProperties.computeIfPresent( - HttpConstants.HttpHeaders.PARTITION_KEY, - (k, v) -> v instanceof String ? v : null); - - if ((epk == null && pkString == null) || epkString == null) { - return lenientFormat( - "expected byte[] value for %s and string value for %s, not (%s, %s)", - WFConstants.BackendHeaders.EFFECTIVE_PARTITION_KEY, - WFConstants.BackendHeaders.EFFECTIVE_PARTITION_KEY_STRING, - epk == null - ? (pkString == null ? "null" : pkString) - : ByteBufUtil.hexDump(epk), epkString == null ? "null" : epkString); - } - - if (operation.getPartitionKey() != null ) { - errorMessage = "partition key and effective partition key may not both be set."; - } - } - } - } - - return errorMessage; - } +class BatchExecUtils { static String getStringOperationType(OperationType operationType) { switch (operationType) { diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/batch/BatchExecutor.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/batch/BatchExecutor.java index ed0d07c294ca..6a38d483b0ef 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/batch/BatchExecutor.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/batch/BatchExecutor.java @@ -9,15 +9,16 @@ import com.azure.cosmos.TransactionalBatch; import com.azure.cosmos.TransactionalBatchRequestOptions; import com.azure.cosmos.TransactionalBatchResponse; -import com.azure.cosmos.implementation.RequestOptions; import reactor.core.publisher.Mono; import java.util.List; +import static com.azure.cosmos.implementation.guava25.base.Preconditions.checkArgument; + public final class BatchExecutor { private final CosmosAsyncContainer container; - private final RequestOptions options; + private final TransactionalBatchRequestOptions options; private final TransactionalBatch transactionalBatch; public BatchExecutor( @@ -27,7 +28,7 @@ public BatchExecutor( this.container = container; this.transactionalBatch = transactionalBatch; - this.options = options.toRequestOptions(); + this.options = options; } /** @@ -37,17 +38,16 @@ public BatchExecutor( */ public final Mono executeAsync() { - List> operations = this.transactionalBatch.getOperations(); - - BatchExecUtils.ensureValid(operations, this.options); + List> operations = BridgeInternal.getOperationsFromTransactionalBatch(this.transactionalBatch); + checkArgument(operations.size() > 0, "Number of operations should be more than 0."); final SinglePartitionKeyServerBatchRequest request = SinglePartitionKeyServerBatchRequest.createAsync( - this.transactionalBatch.getPartitionKey(), + BridgeInternal.getPartitionKeyFromTransactionalBatch(this.transactionalBatch), operations); request.setAtomicBatch(true); request.setShouldContinueOnError(false); return CosmosBridgeInternal.getAsyncDocumentClient(container.getDatabase()) - .executeBatchRequest(BridgeInternal.getLink(container), request, options, false); + .executeBatchRequest(BridgeInternal.getLink(container), request, BridgeInternal.toRequestOptions(options), false); } } diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/batch/BatchResponseParser.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/batch/BatchResponseParser.java index fc39e6a0a21b..e9a93a22b388 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/batch/BatchResponseParser.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/batch/BatchResponseParser.java @@ -3,6 +3,7 @@ package com.azure.cosmos.implementation.batch; +import com.azure.cosmos.BridgeInternal; import com.azure.cosmos.CosmosException; import com.azure.cosmos.TransactionalBatchOperationResult; import com.azure.cosmos.TransactionalBatchResponse; @@ -21,6 +22,7 @@ import java.io.IOException; import java.time.Duration; import java.util.ArrayList; +import java.util.List; import static com.azure.cosmos.implementation.HttpConstants.HttpHeaders.RETRY_AFTER_IN_MILLISECONDS; import static com.azure.cosmos.implementation.HttpConstants.HttpHeaders.SUB_STATUS; @@ -45,15 +47,15 @@ public static Mono fromErrorResponseAsync( if (throwable instanceof CosmosException) { final CosmosException cosmosException = (CosmosException) throwable; - final TransactionalBatchResponse response = new TransactionalBatchResponse( + final TransactionalBatchResponse response = BridgeInternal.createTransactionBatchResponse( cosmosException.getStatusCode(), cosmosException.getSubStatusCode(), - cosmosException.getMessage(), + cosmosException.toString(), cosmosException.getResponseHeaders(), cosmosException.getDiagnostics(), request.getOperations()); - response.createAndPopulateResults(request.getOperations(), 0); + BatchResponseParser.createAndPopulateResults(response, request.getOperations(), cosmosException.getRetryAfterDuration()); return Mono.just(response); } else { return Mono.error(throwable); @@ -75,14 +77,14 @@ public static Mono fromDocumentServiceResponseAsync( final boolean shouldPromoteOperationStatus) { TransactionalBatchResponse response = null; - final String responseContent = documentServiceResponse.getResponseBodyAsString(); + final byte[] responseContent = documentServiceResponse.getResponseBodyAsByteArray(); - if (StringUtils.isNotEmpty(responseContent)) { + if (responseContent != null && responseContent.length > 0) { response = BatchResponseParser.populateFromResponseContentAsync(documentServiceResponse, request, shouldPromoteOperationStatus); if (response == null) { // Convert any payload read failures as InternalServerError - response = new TransactionalBatchResponse( + response = BridgeInternal.createTransactionBatchResponse( HttpResponseStatus.INTERNAL_SERVER_ERROR.code(), 0, "ServerResponseDeserializationFailure", @@ -97,7 +99,7 @@ public static Mono fromDocumentServiceResponseAsync( documentServiceResponse.getResponseHeaders().getOrDefault(SUB_STATUS, String.valueOf(0))); if (response == null) { - response = new TransactionalBatchResponse( + response = BridgeInternal.createTransactionBatchResponse( responseStatusCode, responseSubStatusCode, null, @@ -110,7 +112,7 @@ public static Mono fromDocumentServiceResponseAsync( if (responseStatusCode >= 200 && responseStatusCode <= 299) { // Server should be guaranteeing number of results equal to operations when // batch request is successful - so fail as InternalServerError if this is not the case. - response = new TransactionalBatchResponse( + response = BridgeInternal.createTransactionBatchResponse( HttpResponseStatus.INTERNAL_SERVER_ERROR.code(), 0, "Invalid server response", @@ -133,7 +135,7 @@ public static Mono fromDocumentServiceResponseAsync( } } - response.createAndPopulateResults(request.getOperations(), retryAfterMilliseconds); + BatchResponseParser.createAndPopulateResults(response, request.getOperations(), Duration.ofMillis(retryAfterMilliseconds)); } checkState(response.size() == request.getOperations().size(), @@ -147,15 +149,15 @@ private static TransactionalBatchResponse populateFromResponseContentAsync( final ServerBatchRequest request, final boolean shouldPromoteOperationStatus) { - final ArrayList> results = new ArrayList<>(); - final String responseContent = documentServiceResponse.getResponseBodyAsString(); + final ArrayList> results = new ArrayList<>(request.getOperations().size()); + final byte[] responseContent = documentServiceResponse.getResponseBodyAsByteArray(); - if (responseContent.charAt(0) != HYBRID_V1) { + if (responseContent[0] != (byte)HYBRID_V1) { // Read from a json response body. To enable hybrid row just complete the else part final ObjectMapper mapper = Utils.getSimpleObjectMapper(); try{ - final ObjectNode[] objectNodes = mapper.readValue(documentServiceResponse.getResponseBodyAsString(), ObjectNode[].class); + final ObjectNode[] objectNodes = mapper.readValue(documentServiceResponse.getResponseBodyAsByteArray(), ObjectNode[].class); for (ObjectNode objectInArray : objectNodes) { final TransactionalBatchOperationResult batchOperationResult = BatchResponseParser.createBatchOperationResultFromJson(objectInArray); results.add(batchOperationResult); @@ -172,7 +174,7 @@ private static TransactionalBatchResponse populateFromResponseContentAsync( } int responseStatusCode = documentServiceResponse.getStatusCode(); - int responseSubStatusCode = Integer.parseInt( + Integer responseSubStatusCode = Integer.parseInt( documentServiceResponse.getResponseHeaders().getOrDefault(SUB_STATUS, String.valueOf(HttpConstants.SubStatusCodes.UNKNOWN))); // Status code of the exact operation which failed. @@ -187,7 +189,7 @@ private static TransactionalBatchResponse populateFromResponseContentAsync( } } - final TransactionalBatchResponse response = new TransactionalBatchResponse( + final TransactionalBatchResponse response = BridgeInternal.createTransactionBatchResponse( responseStatusCode, responseSubStatusCode, null, @@ -195,8 +197,7 @@ private static TransactionalBatchResponse populateFromResponseContentAsync( documentServiceResponse.getCosmosDiagnostics(), request.getOperations()); - response.addAll(results); - + BridgeInternal.addTransactionBatchResultInResponse(response, results); return response; } @@ -210,36 +211,47 @@ private static TransactionalBatchResponse populateFromResponseContentAsync( * @return the result */ private static TransactionalBatchOperationResult createBatchOperationResultFromJson(ObjectNode objectNode) { - final TransactionalBatchOperationResult transactionalBatchOperationResult = new TransactionalBatchOperationResult<>(); - final JsonSerializable jsonSerializable = new JsonSerializable(objectNode); - transactionalBatchOperationResult.setResponseStatus(jsonSerializable.getInt(BatchRequestResponseConstant.FIELD_STATUS_CODE)); + final int responseStatusCode = jsonSerializable.getInt(BatchRequestResponseConstant.FIELD_STATUS_CODE); final Integer subStatusCode = jsonSerializable.getInt(BatchRequestResponseConstant.FIELD_SUBSTATUS_CODE); - if(subStatusCode != null) { - transactionalBatchOperationResult.setSubStatusCode(subStatusCode); - } - final Double requestCharge = jsonSerializable.getDouble(BatchRequestResponseConstant.FIELD_REQUEST_CHARGE); - if(requestCharge != null) { - transactionalBatchOperationResult.setRequestCharge(requestCharge); - } - + final String eTag = jsonSerializable.getString(BatchRequestResponseConstant.FIELD_ETAG); + final ObjectNode resourceBody = jsonSerializable.getObject(BatchRequestResponseConstant.FIELD_RESOURCE_BODY); final Integer retryAfterMilliseconds = jsonSerializable.getInt(BatchRequestResponseConstant.FIELD_RETRY_AFTER_MILLISECONDS); - if(retryAfterMilliseconds != null) { - transactionalBatchOperationResult.setRetryAfter(Duration.ofMillis(retryAfterMilliseconds)); - } - final String etag = jsonSerializable.getString(BatchRequestResponseConstant.FIELD_ETAG); - if(etag != null) { - transactionalBatchOperationResult.setETag(etag); - } + return BridgeInternal.createTransactionBatchResult( + eTag, + requestCharge, + resourceBody, + responseStatusCode, + retryAfterMilliseconds != null ? Duration.ofMillis(retryAfterMilliseconds) : null, + subStatusCode); + } - final ObjectNode resourceBody = jsonSerializable.getObject(BatchRequestResponseConstant.FIELD_RESOURCE_BODY); - if(resourceBody != null) { - transactionalBatchOperationResult.setResourceObject(resourceBody); + /** + * Populate results to match number of operations to number of results in case of any error. + * + * @param response The transactionalBatchResponse in which to add the results + * @param operations List of operations for which the wrapper TransactionalBatchResponse is returned. + * @param retryAfterDuration retryAfterDuration. + * */ + private static void createAndPopulateResults(final TransactionalBatchResponse response, + final List> operations, + final Duration retryAfterDuration) { + final ArrayList> results = new ArrayList<>(operations.size()); + for (int i = 0; i < operations.size(); i++) { + results.add( + BridgeInternal.createTransactionBatchResult( + null, + response.getRequestCharge(), + null, + response.getResponseStatus(), + retryAfterDuration, + response.getSubStatusCode() + )); } - return transactionalBatchOperationResult; + BridgeInternal.addTransactionBatchResultInResponse(response, results); } } diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/batch/ItemBatchOperation.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/batch/ItemBatchOperation.java index d5218c244b2f..05e41e6684f6 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/batch/ItemBatchOperation.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/batch/ItemBatchOperation.java @@ -15,7 +15,7 @@ /** * Represents an operation on an item which will be executed as part of a batch request on a container. */ -public final class ItemBatchOperation implements AutoCloseable { +public final class ItemBatchOperation { private TResource resource; private String materialisedResource; @@ -132,22 +132,6 @@ public void setMaterialisedResource(String materialisedResource) { this.materialisedResource = materialisedResource; } - /** - * Closes this {@link ItemBatchOperation}. - */ - public void close() { - try { - if (this.resource instanceof AutoCloseable) { - ((AutoCloseable) this.resource).close(); // assumes an idempotent close implementation - } - this.resource = null; - } catch (Exception ex) { - // - } - - this.materialisedResource = null; - } - public static final class Builder { private final OperationType operationType; diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/directconnectivity/WFConstants.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/directconnectivity/WFConstants.java index 26277ceeebc1..f40dc6706405 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/directconnectivity/WFConstants.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/directconnectivity/WFConstants.java @@ -69,7 +69,6 @@ public static class BackendHeaders { public static final String BINARY_ID = "x-ms-binary-id"; public static final String TIME_TO_LIVE_IN_SECONDS = "x-ms-time-to-live-in-seconds"; public static final String EFFECTIVE_PARTITION_KEY = "x-ms-effective-partition-key"; - public static final String EFFECTIVE_PARTITION_KEY_STRING = "x-ms-effective-partition-key-string"; public static final String BINARY_PASSTHROUGH_REQUEST = "x-ms-binary-passthrough-request"; public static final String FANOUT_OPERATION_STATE = "x-ms-fanout-operation-state"; public static final String CONTENT_SERIALIZATION_FORMAT = "x-ms-documentdb-content-serialization-format"; diff --git a/sdk/cosmos/azure-cosmos/src/test/java/com/azure/cosmos/BatchOperationResultTests.java b/sdk/cosmos/azure-cosmos/src/test/java/com/azure/cosmos/BatchOperationResultTests.java new file mode 100644 index 000000000000..074467d333ca --- /dev/null +++ b/sdk/cosmos/azure-cosmos/src/test/java/com/azure/cosmos/BatchOperationResultTests.java @@ -0,0 +1,67 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +package com.azure.cosmos; + +import com.azure.cosmos.implementation.HttpConstants; +import com.azure.cosmos.implementation.Utils; +import io.netty.handler.codec.http.HttpResponseStatus; +import org.testng.annotations.Test; + +import java.time.Duration; +import static org.assertj.core.api.Assertions.assertThat; + +public class BatchOperationResultTests { + + private static final int TIMEOUT = 40000; + + private TransactionalBatchOperationResult createTestResult() { + TransactionalBatchOperationResult result = BridgeInternal.createTransactionBatchResult( + "TestETag", + 1.4, + Utils.getSimpleObjectMapper().createObjectNode(), + HttpResponseStatus.OK.code(), + Duration.ofMillis(1234), + HttpConstants.SubStatusCodes.NAME_CACHE_IS_STALE + ); + + return result; + } + + @Test(groups = {"unit"}, timeOut = TIMEOUT) + public void propertiesAreSetThroughCopyCtor() { + TransactionalBatchOperationResult other = createTestResult(); + TransactionalBatchOperationResult result = new TransactionalBatchOperationResult(other); + + assertThat(other.getResponseStatus()).isEqualTo(result.getResponseStatus()); + assertThat(other.getSubStatusCode()).isEqualTo(result.getSubStatusCode()); + assertThat(other.getETag()).isEqualTo(result.getETag()); + assertThat(other.getRequestCharge()).isEqualTo(result.getRequestCharge()); + assertThat(other.getRetryAfter()).isEqualTo(result.getRetryAfter()); + assertThat(other.getResourceObject()).isSameAs(result.getResourceObject()); + } + + @Test(groups = {"unit"}, timeOut = TIMEOUT) + public void propertiesAreSetThroughGenericCtor() { + TransactionalBatchOperationResult other = createTestResult(); + Object testObject = new Object(); + TransactionalBatchOperationResult result = new TransactionalBatchOperationResult(other, testObject); + + assertThat(other.getResponseStatus()).isEqualTo(result.getResponseStatus()); + assertThat(other.getSubStatusCode()).isEqualTo(result.getSubStatusCode()); + assertThat(other.getETag()).isEqualTo(result.getETag()); + assertThat(other.getRequestCharge()).isEqualTo(result.getRequestCharge()); + assertThat(other.getRetryAfter()).isEqualTo(result.getRetryAfter()); + assertThat(other.getResourceObject()).isSameAs(result.getResourceObject()); + assertThat(testObject).isSameAs(result.getItem()); + } + + @Test(groups = {"unit"}, timeOut = TIMEOUT) + public void isSuccessStatusCodeTrueFor200To299() { + for (int x = 100; x < 999; ++x) { + TransactionalBatchOperationResult result = new TransactionalBatchOperationResult(x); + boolean success = x >= 200 && x <= 299; + assertThat(result.isSuccessStatusCode()).isEqualTo(success); + } + } +} diff --git a/sdk/cosmos/azure-cosmos/src/test/java/com/azure/cosmos/BatchTestBase.java b/sdk/cosmos/azure-cosmos/src/test/java/com/azure/cosmos/BatchTestBase.java index c2ee856a209a..8a1e5555334c 100644 --- a/sdk/cosmos/azure-cosmos/src/test/java/com/azure/cosmos/BatchTestBase.java +++ b/sdk/cosmos/azure-cosmos/src/test/java/com/azure/cosmos/BatchTestBase.java @@ -6,12 +6,8 @@ import com.azure.cosmos.implementation.ISessionToken; import com.azure.cosmos.implementation.SessionTokenHelper; import com.azure.cosmos.implementation.apachecommons.lang.StringUtils; -import com.azure.cosmos.implementation.directconnectivity.WFConstants; -import com.azure.cosmos.models.CosmosContainerResponse; -import com.azure.cosmos.models.CosmosDatabaseResponse; import com.azure.cosmos.models.CosmosItemResponse; import com.azure.cosmos.models.PartitionKey; -import com.azure.cosmos.models.ThroughputProperties; import com.azure.cosmos.rx.TestSuiteBase; import com.fasterxml.jackson.annotation.JsonProperty; import io.netty.handler.codec.http.HttpResponseStatus; @@ -19,8 +15,7 @@ import java.util.Random; import java.util.UUID; -import static org.testng.Assert.assertEquals; -import static org.testng.Assert.assertTrue; +import static org.assertj.core.api.Assertions.assertThat; public class BatchTestBase extends TestSuiteBase { @@ -37,27 +32,6 @@ public BatchTestBase(CosmosClientBuilder clientBuilder) { super(clientBuilder); } - CosmosAsyncContainer createSharedThroughputContainer(CosmosAsyncClient client) { - CosmosAsyncContainer sharedThroughputContainer = null; - CosmosDatabaseResponse cosmosDatabaseResponse = client.createDatabaseIfNotExists( - "Shared_" + UUID.randomUUID().toString(), - ThroughputProperties.createManualThroughput(12000)).block(); - - CosmosAsyncDatabase db = client.getDatabase(cosmosDatabaseResponse.getProperties().getId()); - - for (int index = 0; index < 5; index++) { - - CosmosContainerResponse cosmosContainerResponse = db.createContainerIfNotExists(getCollectionDefinition()).block(); - assertTrue(Boolean.parseBoolean(cosmosContainerResponse.getResponseHeaders().get(WFConstants.BackendHeaders.SHARE_THROUGHPUT))); - - if (index == 2) { - sharedThroughputContainer = db.getContainer(cosmosContainerResponse.getProperties().getId()); - } - } - - return sharedThroughputContainer; - } - void createJsonTestDocsAsync(CosmosAsyncContainer container) { this.TestDocPk1ExistingA = this.createJsonTestDocAsync(container, this.partitionKey1); this.TestDocPk1ExistingB = this.createJsonTestDocAsync(container, this.partitionKey1); @@ -92,11 +66,11 @@ void verifyByReadAsync(CosmosAsyncContainer container, TestDoc doc, String eTag) CosmosItemResponse response = container.readItem(doc.getId(), partitionKey, TestDoc.class).block(); - assertEquals(HttpResponseStatus.OK.code(), response.getStatusCode()); - assertEquals(doc, response.getItem()); + assertThat(response.getStatusCode()).isEqualTo(HttpResponseStatus.OK.code()); + assertThat(response.getItem()).isEqualTo(doc); if (eTag != null) { - assertEquals(eTag, response.getETag()); + assertThat(response.getETag()).isEqualTo(eTag); } } @@ -108,9 +82,9 @@ void verifyNotFoundAsync(CosmosAsyncContainer container, TestDoc doc) { CosmosItemResponse response = container.readItem(id, partitionKey, TestDoc.class).block(); // Gateway returns response instead of exception - assertEquals(HttpResponseStatus.NOT_FOUND.code(), response.getStatusCode()); + assertThat(response.getStatusCode()).isEqualTo(HttpResponseStatus.NOT_FOUND.code()); } catch (CosmosException ex) { - assertEquals(HttpResponseStatus.NOT_FOUND.code(), ex.getStatusCode()); + assertThat(ex.getStatusCode()).isEqualTo(HttpResponseStatus.NOT_FOUND.code()); } } @@ -125,7 +99,7 @@ private TestDoc createJsonTestDocAsync(CosmosAsyncContainer container, String pa private TestDoc createJsonTestDocAsync(CosmosAsyncContainer container, String partitionKey, int minDesiredSize) { TestDoc doc = this.populateTestDoc(partitionKey, minDesiredSize); CosmosItemResponse createResponse = container.createItem(doc, this.getPartitionKey(partitionKey), null).block(); - assertEquals(HttpResponseStatus.CREATED.code(), createResponse.getStatusCode()); + assertThat(createResponse.getStatusCode()).isEqualTo(HttpResponseStatus.CREATED.code()); return doc; } diff --git a/sdk/cosmos/azure-cosmos/src/test/java/com/azure/cosmos/TransactionalBatchTest.java b/sdk/cosmos/azure-cosmos/src/test/java/com/azure/cosmos/TransactionalBatchTest.java index dfef7a78b31e..93d14d35d7dd 100644 --- a/sdk/cosmos/azure-cosmos/src/test/java/com/azure/cosmos/TransactionalBatchTest.java +++ b/sdk/cosmos/azure-cosmos/src/test/java/com/azure/cosmos/TransactionalBatchTest.java @@ -5,11 +5,10 @@ import com.azure.cosmos.implementation.HttpConstants; import com.azure.cosmos.implementation.ISessionToken; -import com.azure.cosmos.implementation.apachecommons.lang.StringUtils; import com.azure.cosmos.implementation.guava25.base.Function; import com.azure.cosmos.models.CosmosItemResponse; import io.netty.handler.codec.http.HttpResponseStatus; -import org.testng.Assert; +import org.assertj.core.data.Offset; import org.testng.annotations.AfterClass; import org.testng.annotations.BeforeClass; import org.testng.annotations.Factory; @@ -24,7 +23,6 @@ public class TransactionalBatchTest extends BatchTestBase { private CosmosAsyncClient batchClient; private CosmosAsyncContainer batchContainer; - private CosmosAsyncContainer sharedThroughputContainer; @Factory(dataProvider = "simpleClientBuildersWithDirect") public TransactionalBatchTest(CosmosClientBuilder clientBuilder) { @@ -37,18 +35,11 @@ public void before_TransactionalBatchTest() { this.batchClient = getClientBuilder().buildAsyncClient(); CosmosAsyncContainer asyncContainer = getSharedMultiPartitionCosmosContainer(this.batchClient); batchContainer = batchClient.getDatabase(asyncContainer.getDatabase().getId()).getContainer(asyncContainer.getId()); - sharedThroughputContainer = super.createSharedThroughputContainer(this.batchClient); } @AfterClass(groups = {"emulator"}, timeOut = SHUTDOWN_TIMEOUT, alwaysRun = true) public void afterClass() { assertThat(this.batchClient).isNotNull(); - - if (sharedThroughputContainer != null) { - // Delete the database created in this test - sharedThroughputContainer.getDatabase().delete().block(); - } - this.batchClient.close(); } @@ -57,11 +48,6 @@ public void batchCrud() throws Exception { this.runCrudAsync(batchContainer); } - @Test(groups = {"emulator"}, timeOut = TIMEOUT) - public void batchSharedThroughputCrud() throws Exception { - this.runCrudAsync(sharedThroughputContainer); - } - @Test(groups = {"emulator"}, timeOut = TIMEOUT) public void batchOrdered() { CosmosAsyncContainer container = batchContainer; @@ -80,8 +66,8 @@ public void batchOrdered() { this.verifyBatchProcessed(batchResponse, 2); - Assert.assertEquals(HttpResponseStatus.CREATED.code(), batchResponse.get(0).getResponseStatus()); - Assert.assertEquals(HttpResponseStatus.OK.code(), batchResponse.get(1).getResponseStatus()); + assertThat(batchResponse.get(0).getResponseStatus()).isEqualTo(HttpResponseStatus.CREATED.code()); + assertThat(batchResponse.get(1).getResponseStatus()).isEqualTo(HttpResponseStatus.OK.code()); // Ensure that the replace overwrote the doc from the first operation this.verifyByReadAsync(container, replaceDoc); @@ -103,7 +89,7 @@ public void batchItemETagAsync() { this.getPartitionKey(this.partitionKey1), TestDoc.class).block(); - Assert.assertEquals(HttpResponseStatus.OK.code(), response.getStatusCode()); + assertThat(response.getStatusCode()).isEqualTo(HttpResponseStatus.OK.code()); TransactionalBatchItemRequestOptions firstReplaceOptions = new TransactionalBatchItemRequestOptions(); firstReplaceOptions.setIfMatchETag(response.getETag()); @@ -116,8 +102,8 @@ public void batchItemETagAsync() { this.verifyBatchProcessed(batchResponse, 2); - Assert.assertEquals(HttpResponseStatus.CREATED.code(), batchResponse.get(0).getResponseStatus()); - Assert.assertEquals(HttpResponseStatus.OK.code(), batchResponse.get(1).getResponseStatus()); + assertThat(batchResponse.get(0).getResponseStatus()).isEqualTo(HttpResponseStatus.CREATED.code()); + assertThat(batchResponse.get(1).getResponseStatus()).isEqualTo(HttpResponseStatus.OK.code()); // Ensure that the replace overwrote the doc from the first operation this.verifyByReadAsync(container, testDocToCreate, batchResponse.get(0).getETag()); @@ -138,7 +124,7 @@ public void batchItemETagAsync() { this.verifyBatchProcessed(batchResponse, 1, HttpResponseStatus.PRECONDITION_FAILED); - Assert.assertEquals(HttpResponseStatus.PRECONDITION_FAILED.code(), batchResponse.get(0).getResponseStatus()); + assertThat(batchResponse.get(0).getResponseStatus()).isEqualTo(HttpResponseStatus.PRECONDITION_FAILED.code()); // ensure the document was not updated this.verifyByReadAsync(container, this.TestDocPk1ExistingB); @@ -160,7 +146,7 @@ public void batchItemSessionTokenAsync() { this.getPartitionKey(this.partitionKey1), TestDoc.class).block(); - Assert.assertEquals(HttpResponseStatus.OK.code(), readResponse.getStatusCode()); + assertThat(readResponse.getStatusCode()).isEqualTo(HttpResponseStatus.OK.code()); ISessionToken beforeRequestSessionToken = this.getSessionToken(readResponse.getResponseHeaders().get(HttpConstants.HttpHeaders.SESSION_TOKEN)); @@ -172,11 +158,13 @@ public void batchItemSessionTokenAsync() { this.verifyBatchProcessed(batchResponse, 2); - Assert.assertEquals(HttpResponseStatus.CREATED.code(), batchResponse.get(0).getResponseStatus()); - Assert.assertEquals(HttpResponseStatus.OK.code(), batchResponse.get(1).getResponseStatus()); + assertThat(batchResponse.get(0).getResponseStatus()).isEqualTo(HttpResponseStatus.CREATED.code()); + assertThat(batchResponse.get(1).getResponseStatus()).isEqualTo(HttpResponseStatus.OK.code()); ISessionToken afterRequestSessionToken = this.getSessionToken(batchResponse.getResponseHeaders().get(HttpConstants.HttpHeaders.SESSION_TOKEN)); - Assert.assertTrue(afterRequestSessionToken.getLSN() > beforeRequestSessionToken.getLSN(), "Response session token should be more than request session token"); + assertThat(afterRequestSessionToken.getLSN()) + .as("Response session token should be more than request session token") + .isGreaterThan(beforeRequestSessionToken.getLSN()); } @Test(groups = {"emulator"}, timeOut = TIMEOUT) @@ -192,7 +180,7 @@ public void batchWithTooManyOperationsAsync() { } TransactionalBatchResponse batchResponse = container.executeTransactionalBatch(batch).block(); - Assert.assertEquals(HttpResponseStatus.BAD_REQUEST.code(), batchResponse.getResponseStatus()); + assertThat(batchResponse.getResponseStatus()).isEqualTo(HttpResponseStatus.BAD_REQUEST.code()); } @Test(groups = {"emulator"}, timeOut = TIMEOUT) @@ -209,13 +197,13 @@ public void batchReadsOnlyAsync() throws Exception { this.verifyBatchProcessed(batchResponse, 3); - Assert.assertEquals(HttpResponseStatus.OK.code(), batchResponse.get(0).getResponseStatus()); - Assert.assertEquals(HttpResponseStatus.OK.code(), batchResponse.get(1).getResponseStatus()); - Assert.assertEquals(HttpResponseStatus.OK.code(), batchResponse.get(2).getResponseStatus()); + assertThat(batchResponse.get(0).getResponseStatus()).isEqualTo(HttpResponseStatus.OK.code()); + assertThat(batchResponse.get(1).getResponseStatus()).isEqualTo(HttpResponseStatus.OK.code()); + assertThat(batchResponse.get(2).getResponseStatus()).isEqualTo(HttpResponseStatus.OK.code()); - Assert.assertEquals(this.TestDocPk1ExistingA, batchResponse.getOperationResultAtIndex(0, TestDoc.class).getItem()); - Assert.assertEquals(this.TestDocPk1ExistingB, batchResponse.getOperationResultAtIndex(1, TestDoc.class).getItem()); - Assert.assertEquals(this.TestDocPk1ExistingC, batchResponse.getOperationResultAtIndex(2, TestDoc.class).getItem()); + assertThat(batchResponse.getOperationResultAtIndex(0, TestDoc.class).getItem()).isEqualTo(this.TestDocPk1ExistingA); + assertThat(batchResponse.getOperationResultAtIndex(1, TestDoc.class).getItem()).isEqualTo(this.TestDocPk1ExistingB); + assertThat(batchResponse.getOperationResultAtIndex(2, TestDoc.class).getItem()).isEqualTo(this.TestDocPk1ExistingC); } private TransactionalBatchResponse runCrudAsync(CosmosAsyncContainer container) throws Exception { @@ -243,14 +231,14 @@ private TransactionalBatchResponse runCrudAsync(CosmosAsyncContainer container) this.verifyBatchProcessed(batchResponse, 6); - Assert.assertEquals(HttpResponseStatus.CREATED.code(), batchResponse.get(0).getResponseStatus()); - Assert.assertEquals(HttpResponseStatus.OK.code(), batchResponse.get(1).getResponseStatus()); - Assert.assertEquals(HttpResponseStatus.OK.code(), batchResponse.get(2).getResponseStatus()); - Assert.assertEquals(HttpResponseStatus.CREATED.code(), batchResponse.get(3).getResponseStatus()); - Assert.assertEquals(HttpResponseStatus.OK.code(), batchResponse.get(4).getResponseStatus()); - Assert.assertEquals(HttpResponseStatus.NO_CONTENT.code(), batchResponse.get(5).getResponseStatus()); + assertThat(batchResponse.get(0).getResponseStatus()).isEqualTo(HttpResponseStatus.CREATED.code()); + assertThat(batchResponse.get(1).getResponseStatus()).isEqualTo(HttpResponseStatus.OK.code()); + assertThat(batchResponse.get(2).getResponseStatus()).isEqualTo(HttpResponseStatus.OK.code()); + assertThat(batchResponse.get(3).getResponseStatus()).isEqualTo(HttpResponseStatus.CREATED.code()); + assertThat(batchResponse.get(4).getResponseStatus()).isEqualTo(HttpResponseStatus.OK.code()); + assertThat(batchResponse.get(5).getResponseStatus()).isEqualTo(HttpResponseStatus.NO_CONTENT.code()); - Assert.assertEquals(this.TestDocPk1ExistingC, batchResponse.getOperationResultAtIndex(1, TestDoc.class).getItem()); + assertThat(batchResponse.getOperationResultAtIndex(1, TestDoc.class).getItem()).isEqualTo(this.TestDocPk1ExistingC); this.verifyByReadAsync(container, testDocToCreate); this.verifyByReadAsync(container, testDocToReplace); @@ -266,11 +254,6 @@ public void batchWithCreateConflictAsync() { this.runBatchWithCreateConflictAsync(batchContainer); } - @Test(groups = {"emulator"}, timeOut = TIMEOUT) - public void batchWithCreateConflictSharedThroughputAsync() { - this.runBatchWithCreateConflictAsync(this.sharedThroughputContainer); - } - @Test(groups = {"emulator"}, timeOut = TIMEOUT) public void batchWithInvalidCreateAsync() { CosmosAsyncContainer container = batchContainer; @@ -357,9 +340,9 @@ private void runWithErrorAsync( this.verifyBatchProcessed(batchResponse, 3, expectedFailedOperationStatusCode); - Assert.assertEquals(HttpResponseStatus.FAILED_DEPENDENCY.code(), batchResponse.get(0).getResponseStatus()); - Assert.assertEquals(expectedFailedOperationStatusCode.code(), batchResponse.get(1).getResponseStatus()); - Assert.assertEquals(HttpResponseStatus.FAILED_DEPENDENCY.code(), batchResponse.get(2).getResponseStatus()); + assertThat(batchResponse.get(0).getResponseStatus()).isEqualTo(HttpResponseStatus.FAILED_DEPENDENCY.code()); + assertThat(batchResponse.get(1).getResponseStatus()).isEqualTo(expectedFailedOperationStatusCode.code()); + assertThat(batchResponse.get(2).getResponseStatus()).isEqualTo(HttpResponseStatus.FAILED_DEPENDENCY.code()); this.verifyNotFoundAsync(container, testDocToCreate); this.verifyNotFoundAsync(container, anotherTestDocToCreate); @@ -370,21 +353,19 @@ private void verifyBatchProcessed(TransactionalBatchResponse batchResponse, int } private void verifyBatchProcessed(TransactionalBatchResponse batchResponse, int numberOfOperations, HttpResponseStatus expectedStatusCode) { - Assert.assertNotNull(batchResponse); - Assert.assertEquals( - batchResponse.getResponseStatus(), - expectedStatusCode.code(), - "Batch server response had StatusCode {0} instead of {1} expected and had ErrorMessage {2}"); + assertThat(batchResponse).isNotNull(); + assertThat(batchResponse.getResponseStatus()) + .as("Batch server response had StatusCode {0} instead of {1} expected and had ErrorMessage {2}") + .isEqualTo(expectedStatusCode.code()); - Assert.assertEquals(numberOfOperations, batchResponse.size()); - Assert.assertTrue(batchResponse.getRequestCharge() > 0); - Assert.assertTrue(StringUtils.isNotEmpty(batchResponse.getCosmosDiagnostics().toString())); + assertThat(batchResponse.size()).isEqualTo(numberOfOperations); + assertThat(batchResponse.getRequestCharge()).isPositive(); + assertThat(batchResponse.getCosmosDiagnostics().toString()).isNotEmpty(); // Allow a delta since we round both the total charge and the individual operation // charges to 2 decimal places. - Assert.assertEquals( - batchResponse.getRequestCharge(), - batchResponse.stream().mapToDouble(result -> result.getRequestCharge()).sum(), - 0.1); + assertThat(batchResponse.getRequestCharge()) + .isCloseTo(batchResponse.getResults().stream().mapToDouble(TransactionalBatchOperationResult::getRequestCharge).sum(), + Offset.offset(0.1)); } } diff --git a/sdk/cosmos/azure-cosmos/src/test/java/com/azure/cosmos/implementation/batch/BatchOperationResultTests.java b/sdk/cosmos/azure-cosmos/src/test/java/com/azure/cosmos/implementation/batch/BatchOperationResultTests.java deleted file mode 100644 index f9ce0a38668e..000000000000 --- a/sdk/cosmos/azure-cosmos/src/test/java/com/azure/cosmos/implementation/batch/BatchOperationResultTests.java +++ /dev/null @@ -1,68 +0,0 @@ -// Copyright (c) Microsoft Corporation. All rights reserved. -// Licensed under the MIT License. - -package com.azure.cosmos.implementation.batch; - -import com.azure.cosmos.TransactionalBatchOperationResult; -import com.azure.cosmos.implementation.HttpConstants; -import com.azure.cosmos.implementation.Utils; -import io.netty.handler.codec.http.HttpResponseStatus; -import org.testng.annotations.Test; - -import java.time.Duration; - -import static org.testng.Assert.assertEquals; -import static org.testng.Assert.assertSame; - -public class BatchOperationResultTests { - - private static final int TIMEOUT = 40000; - - private TransactionalBatchOperationResult createTestResult() { - TransactionalBatchOperationResult result = new TransactionalBatchOperationResult(HttpResponseStatus.OK.code()); - result.setSubStatusCode(HttpConstants.SubStatusCodes.NAME_CACHE_IS_STALE); - result.setETag("TestETag"); - result.setRequestCharge(1.4); - result.setResourceObject(Utils.getSimpleObjectMapper().createObjectNode()); - result.setRetryAfter(Duration.ofMillis(1234)); - - return result; - } - - @Test(groups = {"unit"}, timeOut = TIMEOUT) - public void propertiesAreSetThroughCopyCtor() { - TransactionalBatchOperationResult other = createTestResult(); - TransactionalBatchOperationResult result = new TransactionalBatchOperationResult(other); - - assertEquals(other.getResponseStatus(), result.getResponseStatus()); - assertEquals(other.getSubStatusCode(), result.getSubStatusCode()); - assertEquals(other.getETag(), result.getETag()); - assertEquals(other.getRequestCharge(), result.getRequestCharge()); - assertEquals(other.getRetryAfter(), result.getRetryAfter()); - assertSame(other.getResourceObject(), result.getResourceObject()); - } - - @Test(groups = {"unit"}, timeOut = TIMEOUT) - public void propertiesAreSetThroughGenericCtor() { - TransactionalBatchOperationResult other = createTestResult(); - Object testObject = new Object(); - TransactionalBatchOperationResult result = new TransactionalBatchOperationResult(other, testObject); - - assertEquals(other.getResponseStatus(), result.getResponseStatus()); - assertEquals(other.getSubStatusCode(), result.getSubStatusCode()); - assertEquals(other.getETag(), result.getETag()); - assertEquals(other.getRequestCharge(), result.getRequestCharge()); - assertEquals(other.getRetryAfter(), result.getRetryAfter()); - assertSame(other.getResourceObject(), result.getResourceObject()); - assertSame(testObject, result.getItem()); - } - - @Test(groups = {"unit"}, timeOut = TIMEOUT) - public void isSuccessStatusCodeTrueFor200To299() { - for (int x = 100; x < 999; ++x) { - TransactionalBatchOperationResult result = new TransactionalBatchOperationResult(x); - boolean success = x >= 200 && x <= 299; - assertEquals(success, result.isSuccessStatusCode()); - } - } -} diff --git a/sdk/cosmos/azure-cosmos/src/test/java/com/azure/cosmos/implementation/batch/BatchResponsePayloadWriter.java b/sdk/cosmos/azure-cosmos/src/test/java/com/azure/cosmos/implementation/batch/BatchResponsePayloadWriter.java index cc89220ae281..bb13e1d9ec2d 100644 --- a/sdk/cosmos/azure-cosmos/src/test/java/com/azure/cosmos/implementation/batch/BatchResponsePayloadWriter.java +++ b/sdk/cosmos/azure-cosmos/src/test/java/com/azure/cosmos/implementation/batch/BatchResponsePayloadWriter.java @@ -10,15 +10,15 @@ import java.util.List; -public class BatchResponsePayloadWriter { +class BatchResponsePayloadWriter { private List> results; - public BatchResponsePayloadWriter(List> results) { + BatchResponsePayloadWriter(List> results) { this.results = results; } - public String generatePayload() { + String generatePayload() { return writeOperationResult().toString(); } diff --git a/sdk/cosmos/azure-cosmos/src/test/java/com/azure/cosmos/implementation/batch/PartitionKeyBatchResponseTests.java b/sdk/cosmos/azure-cosmos/src/test/java/com/azure/cosmos/implementation/batch/PartitionKeyBatchResponseTests.java index c2b691399a45..a097875c8a96 100644 --- a/sdk/cosmos/azure-cosmos/src/test/java/com/azure/cosmos/implementation/batch/PartitionKeyBatchResponseTests.java +++ b/sdk/cosmos/azure-cosmos/src/test/java/com/azure/cosmos/implementation/batch/PartitionKeyBatchResponseTests.java @@ -3,6 +3,7 @@ package com.azure.cosmos.implementation.batch; +import com.azure.cosmos.BridgeInternal; import com.azure.cosmos.TransactionalBatchOperationResult; import com.azure.cosmos.TransactionalBatchResponse; import com.azure.cosmos.implementation.OperationType; @@ -17,7 +18,7 @@ import java.util.Arrays; import java.util.List; -import static org.testng.AssertJUnit.assertEquals; +import static org.assertj.core.api.Assertions.assertThat; public class PartitionKeyBatchResponseTests { @@ -33,8 +34,14 @@ public void statusCodesAreSetThroughResponseAsync() { .id("0") .build(); - TransactionalBatchOperationResult transactionalBatchOperationResult = new TransactionalBatchOperationResult(HttpResponseStatus.OK.code()); - transactionalBatchOperationResult.setETag(operation.getId()); + TransactionalBatchOperationResult transactionalBatchOperationResult = BridgeInternal.createTransactionBatchResult( + operation.getId(), + 0.0, + null, + HttpResponseStatus.OK.code(), + null, + 0 + ); results.add(transactionalBatchOperationResult); @@ -56,6 +63,6 @@ public void statusCodesAreSetThroughResponseAsync() { batchRequest, true).block(); - assertEquals(HttpResponseStatus.OK.code(), batchResponse.getResponseStatus()); + assertThat(batchResponse.getResponseStatus()).isEqualTo(HttpResponseStatus.OK.code()); } } From a8103d833ae50ecc6a58abc408d8855f91fa63e5 Mon Sep 17 00:00:00 2001 From: Rakesh Kumar Date: Thu, 1 Oct 2020 03:21:01 +0530 Subject: [PATCH 03/22] Few more fixes Signed-off-by: Rakesh Kumar --- .../java/com/azure/cosmos/BridgeInternal.java | 4 +- .../azure/cosmos/CosmosAsyncContainer.java | 4 +- .../com/azure/cosmos/CosmosContainer.java | 8 +- .../TransactionalBatchOperationResult.java | 2 + .../cosmos/TransactionalBatchResponse.java | 6 +- .../implementation/AsyncDocumentClient.java | 4 +- .../implementation/RxDocumentClientImpl.java | 42 ++++++--- .../batch/BatchResponseParser.java | 4 +- .../batch/ItemBatchOperation.java | 51 ++++------- .../batch/ServerBatchRequest.java | 20 +---- .../SinglePartitionKeyServerBatchRequest.java | 1 - .../java/com/azure/cosmos/BatchTestBase.java | 23 ++--- .../azure/cosmos/TransactionalBatchTest.java | 90 +++++++------------ ...a => TransactionalBatchResponseTests.java} | 4 +- 14 files changed, 109 insertions(+), 154 deletions(-) rename sdk/cosmos/azure-cosmos/src/test/java/com/azure/cosmos/implementation/batch/{PartitionKeyBatchResponseTests.java => TransactionalBatchResponseTests.java} (95%) diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/BridgeInternal.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/BridgeInternal.java index 55ea1572e23b..ec694fd94c5f 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/BridgeInternal.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/BridgeInternal.java @@ -616,7 +616,7 @@ public static RequestOptions toRequestOptions(TransactionalBatchRequestOptions t } @Warning(value = INTERNAL_USE_ONLY_WARNING) - public static TransactionalBatchOperationResult createTransactionBatchResult( + public static TransactionalBatchOperationResult createTransactionBatchResult( String eTag, Double requestCharge, ObjectNode resourceObject, @@ -624,7 +624,7 @@ public static TransactionalBatchOperationResult createTransactionBatchResult( Duration retryAfter, Integer subStatusCode) { - return new TransactionalBatchOperationResult( + return new TransactionalBatchOperationResult<>( eTag, requestCharge, resourceObject, diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncContainer.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncContainer.java index 0bcbffeebedb..6e1a24152ecc 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncContainer.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncContainer.java @@ -524,7 +524,7 @@ private T transform(Object object, Class classType) { * Use {@link TransactionalBatchResponse#isSuccessStatusCode} on the response returned to ensure that the * transactional batch succeeded. */ - @Beta(Beta.SinceVersion.V4_4_0) + @Beta(Beta.SinceVersion.V4_6_0) public Mono executeTransactionalBatch(TransactionalBatch transactionalBatch) { return executeTransactionalBatch(transactionalBatch, new TransactionalBatchRequestOptions()); } @@ -558,7 +558,7 @@ public Mono executeTransactionalBatch(TransactionalB * Use {@link TransactionalBatchResponse#isSuccessStatusCode} on the response returned to ensure that the * transactional batch succeeded. */ - @Beta(Beta.SinceVersion.V4_4_0) + @Beta(Beta.SinceVersion.V4_6_0) public Mono executeTransactionalBatch( TransactionalBatch transactionalBatch, TransactionalBatchRequestOptions requestOptions) { diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosContainer.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosContainer.java index 789d304cbf09..979d90f349d0 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosContainer.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosContainer.java @@ -471,7 +471,7 @@ public CosmosItemResponse deleteItem(T item, CosmosItemRequestOption * * @param transactionalBatch Batch having list of operation and partition key which will be executed by this container. * - * @return A Mono response which contains details of execution of the transactional batch. + * @return A TransactionalBatchResponse which contains details of execution of the transactional batch. *

* If the transactional batch executes successfully, the value returned by {@link * TransactionalBatchResponse#getResponseStatus} on the response returned will be set to 200}. @@ -494,7 +494,7 @@ public CosmosItemResponse deleteItem(T item, CosmosItemRequestOption * Use {@link TransactionalBatchResponse#isSuccessStatusCode} on the response returned to ensure that the * transactional batch succeeded. */ - @Beta(Beta.SinceVersion.V4_4_0) + @Beta(Beta.SinceVersion.V4_6_0) public TransactionalBatchResponse executeTransactionalBatch(TransactionalBatch transactionalBatch) { return this.blockBatchResponse(asyncContainer.executeTransactionalBatch(transactionalBatch)); } @@ -505,7 +505,7 @@ public TransactionalBatchResponse executeTransactionalBatch(TransactionalBatch t * @param transactionalBatch Batch having list of operation and partition key which will be executed by this container. * @param requestOptions Options that apply specifically to batch request. * - * @return A Mono response which contains details of execution of the transactional batch. + * @return A TransactionalBatchResponse which contains details of execution of the transactional batch. *

* If the transactional batch executes successfully, the value returned by {@link * TransactionalBatchResponse#getResponseStatus} on the response returned will be set to 200}. @@ -528,7 +528,7 @@ public TransactionalBatchResponse executeTransactionalBatch(TransactionalBatch t * Use {@link TransactionalBatchResponse#isSuccessStatusCode} on the response returned to ensure that the * transactional batch succeeded. */ - @Beta(Beta.SinceVersion.V4_4_0) + @Beta(Beta.SinceVersion.V4_6_0) public TransactionalBatchResponse executeTransactionalBatch( TransactionalBatch transactionalBatch, TransactionalBatchRequestOptions requestOptions) { diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/TransactionalBatchOperationResult.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/TransactionalBatchOperationResult.java index 202ca5150753..72cec8b891ff 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/TransactionalBatchOperationResult.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/TransactionalBatchOperationResult.java @@ -69,6 +69,8 @@ public final class TransactionalBatchOperationResult { int responseStatus, Duration retryAfter, Integer subStatusCode) { + checkNotNull(responseStatus, "expected non-null responseStatus"); + this.eTag = eTag; this.requestCharge = requestCharge; this.resourceObject = resourceObject; diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/TransactionalBatchResponse.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/TransactionalBatchResponse.java index d4e03ca55221..7d0537e5dada 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/TransactionalBatchResponse.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/TransactionalBatchResponse.java @@ -11,7 +11,6 @@ import org.slf4j.Logger; import org.slf4j.LoggerFactory; -import java.io.IOException; import java.time.Duration; import java.util.ArrayList; import java.util.Collection; @@ -76,11 +75,10 @@ public class TransactionalBatchResponse { * @param type class type for which deserialization is needed. * * @return TransactionalBatchOperationResult containing the individual result of operation. - * @throws IOException if the body of the resource cannot be read. */ public TransactionalBatchOperationResult getOperationResultAtIndex( final int index, - final Class type) throws IOException { + final Class type) { checkArgument(index >= 0, "expected non-negative index"); checkNotNull(type, "expected non-null type"); @@ -101,6 +99,8 @@ public CosmosDiagnostics getCosmosDiagnostics() { /** * Gets the number of operation results. + * + * @return the number of operations results in this response. */ public int size() { return this.results == null ? 0 : this.results.size(); diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/AsyncDocumentClient.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/AsyncDocumentClient.java index a49e3c049b7d..07618d672f9e 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/AsyncDocumentClient.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/AsyncDocumentClient.java @@ -7,8 +7,6 @@ import com.azure.cosmos.implementation.batch.ServerBatchRequest; import com.azure.cosmos.TransactionalBatchResponse; import com.azure.cosmos.implementation.apachecommons.lang.StringUtils; -import com.azure.cosmos.implementation.caches.RxClientCollectionCache; -import com.azure.cosmos.implementation.caches.RxPartitionKeyRangeCache; import com.azure.cosmos.models.CosmosItemIdentity; import com.azure.cosmos.models.CosmosQueryRequestOptions; import com.azure.cosmos.models.FeedResponse; @@ -806,7 +804,7 @@ Mono executeStoredProcedure(String storedProcedureLink, * @param serverBatchRequest the batch request with the content and flags. * @param options the request options. * @param disableAutomaticIdGeneration the flag for disabling automatic id generation. - * @return a {@link Mono} containing the single resource response with the created document or an error. + * @return a {@link Mono} containing the transactionalBatchResponse response which results of all operations. */ Mono executeBatchRequest(String collectionLink, ServerBatchRequest serverBatchRequest, diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/RxDocumentClientImpl.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/RxDocumentClientImpl.java index 12068708234d..401e6f6777cb 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/RxDocumentClientImpl.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/RxDocumentClientImpl.java @@ -1218,11 +1218,11 @@ private Mono getCreateDocumentRequest(DocumentClientRe return addPartitionKeyInformation(request, content, document, options, collectionObs); } - private RxDocumentServiceRequest getBatchDocumentRequest(DocumentClientRetryPolicy requestRetryPolicy, - String documentCollectionLink, - ServerBatchRequest serverBatchRequest, - RequestOptions options, - boolean disableAutomaticIdGeneration) { + private Mono getBatchDocumentRequest(DocumentClientRetryPolicy requestRetryPolicy, + String documentCollectionLink, + ServerBatchRequest serverBatchRequest, + RequestOptions options, + boolean disableAutomaticIdGeneration) { checkArgument(StringUtils.isNotEmpty(documentCollectionLink), "expected non empty documentCollectionLink"); checkNotNull(serverBatchRequest, "expected non null serverBatchRequest"); @@ -1230,6 +1230,7 @@ private RxDocumentServiceRequest getBatchDocumentRequest(DocumentClientRetryPoli Instant serializationStartTimeUTC = Instant.now(); ByteBuffer content = ByteBuffer.wrap(serverBatchRequest.transferRequestBody().getBytes(StandardCharsets.UTF_8)); Instant serializationEndTimeUTC = Instant.now(); + SerializationDiagnosticsContext.SerializationDiagnostics serializationDiagnostics = new SerializationDiagnosticsContext.SerializationDiagnostics( serializationStartTimeUTC, serializationEndTimeUTC, @@ -1239,6 +1240,7 @@ private RxDocumentServiceRequest getBatchDocumentRequest(DocumentClientRetryPoli Map requestHeaders = this.getRequestHeaders(options, ResourceType.Document, OperationType.Batch); RxDocumentServiceRequest request = RxDocumentServiceRequest.create( + this, OperationType.Batch, ResourceType.Document, path, @@ -1255,14 +1257,32 @@ private RxDocumentServiceRequest getBatchDocumentRequest(DocumentClientRetryPoli serializationDiagnosticsContext.addSerializationDiagnostics(serializationDiagnostics); } - return addBatchHeaders(request, serverBatchRequest); + Mono> collectionObs = + this.collectionCache.resolveCollectionAsync(BridgeInternal.getMetaDataDiagnosticContext(request.requestContext.cosmosDiagnostics), request); + + return collectionObs.map((Utils.ValueHolder collectionValueHolder) -> { + addBatchHeaders(request, serverBatchRequest, collectionValueHolder.v); + return request; + }); } private RxDocumentServiceRequest addBatchHeaders(RxDocumentServiceRequest request, - ServerBatchRequest serverBatchRequest) { + ServerBatchRequest serverBatchRequest, + DocumentCollection collection) { if(serverBatchRequest instanceof SinglePartitionKeyServerBatchRequest) { - PartitionKeyInternal partitionKeyInternal = BridgeInternal.getPartitionKeyInternal(((SinglePartitionKeyServerBatchRequest) serverBatchRequest).getPartitionKey()); + + PartitionKey partitionKey = ((SinglePartitionKeyServerBatchRequest) serverBatchRequest).getPartitionKey(); + PartitionKeyInternal partitionKeyInternal; + + if (partitionKey.equals(PartitionKey.NONE)) { + PartitionKeyDefinition partitionKeyDefinition = collection.getPartitionKey(); + partitionKeyInternal = ModelBridgeInternal.getNonePartitionKey(partitionKeyDefinition); + } else { + // Partition key is always non-null + partitionKeyInternal = BridgeInternal.getPartitionKeyInternal(partitionKey); + } + request.setPartitionKeyInternal(partitionKeyInternal); request.getHeaders().put(HttpConstants.HttpHeaders.PARTITION_KEY, Utils.escapeNonAscii(partitionKeyInternal.toJson())); } else { @@ -2367,8 +2387,10 @@ private Mono executeBatchRequestInternal(String coll try { logger.debug("Executing a Batch request with number of operations {}", serverBatchRequest.getOperations().size()); - RxDocumentServiceRequest documentServiceRequest = getBatchDocumentRequest(requestRetryPolicy, collectionLink, serverBatchRequest, options, disableAutomaticIdGeneration); - Mono responseObservable = create(documentServiceRequest, requestRetryPolicy); + Mono requestObs = getBatchDocumentRequest(requestRetryPolicy, collectionLink, serverBatchRequest, options, disableAutomaticIdGeneration); + Mono responseObservable = requestObs.flatMap(request -> { + return create(request, requestRetryPolicy); + }); return responseObservable .flatMap(serviceResponse -> BatchResponseParser.fromDocumentServiceResponseAsync(serviceResponse, serverBatchRequest, true)) diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/batch/BatchResponseParser.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/batch/BatchResponseParser.java index e9a93a22b388..bc8080cafe60 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/batch/BatchResponseParser.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/batch/BatchResponseParser.java @@ -156,8 +156,8 @@ private static TransactionalBatchResponse populateFromResponseContentAsync( // Read from a json response body. To enable hybrid row just complete the else part final ObjectMapper mapper = Utils.getSimpleObjectMapper(); - try{ - final ObjectNode[] objectNodes = mapper.readValue(documentServiceResponse.getResponseBodyAsByteArray(), ObjectNode[].class); + try { + final ObjectNode[] objectNodes = mapper.readValue(responseContent, ObjectNode[].class); for (ObjectNode objectInArray : objectNodes) { final TransactionalBatchOperationResult batchOperationResult = BatchResponseParser.createBatchOperationResultFromJson(objectInArray); results.add(batchOperationResult); diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/batch/ItemBatchOperation.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/batch/ItemBatchOperation.java index 05e41e6684f6..b9fe695708a4 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/batch/ItemBatchOperation.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/batch/ItemBatchOperation.java @@ -18,7 +18,6 @@ public final class ItemBatchOperation { private TResource resource; - private String materialisedResource; private String id; private int operationIndex; @@ -53,8 +52,7 @@ static JsonSerializable writeOperation(final ItemBatchOperation operation) { jsonSerializable.set(BatchRequestResponseConstant.FIELD_OPERATION_TYPE, BatchExecUtils.getStringOperationType(operation.getOperationType())); if (StringUtils.isNotEmpty(operation.getPartitionKeyJson())) { - // This is set in BatchAsyncContainerExecutor.resolvePartitionKeyRangeIdAsync. For transactional no need to - // pass partition key in operations as batch will have it. + // Used for non transactional batch. jsonSerializable.set(BatchRequestResponseConstant.FIELD_PARTITION_KEY, operation.getPartitionKeyJson()); } @@ -81,55 +79,36 @@ static JsonSerializable writeOperation(final ItemBatchOperation operation) { return jsonSerializable; } - public String getId() { + String getId() { return this.id; } - int getOperationIndex() { - return this.operationIndex; - } - - ItemBatchOperation setOperationIndex(final int value) { - this.operationIndex = value; - return this; - } - - public OperationType getOperationType() { + OperationType getOperationType() { return this.operationType; } - public PartitionKey getPartitionKey() { - return partitionKey; - } - - public ItemBatchOperation setPartitionKey(final PartitionKey value) { - partitionKey = value; - return this; - } - - private String getPartitionKeyJson() { - return partitionKeyJson; + String getPartitionKeyJson() { + return this.partitionKeyJson; } - ItemBatchOperation setPartitionKeyJson(final String value) { - partitionKeyJson = value; - return this; + int getOperationIndex() { + return operationIndex; } - public RequestOptions getRequestOptions() { - return requestOptions; + PartitionKey getPartitionKey() { + return partitionKey; } - public TResource getResource() { - return resource; + void setPartitionKeyJson(String partitionKeyJson) { + this.partitionKeyJson = partitionKeyJson; } - public String getMaterialisedResource() { - return materialisedResource; + RequestOptions getRequestOptions() { + return this.requestOptions; } - public void setMaterialisedResource(String materialisedResource) { - this.materialisedResource = materialisedResource; + TResource getResource() { + return this.resource; } public static final class Builder { diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/batch/ServerBatchRequest.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/batch/ServerBatchRequest.java index 133fa74ab660..55468705ffb6 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/batch/ServerBatchRequest.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/batch/ServerBatchRequest.java @@ -37,20 +37,6 @@ public abstract class ServerBatchRequest { this.maxOperationCount = maxOperationCount; } - /** - * Adds as many operations as possible from the given list of operations. - *

- * Operations are added in order while ensuring the request stream never exceeds {@link #maxBodyLength}. - * - * @param operations Operations to be added; read-only. - * - * @return Any pending operations that were not included in the request. - */ - final List> createBodyStreamAsync( - final List> operations) { - return createBodyStreamAsync(operations, false); - } - /** * Adds as many operations as possible from the given list of operations. * TODO(rakkuma): Similarly for hybrid row, request needs to be parsed to create a request body in any form. @@ -59,14 +45,10 @@ final List> createBodyStreamAsync( * Operations are added in order while ensuring the request body never exceeds {@link #maxBodyLength}. * * @param operations operations to be added; read-only. - * @param ensureContinuousOperationIndexes specifies whether to stop adding operations to the request once there is - * non-continuity in the operation indexes. * * @return Any pending operations that were not included in the request. */ - final List> createBodyStreamAsync( - final List> operations, - final boolean ensureContinuousOperationIndexes) { + final List> createBodyStreamAsync(final List> operations) { checkNotNull(operations, "expected non-null operations"); diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/batch/SinglePartitionKeyServerBatchRequest.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/batch/SinglePartitionKeyServerBatchRequest.java index a499b089b4c9..1d3589e21976 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/batch/SinglePartitionKeyServerBatchRequest.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/batch/SinglePartitionKeyServerBatchRequest.java @@ -4,7 +4,6 @@ package com.azure.cosmos.implementation.batch; import com.azure.cosmos.models.PartitionKey; - import java.util.List; import static com.azure.cosmos.implementation.guava25.base.Preconditions.checkNotNull; diff --git a/sdk/cosmos/azure-cosmos/src/test/java/com/azure/cosmos/BatchTestBase.java b/sdk/cosmos/azure-cosmos/src/test/java/com/azure/cosmos/BatchTestBase.java index 8a1e5555334c..732e8eacf891 100644 --- a/sdk/cosmos/azure-cosmos/src/test/java/com/azure/cosmos/BatchTestBase.java +++ b/sdk/cosmos/azure-cosmos/src/test/java/com/azure/cosmos/BatchTestBase.java @@ -32,7 +32,7 @@ public BatchTestBase(CosmosClientBuilder clientBuilder) { super(clientBuilder); } - void createJsonTestDocsAsync(CosmosAsyncContainer container) { + void createJsonTestDocsAsync(CosmosContainer container) { this.TestDocPk1ExistingA = this.createJsonTestDocAsync(container, this.partitionKey1); this.TestDocPk1ExistingB = this.createJsonTestDocAsync(container, this.partitionKey1); this.TestDocPk1ExistingC = this.createJsonTestDocAsync(container, this.partitionKey1); @@ -48,23 +48,18 @@ TestDoc populateTestDoc(String partitionKey, int minDesiredSize) { return new TestDoc(UUID.randomUUID().toString(), this.random.nextInt(), description, partitionKey); } - public TestDoc populateTestDoc(String id, String partitionKey) { - String description = StringUtils.repeat("x", 20); - return new TestDoc(id, this.random.nextInt(), description, partitionKey); - } - TestDoc getTestDocCopy(TestDoc testDoc) { return new TestDoc(testDoc.getId(), testDoc.getCost(), testDoc.getDescription(), testDoc.getStatus()); } - void verifyByReadAsync(CosmosAsyncContainer container, TestDoc doc) { + void verifyByReadAsync(CosmosContainer container, TestDoc doc) { verifyByReadAsync(container, doc, null); } - void verifyByReadAsync(CosmosAsyncContainer container, TestDoc doc, String eTag) { + void verifyByReadAsync(CosmosContainer container, TestDoc doc, String eTag) { PartitionKey partitionKey = this.getPartitionKey(doc.getStatus()); - CosmosItemResponse response = container.readItem(doc.getId(), partitionKey, TestDoc.class).block(); + CosmosItemResponse response = container.readItem(doc.getId(), partitionKey, TestDoc.class); assertThat(response.getStatusCode()).isEqualTo(HttpResponseStatus.OK.code()); assertThat(response.getItem()).isEqualTo(doc); @@ -74,12 +69,12 @@ void verifyByReadAsync(CosmosAsyncContainer container, TestDoc doc, String eTag) } } - void verifyNotFoundAsync(CosmosAsyncContainer container, TestDoc doc) { + void verifyNotFoundAsync(CosmosContainer container, TestDoc doc) { String id = doc.getId(); PartitionKey partitionKey = this.getPartitionKey(doc.getStatus()); try { - CosmosItemResponse response = container.readItem(id, partitionKey, TestDoc.class).block(); + CosmosItemResponse response = container.readItem(id, partitionKey, TestDoc.class); // Gateway returns response instead of exception assertThat(response.getStatusCode()).isEqualTo(HttpResponseStatus.NOT_FOUND.code()); @@ -92,13 +87,13 @@ PartitionKey getPartitionKey(String partitionKey) { return new PartitionKey(partitionKey); } - private TestDoc createJsonTestDocAsync(CosmosAsyncContainer container, String partitionKey) { + private TestDoc createJsonTestDocAsync(CosmosContainer container, String partitionKey) { return createJsonTestDocAsync(container, partitionKey, 20); } - private TestDoc createJsonTestDocAsync(CosmosAsyncContainer container, String partitionKey, int minDesiredSize) { + private TestDoc createJsonTestDocAsync(CosmosContainer container, String partitionKey, int minDesiredSize) { TestDoc doc = this.populateTestDoc(partitionKey, minDesiredSize); - CosmosItemResponse createResponse = container.createItem(doc, this.getPartitionKey(partitionKey), null).block(); + CosmosItemResponse createResponse = container.createItem(doc, this.getPartitionKey(partitionKey), null); assertThat(createResponse.getStatusCode()).isEqualTo(HttpResponseStatus.CREATED.code()); return doc; } diff --git a/sdk/cosmos/azure-cosmos/src/test/java/com/azure/cosmos/TransactionalBatchTest.java b/sdk/cosmos/azure-cosmos/src/test/java/com/azure/cosmos/TransactionalBatchTest.java index 93d14d35d7dd..1c5e07791659 100644 --- a/sdk/cosmos/azure-cosmos/src/test/java/com/azure/cosmos/TransactionalBatchTest.java +++ b/sdk/cosmos/azure-cosmos/src/test/java/com/azure/cosmos/TransactionalBatchTest.java @@ -21,8 +21,8 @@ public class TransactionalBatchTest extends BatchTestBase { - private CosmosAsyncClient batchClient; - private CosmosAsyncContainer batchContainer; + private CosmosClient batchClient; + private CosmosContainer batchContainer; @Factory(dataProvider = "simpleClientBuildersWithDirect") public TransactionalBatchTest(CosmosClientBuilder clientBuilder) { @@ -32,8 +32,8 @@ public TransactionalBatchTest(CosmosClientBuilder clientBuilder) { @BeforeClass(groups = {"emulator"}, timeOut = SETUP_TIMEOUT) public void before_TransactionalBatchTest() { assertThat(this.batchClient).isNull(); - this.batchClient = getClientBuilder().buildAsyncClient(); - CosmosAsyncContainer asyncContainer = getSharedMultiPartitionCosmosContainer(this.batchClient); + this.batchClient = getClientBuilder().buildClient(); + CosmosAsyncContainer asyncContainer = getSharedMultiPartitionCosmosContainer(this.batchClient.asyncClient()); batchContainer = batchClient.getDatabase(asyncContainer.getDatabase().getId()).getContainer(asyncContainer.getId()); } @@ -43,14 +43,9 @@ public void afterClass() { this.batchClient.close(); } - @Test(groups = {"emulator"}, timeOut = TIMEOUT) - public void batchCrud() throws Exception { - this.runCrudAsync(batchContainer); - } - @Test(groups = {"emulator"}, timeOut = TIMEOUT) public void batchOrdered() { - CosmosAsyncContainer container = batchContainer; + CosmosContainer container = this.batchContainer; this.createJsonTestDocsAsync(container); TestDoc firstDoc = this.populateTestDoc(this.partitionKey1); @@ -61,8 +56,7 @@ public void batchOrdered() { TransactionalBatchResponse batchResponse = container.executeTransactionalBatch( TransactionalBatch.createTransactionalBatch(this.getPartitionKey(this.partitionKey1)) .createItem(firstDoc) - .replaceItem(replaceDoc.getId(), replaceDoc)) - .block(); + .replaceItem(replaceDoc.getId(), replaceDoc)); this.verifyBatchProcessed(batchResponse, 2); @@ -75,7 +69,7 @@ public void batchOrdered() { @Test(groups = {"emulator"}, timeOut = TIMEOUT) public void batchItemETagAsync() { - CosmosAsyncContainer container = batchContainer; + CosmosContainer container = batchContainer; this.createJsonTestDocsAsync(container); { @@ -87,7 +81,7 @@ public void batchItemETagAsync() { CosmosItemResponse response = container.readItem( this.TestDocPk1ExistingA.getId(), this.getPartitionKey(this.partitionKey1), - TestDoc.class).block(); + TestDoc.class); assertThat(response.getStatusCode()).isEqualTo(HttpResponseStatus.OK.code()); @@ -97,8 +91,7 @@ public void batchItemETagAsync() { TransactionalBatchResponse batchResponse = container.executeTransactionalBatch( TransactionalBatch.createTransactionalBatch(this.getPartitionKey(this.partitionKey1)) .createItem(testDocToCreate) - .replaceItem(testDocToReplace.getId(), testDocToReplace, firstReplaceOptions)) - .block(); + .replaceItem(testDocToReplace.getId(), testDocToReplace, firstReplaceOptions)); this.verifyBatchProcessed(batchResponse, 2); @@ -119,8 +112,7 @@ public void batchItemETagAsync() { TransactionalBatchResponse batchResponse = container.executeTransactionalBatch( TransactionalBatch.createTransactionalBatch(this.getPartitionKey(this.partitionKey1)) - .replaceItem(testDocToReplace.getId(), testDocToReplace, replaceOptions)) - .block(); + .replaceItem(testDocToReplace.getId(), testDocToReplace, replaceOptions)); this.verifyBatchProcessed(batchResponse, 1, HttpResponseStatus.PRECONDITION_FAILED); @@ -133,7 +125,7 @@ public void batchItemETagAsync() { @Test(groups = {"emulator"}, timeOut = TIMEOUT) public void batchItemSessionTokenAsync() { - CosmosAsyncContainer container = batchContainer; + CosmosContainer container = batchContainer; this.createJsonTestDocsAsync(container); TestDoc testDocToCreate = this.populateTestDoc(this.partitionKey1); @@ -144,7 +136,7 @@ public void batchItemSessionTokenAsync() { CosmosItemResponse readResponse = container.readItem( this.TestDocPk1ExistingA.getId(), this.getPartitionKey(this.partitionKey1), - TestDoc.class).block(); + TestDoc.class); assertThat(readResponse.getStatusCode()).isEqualTo(HttpResponseStatus.OK.code()); @@ -153,8 +145,7 @@ public void batchItemSessionTokenAsync() { TransactionalBatchResponse batchResponse = container.executeTransactionalBatch( TransactionalBatch.createTransactionalBatch(this.getPartitionKey(this.partitionKey1)) .createItem(testDocToCreate) - .replaceItem(testDocToReplace.getId(), testDocToReplace)) - .block(); + .replaceItem(testDocToReplace.getId(), testDocToReplace)); this.verifyBatchProcessed(batchResponse, 2); @@ -169,7 +160,7 @@ public void batchItemSessionTokenAsync() { @Test(groups = {"emulator"}, timeOut = TIMEOUT) public void batchWithTooManyOperationsAsync() { - CosmosAsyncContainer container = batchContainer; + CosmosContainer container = batchContainer; int operationCount = MAX_OPERATIONS_IN_DIRECT_MODE_BATCH_REQUEST + 1; // Increase the doc size by a bit so all docs won't fit in one server request. @@ -179,21 +170,20 @@ public void batchWithTooManyOperationsAsync() { batch.readItem("someId"); } - TransactionalBatchResponse batchResponse = container.executeTransactionalBatch(batch).block(); + TransactionalBatchResponse batchResponse = container.executeTransactionalBatch(batch); assertThat(batchResponse.getResponseStatus()).isEqualTo(HttpResponseStatus.BAD_REQUEST.code()); } @Test(groups = {"emulator"}, timeOut = TIMEOUT) public void batchReadsOnlyAsync() throws Exception { - CosmosAsyncContainer container = batchContainer; + CosmosContainer container = batchContainer; this.createJsonTestDocsAsync(container); TransactionalBatchResponse batchResponse = container.executeTransactionalBatch( TransactionalBatch.createTransactionalBatch(this.getPartitionKey(this.partitionKey1)) .readItem(this.TestDocPk1ExistingA.getId()) .readItem(this.TestDocPk1ExistingB.getId()) - .readItem(this.TestDocPk1ExistingC.getId())) - .block(); + .readItem(this.TestDocPk1ExistingC.getId())); this.verifyBatchProcessed(batchResponse, 3); @@ -206,7 +196,9 @@ public void batchReadsOnlyAsync() throws Exception { assertThat(batchResponse.getOperationResultAtIndex(2, TestDoc.class).getItem()).isEqualTo(this.TestDocPk1ExistingC); } - private TransactionalBatchResponse runCrudAsync(CosmosAsyncContainer container) throws Exception { + @Test(groups = {"emulator"}, timeOut = TIMEOUT) + public void batchCrud() throws Exception { + CosmosContainer container = batchContainer; this.createJsonTestDocsAsync(container); BatchTestBase.TestDoc testDocToCreate = this.populateTestDoc(this.partitionKey1); @@ -226,8 +218,7 @@ private TransactionalBatchResponse runCrudAsync(CosmosAsyncContainer container) .replaceItem(testDocToReplace.getId(), testDocToReplace) .upsertItem(testDocToUpsert) .upsertItem(anotherTestDocToUpsert) - .deleteItem(this.TestDocPk1ExistingD.getId())) - .block(); + .deleteItem(this.TestDocPk1ExistingD.getId())); this.verifyBatchProcessed(batchResponse, 6); @@ -245,39 +236,28 @@ private TransactionalBatchResponse runCrudAsync(CosmosAsyncContainer container) this.verifyByReadAsync(container, testDocToUpsert); this.verifyByReadAsync(container, anotherTestDocToUpsert); this.verifyNotFoundAsync(container, this.TestDocPk1ExistingD); - - return batchResponse; - } - - @Test(groups = {"emulator"}, timeOut = TIMEOUT) - public void batchWithCreateConflictAsync() { - this.runBatchWithCreateConflictAsync(batchContainer); } @Test(groups = {"emulator"}, timeOut = TIMEOUT) public void batchWithInvalidCreateAsync() { - CosmosAsyncContainer container = batchContainer; - // partition key mismatch between doc and and value passed in to the operation this.runWithErrorAsync( - container, + batchContainer, batch -> batch.createItem(this.populateTestDoc(UUID.randomUUID().toString())), HttpResponseStatus.BAD_REQUEST); } @Test(groups = {"emulator"}, timeOut = TIMEOUT) public void batchWithReadOfNonExistentEntityAsync() { - CosmosAsyncContainer container = batchContainer; this.runWithErrorAsync( - container, + batchContainer, batch -> batch.readItem(UUID.randomUUID().toString()), HttpResponseStatus.NOT_FOUND); } @Test(groups = {"emulator"}, timeOut = TIMEOUT) public void batchWithReplaceOfStaleEntityAsync() { - CosmosAsyncContainer container = batchContainer; - this.createJsonTestDocsAsync(container); + this.createJsonTestDocsAsync(batchContainer); TestDoc staleTestDocToReplace = this.getTestDocCopy(this.TestDocPk1ExistingA); staleTestDocToReplace.setCost(staleTestDocToReplace.getCost() + 1); @@ -286,43 +266,42 @@ public void batchWithReplaceOfStaleEntityAsync() { staleReplaceOptions.setIfMatchETag(UUID.randomUUID().toString()); this.runWithErrorAsync( - container, + batchContainer, batch -> batch.replaceItem(staleTestDocToReplace.getId(), staleTestDocToReplace, staleReplaceOptions), HttpResponseStatus.PRECONDITION_FAILED); // make sure the stale doc hasn't changed - this.verifyByReadAsync(container, this.TestDocPk1ExistingA); + this.verifyByReadAsync(batchContainer, this.TestDocPk1ExistingA); } @Test(groups = {"emulator"}, timeOut = TIMEOUT) public void batchWithDeleteOfNonExistentEntityAsync() { - CosmosAsyncContainer container = batchContainer; - this.runWithErrorAsync( - container, + batchContainer, batch -> batch.deleteItem(UUID.randomUUID().toString()), HttpResponseStatus.NOT_FOUND); } - private void runBatchWithCreateConflictAsync(CosmosAsyncContainer container) { - this.createJsonTestDocsAsync(container); + @Test(groups = {"emulator"}, timeOut = TIMEOUT) + public void batchWithCreateConflictAsync() { + this.createJsonTestDocsAsync(batchContainer); // try to create a doc with id that already exists (should return a Conflict) TestDoc conflictingTestDocToCreate = this.getTestDocCopy(this.TestDocPk1ExistingA); conflictingTestDocToCreate.setCost(conflictingTestDocToCreate.getCost()); this.runWithErrorAsync( - container, + batchContainer, batch -> batch.createItem(conflictingTestDocToCreate), HttpResponseStatus.CONFLICT); // make sure the conflicted doc hasn't changed - this.verifyByReadAsync(container, this.TestDocPk1ExistingA); + this.verifyByReadAsync(batchContainer, this.TestDocPk1ExistingA); } private void runWithErrorAsync( - CosmosAsyncContainer container, + CosmosContainer container, Function appendOperation, HttpResponseStatus expectedFailedOperationStatusCode) { @@ -335,8 +314,7 @@ private void runWithErrorAsync( appendOperation.apply(batch); TransactionalBatchResponse batchResponse = container.executeTransactionalBatch( - batch.createItem(anotherTestDocToCreate)) - .block(); + batch.createItem(anotherTestDocToCreate)); this.verifyBatchProcessed(batchResponse, 3, expectedFailedOperationStatusCode); diff --git a/sdk/cosmos/azure-cosmos/src/test/java/com/azure/cosmos/implementation/batch/PartitionKeyBatchResponseTests.java b/sdk/cosmos/azure-cosmos/src/test/java/com/azure/cosmos/implementation/batch/TransactionalBatchResponseTests.java similarity index 95% rename from sdk/cosmos/azure-cosmos/src/test/java/com/azure/cosmos/implementation/batch/PartitionKeyBatchResponseTests.java rename to sdk/cosmos/azure-cosmos/src/test/java/com/azure/cosmos/implementation/batch/TransactionalBatchResponseTests.java index a097875c8a96..eb44f01fd3ee 100644 --- a/sdk/cosmos/azure-cosmos/src/test/java/com/azure/cosmos/implementation/batch/PartitionKeyBatchResponseTests.java +++ b/sdk/cosmos/azure-cosmos/src/test/java/com/azure/cosmos/implementation/batch/TransactionalBatchResponseTests.java @@ -20,7 +20,7 @@ import static org.assertj.core.api.Assertions.assertThat; -public class PartitionKeyBatchResponseTests { +public class TransactionalBatchResponseTests { private static final int TIMEOUT = 40000; @@ -59,7 +59,7 @@ public void statusCodesAreSetThroughResponseAsync() { responseContent.getBytes(StandardCharsets.UTF_8)); TransactionalBatchResponse batchResponse = BatchResponseParser.fromDocumentServiceResponseAsync( - new RxDocumentServiceResponse(storeResponse), + new RxDocumentServiceResponse(null, storeResponse), batchRequest, true).block(); From 55ac0c568a941cadb6e9c93cfd98ba823310c7af Mon Sep 17 00:00:00 2001 From: Rakesh Kumar Date: Thu, 1 Oct 2020 03:41:51 +0530 Subject: [PATCH 04/22] Fix Signed-off-by: Rakesh Kumar --- .../com/azure/cosmos/implementation/RxDocumentClientImpl.java | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/RxDocumentClientImpl.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/RxDocumentClientImpl.java index d58e3a89512e..4db0830865f8 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/RxDocumentClientImpl.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/RxDocumentClientImpl.java @@ -1340,7 +1340,7 @@ private RxDocumentServiceRequest addBatchHeaders(RxDocumentServiceRequest reques return request; } - private void populateHeaders(RxDocumentServiceRequest request, RequestVerb httpMethod) { + private Mono populateHeaders(RxDocumentServiceRequest request, RequestVerb httpMethod) { request.getHeaders().put(HttpConstants.HttpHeaders.X_DATE, Utils.nowAsRFC1123()); if (this.masterKeyOrResourceToken != null || this.resourceTokensMap != null || this.cosmosAuthorizationTokenResolver != null || this.credential != null) { From 136c20cb3af818cb5bcef55dd07e4281cd21e26c Mon Sep 17 00:00:00 2001 From: Rakesh Kumar Date: Thu, 1 Oct 2020 13:05:44 +0530 Subject: [PATCH 05/22] fix Signed-off-by: Rakesh Kumar --- .../azure/cosmos/TransactionalBatchTest.java | 29 +++++++++---------- 1 file changed, 14 insertions(+), 15 deletions(-) diff --git a/sdk/cosmos/azure-cosmos/src/test/java/com/azure/cosmos/TransactionalBatchTest.java b/sdk/cosmos/azure-cosmos/src/test/java/com/azure/cosmos/TransactionalBatchTest.java index 1c5e07791659..6dd51473ae0d 100644 --- a/sdk/cosmos/azure-cosmos/src/test/java/com/azure/cosmos/TransactionalBatchTest.java +++ b/sdk/cosmos/azure-cosmos/src/test/java/com/azure/cosmos/TransactionalBatchTest.java @@ -29,7 +29,7 @@ public TransactionalBatchTest(CosmosClientBuilder clientBuilder) { super(clientBuilder); } - @BeforeClass(groups = {"emulator"}, timeOut = SETUP_TIMEOUT) + @BeforeClass(groups = {"simple"}, timeOut = SETUP_TIMEOUT) public void before_TransactionalBatchTest() { assertThat(this.batchClient).isNull(); this.batchClient = getClientBuilder().buildClient(); @@ -37,13 +37,12 @@ public void before_TransactionalBatchTest() { batchContainer = batchClient.getDatabase(asyncContainer.getDatabase().getId()).getContainer(asyncContainer.getId()); } - @AfterClass(groups = {"emulator"}, timeOut = SHUTDOWN_TIMEOUT, alwaysRun = true) + @AfterClass(groups = {"simple"}, timeOut = SHUTDOWN_TIMEOUT, alwaysRun = true) public void afterClass() { - assertThat(this.batchClient).isNotNull(); - this.batchClient.close(); + safeCloseSyncClient(this.batchClient); } - @Test(groups = {"emulator"}, timeOut = TIMEOUT) + @Test(groups = {"simple"}, timeOut = TIMEOUT) public void batchOrdered() { CosmosContainer container = this.batchContainer; this.createJsonTestDocsAsync(container); @@ -67,7 +66,7 @@ public void batchOrdered() { this.verifyByReadAsync(container, replaceDoc); } - @Test(groups = {"emulator"}, timeOut = TIMEOUT) + @Test(groups = {"simple"}, timeOut = TIMEOUT) public void batchItemETagAsync() { CosmosContainer container = batchContainer; this.createJsonTestDocsAsync(container); @@ -123,7 +122,7 @@ public void batchItemETagAsync() { } } - @Test(groups = {"emulator"}, timeOut = TIMEOUT) + @Test(groups = {"simple"}, timeOut = TIMEOUT) public void batchItemSessionTokenAsync() { CosmosContainer container = batchContainer; this.createJsonTestDocsAsync(container); @@ -158,7 +157,7 @@ public void batchItemSessionTokenAsync() { .isGreaterThan(beforeRequestSessionToken.getLSN()); } - @Test(groups = {"emulator"}, timeOut = TIMEOUT) + @Test(groups = {"simple"}, timeOut = TIMEOUT) public void batchWithTooManyOperationsAsync() { CosmosContainer container = batchContainer; int operationCount = MAX_OPERATIONS_IN_DIRECT_MODE_BATCH_REQUEST + 1; @@ -174,7 +173,7 @@ public void batchWithTooManyOperationsAsync() { assertThat(batchResponse.getResponseStatus()).isEqualTo(HttpResponseStatus.BAD_REQUEST.code()); } - @Test(groups = {"emulator"}, timeOut = TIMEOUT) + @Test(groups = {"simple"}, timeOut = TIMEOUT) public void batchReadsOnlyAsync() throws Exception { CosmosContainer container = batchContainer; this.createJsonTestDocsAsync(container); @@ -196,7 +195,7 @@ public void batchReadsOnlyAsync() throws Exception { assertThat(batchResponse.getOperationResultAtIndex(2, TestDoc.class).getItem()).isEqualTo(this.TestDocPk1ExistingC); } - @Test(groups = {"emulator"}, timeOut = TIMEOUT) + @Test(groups = {"simple"}, timeOut = TIMEOUT) public void batchCrud() throws Exception { CosmosContainer container = batchContainer; this.createJsonTestDocsAsync(container); @@ -238,7 +237,7 @@ public void batchCrud() throws Exception { this.verifyNotFoundAsync(container, this.TestDocPk1ExistingD); } - @Test(groups = {"emulator"}, timeOut = TIMEOUT) + @Test(groups = {"simple"}, timeOut = TIMEOUT) public void batchWithInvalidCreateAsync() { // partition key mismatch between doc and and value passed in to the operation this.runWithErrorAsync( @@ -247,7 +246,7 @@ public void batchWithInvalidCreateAsync() { HttpResponseStatus.BAD_REQUEST); } - @Test(groups = {"emulator"}, timeOut = TIMEOUT) + @Test(groups = {"simple"}, timeOut = TIMEOUT) public void batchWithReadOfNonExistentEntityAsync() { this.runWithErrorAsync( batchContainer, @@ -255,7 +254,7 @@ public void batchWithReadOfNonExistentEntityAsync() { HttpResponseStatus.NOT_FOUND); } - @Test(groups = {"emulator"}, timeOut = TIMEOUT) + @Test(groups = {"simple"}, timeOut = TIMEOUT) public void batchWithReplaceOfStaleEntityAsync() { this.createJsonTestDocsAsync(batchContainer); @@ -274,7 +273,7 @@ public void batchWithReplaceOfStaleEntityAsync() { this.verifyByReadAsync(batchContainer, this.TestDocPk1ExistingA); } - @Test(groups = {"emulator"}, timeOut = TIMEOUT) + @Test(groups = {"simple"}, timeOut = TIMEOUT) public void batchWithDeleteOfNonExistentEntityAsync() { this.runWithErrorAsync( batchContainer, @@ -282,7 +281,7 @@ public void batchWithDeleteOfNonExistentEntityAsync() { HttpResponseStatus.NOT_FOUND); } - @Test(groups = {"emulator"}, timeOut = TIMEOUT) + @Test(groups = {"simple"}, timeOut = TIMEOUT) public void batchWithCreateConflictAsync() { this.createJsonTestDocsAsync(batchContainer); From 24633c980f08de87372b7ed8fa401d29f2d1a6ae Mon Sep 17 00:00:00 2001 From: Rakesh Kumar Date: Thu, 1 Oct 2020 15:37:02 +0530 Subject: [PATCH 06/22] Added issue link Signed-off-by: Rakesh Kumar --- .../azure/cosmos/implementation/batch/BatchResponseParser.java | 2 ++ .../azure/cosmos/implementation/batch/ItemBatchOperation.java | 1 + .../azure/cosmos/implementation/batch/ServerBatchRequest.java | 2 +- 3 files changed, 4 insertions(+), 1 deletion(-) diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/batch/BatchResponseParser.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/batch/BatchResponseParser.java index bc8080cafe60..aa872d46e98e 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/batch/BatchResponseParser.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/batch/BatchResponseParser.java @@ -169,6 +169,7 @@ private static TransactionalBatchResponse populateFromResponseContentAsync( } else { // TODO(rakkuma): Implement hybrid row response parsing logic here. Parse the response hybrid row buffer // into array list of TransactionalBatchOperationResult. Remaining part is taken care from the caller function. + // Issue: https://github.com/Azure/azure-sdk-for-java/issues/15856 logger.error("Hybrid row is not implemented right now"); return null; } @@ -205,6 +206,7 @@ private static TransactionalBatchResponse populateFromResponseContentAsync( * Read batch operation result result. * * TODO(rakkuma): Similarly hybrid row result needs to be parsed. + * Issue: https://github.com/Azure/azure-sdk-for-java/issues/15856 * * @param objectNode having response for a single operation. * diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/batch/ItemBatchOperation.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/batch/ItemBatchOperation.java index b9fe695708a4..eda92068dcd4 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/batch/ItemBatchOperation.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/batch/ItemBatchOperation.java @@ -46,6 +46,7 @@ private ItemBatchOperation( } // TODO(rakkuma): Similarly for hybrid row, operation needs to be written in Hybrid row. + // Issue: https://github.com/Azure/azure-sdk-for-java/issues/15856 static JsonSerializable writeOperation(final ItemBatchOperation operation) { final JsonSerializable jsonSerializable = new JsonSerializable(); diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/batch/ServerBatchRequest.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/batch/ServerBatchRequest.java index 55468705ffb6..14e78602c173 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/batch/ServerBatchRequest.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/batch/ServerBatchRequest.java @@ -40,7 +40,7 @@ public abstract class ServerBatchRequest { /** * Adds as many operations as possible from the given list of operations. * TODO(rakkuma): Similarly for hybrid row, request needs to be parsed to create a request body in any form. - * + * Issue: https://github.com/Azure/azure-sdk-for-java/issues/15856 *

* Operations are added in order while ensuring the request body never exceeds {@link #maxBodyLength}. * From 6cd936358e28cff3de85b8dccb0d21e7d7eae620 Mon Sep 17 00:00:00 2001 From: Rakesh Kumar Date: Thu, 1 Oct 2020 15:39:59 +0530 Subject: [PATCH 07/22] beta version change Signed-off-by: Rakesh Kumar --- .../src/main/java/com/azure/cosmos/CosmosAsyncContainer.java | 4 ++-- .../src/main/java/com/azure/cosmos/CosmosContainer.java | 4 ++-- .../src/main/java/com/azure/cosmos/util/Beta.java | 4 +++- 3 files changed, 7 insertions(+), 5 deletions(-) diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncContainer.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncContainer.java index 6e1a24152ecc..064fc5b9c26c 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncContainer.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncContainer.java @@ -524,7 +524,7 @@ private T transform(Object object, Class classType) { * Use {@link TransactionalBatchResponse#isSuccessStatusCode} on the response returned to ensure that the * transactional batch succeeded. */ - @Beta(Beta.SinceVersion.V4_6_0) + @Beta(Beta.SinceVersion.V4_7_0) public Mono executeTransactionalBatch(TransactionalBatch transactionalBatch) { return executeTransactionalBatch(transactionalBatch, new TransactionalBatchRequestOptions()); } @@ -558,7 +558,7 @@ public Mono executeTransactionalBatch(TransactionalB * Use {@link TransactionalBatchResponse#isSuccessStatusCode} on the response returned to ensure that the * transactional batch succeeded. */ - @Beta(Beta.SinceVersion.V4_6_0) + @Beta(Beta.SinceVersion.V4_7_0) public Mono executeTransactionalBatch( TransactionalBatch transactionalBatch, TransactionalBatchRequestOptions requestOptions) { diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosContainer.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosContainer.java index 979d90f349d0..7dbc66d9f3c9 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosContainer.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosContainer.java @@ -494,7 +494,7 @@ public CosmosItemResponse deleteItem(T item, CosmosItemRequestOption * Use {@link TransactionalBatchResponse#isSuccessStatusCode} on the response returned to ensure that the * transactional batch succeeded. */ - @Beta(Beta.SinceVersion.V4_6_0) + @Beta(Beta.SinceVersion.V4_7_0) public TransactionalBatchResponse executeTransactionalBatch(TransactionalBatch transactionalBatch) { return this.blockBatchResponse(asyncContainer.executeTransactionalBatch(transactionalBatch)); } @@ -528,7 +528,7 @@ public TransactionalBatchResponse executeTransactionalBatch(TransactionalBatch t * Use {@link TransactionalBatchResponse#isSuccessStatusCode} on the response returned to ensure that the * transactional batch succeeded. */ - @Beta(Beta.SinceVersion.V4_6_0) + @Beta(Beta.SinceVersion.V4_7_0) public TransactionalBatchResponse executeTransactionalBatch( TransactionalBatch transactionalBatch, TransactionalBatchRequestOptions requestOptions) { diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/util/Beta.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/util/Beta.java index b585233ba6b0..800578a4e673 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/util/Beta.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/util/Beta.java @@ -43,6 +43,8 @@ public enum SinceVersion { /** v4.5.1 */ V4_5_1, /** v4.6.0 */ - V4_6_0 + V4_6_0, + /** v4.7.0 */ + V4_7_0 } } From d9500e05e9e241a658b4d1ec572792dabbd48a7f Mon Sep 17 00:00:00 2001 From: Rakesh Kumar Date: Thu, 1 Oct 2020 21:51:00 +0530 Subject: [PATCH 08/22] Add Beta change to public classes Signed-off-by: Rakesh Kumar --- .../azure/cosmos/CosmosAsyncContainer.java | 23 +++++++----- .../com/azure/cosmos/TransactionalBatch.java | 2 ++ .../TransactionalBatchItemRequestOptions.java | 5 +++ .../TransactionalBatchOperationResult.java | 2 ++ .../TransactionalBatchRequestOptions.java | 5 +++ .../cosmos/TransactionalBatchResponse.java | 2 ++ .../azure/cosmos/TransactionalBatchTest.java | 35 +++++++++++++++++-- 7 files changed, 63 insertions(+), 11 deletions(-) diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncContainer.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncContainer.java index 064fc5b9c26c..71a595c7481e 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncContainer.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncContainer.java @@ -567,15 +567,20 @@ public Mono executeTransactionalBatch( requestOptions = new TransactionalBatchRequestOptions(); } - final BatchExecutor executor = new BatchExecutor(this, transactionalBatch, requestOptions); - final Mono responseMono = executor.executeAsync(); - - return withContext(context -> database.getClient().getTracerProvider(). - traceEnabledBatchResponsePublisher(responseMono, - context, - this.batchSpanName, - database.getId(), - database.getClient().getServiceEndpoint())); + final TransactionalBatchRequestOptions transactionalBatchRequestOptions = requestOptions; + + return withContext(context -> { + final BatchExecutor executor = new BatchExecutor(this, transactionalBatch, transactionalBatchRequestOptions); + final Mono responseMono = executor.executeAsync(); + + return database.getClient().getTracerProvider(). + traceEnabledBatchResponsePublisher( + responseMono, + context, + this.batchSpanName, + database.getId(), + database.getClient().getServiceEndpoint()); + }); } /** diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/TransactionalBatch.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/TransactionalBatch.java index 6bec92566fe2..bee2fe8a1905 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/TransactionalBatch.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/TransactionalBatch.java @@ -7,6 +7,7 @@ import com.azure.cosmos.implementation.apachecommons.collections.list.UnmodifiableList; import com.azure.cosmos.implementation.batch.ItemBatchOperation; import com.azure.cosmos.models.PartitionKey; +import com.azure.cosmos.util.Beta; import java.util.ArrayList; import java.util.List; @@ -85,6 +86,7 @@ * See: * Limits on TransactionalBatch requests. */ +@Beta(Beta.SinceVersion.V4_7_0) public final class TransactionalBatch { private final ArrayList> operations; diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/TransactionalBatchItemRequestOptions.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/TransactionalBatchItemRequestOptions.java index bde6b127bfda..f6be46188233 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/TransactionalBatchItemRequestOptions.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/TransactionalBatchItemRequestOptions.java @@ -4,7 +4,12 @@ package com.azure.cosmos; import com.azure.cosmos.implementation.RequestOptions; +import com.azure.cosmos.util.Beta; +/** + * Encapsulates options that can be specified for an operation within a {@link TransactionalBatch}. + */ +@Beta(Beta.SinceVersion.V4_7_0) public final class TransactionalBatchItemRequestOptions { private String ifMatchETag; private String ifNoneMatchETag; diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/TransactionalBatchOperationResult.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/TransactionalBatchOperationResult.java index 72cec8b891ff..b0a09537bca9 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/TransactionalBatchOperationResult.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/TransactionalBatchOperationResult.java @@ -3,6 +3,7 @@ package com.azure.cosmos; +import com.azure.cosmos.util.Beta; import com.fasterxml.jackson.databind.node.ObjectNode; import java.time.Duration; @@ -13,6 +14,7 @@ * * @param the type parameter */ +@Beta(Beta.SinceVersion.V4_7_0) public final class TransactionalBatchOperationResult { private String eTag; diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/TransactionalBatchRequestOptions.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/TransactionalBatchRequestOptions.java index c15bf7abe827..0df56e7989c6 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/TransactionalBatchRequestOptions.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/TransactionalBatchRequestOptions.java @@ -4,7 +4,12 @@ package com.azure.cosmos; import com.azure.cosmos.implementation.RequestOptions; +import com.azure.cosmos.util.Beta; +/** + * Encapsulates options that can be specified for a {@link TransactionalBatch}. + */ +@Beta(Beta.SinceVersion.V4_7_0) public final class TransactionalBatchRequestOptions { private ConsistencyLevel consistencyLevel; private String sessionToken; diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/TransactionalBatchResponse.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/TransactionalBatchResponse.java index 7d0537e5dada..ff9afce9e272 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/TransactionalBatchResponse.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/TransactionalBatchResponse.java @@ -8,6 +8,7 @@ import com.azure.cosmos.implementation.apachecommons.collections.list.UnmodifiableList; import com.azure.cosmos.implementation.apachecommons.lang.StringUtils; import com.azure.cosmos.implementation.batch.ItemBatchOperation; +import com.azure.cosmos.util.Beta; import org.slf4j.Logger; import org.slf4j.LoggerFactory; @@ -23,6 +24,7 @@ /** * Response of a {@link TransactionalBatch} request. */ +@Beta(Beta.SinceVersion.V4_7_0) public class TransactionalBatchResponse { private final static Logger logger = LoggerFactory.getLogger(TransactionalBatchResponse.class); diff --git a/sdk/cosmos/azure-cosmos/src/test/java/com/azure/cosmos/TransactionalBatchTest.java b/sdk/cosmos/azure-cosmos/src/test/java/com/azure/cosmos/TransactionalBatchTest.java index 6dd51473ae0d..31f31ede1993 100644 --- a/sdk/cosmos/azure-cosmos/src/test/java/com/azure/cosmos/TransactionalBatchTest.java +++ b/sdk/cosmos/azure-cosmos/src/test/java/com/azure/cosmos/TransactionalBatchTest.java @@ -13,6 +13,7 @@ import org.testng.annotations.BeforeClass; import org.testng.annotations.Factory; import org.testng.annotations.Test; +import reactor.core.publisher.Mono; import java.util.UUID; @@ -23,6 +24,7 @@ public class TransactionalBatchTest extends BatchTestBase { private CosmosClient batchClient; private CosmosContainer batchContainer; + private CosmosAsyncContainer batchAsyncContainer; @Factory(dataProvider = "simpleClientBuildersWithDirect") public TransactionalBatchTest(CosmosClientBuilder clientBuilder) { @@ -33,8 +35,8 @@ public TransactionalBatchTest(CosmosClientBuilder clientBuilder) { public void before_TransactionalBatchTest() { assertThat(this.batchClient).isNull(); this.batchClient = getClientBuilder().buildClient(); - CosmosAsyncContainer asyncContainer = getSharedMultiPartitionCosmosContainer(this.batchClient.asyncClient()); - batchContainer = batchClient.getDatabase(asyncContainer.getDatabase().getId()).getContainer(asyncContainer.getId()); + batchAsyncContainer = getSharedMultiPartitionCosmosContainer(this.batchClient.asyncClient()); + batchContainer = batchClient.getDatabase(batchAsyncContainer.getDatabase().getId()).getContainer(batchAsyncContainer.getId()); } @AfterClass(groups = {"simple"}, timeOut = SHUTDOWN_TIMEOUT, alwaysRun = true) @@ -42,6 +44,35 @@ public void afterClass() { safeCloseSyncClient(this.batchClient); } + @Test(groups = {"simple"}, timeOut = TIMEOUT) + public void batchExecutionRepeat() { + CosmosContainer container = this.batchContainer; + this.createJsonTestDocsAsync(container); + + TestDoc firstDoc = this.populateTestDoc(this.partitionKey1); + TestDoc replaceDoc = this.getTestDocCopy(firstDoc); + replaceDoc.setCost(replaceDoc.getCost() + 1); + + Mono batchResponseMono = batchAsyncContainer.executeTransactionalBatch( + TransactionalBatch.createTransactionalBatch(this.getPartitionKey(this.partitionKey1)) + .createItem(firstDoc) + .replaceItem(replaceDoc.getId(), replaceDoc)); + + TransactionalBatchResponse batchResponse1 = batchResponseMono.block(); + this.verifyBatchProcessed(batchResponse1, 2); + + assertThat(batchResponse1.get(0).getResponseStatus()).isEqualTo(HttpResponseStatus.CREATED.code()); + assertThat(batchResponse1.get(1).getResponseStatus()).isEqualTo(HttpResponseStatus.OK.code()); + + // Block again. + TransactionalBatchResponse batchResponse2 = batchResponseMono.block(); + assertThat(batchResponse2.get(0).getResponseStatus()).isEqualTo(HttpResponseStatus.CONFLICT.code()); + assertThat(batchResponse2.get(1).getResponseStatus()).isEqualTo(HttpResponseStatus.FAILED_DEPENDENCY.code()); + + // Ensure that the replace overwrote the doc from the first operation + this.verifyByReadAsync(container, replaceDoc); + } + @Test(groups = {"simple"}, timeOut = TIMEOUT) public void batchOrdered() { CosmosContainer container = this.batchContainer; From 5afc339455bda424646f29143d608f83436c9a90 Mon Sep 17 00:00:00 2001 From: Rakesh Kumar Date: Thu, 1 Oct 2020 23:52:36 +0530 Subject: [PATCH 09/22] Minor fix Signed-off-by: Rakesh Kumar --- .../azure/cosmos/CosmosAsyncContainer.java | 4 +-- .../com/azure/cosmos/CosmosContainer.java | 4 +-- .../com/azure/cosmos/TransactionalBatch.java | 33 ++++++++++--------- .../TransactionalBatchItemRequestOptions.java | 8 ++--- .../TransactionalBatchOperationResult.java | 23 ++++--------- .../TransactionalBatchRequestOptions.java | 12 ++----- .../batch/ServerBatchRequest.java | 2 +- .../cosmos/BatchOperationResultTests.java | 10 +++++- .../azure/cosmos/TransactionalBatchTest.java | 4 +-- 9 files changed, 47 insertions(+), 53 deletions(-) diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncContainer.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncContainer.java index 71a595c7481e..eeae29310b67 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncContainer.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncContainer.java @@ -497,7 +497,7 @@ private T transform(Object object, Class classType) { } /** - * Executes the transactional batch at the Azure Cosmos service as an asynchronous operation. + * Executes the transactional batch. * * @param transactionalBatch Batch having list of operation and partition key which will be executed by this container. * @@ -530,7 +530,7 @@ public Mono executeTransactionalBatch(TransactionalB } /** - * Executes the transactional batch at the Azure Cosmos service as an asynchronous operation. + * Executes the transactional batch. * * @param transactionalBatch Batch having list of operation and partition key which will be executed by this container. * @param requestOptions Options that apply specifically to batch request. diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosContainer.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosContainer.java index 7dbc66d9f3c9..9199ada4fda1 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosContainer.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosContainer.java @@ -467,7 +467,7 @@ public CosmosItemResponse deleteItem(T item, CosmosItemRequestOption } /** - * Executes the transactional batch at the Azure Cosmos service as an asynchronous operation. + * Executes the transactional batch. * * @param transactionalBatch Batch having list of operation and partition key which will be executed by this container. * @@ -500,7 +500,7 @@ public TransactionalBatchResponse executeTransactionalBatch(TransactionalBatch t } /** - * Executes the transactional batch at the Azure Cosmos service as an asynchronous operation. + * Executes the transactional batch. * * @param transactionalBatch Batch having list of operation and partition key which will be executed by this container. * @param requestOptions Options that apply specifically to batch request. diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/TransactionalBatch.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/TransactionalBatch.java index bee2fe8a1905..b1a31188f88b 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/TransactionalBatch.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/TransactionalBatch.java @@ -46,18 +46,18 @@ * .upsertItem(test3) * .deleteItem("reading"); * - * try (TransactionalBatchResponse response = container.executeTransactionalBatch(batch) { + * TransactionalBatchResponse response = container.executeTransactionalBatch(batch); * - * if (!response.IsSuccessStatusCode) { - * // Handle and log exception - * return; - * } + * if (!response.IsSuccessStatusCode) { + * // Handle and log exception + * return; + * } * - * // Look up interested results - e.g., via typed access on operation results + * // Look up interested results - e.g., via typed access on operation results + * + * TransactionalBatchOperationResult result = response.getOperationResultAtIndex(0, ToDoActivity.class); + * ToDoActivity readActivity = result.getItem(); * - * TransactionalBatchOperationResult result = response.getOperationResultAtIndex(0, ToDoActivity.class); - * ToDoActivity readActivity = result.getItem(); - * } * } * * Example @@ -71,16 +71,16 @@ * .readItem("jogging") * .readItem("running") * - * try (TransactionalBatchResponse response = container.executeTransactionalBatch(batch) { + * TransactionalBatchResponse response = container.executeTransactionalBatch(batch); * - * // Look up interested results - eg. via direct access to operation result stream + * // Look up interested results - eg. via direct access to operation result stream * - * List resultItems = new ArrayList(); + * List resultItems = new ArrayList(); * - * for (TransactionalBatchOperationResult result : response) { - * resultItems.add(result.getResourceObject().toString()) - * } + * for (TransactionalBatchOperationResult result : response) { + * resultItems.add(result.getResourceObject().toString()) * } + * * } *

* See: @@ -132,6 +132,7 @@ public TransactionalBatch createItem(T item) { * * @param item A JSON serializable object that must contain an id property. * @param requestOptions The options for the item request. + * * @return The transactional batch instance with the operation added. */ public TransactionalBatch createItem(T item, TransactionalBatchItemRequestOptions requestOptions) { @@ -245,6 +246,7 @@ public TransactionalBatch replaceItem(String id, T item) { * @param id The unique id of the item. * @param item A JSON serializable object that must contain an id property. * @param requestOptions The options for the item request. + * * @return The transactional batch instance with the operation added. */ public TransactionalBatch replaceItem( @@ -285,6 +287,7 @@ public TransactionalBatch upsertItem(T item) { * * @param item A JSON serializable object that must contain an id property. * @param requestOptions The options for the item request. + * * @return The transactional batch instance with the operation added. */ public TransactionalBatch upsertItem(T item, TransactionalBatchItemRequestOptions requestOptions) { diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/TransactionalBatchItemRequestOptions.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/TransactionalBatchItemRequestOptions.java index f6be46188233..dde32da47d0e 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/TransactionalBatchItemRequestOptions.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/TransactionalBatchItemRequestOptions.java @@ -15,7 +15,7 @@ public final class TransactionalBatchItemRequestOptions { private String ifNoneMatchETag; /** - * Gets the If-Match (ETag) associated with the request in the Azure Cosmos DB service. + * Gets the If-Match (ETag) associated with the operation in TransactionalBatch. * * @return ifMatchETag the ifMatchETag associated with the request. */ @@ -24,7 +24,7 @@ public String getIfMatchETag() { } /** - * Sets the If-Match (ETag) associated with the request in the Azure Cosmos DB service. + * Sets the If-Match (ETag) associated with the operation in TransactionalBatch. * * @param ifMatchETag the ifMatchETag associated with the request. * @return the current request options @@ -35,7 +35,7 @@ public TransactionalBatchItemRequestOptions setIfMatchETag(final String ifMatchE } /** - * Gets the If-None-Match (ETag) associated with the request in the Azure Cosmos DB service. + * Gets the If-None-Match (ETag) associated with the request in operation in TransactionalBatch. * * @return the ifNoneMatchETag associated with the request. */ @@ -44,7 +44,7 @@ public String getIfNoneMatchETag() { } /** - * Sets the If-None-Match (ETag) associated with the request in the Azure Cosmos DB service. + * Sets the If-None-Match (ETag) associated with the request in operation in TransactionalBatch. * * @param ifNoneMatchEtag the ifNoneMatchETag associated with the request. * @return the current request options diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/TransactionalBatchOperationResult.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/TransactionalBatchOperationResult.java index b0a09537bca9..1a83cdbeae9b 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/TransactionalBatchOperationResult.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/TransactionalBatchOperationResult.java @@ -26,18 +26,9 @@ public final class TransactionalBatchOperationResult { private Integer subStatusCode; /** - * Instantiates a new Transactional batch operation result. + * Instantiates a new Transactional batch operation result using a TransactionalBatchOperationResult. * - * @param responseStatus the response status - */ - TransactionalBatchOperationResult(final int responseStatus) { - this.responseStatus = responseStatus; - } - - /** - * Instantiates a new Transactional batch operation result. - * - * @param other the other + * @param other the other TransactionalBatchOperationResult. */ TransactionalBatchOperationResult(final TransactionalBatchOperationResult other) { @@ -52,10 +43,10 @@ public final class TransactionalBatchOperationResult { } /** - * Instantiates a new Transactional batch operation result. + * Instantiates a new Transactional batch operation result using a TransactionalBatchOperationResult and item. * - * @param result the result - * @param item the item + * @param result the TransactionalBatchOperationResult. + * @param item the item. */ TransactionalBatchOperationResult(TransactionalBatchOperationResult result, TResource item) { this(result); @@ -83,7 +74,7 @@ public final class TransactionalBatchOperationResult { /** * Gets the entity tag associated with the current item. - *

+ * * ETags are used for concurrency checking when updating resources. * * @return Entity tag associated with the current item. @@ -104,7 +95,7 @@ public Double getRequestCharge() { /** * Gets the item associated with the current result. * - * @return Resource associated with the current result. + * @return item associated with the current result. */ public TResource getItem() { return this.item; diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/TransactionalBatchRequestOptions.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/TransactionalBatchRequestOptions.java index 0df56e7989c6..eb1b9b507a72 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/TransactionalBatchRequestOptions.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/TransactionalBatchRequestOptions.java @@ -14,19 +14,11 @@ public final class TransactionalBatchRequestOptions { private ConsistencyLevel consistencyLevel; private String sessionToken; - /** - * Constructor - */ - public TransactionalBatchRequestOptions() { - super(); - } - /** * Gets the consistency level required for the request. * * @return the consistency level. */ - public ConsistencyLevel getConsistencyLevel() { return consistencyLevel; } @@ -35,7 +27,7 @@ public ConsistencyLevel getConsistencyLevel() { * Sets the consistency level required for the request. * * @param consistencyLevel the consistency level. - * @return the CosmosItemRequestOptions. + * @return the TransactionalBatchRequestOptions. */ TransactionalBatchRequestOptions setConsistencyLevel(ConsistencyLevel consistencyLevel) { this.consistencyLevel = consistencyLevel; @@ -55,7 +47,7 @@ public String getSessionToken() { * Sets the token for use with session consistency. * * @param sessionToken the session token. - * @return the CosmosItemRequestOptions. + * @return the TransactionalBatchRequestOptions. */ public TransactionalBatchRequestOptions setSessionToken(String sessionToken) { this.sessionToken = sessionToken; diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/batch/ServerBatchRequest.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/batch/ServerBatchRequest.java index 14e78602c173..ff3f90ed582f 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/batch/ServerBatchRequest.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/batch/ServerBatchRequest.java @@ -41,7 +41,7 @@ public abstract class ServerBatchRequest { * Adds as many operations as possible from the given list of operations. * TODO(rakkuma): Similarly for hybrid row, request needs to be parsed to create a request body in any form. * Issue: https://github.com/Azure/azure-sdk-for-java/issues/15856 - *

+ * * Operations are added in order while ensuring the request body never exceeds {@link #maxBodyLength}. * * @param operations operations to be added; read-only. diff --git a/sdk/cosmos/azure-cosmos/src/test/java/com/azure/cosmos/BatchOperationResultTests.java b/sdk/cosmos/azure-cosmos/src/test/java/com/azure/cosmos/BatchOperationResultTests.java index 074467d333ca..29b376ec607f 100644 --- a/sdk/cosmos/azure-cosmos/src/test/java/com/azure/cosmos/BatchOperationResultTests.java +++ b/sdk/cosmos/azure-cosmos/src/test/java/com/azure/cosmos/BatchOperationResultTests.java @@ -59,7 +59,15 @@ public void propertiesAreSetThroughGenericCtor() { @Test(groups = {"unit"}, timeOut = TIMEOUT) public void isSuccessStatusCodeTrueFor200To299() { for (int x = 100; x < 999; ++x) { - TransactionalBatchOperationResult result = new TransactionalBatchOperationResult(x); + TransactionalBatchOperationResult result = BridgeInternal.createTransactionBatchResult( + null, + null, + null, + x, + null, + null + ); + boolean success = x >= 200 && x <= 299; assertThat(result.isSuccessStatusCode()).isEqualTo(success); } diff --git a/sdk/cosmos/azure-cosmos/src/test/java/com/azure/cosmos/TransactionalBatchTest.java b/sdk/cosmos/azure-cosmos/src/test/java/com/azure/cosmos/TransactionalBatchTest.java index 31f31ede1993..b6151b5b74bc 100644 --- a/sdk/cosmos/azure-cosmos/src/test/java/com/azure/cosmos/TransactionalBatchTest.java +++ b/sdk/cosmos/azure-cosmos/src/test/java/com/azure/cosmos/TransactionalBatchTest.java @@ -205,7 +205,7 @@ public void batchWithTooManyOperationsAsync() { } @Test(groups = {"simple"}, timeOut = TIMEOUT) - public void batchReadsOnlyAsync() throws Exception { + public void batchReadsOnlyAsync() { CosmosContainer container = batchContainer; this.createJsonTestDocsAsync(container); @@ -227,7 +227,7 @@ public void batchReadsOnlyAsync() throws Exception { } @Test(groups = {"simple"}, timeOut = TIMEOUT) - public void batchCrud() throws Exception { + public void batchCrud() { CosmosContainer container = batchContainer; this.createJsonTestDocsAsync(container); From 47353b5128286eb4788a3c99eab7286d34e10fe0 Mon Sep 17 00:00:00 2001 From: Rakesh Kumar Date: Mon, 5 Oct 2020 20:52:49 +0530 Subject: [PATCH 10/22] Code review Signed-off-by: Rakesh Kumar --- .../java/com/azure/cosmos/BridgeInternal.java | 16 +- .../azure/cosmos/CosmosAsyncContainer.java | 16 +- .../com/azure/cosmos/CosmosContainer.java | 16 +- .../com/azure/cosmos/TransactionalBatch.java | 21 +- .../TransactionalBatchOperationResult.java | 43 ++-- .../TransactionalBatchRequestOptions.java | 2 +- .../cosmos/TransactionalBatchResponse.java | 113 ++++----- .../implementation/RxDocumentClientImpl.java | 5 +- .../cosmos/implementation/TracerProvider.java | 2 +- .../implementation/batch/BatchExecUtils.java | 80 ++++++- .../implementation/batch/BatchExecutor.java | 2 +- .../batch/BatchResponseParser.java | 103 +++------ .../batch/ServerBatchRequest.java | 2 +- .../SinglePartitionKeyServerBatchRequest.java | 4 +- .../cosmos/BatchOperationResultTests.java | 24 +- .../java/com/azure/cosmos/BatchTestBase.java | 108 +++++++-- .../azure/cosmos/TransactionalBatchTest.java | 216 +++++++++++------- .../batch/BatchResponsePayloadWriter.java | 8 +- .../TransactionalBatchResponseTests.java | 100 +++++++- 19 files changed, 578 insertions(+), 303 deletions(-) diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/BridgeInternal.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/BridgeInternal.java index ec694fd94c5f..23468e83bd19 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/BridgeInternal.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/BridgeInternal.java @@ -618,17 +618,17 @@ public static RequestOptions toRequestOptions(TransactionalBatchRequestOptions t @Warning(value = INTERNAL_USE_ONLY_WARNING) public static TransactionalBatchOperationResult createTransactionBatchResult( String eTag, - Double requestCharge, + double requestCharge, ObjectNode resourceObject, - int responseStatus, + int statusCode, Duration retryAfter, - Integer subStatusCode) { + int subStatusCode) { return new TransactionalBatchOperationResult<>( eTag, requestCharge, resourceObject, - responseStatus, + statusCode, retryAfter, subStatusCode); } @@ -636,19 +636,17 @@ public static TransactionalBatchOperationResult createTransactionBatchResult( @Warning(value = INTERNAL_USE_ONLY_WARNING) public static TransactionalBatchResponse createTransactionBatchResponse( int responseStatusCode, - Integer responseSubStatusCode, + int responseSubStatusCode, String errorMessage, Map responseHeaders, - CosmosDiagnostics cosmosDiagnostics, - List> operations) { + CosmosDiagnostics cosmosDiagnostics) { return new TransactionalBatchResponse( responseStatusCode, responseSubStatusCode, errorMessage, responseHeaders, - cosmosDiagnostics, - operations); + cosmosDiagnostics); } @Warning(value = INTERNAL_USE_ONLY_WARNING) diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncContainer.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncContainer.java index eeae29310b67..a53e23fa5bdd 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncContainer.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncContainer.java @@ -504,21 +504,21 @@ private T transform(Object object, Class classType) { * @return A Mono response which contains details of execution of the transactional batch. *

* If the transactional batch executes successfully, the value returned by {@link - * TransactionalBatchResponse#getResponseStatus} on the response returned will be set to 200}. + * TransactionalBatchResponse#getStatusCode} on the response returned will be set to 200}. *

* If an operation within the transactional batch fails during execution, no changes from the batch will be * committed and the status of the failing operation is made available by {@link - * TransactionalBatchResponse#getResponseStatus}. To obtain information about the operations that failed, the + * TransactionalBatchResponse#getStatusCode}. To obtain information about the operations that failed, the * response can be enumerated. This returns {@link TransactionalBatchOperationResult} instances corresponding to * each operation in the transactional batch in the order they were added to the transactional batch. For a result * corresponding to an operation within the transactional batch, use - * {@link TransactionalBatchOperationResult#getResponseStatus} + * {@link TransactionalBatchOperationResult#getStatusCode} * to access the status of the operation. If the operation was not executed or it was aborted due to the failure of * another operation within the transactional batch, the value of this field will be 424; * for the operation that caused the batch to abort, the value of this field * will indicate the cause of failure. *

- * The value returned by {@link TransactionalBatchResponse#getResponseStatus} on the response returned may also have + * The value returned by {@link TransactionalBatchResponse#getStatusCode} on the response returned may also have * values such as 500 in case of server errors and 429. *

* Use {@link TransactionalBatchResponse#isSuccessStatusCode} on the response returned to ensure that the @@ -538,21 +538,21 @@ public Mono executeTransactionalBatch(TransactionalB * @return A Mono response which contains details of execution of the transactional batch. *

* If the transactional batch executes successfully, the value returned by {@link - * TransactionalBatchResponse#getResponseStatus} on the response returned will be set to 200}. + * TransactionalBatchResponse#getStatusCode} on the response returned will be set to 200}. *

* If an operation within the transactional batch fails during execution, no changes from the batch will be * committed and the status of the failing operation is made available by {@link - * TransactionalBatchResponse#getResponseStatus}. To obtain information about the operations that failed, the + * TransactionalBatchResponse#getStatusCode}. To obtain information about the operations that failed, the * response can be enumerated. This returns {@link TransactionalBatchOperationResult} instances corresponding to * each operation in the transactional batch in the order they were added to the transactional batch. For a result * corresponding to an operation within the transactional batch, use - * {@link TransactionalBatchOperationResult#getResponseStatus} + * {@link TransactionalBatchOperationResult#getStatusCode} * to access the status of the operation. If the operation was not executed or it was aborted due to the failure of * another operation within the transactional batch, the value of this field will be 424; * for the operation that caused the batch to abort, the value of this field * will indicate the cause of failure. *

- * The value returned by {@link TransactionalBatchResponse#getResponseStatus} on the response returned may also have + * The value returned by {@link TransactionalBatchResponse#getStatusCode} on the response returned may also have * values such as 500 in case of server errors and 429. *

* Use {@link TransactionalBatchResponse#isSuccessStatusCode} on the response returned to ensure that the diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosContainer.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosContainer.java index 9199ada4fda1..55da0ee93db0 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosContainer.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosContainer.java @@ -474,21 +474,21 @@ public CosmosItemResponse deleteItem(T item, CosmosItemRequestOption * @return A TransactionalBatchResponse which contains details of execution of the transactional batch. *

* If the transactional batch executes successfully, the value returned by {@link - * TransactionalBatchResponse#getResponseStatus} on the response returned will be set to 200}. + * TransactionalBatchResponse#getStatusCode} on the response returned will be set to 200}. *

* If an operation within the transactional batch fails during execution, no changes from the batch will be * committed and the status of the failing operation is made available by {@link - * TransactionalBatchResponse#getResponseStatus}. To obtain information about the operations that failed, the + * TransactionalBatchResponse#getStatusCode}. To obtain information about the operations that failed, the * response can be enumerated. This returns {@link TransactionalBatchOperationResult} instances corresponding to * each operation in the transactional batch in the order they were added to the transactional batch. For a result * corresponding to an operation within the transactional batch, use - * {@link TransactionalBatchOperationResult#getResponseStatus} + * {@link TransactionalBatchOperationResult#getStatusCode} * to access the status of the operation. If the operation was not executed or it was aborted due to the failure of * another operation within the transactional batch, the value of this field will be 424; * for the operation that caused the batch to abort, the value of this field * will indicate the cause of failure. *

- * The value returned by {@link TransactionalBatchResponse#getResponseStatus} on the response returned may also have + * The value returned by {@link TransactionalBatchResponse#getStatusCode} on the response returned may also have * values such as 500 in case of server errors and 429. *

* Use {@link TransactionalBatchResponse#isSuccessStatusCode} on the response returned to ensure that the @@ -508,21 +508,21 @@ public TransactionalBatchResponse executeTransactionalBatch(TransactionalBatch t * @return A TransactionalBatchResponse which contains details of execution of the transactional batch. *

* If the transactional batch executes successfully, the value returned by {@link - * TransactionalBatchResponse#getResponseStatus} on the response returned will be set to 200}. + * TransactionalBatchResponse#getStatusCode} on the response returned will be set to 200}. *

* If an operation within the transactional batch fails during execution, no changes from the batch will be * committed and the status of the failing operation is made available by {@link - * TransactionalBatchResponse#getResponseStatus}. To obtain information about the operations that failed, the + * TransactionalBatchResponse#getStatusCode}. To obtain information about the operations that failed, the * response can be enumerated. This returns {@link TransactionalBatchOperationResult} instances corresponding to * each operation in the transactional batch in the order they were added to the transactional batch. For a result * corresponding to an operation within the transactional batch, use - * {@link TransactionalBatchOperationResult#getResponseStatus} + * {@link TransactionalBatchOperationResult#getStatusCode} * to access the status of the operation. If the operation was not executed or it was aborted due to the failure of * another operation within the transactional batch, the value of this field will be 424; * for the operation that caused the batch to abort, the value of this field * will indicate the cause of failure. *

- * The value returned by {@link TransactionalBatchResponse#getResponseStatus} on the response returned may also have + * The value returned by {@link TransactionalBatchResponse#getStatusCode} on the response returned may also have * values such as 500 in case of server errors and 429. *

* Use {@link TransactionalBatchResponse#isSuccessStatusCode} on the response returned to ensure that the diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/TransactionalBatch.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/TransactionalBatch.java index b1a31188f88b..cd2a05d24988 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/TransactionalBatch.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/TransactionalBatch.java @@ -18,10 +18,9 @@ * Represents a batch of operations against items with the same {@link PartitionKey} in a container that will be performed * in a transactional manner at the Azure Cosmos DB service. *

- * Use {@link TransactionalBatch#createTransactionalBatch(PartitionKey)} or new {@link #TransactionalBatch(PartitionKey)} - * to create an instance of TransactionalBatch + * Use {@link TransactionalBatch#createTransactionalBatch(PartitionKey)} to create an instance of TransactionalBatch. * Example - * This example atomically modifies a set of documents as a batch. + * This example atomically modifies a set of items as a batch. * 

{@code
  * public class ToDoActivity {
  *     public final String type;
@@ -61,7 +60,7 @@
  * }
* * Example - *

This example atomically reads a set of documents as a batch. + *

This example atomically reads a set of items as a batch. * 

{@code
  * String activityType = "personal";
  *
@@ -72,13 +71,11 @@
  *     .readItem("running")
  *
  * TransactionalBatchResponse response = container.executeTransactionalBatch(batch);
+ * List resultItems = new ArrayList();
  *
- * // Look up interested results - eg. via direct access to operation result stream
- *
- * List resultItems = new ArrayList();
- *
- * for (TransactionalBatchOperationResult result : response) {
- *     resultItems.add(result.getResourceObject().toString())
+ * for (int i = 0; i < response.size(); i++) {
+ *     TransactionalBatchOperationResult result = response.getOperationResultAtIndex(0, ToDoActivity.class);
+ *     resultItems.add(result.getItem());
  * }
  *
  * }
@@ -89,10 +86,10 @@ @Beta(Beta.SinceVersion.V4_7_0) public final class TransactionalBatch { - private final ArrayList> operations; + private final List> operations; private final PartitionKey partitionKey; - public TransactionalBatch(PartitionKey partitionKey) { + TransactionalBatch(PartitionKey partitionKey) { checkNotNull(partitionKey, "expected non-null partitionKey"); this.operations = new ArrayList<>(); diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/TransactionalBatchOperationResult.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/TransactionalBatchOperationResult.java index 1a83cdbeae9b..127f1ab0ac41 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/TransactionalBatchOperationResult.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/TransactionalBatchOperationResult.java @@ -18,12 +18,12 @@ public final class TransactionalBatchOperationResult { private String eTag; - private Double requestCharge; + private double requestCharge; private TResource item; private ObjectNode resourceObject; - private int responseStatus; + private int statusCode; private Duration retryAfter; - private Integer subStatusCode; + private int subStatusCode; /** * Instantiates a new Transactional batch operation result using a TransactionalBatchOperationResult. @@ -34,7 +34,7 @@ public final class TransactionalBatchOperationResult { checkNotNull(other, "expected non-null other"); - this.responseStatus = other.responseStatus; + this.statusCode = other.statusCode; this.subStatusCode = other.subStatusCode; this.eTag = other.eTag; this.requestCharge = other.requestCharge; @@ -57,17 +57,17 @@ public final class TransactionalBatchOperationResult { * Initializes a new instance of the {@link TransactionalBatchOperationResult} class. */ TransactionalBatchOperationResult(String eTag, - Double requestCharge, + double requestCharge, ObjectNode resourceObject, - int responseStatus, + int statusCode, Duration retryAfter, - Integer subStatusCode) { - checkNotNull(responseStatus, "expected non-null responseStatus"); + int subStatusCode) { + checkNotNull(statusCode, "expected non-null statusCode"); this.eTag = eTag; this.requestCharge = requestCharge; this.resourceObject = resourceObject; - this.responseStatus = responseStatus; + this.statusCode = statusCode; this.retryAfter = retryAfter; this.subStatusCode = subStatusCode; } @@ -84,11 +84,14 @@ public String getETag() { } /** - * Gets the request charge in request units for the current operation. + * Gets the request charge as request units (RU) consumed by the current operation. + *

+ * For more information about the RU and factors that can impact the effective charges please visit + * Request Units in Azure Cosmos DB * - * @return Request charge in request units for the current operation. + * @return the request charge. */ - public Double getRequestCharge() { + public double getRequestCharge() { return this.requestCharge; } @@ -111,11 +114,11 @@ public Duration getRetryAfter() { } /** - * Gets sub status code. + * Gets sub status code associated with the current result. * * @return the sub status code */ - public Integer getSubStatusCode() { + public int getSubStatusCode() { return this.subStatusCode; } @@ -125,19 +128,19 @@ public Integer getSubStatusCode() { * @return {@code true} if the current operation completed successfully; {@code false} otherwise. */ public boolean isSuccessStatusCode() { - return 200 <= this.responseStatus && this.responseStatus <= 299; + return 200 <= this.statusCode && this.statusCode <= 299; } /** - * Gets response status. + * Gets the HTTP status code associated with the current result. * - * @return the response status + * @return the status code. */ - public int getResponseStatus() { - return this.responseStatus; + public int getStatusCode() { + return this.statusCode; } - public ObjectNode getResourceObject() { + ObjectNode getResourceObject() { return resourceObject; } } diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/TransactionalBatchRequestOptions.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/TransactionalBatchRequestOptions.java index eb1b9b507a72..358e4bc594e5 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/TransactionalBatchRequestOptions.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/TransactionalBatchRequestOptions.java @@ -19,7 +19,7 @@ public final class TransactionalBatchRequestOptions { * * @return the consistency level. */ - public ConsistencyLevel getConsistencyLevel() { + ConsistencyLevel getConsistencyLevel() { return consistencyLevel; } diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/TransactionalBatchResponse.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/TransactionalBatchResponse.java index ff9afce9e272..22d5391b1147 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/TransactionalBatchResponse.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/TransactionalBatchResponse.java @@ -3,14 +3,9 @@ package com.azure.cosmos; -import com.azure.cosmos.implementation.HttpConstants; import com.azure.cosmos.implementation.JsonSerializable; -import com.azure.cosmos.implementation.apachecommons.collections.list.UnmodifiableList; -import com.azure.cosmos.implementation.apachecommons.lang.StringUtils; -import com.azure.cosmos.implementation.batch.ItemBatchOperation; +import com.azure.cosmos.implementation.batch.BatchExecUtils; import com.azure.cosmos.util.Beta; -import org.slf4j.Logger; -import org.slf4j.LoggerFactory; import java.time.Duration; import java.util.ArrayList; @@ -27,44 +22,37 @@ @Beta(Beta.SinceVersion.V4_7_0) public class TransactionalBatchResponse { - private final static Logger logger = LoggerFactory.getLogger(TransactionalBatchResponse.class); - - private Map responseHeaders; - private final int responseStatus; - private String errorMessage; - private List> results; - private Integer subStatusCode; - private List> operations; - private CosmosDiagnostics cosmosDiagnostics; + private final Map responseHeaders; + private final int statusCode; + private final String errorMessage; + private final List> results; + private final int subStatusCode; + private final CosmosDiagnostics cosmosDiagnostics; /** * Initializes a new instance of the {@link TransactionalBatchResponse} class. * - * @param responseStatus the response status. + * @param statusCode the response status code. * @param subStatusCode the response sub-status code. * @param errorMessage an error message or {@code null}. * @param responseHeaders the response http headers * @param cosmosDiagnostics the diagnostic - * @param operations a {@link List list} of {@link ItemBatchOperation batch operations}. */ TransactionalBatchResponse( - final int responseStatus, - final Integer subStatusCode, + final int statusCode, + final int subStatusCode, final String errorMessage, final Map responseHeaders, - final CosmosDiagnostics cosmosDiagnostics, - final List> operations) { + final CosmosDiagnostics cosmosDiagnostics) { - checkNotNull(responseStatus, "expected non-null responseStatus"); + checkNotNull(statusCode, "expected non-null statusCode"); checkNotNull(responseHeaders, "expected non-null responseHeaders"); - checkNotNull(operations, "expected non-null operations"); - this.responseStatus = responseStatus; + this.statusCode = statusCode; this.subStatusCode = subStatusCode; this.errorMessage = errorMessage; this.responseHeaders = responseHeaders; this.cosmosDiagnostics = cosmosDiagnostics; - this.operations = UnmodifiableList.unmodifiableList(operations); this.results = new ArrayList<>(); } @@ -95,7 +83,12 @@ public TransactionalBatchOperationResult getOperationResultAtIndex( return new TransactionalBatchOperationResult(result, item); } - public CosmosDiagnostics getCosmosDiagnostics() { + /** + * Gets the diagnostics information for the current request to Azure Cosmos DB service. + * + * @return diagnostics information for the current request to Azure Cosmos DB service. + */ + public CosmosDiagnostics getDiagnostics() { return cosmosDiagnostics; } @@ -114,7 +107,7 @@ public int size() { * @return a value indicating whether the batch was successfully processed. */ public boolean isSuccessStatusCode() { - return this.responseStatus >= 200 && this.responseStatus <= 299; + return this.statusCode >= 200 && this.statusCode <= 299; } /** @@ -123,7 +116,7 @@ public boolean isSuccessStatusCode() { * @return the activity ID that identifies the server request made to execute the batch. */ public String getActivityId() { - return this.responseHeaders.get(HttpConstants.HttpHeaders.ACTIVITY_ID); + return BatchExecUtils.getActivityId(this.responseHeaders); } /** @@ -136,37 +129,39 @@ public String getErrorMessage() { } /** - * Gets the request charge for the batch request. + * Gets the request charge as request units (RU) consumed by the batch operation. + *

+ * For more information about the RU and factors that can impact the effective charges please visit + * Request Units in Azure Cosmos DB * - * @return the request charge measured in request units. + * @return the request charge. */ public double getRequestCharge() { - final String value = this.responseHeaders.get(HttpConstants.HttpHeaders.REQUEST_CHARGE); - if (StringUtils.isEmpty(value)) { - return 0; - } + return BatchExecUtils.getRequestCharge(this.responseHeaders); + } - try { - return Double.valueOf(value); - } catch (NumberFormatException e) { - logger.warn("INVALID x-ms-request-charge value {}.", value); - return 0; - } + /** + * Gets the HTTP status code associated with the response. + * + * @return the status code. + */ + public int getStatusCode() { + return this.statusCode; } /** - * Gets the response status code of the batch request. + * Gets the token used for managing client's consistency requirements. * - * @return the response status code of the batch request. + * @return the session token. */ - public int getResponseStatus() { - return this.responseStatus; + public String getSessionToken() { + return BatchExecUtils.getSessionToken(this.responseHeaders); } /** - * Gets the response headers. + * Gets the headers associated with the response. * - * @return the response header map. + * @return the response headers. */ public Map getResponseHeaders() { return this.responseHeaders; @@ -178,14 +173,15 @@ public Map getResponseHeaders() { * @return the amount of time to wait before retrying this or any other request due to throttling. */ public Duration getRetryAfter() { - if (this.responseHeaders.containsKey(HttpConstants.HttpHeaders.RETRY_AFTER)) { - return Duration.parse(this.responseHeaders.get(HttpConstants.HttpHeaders.RETRY_AFTER)); - } - - return null; + return BatchExecUtils.getRetryAfter(this.responseHeaders); } - public Integer getSubStatusCode() { + /** + * Gets the HTTP sub status code associated with the response. + * + * @return the sub status code. + */ + public int getSubStatusCode() { return this.subStatusCode; } @@ -209,8 +205,17 @@ public TransactionalBatchOperationResult get(int index) { return this.results.get(index); } - public boolean isEmpty() { - return this.results.isEmpty(); + /** + * Gets the end-to-end request latency for the current request to Azure Cosmos DB service. + * + * @return end-to-end request latency for the current request to Azure Cosmos DB service. + */ + public Duration getDuration() { + if (cosmosDiagnostics == null) { + return Duration.ZERO; + } + + return this.cosmosDiagnostics.getDuration(); } boolean addAll(Collection> collection) { diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/RxDocumentClientImpl.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/RxDocumentClientImpl.java index 4db0830865f8..830f64221d38 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/RxDocumentClientImpl.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/RxDocumentClientImpl.java @@ -2484,12 +2484,11 @@ private Mono executeBatchRequestInternal(String coll }); return responseObservable - .flatMap(serviceResponse -> BatchResponseParser.fromDocumentServiceResponseAsync(serviceResponse, serverBatchRequest, true)) - .onErrorResume(throwable -> BatchResponseParser.fromErrorResponseAsync(throwable, serverBatchRequest)); + .flatMap(serviceResponse -> BatchResponseParser.fromDocumentServiceResponseAsync(serviceResponse, serverBatchRequest, true)); } catch (Exception ex) { logger.debug("Failure in executing a batch due to [{}]", ex.getMessage(), ex); - return BatchResponseParser.fromErrorResponseAsync(ex, serverBatchRequest); + return Mono.error(ex); } } diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/TracerProvider.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/TracerProvider.java index b15fcb29b76f..3887ffc7557b 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/TracerProvider.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/TracerProvider.java @@ -115,7 +115,7 @@ public Mono traceEnabledBatchResponsePublisher(Mono< String databaseId, String endpoint) { return traceEnabledPublisher(resultPublisher, context, spanName, databaseId, endpoint, - TransactionalBatchResponse::getResponseStatus); + TransactionalBatchResponse::getStatusCode); } public Mono> traceEnabledCosmosItemResponsePublisher(Mono> resultPublisher, diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/batch/BatchExecUtils.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/batch/BatchExecUtils.java index 18b2e8f1a5c7..7767325f4ffe 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/batch/BatchExecUtils.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/batch/BatchExecUtils.java @@ -3,7 +3,14 @@ package com.azure.cosmos.implementation.batch; +import com.azure.cosmos.implementation.HttpConstants; import com.azure.cosmos.implementation.OperationType; +import com.azure.cosmos.implementation.apachecommons.lang.StringUtils; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +import java.time.Duration; +import java.util.Map; import static com.azure.cosmos.implementation.batch.BatchRequestResponseConstant.OPERATION_CREATE; import static com.azure.cosmos.implementation.batch.BatchRequestResponseConstant.OPERATION_DELETE; @@ -14,7 +21,9 @@ /** * Util methods for batch requests/response. */ -class BatchExecUtils { +public final class BatchExecUtils { + + private final static Logger logger = LoggerFactory.getLogger(BatchExecUtils.class); static String getStringOperationType(OperationType operationType) { switch (operationType) { @@ -32,4 +41,73 @@ static String getStringOperationType(OperationType operationType) { return null; } + + public static Duration getRetryAfter(Map responseHeaders) { + long retryIntervalInMilliseconds = 0; + + if (responseHeaders != null) { + String header = responseHeaders.get(HttpConstants.HttpHeaders.RETRY_AFTER_IN_MILLISECONDS); + + if (StringUtils.isNotEmpty(header)) { + try { + retryIntervalInMilliseconds = Long.parseLong(header); + } catch (NumberFormatException e) { + // If the value cannot be parsed as long, return 0. + } + } + } + + return Duration.ofMillis(retryIntervalInMilliseconds); + } + + public static String getSessionToken(Map responseHeaders) { + if (responseHeaders != null) { + return responseHeaders.get(HttpConstants.HttpHeaders.SESSION_TOKEN); + } + + return null; + } + + public static String getActivityId(Map responseHeaders) { + if (responseHeaders != null) { + return responseHeaders.get(HttpConstants.HttpHeaders.ACTIVITY_ID); + } + + return null; + } + + public static double getRequestCharge(Map responseHeaders) { + if (responseHeaders == null) { + return 0; + } + + final String value = responseHeaders.get(HttpConstants.HttpHeaders.REQUEST_CHARGE); + if (StringUtils.isEmpty(value)) { + return 0; + } + + try { + return Double.valueOf(value); + } catch (NumberFormatException e) { + logger.warn("INVALID x-ms-request-charge value {}.", value); + return 0; + } + } + + public static int getSubStatusCode(Map responseHeaders) { + int code = HttpConstants.SubStatusCodes.UNKNOWN; + + if (responseHeaders != null) { + String subStatusString = responseHeaders.get(HttpConstants.HttpHeaders.SUB_STATUS); + if (StringUtils.isNotEmpty(subStatusString)) { + try { + code = Integer.parseInt(subStatusString); + } catch (NumberFormatException e) { + // If value cannot be parsed as Integer, return Unknown. + } + } + } + + return code; + } } diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/batch/BatchExecutor.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/batch/BatchExecutor.java index 6a38d483b0ef..b3160e4ac317 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/batch/BatchExecutor.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/batch/BatchExecutor.java @@ -41,7 +41,7 @@ public final Mono executeAsync() { List> operations = BridgeInternal.getOperationsFromTransactionalBatch(this.transactionalBatch); checkArgument(operations.size() > 0, "Number of operations should be more than 0."); - final SinglePartitionKeyServerBatchRequest request = SinglePartitionKeyServerBatchRequest.createAsync( + final SinglePartitionKeyServerBatchRequest request = SinglePartitionKeyServerBatchRequest.createBatchRequest( BridgeInternal.getPartitionKeyFromTransactionalBatch(this.transactionalBatch), operations); request.setAtomicBatch(true); diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/batch/BatchResponseParser.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/batch/BatchResponseParser.java index aa872d46e98e..bbf50ba96c2b 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/batch/BatchResponseParser.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/batch/BatchResponseParser.java @@ -4,14 +4,12 @@ package com.azure.cosmos.implementation.batch; import com.azure.cosmos.BridgeInternal; -import com.azure.cosmos.CosmosException; import com.azure.cosmos.TransactionalBatchOperationResult; import com.azure.cosmos.TransactionalBatchResponse; import com.azure.cosmos.implementation.HttpConstants; import com.azure.cosmos.implementation.JsonSerializable; import com.azure.cosmos.implementation.RxDocumentServiceResponse; import com.azure.cosmos.implementation.Utils; -import com.azure.cosmos.implementation.apachecommons.lang.StringUtils; import com.fasterxml.jackson.databind.ObjectMapper; import com.fasterxml.jackson.databind.node.ObjectNode; import io.netty.handler.codec.http.HttpResponseStatus; @@ -24,8 +22,6 @@ import java.util.ArrayList; import java.util.List; -import static com.azure.cosmos.implementation.HttpConstants.HttpHeaders.RETRY_AFTER_IN_MILLISECONDS; -import static com.azure.cosmos.implementation.HttpConstants.HttpHeaders.SUB_STATUS; import static com.azure.cosmos.implementation.guava25.base.Preconditions.checkState; public final class BatchResponseParser { @@ -33,35 +29,6 @@ public final class BatchResponseParser { private final static Logger logger = LoggerFactory.getLogger(BatchResponseParser.class); private final static char HYBRID_V1 = 129; - /** Creates a transactional batch response} from a exception - * - * @param throwable the {@link Throwable error}. - * @param request the {@link ServerBatchRequest batch request} that produced {@code message}. - * - * @return a Mono that provides the {@link TransactionalBatchResponse transactional batch response} created - * from {@link TransactionalBatchResponse message} when the asynchronous operation completes. - */ - public static Mono fromErrorResponseAsync( - final Throwable throwable, - final ServerBatchRequest request) { - - if (throwable instanceof CosmosException) { - final CosmosException cosmosException = (CosmosException) throwable; - final TransactionalBatchResponse response = BridgeInternal.createTransactionBatchResponse( - cosmosException.getStatusCode(), - cosmosException.getSubStatusCode(), - cosmosException.toString(), - cosmosException.getResponseHeaders(), - cosmosException.getDiagnostics(), - request.getOperations()); - - BatchResponseParser.createAndPopulateResults(response, request.getOperations(), cosmosException.getRetryAfterDuration()); - return Mono.just(response); - } else { - return Mono.error(throwable); - } - } - /** Creates a transactional batch response} from a response message * * @param documentServiceResponse the {@link RxDocumentServiceResponse response message}. @@ -69,7 +36,7 @@ public static Mono fromErrorResponseAsync( * @param shouldPromoteOperationStatus indicates whether the operation status should be promoted. * * @return a Mono that provides the {@link TransactionalBatchResponse transactional batch response} created - * from {@link RxDocumentServiceResponse message} when the asynchronous operation completes. + * from {@link RxDocumentServiceResponse message} when the batch operation completes. */ public static Mono fromDocumentServiceResponseAsync( final RxDocumentServiceResponse documentServiceResponse, @@ -86,17 +53,15 @@ public static Mono fromDocumentServiceResponseAsync( // Convert any payload read failures as InternalServerError response = BridgeInternal.createTransactionBatchResponse( HttpResponseStatus.INTERNAL_SERVER_ERROR.code(), - 0, + HttpConstants.SubStatusCodes.UNKNOWN, "ServerResponseDeserializationFailure", documentServiceResponse.getResponseHeaders(), - documentServiceResponse.getCosmosDiagnostics(), - request.getOperations()); + documentServiceResponse.getCosmosDiagnostics()); } } int responseStatusCode = documentServiceResponse.getStatusCode(); - int responseSubStatusCode = Integer.parseInt( - documentServiceResponse.getResponseHeaders().getOrDefault(SUB_STATUS, String.valueOf(0))); + int responseSubStatusCode = BatchExecUtils.getSubStatusCode(documentServiceResponse.getResponseHeaders()); if (response == null) { response = BridgeInternal.createTransactionBatchResponse( @@ -104,8 +69,7 @@ public static Mono fromDocumentServiceResponseAsync( responseSubStatusCode, null, documentServiceResponse.getResponseHeaders(), - documentServiceResponse.getCosmosDiagnostics(), - request.getOperations()); + documentServiceResponse.getCosmosDiagnostics()); } if (response.size() != request.getOperations().size()) { @@ -114,28 +78,19 @@ public static Mono fromDocumentServiceResponseAsync( // batch request is successful - so fail as InternalServerError if this is not the case. response = BridgeInternal.createTransactionBatchResponse( HttpResponseStatus.INTERNAL_SERVER_ERROR.code(), - 0, + HttpConstants.SubStatusCodes.UNKNOWN, "Invalid server response", documentServiceResponse.getResponseHeaders(), - documentServiceResponse.getCosmosDiagnostics(), - request.getOperations()); + documentServiceResponse.getCosmosDiagnostics()); } // When the overall response status code is TooManyRequests, propagate the RetryAfter into the individual operations. - int retryAfterMilliseconds = 0; - + Duration retryAfterDuration = Duration.ZERO; if (responseStatusCode == HttpResponseStatus.TOO_MANY_REQUESTS.code()) { - String retryResponseValue = documentServiceResponse.getResponseHeaders().getOrDefault(RETRY_AFTER_IN_MILLISECONDS, null); - if (StringUtils.isNotEmpty(retryResponseValue)) { - try { - retryAfterMilliseconds = Integer.parseInt(retryResponseValue); - } catch (NumberFormatException ex) { - // Do nothing. It's number format exception - } - } + retryAfterDuration = BatchExecUtils.getRetryAfter(documentServiceResponse.getResponseHeaders()); } - BatchResponseParser.createAndPopulateResults(response, request.getOperations(), Duration.ofMillis(retryAfterMilliseconds)); + BatchResponseParser.createAndPopulateResults(response, request.getOperations(), retryAfterDuration); } checkState(response.size() == request.getOperations().size(), @@ -167,23 +122,20 @@ private static TransactionalBatchResponse populateFromResponseContentAsync( } } else { - // TODO(rakkuma): Implement hybrid row response parsing logic here. Parse the response hybrid row buffer - // into array list of TransactionalBatchOperationResult. Remaining part is taken care from the caller function. + // TODO(rakkuma): Implement hybrid row response parsing logic here. // Issue: https://github.com/Azure/azure-sdk-for-java/issues/15856 logger.error("Hybrid row is not implemented right now"); return null; } int responseStatusCode = documentServiceResponse.getStatusCode(); - Integer responseSubStatusCode = Integer.parseInt( - documentServiceResponse.getResponseHeaders().getOrDefault(SUB_STATUS, String.valueOf(HttpConstants.SubStatusCodes.UNKNOWN))); + int responseSubStatusCode = BatchExecUtils.getSubStatusCode(documentServiceResponse.getResponseHeaders()); // Status code of the exact operation which failed. - if (responseStatusCode == HttpResponseStatus.MULTI_STATUS.code() - && shouldPromoteOperationStatus) { + if (responseStatusCode == HttpResponseStatus.MULTI_STATUS.code() && shouldPromoteOperationStatus) { for (TransactionalBatchOperationResult result : results) { - if (result.getResponseStatus()!= HttpResponseStatus.FAILED_DEPENDENCY.code()) { - responseStatusCode = result.getResponseStatus(); + if (result.getStatusCode()!= HttpResponseStatus.FAILED_DEPENDENCY.code()) { + responseStatusCode = result.getStatusCode(); responseSubStatusCode = result.getSubStatusCode(); break; } @@ -195,8 +147,7 @@ private static TransactionalBatchResponse populateFromResponseContentAsync( responseSubStatusCode, null, documentServiceResponse.getResponseHeaders(), - documentServiceResponse.getCosmosDiagnostics(), - request.getOperations()); + documentServiceResponse.getCosmosDiagnostics()); BridgeInternal.addTransactionBatchResultInResponse(response, results); return response; @@ -215,9 +166,17 @@ private static TransactionalBatchResponse populateFromResponseContentAsync( private static TransactionalBatchOperationResult createBatchOperationResultFromJson(ObjectNode objectNode) { final JsonSerializable jsonSerializable = new JsonSerializable(objectNode); - final int responseStatusCode = jsonSerializable.getInt(BatchRequestResponseConstant.FIELD_STATUS_CODE); - final Integer subStatusCode = jsonSerializable.getInt(BatchRequestResponseConstant.FIELD_SUBSTATUS_CODE); - final Double requestCharge = jsonSerializable.getDouble(BatchRequestResponseConstant.FIELD_REQUEST_CHARGE); + final int statusCode = jsonSerializable.getInt(BatchRequestResponseConstant.FIELD_STATUS_CODE); + Integer subStatusCode = jsonSerializable.getInt(BatchRequestResponseConstant.FIELD_SUBSTATUS_CODE); + if (subStatusCode == null) { + subStatusCode = HttpConstants.SubStatusCodes.UNKNOWN; + } + + Double requestCharge = jsonSerializable.getDouble(BatchRequestResponseConstant.FIELD_REQUEST_CHARGE); + if (requestCharge == null) { + requestCharge = (double) 0; + } + final String eTag = jsonSerializable.getString(BatchRequestResponseConstant.FIELD_ETAG); final ObjectNode resourceBody = jsonSerializable.getObject(BatchRequestResponseConstant.FIELD_RESOURCE_BODY); final Integer retryAfterMilliseconds = jsonSerializable.getInt(BatchRequestResponseConstant.FIELD_RETRY_AFTER_MILLISECONDS); @@ -226,8 +185,8 @@ private static TransactionalBatchOperationResult createBatchOperationResultFr eTag, requestCharge, resourceBody, - responseStatusCode, - retryAfterMilliseconds != null ? Duration.ofMillis(retryAfterMilliseconds) : null, + statusCode, + retryAfterMilliseconds != null ? Duration.ofMillis(retryAfterMilliseconds) : Duration.ZERO, subStatusCode); } @@ -241,14 +200,14 @@ private static TransactionalBatchOperationResult createBatchOperationResultFr private static void createAndPopulateResults(final TransactionalBatchResponse response, final List> operations, final Duration retryAfterDuration) { - final ArrayList> results = new ArrayList<>(operations.size()); + final List> results = new ArrayList<>(operations.size()); for (int i = 0; i < operations.size(); i++) { results.add( BridgeInternal.createTransactionBatchResult( null, response.getRequestCharge(), null, - response.getResponseStatus(), + response.getStatusCode(), retryAfterDuration, response.getSubStatusCode() )); diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/batch/ServerBatchRequest.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/batch/ServerBatchRequest.java index ff3f90ed582f..c4f40aa6d8c7 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/batch/ServerBatchRequest.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/batch/ServerBatchRequest.java @@ -48,7 +48,7 @@ public abstract class ServerBatchRequest { * * @return Any pending operations that were not included in the request. */ - final List> createBodyStreamAsync(final List> operations) { + final List> createBodyOfBatchRequest(final List> operations) { checkNotNull(operations, "expected non-null operations"); diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/batch/SinglePartitionKeyServerBatchRequest.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/batch/SinglePartitionKeyServerBatchRequest.java index 1d3589e21976..3d6e11ad5c41 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/batch/SinglePartitionKeyServerBatchRequest.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/batch/SinglePartitionKeyServerBatchRequest.java @@ -32,7 +32,7 @@ private SinglePartitionKeyServerBatchRequest(final PartitionKey partitionKey) { * * @return A newly created instance of {@link SinglePartitionKeyServerBatchRequest}. */ - static SinglePartitionKeyServerBatchRequest createAsync( + static SinglePartitionKeyServerBatchRequest createBatchRequest( final PartitionKey partitionKey, final List> operations) { @@ -40,7 +40,7 @@ static SinglePartitionKeyServerBatchRequest createAsync( checkNotNull(operations, "expected non-null operations"); final SinglePartitionKeyServerBatchRequest request = new SinglePartitionKeyServerBatchRequest(partitionKey); - request.createBodyStreamAsync(operations); + request.createBodyOfBatchRequest(operations); return request; } diff --git a/sdk/cosmos/azure-cosmos/src/test/java/com/azure/cosmos/BatchOperationResultTests.java b/sdk/cosmos/azure-cosmos/src/test/java/com/azure/cosmos/BatchOperationResultTests.java index 29b376ec607f..f5d0e143f4e0 100644 --- a/sdk/cosmos/azure-cosmos/src/test/java/com/azure/cosmos/BatchOperationResultTests.java +++ b/sdk/cosmos/azure-cosmos/src/test/java/com/azure/cosmos/BatchOperationResultTests.java @@ -5,6 +5,7 @@ import com.azure.cosmos.implementation.HttpConstants; import com.azure.cosmos.implementation.Utils; +import com.fasterxml.jackson.databind.node.ObjectNode; import io.netty.handler.codec.http.HttpResponseStatus; import org.testng.annotations.Test; @@ -14,12 +15,13 @@ public class BatchOperationResultTests { private static final int TIMEOUT = 40000; + private ObjectNode objectNode = Utils.getSimpleObjectMapper().createObjectNode(); private TransactionalBatchOperationResult createTestResult() { TransactionalBatchOperationResult result = BridgeInternal.createTransactionBatchResult( "TestETag", 1.4, - Utils.getSimpleObjectMapper().createObjectNode(), + objectNode, HttpResponseStatus.OK.code(), Duration.ofMillis(1234), HttpConstants.SubStatusCodes.NAME_CACHE_IS_STALE @@ -28,12 +30,24 @@ private TransactionalBatchOperationResult createTestResult() { return result; } + @Test(groups = {"unit"}, timeOut = TIMEOUT) + public void propertiesAreSetThroughCtor() { + TransactionalBatchOperationResult result = createTestResult(); + + assertThat(result.getStatusCode()).isEqualTo(HttpResponseStatus.OK.code()); + assertThat(result.getSubStatusCode()).isEqualTo(HttpConstants.SubStatusCodes.NAME_CACHE_IS_STALE); + assertThat(result.getETag()).isEqualTo("TestETag"); + assertThat(result.getRequestCharge()).isEqualTo(1.4); + assertThat(result.getRetryAfter()).isEqualTo(Duration.ofMillis(1234)); + assertThat(result.getResourceObject()).isSameAs(objectNode); + } + @Test(groups = {"unit"}, timeOut = TIMEOUT) public void propertiesAreSetThroughCopyCtor() { TransactionalBatchOperationResult other = createTestResult(); TransactionalBatchOperationResult result = new TransactionalBatchOperationResult(other); - assertThat(other.getResponseStatus()).isEqualTo(result.getResponseStatus()); + assertThat(other.getStatusCode()).isEqualTo(result.getStatusCode()); assertThat(other.getSubStatusCode()).isEqualTo(result.getSubStatusCode()); assertThat(other.getETag()).isEqualTo(result.getETag()); assertThat(other.getRequestCharge()).isEqualTo(result.getRequestCharge()); @@ -47,7 +61,7 @@ public void propertiesAreSetThroughGenericCtor() { Object testObject = new Object(); TransactionalBatchOperationResult result = new TransactionalBatchOperationResult(other, testObject); - assertThat(other.getResponseStatus()).isEqualTo(result.getResponseStatus()); + assertThat(other.getStatusCode()).isEqualTo(result.getStatusCode()); assertThat(other.getSubStatusCode()).isEqualTo(result.getSubStatusCode()); assertThat(other.getETag()).isEqualTo(result.getETag()); assertThat(other.getRequestCharge()).isEqualTo(result.getRequestCharge()); @@ -61,11 +75,11 @@ public void isSuccessStatusCodeTrueFor200To299() { for (int x = 100; x < 999; ++x) { TransactionalBatchOperationResult result = BridgeInternal.createTransactionBatchResult( null, - null, + 0.0, null, x, null, - null + 0 ); boolean success = x >= 200 && x <= 299; diff --git a/sdk/cosmos/azure-cosmos/src/test/java/com/azure/cosmos/BatchTestBase.java b/sdk/cosmos/azure-cosmos/src/test/java/com/azure/cosmos/BatchTestBase.java index 732e8eacf891..9aa1f3701622 100644 --- a/sdk/cosmos/azure-cosmos/src/test/java/com/azure/cosmos/BatchTestBase.java +++ b/sdk/cosmos/azure-cosmos/src/test/java/com/azure/cosmos/BatchTestBase.java @@ -12,6 +12,7 @@ import com.fasterxml.jackson.annotation.JsonProperty; import io.netty.handler.codec.http.HttpResponseStatus; +import java.util.Objects; import java.util.Random; import java.util.UUID; @@ -22,7 +23,7 @@ public class BatchTestBase extends TestSuiteBase { private Random random = new Random(); String partitionKey1 = "TBD1"; - // Documents in partitionKey1 + // items in partitionKey1 TestDoc TestDocPk1ExistingA; TestDoc TestDocPk1ExistingB ; TestDoc TestDocPk1ExistingC; @@ -32,11 +33,11 @@ public BatchTestBase(CosmosClientBuilder clientBuilder) { super(clientBuilder); } - void createJsonTestDocsAsync(CosmosContainer container) { - this.TestDocPk1ExistingA = this.createJsonTestDocAsync(container, this.partitionKey1); - this.TestDocPk1ExistingB = this.createJsonTestDocAsync(container, this.partitionKey1); - this.TestDocPk1ExistingC = this.createJsonTestDocAsync(container, this.partitionKey1); - this.TestDocPk1ExistingD = this.createJsonTestDocAsync(container, this.partitionKey1); + void createJsonTestDocs(CosmosContainer container) { + this.TestDocPk1ExistingA = this.createJsonTestDoc(container, this.partitionKey1); + this.TestDocPk1ExistingB = this.createJsonTestDoc(container, this.partitionKey1); + this.TestDocPk1ExistingC = this.createJsonTestDoc(container, this.partitionKey1); + this.TestDocPk1ExistingD = this.createJsonTestDoc(container, this.partitionKey1); } TestDoc populateTestDoc(String partitionKey) { @@ -52,11 +53,11 @@ TestDoc getTestDocCopy(TestDoc testDoc) { return new TestDoc(testDoc.getId(), testDoc.getCost(), testDoc.getDescription(), testDoc.getStatus()); } - void verifyByReadAsync(CosmosContainer container, TestDoc doc) { - verifyByReadAsync(container, doc, null); + void verifyByRead(CosmosContainer container, TestDoc doc) { + verifyByRead(container, doc, null); } - void verifyByReadAsync(CosmosContainer container, TestDoc doc, String eTag) { + void verifyByRead(CosmosContainer container, TestDoc doc, String eTag) { PartitionKey partitionKey = this.getPartitionKey(doc.getStatus()); CosmosItemResponse response = container.readItem(doc.getId(), partitionKey, TestDoc.class); @@ -69,7 +70,7 @@ void verifyByReadAsync(CosmosContainer container, TestDoc doc, String eTag) { } } - void verifyNotFoundAsync(CosmosContainer container, TestDoc doc) { + void verifyNotFound(CosmosContainer container, TestDoc doc) { String id = doc.getId(); PartitionKey partitionKey = this.getPartitionKey(doc.getStatus()); @@ -87,11 +88,11 @@ PartitionKey getPartitionKey(String partitionKey) { return new PartitionKey(partitionKey); } - private TestDoc createJsonTestDocAsync(CosmosContainer container, String partitionKey) { - return createJsonTestDocAsync(container, partitionKey, 20); + private TestDoc createJsonTestDoc(CosmosContainer container, String partitionKey) { + return createJsonTestDoc(container, partitionKey, 20); } - private TestDoc createJsonTestDocAsync(CosmosContainer container, String partitionKey, int minDesiredSize) { + TestDoc createJsonTestDoc(CosmosContainer container, String partitionKey, int minDesiredSize) { TestDoc doc = this.populateTestDoc(partitionKey, minDesiredSize); CosmosItemResponse createResponse = container.createItem(doc, this.getPartitionKey(partitionKey), null); assertThat(createResponse.getStatusCode()).isEqualTo(HttpResponseStatus.CREATED.code()); @@ -182,4 +183,85 @@ public void setStatus(String status) { this.status = status; } } + + public static class EventDoc { + + public String id; + int clicks; + int views; + String type; + + @JsonProperty("mypk") + public String partitionKey; + + + public EventDoc() { + + } + + public EventDoc(String id, int clicks, int views, String type, String partitionKey) { + this.id = id; + this.clicks = clicks; + this.views = views; + this.type = type; + this.partitionKey = partitionKey; + } + + public String getId() { + return id; + } + + public int getClicks() { + return clicks; + } + + public int getViews() { + return views; + } + + public String getType() { + return type; + } + + public String getPartitionKey() { + return partitionKey; + } + + public void setId(String id) { + this.id = id; + } + + public void setClicks(int clicks) { + this.clicks = clicks; + } + + public void setViews(int views) { + this.views = views; + } + + public void setType(String type) { + this.type = type; + } + + public void setPartitionKey(String partitionKey) { + this.partitionKey = partitionKey; + } + + @Override + public boolean equals(Object o) { + if (this == o) return true; + if (o == null || getClass() != o.getClass()) return false; + EventDoc eventDoc = (EventDoc) o; + return clicks == eventDoc.clicks && + views == eventDoc.views && + Objects.equals(id, eventDoc.id) && + Objects.equals(type, eventDoc.type) && + Objects.equals(partitionKey, eventDoc.partitionKey); + } + + @Override + public int hashCode() { + return Objects.hash(id, clicks, views, type, partitionKey); + } + } } diff --git a/sdk/cosmos/azure-cosmos/src/test/java/com/azure/cosmos/TransactionalBatchTest.java b/sdk/cosmos/azure-cosmos/src/test/java/com/azure/cosmos/TransactionalBatchTest.java index b6151b5b74bc..d4694b4ec5fe 100644 --- a/sdk/cosmos/azure-cosmos/src/test/java/com/azure/cosmos/TransactionalBatchTest.java +++ b/sdk/cosmos/azure-cosmos/src/test/java/com/azure/cosmos/TransactionalBatchTest.java @@ -8,15 +8,16 @@ import com.azure.cosmos.implementation.guava25.base.Function; import com.azure.cosmos.models.CosmosItemResponse; import io.netty.handler.codec.http.HttpResponseStatus; +import org.assertj.core.api.Assertions; import org.assertj.core.data.Offset; import org.testng.annotations.AfterClass; import org.testng.annotations.BeforeClass; import org.testng.annotations.Factory; import org.testng.annotations.Test; -import reactor.core.publisher.Mono; import java.util.UUID; +import static com.azure.cosmos.implementation.batch.BatchRequestResponseConstant.MAX_DIRECT_MODE_BATCH_REQUEST_BODY_SIZE_IN_BYTES; import static com.azure.cosmos.implementation.batch.BatchRequestResponseConstant.MAX_OPERATIONS_IN_DIRECT_MODE_BATCH_REQUEST; import static org.assertj.core.api.Assertions.assertThat; @@ -24,7 +25,6 @@ public class TransactionalBatchTest extends BatchTestBase { private CosmosClient batchClient; private CosmosContainer batchContainer; - private CosmosAsyncContainer batchAsyncContainer; @Factory(dataProvider = "simpleClientBuildersWithDirect") public TransactionalBatchTest(CosmosClientBuilder clientBuilder) { @@ -35,7 +35,7 @@ public TransactionalBatchTest(CosmosClientBuilder clientBuilder) { public void before_TransactionalBatchTest() { assertThat(this.batchClient).isNull(); this.batchClient = getClientBuilder().buildClient(); - batchAsyncContainer = getSharedMultiPartitionCosmosContainer(this.batchClient.asyncClient()); + CosmosAsyncContainer batchAsyncContainer = getSharedMultiPartitionCosmosContainer(this.batchClient.asyncClient()); batchContainer = batchClient.getDatabase(batchAsyncContainer.getDatabase().getId()).getContainer(batchAsyncContainer.getId()); } @@ -45,62 +45,73 @@ public void afterClass() { } @Test(groups = {"simple"}, timeOut = TIMEOUT) - public void batchExecutionRepeat() { + public void batchOrdered() { CosmosContainer container = this.batchContainer; - this.createJsonTestDocsAsync(container); + this.createJsonTestDocs(container); TestDoc firstDoc = this.populateTestDoc(this.partitionKey1); + TestDoc replaceDoc = this.getTestDocCopy(firstDoc); replaceDoc.setCost(replaceDoc.getCost() + 1); - Mono batchResponseMono = batchAsyncContainer.executeTransactionalBatch( + TransactionalBatchResponse batchResponse = container.executeTransactionalBatch( TransactionalBatch.createTransactionalBatch(this.getPartitionKey(this.partitionKey1)) .createItem(firstDoc) .replaceItem(replaceDoc.getId(), replaceDoc)); - TransactionalBatchResponse batchResponse1 = batchResponseMono.block(); - this.verifyBatchProcessed(batchResponse1, 2); - - assertThat(batchResponse1.get(0).getResponseStatus()).isEqualTo(HttpResponseStatus.CREATED.code()); - assertThat(batchResponse1.get(1).getResponseStatus()).isEqualTo(HttpResponseStatus.OK.code()); + this.verifyBatchProcessed(batchResponse, 2); - // Block again. - TransactionalBatchResponse batchResponse2 = batchResponseMono.block(); - assertThat(batchResponse2.get(0).getResponseStatus()).isEqualTo(HttpResponseStatus.CONFLICT.code()); - assertThat(batchResponse2.get(1).getResponseStatus()).isEqualTo(HttpResponseStatus.FAILED_DEPENDENCY.code()); + assertThat(batchResponse.get(0).getStatusCode()).isEqualTo(HttpResponseStatus.CREATED.code()); + assertThat(batchResponse.get(1).getStatusCode()).isEqualTo(HttpResponseStatus.OK.code()); // Ensure that the replace overwrote the doc from the first operation - this.verifyByReadAsync(container, replaceDoc); + this.verifyByRead(container, replaceDoc); } + @Test(groups = {"simple"}, timeOut = TIMEOUT) - public void batchOrdered() { + public void batchMultipleItemExecution() { CosmosContainer container = this.batchContainer; - this.createJsonTestDocsAsync(container); + this.createJsonTestDocs(container); TestDoc firstDoc = this.populateTestDoc(this.partitionKey1); - TestDoc replaceDoc = this.getTestDocCopy(firstDoc); replaceDoc.setCost(replaceDoc.getCost() + 1); + EventDoc eventDoc1 = new EventDoc(UUID.randomUUID().toString(), 2, 4, "type1", this.partitionKey1); + EventDoc readEventDoc = new EventDoc(UUID.randomUUID().toString(), 6, 14, "type2", this.partitionKey1); + CosmosItemResponse createResponse = container.createItem(readEventDoc, this.getPartitionKey(this.partitionKey1), null); + assertThat(createResponse.getStatusCode()).isEqualTo(HttpResponseStatus.CREATED.code()); + TransactionalBatchResponse batchResponse = container.executeTransactionalBatch( TransactionalBatch.createTransactionalBatch(this.getPartitionKey(this.partitionKey1)) .createItem(firstDoc) - .replaceItem(replaceDoc.getId(), replaceDoc)); + .createItem(eventDoc1) + .replaceItem(replaceDoc.getId(), replaceDoc) + .readItem(readEventDoc.getId())); - this.verifyBatchProcessed(batchResponse, 2); + this.verifyBatchProcessed(batchResponse, 4); + + assertThat(batchResponse.get(0).getStatusCode()).isEqualTo(HttpResponseStatus.CREATED.code()); + assertThat(batchResponse.getOperationResultAtIndex(0, TestDoc.class).getItem()).isEqualTo(firstDoc); + + assertThat(batchResponse.get(1).getStatusCode()).isEqualTo(HttpResponseStatus.CREATED.code()); + assertThat(batchResponse.getOperationResultAtIndex(1, EventDoc.class).getItem()).isEqualTo(eventDoc1); - assertThat(batchResponse.get(0).getResponseStatus()).isEqualTo(HttpResponseStatus.CREATED.code()); - assertThat(batchResponse.get(1).getResponseStatus()).isEqualTo(HttpResponseStatus.OK.code()); + assertThat(batchResponse.get(2).getStatusCode()).isEqualTo(HttpResponseStatus.OK.code()); + assertThat(batchResponse.getOperationResultAtIndex(2, TestDoc.class).getItem()).isEqualTo(replaceDoc); + + assertThat(batchResponse.get(3).getStatusCode()).isEqualTo(HttpResponseStatus.OK.code()); + assertThat(batchResponse.getOperationResultAtIndex(3, EventDoc.class).getItem()).isEqualTo(readEventDoc); // Ensure that the replace overwrote the doc from the first operation - this.verifyByReadAsync(container, replaceDoc); + this.verifyByRead(container, replaceDoc); } @Test(groups = {"simple"}, timeOut = TIMEOUT) - public void batchItemETagAsync() { + public void batchItemETagTest() { CosmosContainer container = batchContainer; - this.createJsonTestDocsAsync(container); + this.createJsonTestDocs(container); { BatchTestBase.TestDoc testDocToCreate = this.populateTestDoc(this.partitionKey1); @@ -125,12 +136,12 @@ public void batchItemETagAsync() { this.verifyBatchProcessed(batchResponse, 2); - assertThat(batchResponse.get(0).getResponseStatus()).isEqualTo(HttpResponseStatus.CREATED.code()); - assertThat(batchResponse.get(1).getResponseStatus()).isEqualTo(HttpResponseStatus.OK.code()); + assertThat(batchResponse.get(0).getStatusCode()).isEqualTo(HttpResponseStatus.CREATED.code()); + assertThat(batchResponse.get(1).getStatusCode()).isEqualTo(HttpResponseStatus.OK.code()); // Ensure that the replace overwrote the doc from the first operation - this.verifyByReadAsync(container, testDocToCreate, batchResponse.get(0).getETag()); - this.verifyByReadAsync(container, testDocToReplace, batchResponse.get(1).getETag()); + this.verifyByRead(container, testDocToCreate, batchResponse.get(0).getETag()); + this.verifyByRead(container, testDocToReplace, batchResponse.get(1).getETag()); } { @@ -146,17 +157,17 @@ public void batchItemETagAsync() { this.verifyBatchProcessed(batchResponse, 1, HttpResponseStatus.PRECONDITION_FAILED); - assertThat(batchResponse.get(0).getResponseStatus()).isEqualTo(HttpResponseStatus.PRECONDITION_FAILED.code()); + assertThat(batchResponse.get(0).getStatusCode()).isEqualTo(HttpResponseStatus.PRECONDITION_FAILED.code()); - // ensure the document was not updated - this.verifyByReadAsync(container, this.TestDocPk1ExistingB); + // ensure the item was not updated + this.verifyByRead(container, this.TestDocPk1ExistingB); } } @Test(groups = {"simple"}, timeOut = TIMEOUT) - public void batchItemSessionTokenAsync() { + public void batchItemSessionTokenTest() { CosmosContainer container = batchContainer; - this.createJsonTestDocsAsync(container); + this.createJsonTestDocs(container); TestDoc testDocToCreate = this.populateTestDoc(this.partitionKey1); @@ -179,8 +190,8 @@ public void batchItemSessionTokenAsync() { this.verifyBatchProcessed(batchResponse, 2); - assertThat(batchResponse.get(0).getResponseStatus()).isEqualTo(HttpResponseStatus.CREATED.code()); - assertThat(batchResponse.get(1).getResponseStatus()).isEqualTo(HttpResponseStatus.OK.code()); + assertThat(batchResponse.get(0).getStatusCode()).isEqualTo(HttpResponseStatus.CREATED.code()); + assertThat(batchResponse.get(1).getStatusCode()).isEqualTo(HttpResponseStatus.OK.code()); ISessionToken afterRequestSessionToken = this.getSessionToken(batchResponse.getResponseHeaders().get(HttpConstants.HttpHeaders.SESSION_TOKEN)); assertThat(afterRequestSessionToken.getLSN()) @@ -189,8 +200,7 @@ public void batchItemSessionTokenAsync() { } @Test(groups = {"simple"}, timeOut = TIMEOUT) - public void batchWithTooManyOperationsAsync() { - CosmosContainer container = batchContainer; + public void batchWithTooManyOperationsTest() { int operationCount = MAX_OPERATIONS_IN_DIRECT_MODE_BATCH_REQUEST + 1; // Increase the doc size by a bit so all docs won't fit in one server request. @@ -200,14 +210,57 @@ public void batchWithTooManyOperationsAsync() { batch.readItem("someId"); } - TransactionalBatchResponse batchResponse = container.executeTransactionalBatch(batch); - assertThat(batchResponse.getResponseStatus()).isEqualTo(HttpResponseStatus.BAD_REQUEST.code()); + try { + batchContainer.executeTransactionalBatch(batch); + Assertions.fail("Should throw bad request exception"); + } catch (CosmosException ex) { + assertThat(ex.getStatusCode()).isEqualTo(HttpResponseStatus.BAD_REQUEST.code()); + } + } + + @Test(groups = {"emulator"}, timeOut = TIMEOUT * 10) + public void batchLargerThanServerRequest() { + int operationCount = 20; + int appxDocSize = (MAX_DIRECT_MODE_BATCH_REQUEST_BODY_SIZE_IN_BYTES * 11) / operationCount; + + // Increase the doc size by a bit so all docs won't fit in one server request. + appxDocSize = (int)(appxDocSize * 1.05); + TransactionalBatch batch = TransactionalBatch.createTransactionalBatch(this.getPartitionKey(this.partitionKey1)); + + for (int i = 0; i < operationCount; i++) { + TestDoc doc = this.populateTestDoc(this.partitionKey1, appxDocSize); + batch.createItem(doc); + } + + try { + batchContainer.executeTransactionalBatch(batch); + Assertions.fail("Should throw REQUEST_ENTITY_TOO_LARGE exception"); + } catch (CosmosException ex) { + assertThat(ex.getStatusCode()).isEqualTo(HttpResponseStatus.REQUEST_ENTITY_TOO_LARGE.code()); + } + } + + @Test(groups = {"emulator"}, timeOut = TIMEOUT * 10) + public void batchServerResponseTooLarge() { + int operationCount = 10; + int appxDocSizeInBytes = 1 * 1024 * 1024; + + TestDoc doc = this.createJsonTestDoc(batchContainer, this.partitionKey1, appxDocSizeInBytes); + + TransactionalBatch batch = TransactionalBatch.createTransactionalBatch(this.getPartitionKey(this.partitionKey1)); + for (int i = 0; i < operationCount; i++) { + batch.readItem(doc.getId()); + } + + TransactionalBatchResponse batchResponse = batchContainer.executeTransactionalBatch(batch); + assertThat(batchResponse.getStatusCode()).isEqualTo(HttpResponseStatus.REQUEST_ENTITY_TOO_LARGE.code()); + assertThat(batchResponse.get(1).getStatusCode()).isEqualTo(HttpResponseStatus.FAILED_DEPENDENCY.code()); } @Test(groups = {"simple"}, timeOut = TIMEOUT) - public void batchReadsOnlyAsync() { + public void batchReadsOnlyTest() { CosmosContainer container = batchContainer; - this.createJsonTestDocsAsync(container); + this.createJsonTestDocs(container); TransactionalBatchResponse batchResponse = container.executeTransactionalBatch( TransactionalBatch.createTransactionalBatch(this.getPartitionKey(this.partitionKey1)) @@ -217,9 +270,9 @@ public void batchReadsOnlyAsync() { this.verifyBatchProcessed(batchResponse, 3); - assertThat(batchResponse.get(0).getResponseStatus()).isEqualTo(HttpResponseStatus.OK.code()); - assertThat(batchResponse.get(1).getResponseStatus()).isEqualTo(HttpResponseStatus.OK.code()); - assertThat(batchResponse.get(2).getResponseStatus()).isEqualTo(HttpResponseStatus.OK.code()); + assertThat(batchResponse.get(0).getStatusCode()).isEqualTo(HttpResponseStatus.OK.code()); + assertThat(batchResponse.get(1).getStatusCode()).isEqualTo(HttpResponseStatus.OK.code()); + assertThat(batchResponse.get(2).getStatusCode()).isEqualTo(HttpResponseStatus.OK.code()); assertThat(batchResponse.getOperationResultAtIndex(0, TestDoc.class).getItem()).isEqualTo(this.TestDocPk1ExistingA); assertThat(batchResponse.getOperationResultAtIndex(1, TestDoc.class).getItem()).isEqualTo(this.TestDocPk1ExistingB); @@ -229,7 +282,7 @@ public void batchReadsOnlyAsync() { @Test(groups = {"simple"}, timeOut = TIMEOUT) public void batchCrud() { CosmosContainer container = batchContainer; - this.createJsonTestDocsAsync(container); + this.createJsonTestDocs(container); BatchTestBase.TestDoc testDocToCreate = this.populateTestDoc(this.partitionKey1); BatchTestBase.TestDoc testDocToUpsert = this.populateTestDoc(this.partitionKey1); @@ -252,42 +305,42 @@ public void batchCrud() { this.verifyBatchProcessed(batchResponse, 6); - assertThat(batchResponse.get(0).getResponseStatus()).isEqualTo(HttpResponseStatus.CREATED.code()); - assertThat(batchResponse.get(1).getResponseStatus()).isEqualTo(HttpResponseStatus.OK.code()); - assertThat(batchResponse.get(2).getResponseStatus()).isEqualTo(HttpResponseStatus.OK.code()); - assertThat(batchResponse.get(3).getResponseStatus()).isEqualTo(HttpResponseStatus.CREATED.code()); - assertThat(batchResponse.get(4).getResponseStatus()).isEqualTo(HttpResponseStatus.OK.code()); - assertThat(batchResponse.get(5).getResponseStatus()).isEqualTo(HttpResponseStatus.NO_CONTENT.code()); + assertThat(batchResponse.get(0).getStatusCode()).isEqualTo(HttpResponseStatus.CREATED.code()); + assertThat(batchResponse.get(1).getStatusCode()).isEqualTo(HttpResponseStatus.OK.code()); + assertThat(batchResponse.get(2).getStatusCode()).isEqualTo(HttpResponseStatus.OK.code()); + assertThat(batchResponse.get(3).getStatusCode()).isEqualTo(HttpResponseStatus.CREATED.code()); + assertThat(batchResponse.get(4).getStatusCode()).isEqualTo(HttpResponseStatus.OK.code()); + assertThat(batchResponse.get(5).getStatusCode()).isEqualTo(HttpResponseStatus.NO_CONTENT.code()); assertThat(batchResponse.getOperationResultAtIndex(1, TestDoc.class).getItem()).isEqualTo(this.TestDocPk1ExistingC); - this.verifyByReadAsync(container, testDocToCreate); - this.verifyByReadAsync(container, testDocToReplace); - this.verifyByReadAsync(container, testDocToUpsert); - this.verifyByReadAsync(container, anotherTestDocToUpsert); - this.verifyNotFoundAsync(container, this.TestDocPk1ExistingD); + this.verifyByRead(container, testDocToCreate); + this.verifyByRead(container, testDocToReplace); + this.verifyByRead(container, testDocToUpsert); + this.verifyByRead(container, anotherTestDocToUpsert); + this.verifyNotFound(container, this.TestDocPk1ExistingD); } @Test(groups = {"simple"}, timeOut = TIMEOUT) - public void batchWithInvalidCreateAsync() { + public void batchWithInvalidCreateTest() { // partition key mismatch between doc and and value passed in to the operation - this.runWithErrorAsync( + this.runWithError( batchContainer, batch -> batch.createItem(this.populateTestDoc(UUID.randomUUID().toString())), HttpResponseStatus.BAD_REQUEST); } @Test(groups = {"simple"}, timeOut = TIMEOUT) - public void batchWithReadOfNonExistentEntityAsync() { - this.runWithErrorAsync( + public void batchWithReadOfNonExistentEntityTest() { + this.runWithError( batchContainer, batch -> batch.readItem(UUID.randomUUID().toString()), HttpResponseStatus.NOT_FOUND); } @Test(groups = {"simple"}, timeOut = TIMEOUT) - public void batchWithReplaceOfStaleEntityAsync() { - this.createJsonTestDocsAsync(batchContainer); + public void batchWithReplaceOfStaleEntity() { + this.createJsonTestDocs(batchContainer); TestDoc staleTestDocToReplace = this.getTestDocCopy(this.TestDocPk1ExistingA); staleTestDocToReplace.setCost(staleTestDocToReplace.getCost() + 1); @@ -295,42 +348,42 @@ public void batchWithReplaceOfStaleEntityAsync() { TransactionalBatchItemRequestOptions staleReplaceOptions = new TransactionalBatchItemRequestOptions(); staleReplaceOptions.setIfMatchETag(UUID.randomUUID().toString()); - this.runWithErrorAsync( + this.runWithError( batchContainer, batch -> batch.replaceItem(staleTestDocToReplace.getId(), staleTestDocToReplace, staleReplaceOptions), HttpResponseStatus.PRECONDITION_FAILED); // make sure the stale doc hasn't changed - this.verifyByReadAsync(batchContainer, this.TestDocPk1ExistingA); + this.verifyByRead(batchContainer, this.TestDocPk1ExistingA); } @Test(groups = {"simple"}, timeOut = TIMEOUT) - public void batchWithDeleteOfNonExistentEntityAsync() { - this.runWithErrorAsync( + public void batchWithDeleteOfNonExistentEntity() { + this.runWithError( batchContainer, batch -> batch.deleteItem(UUID.randomUUID().toString()), HttpResponseStatus.NOT_FOUND); } @Test(groups = {"simple"}, timeOut = TIMEOUT) - public void batchWithCreateConflictAsync() { - this.createJsonTestDocsAsync(batchContainer); + public void batchWithCreateConflict() { + this.createJsonTestDocs(batchContainer); // try to create a doc with id that already exists (should return a Conflict) TestDoc conflictingTestDocToCreate = this.getTestDocCopy(this.TestDocPk1ExistingA); conflictingTestDocToCreate.setCost(conflictingTestDocToCreate.getCost()); - this.runWithErrorAsync( + this.runWithError( batchContainer, batch -> batch.createItem(conflictingTestDocToCreate), HttpResponseStatus.CONFLICT); // make sure the conflicted doc hasn't changed - this.verifyByReadAsync(batchContainer, this.TestDocPk1ExistingA); + this.verifyByRead(batchContainer, this.TestDocPk1ExistingA); } - private void runWithErrorAsync( + private void runWithError( CosmosContainer container, Function appendOperation, HttpResponseStatus expectedFailedOperationStatusCode) { @@ -348,12 +401,12 @@ private void runWithErrorAsync( this.verifyBatchProcessed(batchResponse, 3, expectedFailedOperationStatusCode); - assertThat(batchResponse.get(0).getResponseStatus()).isEqualTo(HttpResponseStatus.FAILED_DEPENDENCY.code()); - assertThat(batchResponse.get(1).getResponseStatus()).isEqualTo(expectedFailedOperationStatusCode.code()); - assertThat(batchResponse.get(2).getResponseStatus()).isEqualTo(HttpResponseStatus.FAILED_DEPENDENCY.code()); + assertThat(batchResponse.get(0).getStatusCode()).isEqualTo(HttpResponseStatus.FAILED_DEPENDENCY.code()); + assertThat(batchResponse.get(1).getStatusCode()).isEqualTo(expectedFailedOperationStatusCode.code()); + assertThat(batchResponse.get(2).getStatusCode()).isEqualTo(HttpResponseStatus.FAILED_DEPENDENCY.code()); - this.verifyNotFoundAsync(container, testDocToCreate); - this.verifyNotFoundAsync(container, anotherTestDocToCreate); + this.verifyNotFound(container, testDocToCreate); + this.verifyNotFound(container, anotherTestDocToCreate); } private void verifyBatchProcessed(TransactionalBatchResponse batchResponse, int numberOfOperations) { @@ -362,13 +415,14 @@ private void verifyBatchProcessed(TransactionalBatchResponse batchResponse, int private void verifyBatchProcessed(TransactionalBatchResponse batchResponse, int numberOfOperations, HttpResponseStatus expectedStatusCode) { assertThat(batchResponse).isNotNull(); - assertThat(batchResponse.getResponseStatus()) - .as("Batch server response had StatusCode {0} instead of {1} expected and had ErrorMessage {2}") + assertThat(batchResponse.getStatusCode()) + .as("Batch server response had StatusCode {0} instead of {1} expected and had ErrorMessage {2}", + batchResponse.getStatusCode(), expectedStatusCode.code()) .isEqualTo(expectedStatusCode.code()); assertThat(batchResponse.size()).isEqualTo(numberOfOperations); assertThat(batchResponse.getRequestCharge()).isPositive(); - assertThat(batchResponse.getCosmosDiagnostics().toString()).isNotEmpty(); + assertThat(batchResponse.getDiagnostics().toString()).isNotEmpty(); // Allow a delta since we round both the total charge and the individual operation // charges to 2 decimal places. diff --git a/sdk/cosmos/azure-cosmos/src/test/java/com/azure/cosmos/implementation/batch/BatchResponsePayloadWriter.java b/sdk/cosmos/azure-cosmos/src/test/java/com/azure/cosmos/implementation/batch/BatchResponsePayloadWriter.java index bb13e1d9ec2d..eb29c2b61743 100644 --- a/sdk/cosmos/azure-cosmos/src/test/java/com/azure/cosmos/implementation/batch/BatchResponsePayloadWriter.java +++ b/sdk/cosmos/azure-cosmos/src/test/java/com/azure/cosmos/implementation/batch/BatchResponsePayloadWriter.java @@ -36,10 +36,14 @@ private ArrayNode writeOperationResult() { private JsonSerializable writeResult(TransactionalBatchOperationResult result) { JsonSerializable jsonSerializable = new JsonSerializable(); - jsonSerializable.set(BatchRequestResponseConstant.FIELD_STATUS_CODE, result.getResponseStatus()); + jsonSerializable.set(BatchRequestResponseConstant.FIELD_STATUS_CODE, result.getStatusCode()); jsonSerializable.set(BatchRequestResponseConstant.FIELD_SUBSTATUS_CODE, result.getSubStatusCode()); jsonSerializable.set(BatchRequestResponseConstant.FIELD_ETAG, result.getETag()); - jsonSerializable.set(BatchRequestResponseConstant.FIELD_RESOURCE_BODY, result.getResourceObject()); + jsonSerializable.set(BatchRequestResponseConstant.FIELD_REQUEST_CHARGE, result.getRequestCharge()); + + if(result.getRetryAfter() != null) { + jsonSerializable.set(BatchRequestResponseConstant.FIELD_RETRY_AFTER_MILLISECONDS, result.getRetryAfter().toMillis()); + } return jsonSerializable; } diff --git a/sdk/cosmos/azure-cosmos/src/test/java/com/azure/cosmos/implementation/batch/TransactionalBatchResponseTests.java b/sdk/cosmos/azure-cosmos/src/test/java/com/azure/cosmos/implementation/batch/TransactionalBatchResponseTests.java index eb44f01fd3ee..1a547b33eb3f 100644 --- a/sdk/cosmos/azure-cosmos/src/test/java/com/azure/cosmos/implementation/batch/TransactionalBatchResponseTests.java +++ b/sdk/cosmos/azure-cosmos/src/test/java/com/azure/cosmos/implementation/batch/TransactionalBatchResponseTests.java @@ -6,6 +6,7 @@ import com.azure.cosmos.BridgeInternal; import com.azure.cosmos.TransactionalBatchOperationResult; import com.azure.cosmos.TransactionalBatchResponse; +import com.azure.cosmos.implementation.HttpConstants; import com.azure.cosmos.implementation.OperationType; import com.azure.cosmos.implementation.RxDocumentServiceResponse; import com.azure.cosmos.implementation.directconnectivity.StoreResponse; @@ -14,9 +15,13 @@ import org.testng.annotations.Test; import java.nio.charset.StandardCharsets; +import java.time.Duration; import java.util.ArrayList; import java.util.Arrays; +import java.util.HashMap; import java.util.List; +import java.util.Map; +import java.util.UUID; import static org.assertj.core.api.Assertions.assertThat; @@ -25,7 +30,7 @@ public class TransactionalBatchResponseTests { private static final int TIMEOUT = 40000; @Test(groups = {"unit"}, timeOut = TIMEOUT) - public void statusCodesAreSetThroughResponseAsync() { + public void validateAllSetValuesInResponse() { List> results = new ArrayList<>(); ItemBatchOperation[] arrayOperations = new ItemBatchOperation[1]; @@ -34,25 +39,88 @@ public void statusCodesAreSetThroughResponseAsync() { .id("0") .build(); + arrayOperations[0] = operation; + SinglePartitionKeyServerBatchRequest batchRequest = SinglePartitionKeyServerBatchRequest.createBatchRequest( + PartitionKey.NONE, + Arrays.asList(arrayOperations)); + + // Create dummy result TransactionalBatchOperationResult transactionalBatchOperationResult = BridgeInternal.createTransactionBatchResult( operation.getId(), - 0.0, - null, - HttpResponseStatus.OK.code(), + 5.0, null, - 0 + HttpResponseStatus.NOT_MODIFIED.code(), + Duration.ofMillis(100), + HttpConstants.SubStatusCodes.PARTITION_KEY_MISMATCH ); results.add(transactionalBatchOperationResult); + String responseContent = new BatchResponsePayloadWriter(results).generatePayload(); - arrayOperations[0] = operation; + // TransactionalBatchResponse headers + String activityId = UUID.randomUUID().toString(); + Map headers = new HashMap<>(); + headers.put(HttpConstants.HttpHeaders.ACTIVITY_ID, activityId); + headers.put(HttpConstants.HttpHeaders.REQUEST_CHARGE, "4.5"); + headers.put(HttpConstants.HttpHeaders.SESSION_TOKEN, "token123"); + headers.put(HttpConstants.HttpHeaders.RETRY_AFTER_IN_MILLISECONDS, "1234"); + headers.put(HttpConstants.HttpHeaders.SUB_STATUS, String.valueOf(HttpConstants.SubStatusCodes.PARTITION_KEY_RANGE_GONE)); - String responseContent = new BatchResponsePayloadWriter(results).generatePayload(); + StoreResponse storeResponse = new StoreResponse( + HttpResponseStatus.OK.code(), + new ArrayList<>(headers.entrySet()), + responseContent.getBytes(StandardCharsets.UTF_8)); + + TransactionalBatchResponse batchResponse = BatchResponseParser.fromDocumentServiceResponseAsync( + new RxDocumentServiceResponse(null, storeResponse), + batchRequest, + true).block(); + + // Validate response fields + assertThat(batchResponse.getActivityId()).isEqualTo(activityId); + assertThat(batchResponse.getRequestCharge()).isEqualTo(4.5); + assertThat(batchResponse.getStatusCode()).isEqualTo(HttpResponseStatus.OK.code()); + assertThat(batchResponse.getSessionToken()).isEqualTo("token123"); + assertThat(batchResponse.getResponseHeaders()).isEqualTo(headers); + assertThat(batchResponse.getRetryAfter()).isEqualTo(Duration.ofMillis(1234)); + assertThat(batchResponse.getSubStatusCode()).isEqualTo(HttpConstants.SubStatusCodes.PARTITION_KEY_RANGE_GONE); + + // Validate result fields + assertThat(batchResponse.get(0).getETag()).isEqualTo(operation.getId()); + assertThat(batchResponse.get(0).getRequestCharge()).isEqualTo(5.0); + assertThat(batchResponse.get(0).getRetryAfter()).isEqualTo(Duration.ofMillis(100)); + assertThat(batchResponse.get(0).getSubStatusCode()).isEqualTo(HttpConstants.SubStatusCodes.PARTITION_KEY_MISMATCH); + assertThat(batchResponse.get(0).getStatusCode()).isEqualTo(HttpResponseStatus.NOT_MODIFIED.code()); + } - SinglePartitionKeyServerBatchRequest batchRequest = SinglePartitionKeyServerBatchRequest.createAsync( + @Test(groups = {"unit"}, timeOut = TIMEOUT) + public void validateEmptyHeaderInResponse() { + List> results = new ArrayList<>(); + ItemBatchOperation[] arrayOperations = new ItemBatchOperation[1]; + + ItemBatchOperation operation = new ItemBatchOperation.Builder(OperationType.Read,0) + .partitionKey(PartitionKey.NONE) + .id("0") + .build(); + + arrayOperations[0] = operation; + SinglePartitionKeyServerBatchRequest batchRequest = SinglePartitionKeyServerBatchRequest.createBatchRequest( PartitionKey.NONE, Arrays.asList(arrayOperations)); + // Create dummy result + TransactionalBatchOperationResult transactionalBatchOperationResult = BridgeInternal.createTransactionBatchResult( + null, + 5.0, + null, + HttpResponseStatus.NOT_MODIFIED.code(), + null, + 0 + ); + + results.add(transactionalBatchOperationResult); + String responseContent = new BatchResponsePayloadWriter(results).generatePayload(); + StoreResponse storeResponse = new StoreResponse( HttpResponseStatus.OK.code(), new ArrayList<>(), @@ -63,6 +131,20 @@ public void statusCodesAreSetThroughResponseAsync() { batchRequest, true).block(); - assertThat(batchResponse.getResponseStatus()).isEqualTo(HttpResponseStatus.OK.code()); + // Validate response fields + assertThat(batchResponse.getStatusCode()).isEqualTo(HttpResponseStatus.OK.code()); + assertThat(batchResponse.getActivityId()).isNull(); + assertThat(batchResponse.getRequestCharge()).isEqualTo(0); + assertThat(batchResponse.getSessionToken()).isNull(); + assertThat(batchResponse.getResponseHeaders()).isEmpty(); + assertThat(batchResponse.getRetryAfter()).isEqualTo(Duration.ZERO); + assertThat(batchResponse.getSubStatusCode()).isEqualTo(0); + + // Validate result fields + assertThat(batchResponse.get(0).getETag()).isNull(); + assertThat(batchResponse.get(0).getRequestCharge()).isEqualTo(5.0); + assertThat(batchResponse.get(0).getRetryAfter()).isEqualTo(Duration.ZERO); + assertThat(batchResponse.get(0).getSubStatusCode()).isEqualTo(0); + assertThat(batchResponse.get(0).getStatusCode()).isEqualTo(HttpResponseStatus.NOT_MODIFIED.code()); } } From 439604677fe9d76df317810916f6910bd22edbf6 Mon Sep 17 00:00:00 2001 From: Rakesh Kumar Date: Mon, 5 Oct 2020 22:33:00 +0530 Subject: [PATCH 11/22] Fix Signed-off-by: Rakesh Kumar --- .../azure/cosmos/CosmosAsyncContainer.java | 24 ++++++++----------- .../com/azure/cosmos/CosmosContainer.java | 24 ++++++++----------- .../com/azure/cosmos/TransactionalBatch.java | 2 +- .../TransactionalBatchOperationResult.java | 12 +++++----- .../batch/ItemBatchOperation.java | 8 +++---- .../java/com/azure/cosmos/BatchTestBase.java | 5 ++-- .../azure/cosmos/TransactionalBatchTest.java | 5 ++-- 7 files changed, 35 insertions(+), 45 deletions(-) diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncContainer.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncContainer.java index a53e23fa5bdd..8a2667e70853 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncContainer.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncContainer.java @@ -508,19 +508,17 @@ private T transform(Object object, Class classType) { *

* If an operation within the transactional batch fails during execution, no changes from the batch will be * committed and the status of the failing operation is made available by {@link - * TransactionalBatchResponse#getStatusCode}. To obtain information about the operations that failed, the - * response can be enumerated. This returns {@link TransactionalBatchOperationResult} instances corresponding to - * each operation in the transactional batch in the order they were added to the transactional batch. For a result - * corresponding to an operation within the transactional batch, use + * TransactionalBatchResponse#getStatusCode} or by the exception. To obtain information about the operations + * that failed in case of some user error like conflict, not found etc, the response can be enumerated. + * This returns {@link TransactionalBatchOperationResult} instances corresponding to each operation in the + * transactional batch in the order they were added to the transactional batch. + * For a result corresponding to an operation within the transactional batch, use * {@link TransactionalBatchOperationResult#getStatusCode} * to access the status of the operation. If the operation was not executed or it was aborted due to the failure of * another operation within the transactional batch, the value of this field will be 424; * for the operation that caused the batch to abort, the value of this field * will indicate the cause of failure. *

- * The value returned by {@link TransactionalBatchResponse#getStatusCode} on the response returned may also have - * values such as 500 in case of server errors and 429. - *

* Use {@link TransactionalBatchResponse#isSuccessStatusCode} on the response returned to ensure that the * transactional batch succeeded. */ @@ -542,19 +540,17 @@ public Mono executeTransactionalBatch(TransactionalB *

* If an operation within the transactional batch fails during execution, no changes from the batch will be * committed and the status of the failing operation is made available by {@link - * TransactionalBatchResponse#getStatusCode}. To obtain information about the operations that failed, the - * response can be enumerated. This returns {@link TransactionalBatchOperationResult} instances corresponding to - * each operation in the transactional batch in the order they were added to the transactional batch. For a result - * corresponding to an operation within the transactional batch, use + * TransactionalBatchResponse#getStatusCode} or by the exception. To obtain information about the operations + * that failed in case of some user error like conflict, not found etc, the response can be enumerated. + * This returns {@link TransactionalBatchOperationResult} instances corresponding to each operation in the + * transactional batch in the order they were added to the transactional batch. + * For a result corresponding to an operation within the transactional batch, use * {@link TransactionalBatchOperationResult#getStatusCode} * to access the status of the operation. If the operation was not executed or it was aborted due to the failure of * another operation within the transactional batch, the value of this field will be 424; * for the operation that caused the batch to abort, the value of this field * will indicate the cause of failure. *

- * The value returned by {@link TransactionalBatchResponse#getStatusCode} on the response returned may also have - * values such as 500 in case of server errors and 429. - *

* Use {@link TransactionalBatchResponse#isSuccessStatusCode} on the response returned to ensure that the * transactional batch succeeded. */ diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosContainer.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosContainer.java index 55da0ee93db0..fa991d5f2758 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosContainer.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosContainer.java @@ -478,19 +478,17 @@ public CosmosItemResponse deleteItem(T item, CosmosItemRequestOption *

* If an operation within the transactional batch fails during execution, no changes from the batch will be * committed and the status of the failing operation is made available by {@link - * TransactionalBatchResponse#getStatusCode}. To obtain information about the operations that failed, the - * response can be enumerated. This returns {@link TransactionalBatchOperationResult} instances corresponding to - * each operation in the transactional batch in the order they were added to the transactional batch. For a result - * corresponding to an operation within the transactional batch, use + * TransactionalBatchResponse#getStatusCode} or by the exception. To obtain information about the operations + * that failed in case of some user error like conflict, not found etc, the response can be enumerated. + * This returns {@link TransactionalBatchOperationResult} instances corresponding to each operation in the + * transactional batch in the order they were added to the transactional batch. + * For a result corresponding to an operation within the transactional batch, use * {@link TransactionalBatchOperationResult#getStatusCode} * to access the status of the operation. If the operation was not executed or it was aborted due to the failure of * another operation within the transactional batch, the value of this field will be 424; * for the operation that caused the batch to abort, the value of this field * will indicate the cause of failure. *

- * The value returned by {@link TransactionalBatchResponse#getStatusCode} on the response returned may also have - * values such as 500 in case of server errors and 429. - *

* Use {@link TransactionalBatchResponse#isSuccessStatusCode} on the response returned to ensure that the * transactional batch succeeded. */ @@ -512,19 +510,17 @@ public TransactionalBatchResponse executeTransactionalBatch(TransactionalBatch t *

* If an operation within the transactional batch fails during execution, no changes from the batch will be * committed and the status of the failing operation is made available by {@link - * TransactionalBatchResponse#getStatusCode}. To obtain information about the operations that failed, the - * response can be enumerated. This returns {@link TransactionalBatchOperationResult} instances corresponding to - * each operation in the transactional batch in the order they were added to the transactional batch. For a result - * corresponding to an operation within the transactional batch, use + * TransactionalBatchResponse#getStatusCode} or by the exception. To obtain information about the operations + * that failed in case of some user error like conflict, not found etc, the response can be enumerated. + * This returns {@link TransactionalBatchOperationResult} instances corresponding to each operation in the + * transactional batch in the order they were added to the transactional batch. + * For a result corresponding to an operation within the transactional batch, use * {@link TransactionalBatchOperationResult#getStatusCode} * to access the status of the operation. If the operation was not executed or it was aborted due to the failure of * another operation within the transactional batch, the value of this field will be 424; * for the operation that caused the batch to abort, the value of this field * will indicate the cause of failure. *

- * The value returned by {@link TransactionalBatchResponse#getStatusCode} on the response returned may also have - * values such as 500 in case of server errors and 429. - *

* Use {@link TransactionalBatchResponse#isSuccessStatusCode} on the response returned to ensure that the * transactional batch succeeded. */ diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/TransactionalBatch.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/TransactionalBatch.java index cd2a05d24988..059fdd5eb856 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/TransactionalBatch.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/TransactionalBatch.java @@ -47,7 +47,7 @@ * * TransactionalBatchResponse response = container.executeTransactionalBatch(batch); * - * if (!response.IsSuccessStatusCode) { + * if (!response.isSuccessStatusCode()) { * // Handle and log exception * return; * } diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/TransactionalBatchOperationResult.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/TransactionalBatchOperationResult.java index 127f1ab0ac41..13eae1b26758 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/TransactionalBatchOperationResult.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/TransactionalBatchOperationResult.java @@ -17,13 +17,13 @@ @Beta(Beta.SinceVersion.V4_7_0) public final class TransactionalBatchOperationResult { - private String eTag; - private double requestCharge; + private final String eTag; + private final double requestCharge; private TResource item; - private ObjectNode resourceObject; - private int statusCode; - private Duration retryAfter; - private int subStatusCode; + private final ObjectNode resourceObject; + private final int statusCode; + private final Duration retryAfter; + private final int subStatusCode; /** * Instantiates a new Transactional batch operation result using a TransactionalBatchOperationResult. diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/batch/ItemBatchOperation.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/batch/ItemBatchOperation.java index eda92068dcd4..0db9ba62c406 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/batch/ItemBatchOperation.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/batch/ItemBatchOperation.java @@ -19,12 +19,12 @@ public final class ItemBatchOperation { private TResource resource; - private String id; - private int operationIndex; - private PartitionKey partitionKey; + private final String id; + private final int operationIndex; + private final PartitionKey partitionKey; private String partitionKeyJson; private final OperationType operationType; - private RequestOptions requestOptions; + private final RequestOptions requestOptions; private ItemBatchOperation( final OperationType operationType, diff --git a/sdk/cosmos/azure-cosmos/src/test/java/com/azure/cosmos/BatchTestBase.java b/sdk/cosmos/azure-cosmos/src/test/java/com/azure/cosmos/BatchTestBase.java index 9aa1f3701622..e07ed3fa01f5 100644 --- a/sdk/cosmos/azure-cosmos/src/test/java/com/azure/cosmos/BatchTestBase.java +++ b/sdk/cosmos/azure-cosmos/src/test/java/com/azure/cosmos/BatchTestBase.java @@ -11,6 +11,7 @@ import com.azure.cosmos.rx.TestSuiteBase; import com.fasterxml.jackson.annotation.JsonProperty; import io.netty.handler.codec.http.HttpResponseStatus; +import org.assertj.core.api.Assertions; import java.util.Objects; import java.util.Random; @@ -76,9 +77,7 @@ void verifyNotFound(CosmosContainer container, TestDoc doc) { try { CosmosItemResponse response = container.readItem(id, partitionKey, TestDoc.class); - - // Gateway returns response instead of exception - assertThat(response.getStatusCode()).isEqualTo(HttpResponseStatus.NOT_FOUND.code()); + Assertions.fail("Should throw NOT_FOUND exception"); } catch (CosmosException ex) { assertThat(ex.getStatusCode()).isEqualTo(HttpResponseStatus.NOT_FOUND.code()); } diff --git a/sdk/cosmos/azure-cosmos/src/test/java/com/azure/cosmos/TransactionalBatchTest.java b/sdk/cosmos/azure-cosmos/src/test/java/com/azure/cosmos/TransactionalBatchTest.java index d4694b4ec5fe..01a950241a93 100644 --- a/sdk/cosmos/azure-cosmos/src/test/java/com/azure/cosmos/TransactionalBatchTest.java +++ b/sdk/cosmos/azure-cosmos/src/test/java/com/azure/cosmos/TransactionalBatchTest.java @@ -68,7 +68,6 @@ public void batchOrdered() { this.verifyByRead(container, replaceDoc); } - @Test(groups = {"simple"}, timeOut = TIMEOUT) public void batchMultipleItemExecution() { CosmosContainer container = this.batchContainer; @@ -218,7 +217,7 @@ public void batchWithTooManyOperationsTest() { } } - @Test(groups = {"emulator"}, timeOut = TIMEOUT * 10) + @Test(groups = {"simple"}, timeOut = TIMEOUT * 10) public void batchLargerThanServerRequest() { int operationCount = 20; int appxDocSize = (MAX_DIRECT_MODE_BATCH_REQUEST_BODY_SIZE_IN_BYTES * 11) / operationCount; @@ -240,7 +239,7 @@ public void batchLargerThanServerRequest() { } } - @Test(groups = {"emulator"}, timeOut = TIMEOUT * 10) + @Test(groups = {"simple"}, timeOut = TIMEOUT * 10) public void batchServerResponseTooLarge() { int operationCount = 10; int appxDocSizeInBytes = 1 * 1024 * 1024; From 3e8dbc18e8e0707ef90769d2521957e3fa4f18ab Mon Sep 17 00:00:00 2001 From: Rakesh Kumar Date: Tue, 6 Oct 2020 02:16:23 +0530 Subject: [PATCH 12/22] Minor fix Signed-off-by: Rakesh Kumar --- .../cosmos/implementation/RxDocumentClientImpl.java | 2 +- .../implementation/batch/BatchResponseParser.java | 13 ++++++------- .../batch/TransactionalBatchResponseTests.java | 8 ++++---- 3 files changed, 11 insertions(+), 12 deletions(-) diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/RxDocumentClientImpl.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/RxDocumentClientImpl.java index 830f64221d38..25f83a54db86 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/RxDocumentClientImpl.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/RxDocumentClientImpl.java @@ -2484,7 +2484,7 @@ private Mono executeBatchRequestInternal(String coll }); return responseObservable - .flatMap(serviceResponse -> BatchResponseParser.fromDocumentServiceResponseAsync(serviceResponse, serverBatchRequest, true)); + .map(serviceResponse -> BatchResponseParser.fromDocumentServiceResponse(serviceResponse, serverBatchRequest, true)); } catch (Exception ex) { logger.debug("Failure in executing a batch due to [{}]", ex.getMessage(), ex); diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/batch/BatchResponseParser.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/batch/BatchResponseParser.java index bbf50ba96c2b..5bcb5a8e601f 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/batch/BatchResponseParser.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/batch/BatchResponseParser.java @@ -15,7 +15,6 @@ import io.netty.handler.codec.http.HttpResponseStatus; import org.slf4j.Logger; import org.slf4j.LoggerFactory; -import reactor.core.publisher.Mono; import java.io.IOException; import java.time.Duration; @@ -29,16 +28,16 @@ public final class BatchResponseParser { private final static Logger logger = LoggerFactory.getLogger(BatchResponseParser.class); private final static char HYBRID_V1 = 129; - /** Creates a transactional batch response} from a response message + /** Creates a transactional batch response from a documentServiceResponse. * * @param documentServiceResponse the {@link RxDocumentServiceResponse response message}. * @param request the {@link ServerBatchRequest batch request} that produced {@code message}. * @param shouldPromoteOperationStatus indicates whether the operation status should be promoted. * - * @return a Mono that provides the {@link TransactionalBatchResponse transactional batch response} created + * @return the {@link TransactionalBatchResponse transactional batch response} created * from {@link RxDocumentServiceResponse message} when the batch operation completes. */ - public static Mono fromDocumentServiceResponseAsync( + public static TransactionalBatchResponse fromDocumentServiceResponse( final RxDocumentServiceResponse documentServiceResponse, final ServerBatchRequest request, final boolean shouldPromoteOperationStatus) { @@ -47,7 +46,7 @@ public static Mono fromDocumentServiceResponseAsync( final byte[] responseContent = documentServiceResponse.getResponseBodyAsByteArray(); if (responseContent != null && responseContent.length > 0) { - response = BatchResponseParser.populateFromResponseContentAsync(documentServiceResponse, request, shouldPromoteOperationStatus); + response = BatchResponseParser.populateFromResponseContent(documentServiceResponse, request, shouldPromoteOperationStatus); if (response == null) { // Convert any payload read failures as InternalServerError @@ -96,10 +95,10 @@ public static Mono fromDocumentServiceResponseAsync( checkState(response.size() == request.getOperations().size(), "Number of responses should be equal to number of operations in request."); - return Mono.just(response); + return response; } - private static TransactionalBatchResponse populateFromResponseContentAsync( + private static TransactionalBatchResponse populateFromResponseContent( final RxDocumentServiceResponse documentServiceResponse, final ServerBatchRequest request, final boolean shouldPromoteOperationStatus) { diff --git a/sdk/cosmos/azure-cosmos/src/test/java/com/azure/cosmos/implementation/batch/TransactionalBatchResponseTests.java b/sdk/cosmos/azure-cosmos/src/test/java/com/azure/cosmos/implementation/batch/TransactionalBatchResponseTests.java index 1a547b33eb3f..5189fb975e17 100644 --- a/sdk/cosmos/azure-cosmos/src/test/java/com/azure/cosmos/implementation/batch/TransactionalBatchResponseTests.java +++ b/sdk/cosmos/azure-cosmos/src/test/java/com/azure/cosmos/implementation/batch/TransactionalBatchResponseTests.java @@ -71,10 +71,10 @@ public void validateAllSetValuesInResponse() { new ArrayList<>(headers.entrySet()), responseContent.getBytes(StandardCharsets.UTF_8)); - TransactionalBatchResponse batchResponse = BatchResponseParser.fromDocumentServiceResponseAsync( + TransactionalBatchResponse batchResponse = BatchResponseParser.fromDocumentServiceResponse( new RxDocumentServiceResponse(null, storeResponse), batchRequest, - true).block(); + true); // Validate response fields assertThat(batchResponse.getActivityId()).isEqualTo(activityId); @@ -126,10 +126,10 @@ public void validateEmptyHeaderInResponse() { new ArrayList<>(), responseContent.getBytes(StandardCharsets.UTF_8)); - TransactionalBatchResponse batchResponse = BatchResponseParser.fromDocumentServiceResponseAsync( + TransactionalBatchResponse batchResponse = BatchResponseParser.fromDocumentServiceResponse( new RxDocumentServiceResponse(null, storeResponse), batchRequest, - true).block(); + true); // Validate response fields assertThat(batchResponse.getStatusCode()).isEqualTo(HttpResponseStatus.OK.code()); From ed02b3eb706345d4a433016c11854d28d823996b Mon Sep 17 00:00:00 2001 From: Rakesh Kumar Date: Tue, 6 Oct 2020 18:55:50 +0530 Subject: [PATCH 13/22] Session token fix and test cases Signed-off-by: Rakesh Kumar --- .../implementation/RxDocumentClientImpl.java | 2 +- .../batch/BatchResponseParser.java | 3 +- .../batch/ServerBatchRequest.java | 8 +- .../directconnectivity/ConsistencyWriter.java | 6 +- .../java/com/azure/cosmos/BatchTestBase.java | 76 +++++- .../TransactionalBatchAsyncContainerTest.java | 125 +++++++++ .../azure/cosmos/TransactionalBatchTest.java | 239 +++++++++++++++--- 7 files changed, 409 insertions(+), 50 deletions(-) create mode 100644 sdk/cosmos/azure-cosmos/src/test/java/com/azure/cosmos/TransactionalBatchAsyncContainerTest.java diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/RxDocumentClientImpl.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/RxDocumentClientImpl.java index 25f83a54db86..a03f17c1f240 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/RxDocumentClientImpl.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/RxDocumentClientImpl.java @@ -1272,7 +1272,7 @@ private Mono getBatchDocumentRequest(DocumentClientRet checkNotNull(serverBatchRequest, "expected non null serverBatchRequest"); Instant serializationStartTimeUTC = Instant.now(); - ByteBuffer content = ByteBuffer.wrap(serverBatchRequest.transferRequestBody().getBytes(StandardCharsets.UTF_8)); + ByteBuffer content = ByteBuffer.wrap(serverBatchRequest.getRequestBody().getBytes(StandardCharsets.UTF_8)); Instant serializationEndTimeUTC = Instant.now(); SerializationDiagnosticsContext.SerializationDiagnostics serializationDiagnostics = new SerializationDiagnosticsContext.SerializationDiagnostics( diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/batch/BatchResponseParser.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/batch/BatchResponseParser.java index 5bcb5a8e601f..5162fb264e05 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/batch/BatchResponseParser.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/batch/BatchResponseParser.java @@ -133,7 +133,8 @@ private static TransactionalBatchResponse populateFromResponseContent( // Status code of the exact operation which failed. if (responseStatusCode == HttpResponseStatus.MULTI_STATUS.code() && shouldPromoteOperationStatus) { for (TransactionalBatchOperationResult result : results) { - if (result.getStatusCode()!= HttpResponseStatus.FAILED_DEPENDENCY.code()) { + if (result.getStatusCode() != HttpResponseStatus.FAILED_DEPENDENCY.code() && + result.getStatusCode() >= 400) { responseStatusCode = result.getStatusCode(); responseSubStatusCode = result.getSubStatusCode(); break; diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/batch/ServerBatchRequest.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/batch/ServerBatchRequest.java index c4f40aa6d8c7..41ac86a20363 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/batch/ServerBatchRequest.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/batch/ServerBatchRequest.java @@ -80,14 +80,10 @@ final List> createBodyOfBatchRequest(final List writePrivateAsync( }).flatMap(primaryUri -> { try { primaryURI.set(primaryUri); - if (this.useMultipleWriteLocations && + if ((this.useMultipleWriteLocations || request.getOperationType() == OperationType.Batch) && RequestHelper.getConsistencyLevelToUse(this.serviceConfigReader, request) == ConsistencyLevel.SESSION) { // Set session token to ensure session consistency for write requests - // when writes can be issued to multiple locations + // 1. when writes can be issued to multiple locations + // 2. When we have Batch requests, since it can have Reads in it. SessionTokenHelper.setPartitionLocalSessionToken(request, this.sessionContainer); } else { // When writes can only go to single location, there is no reason diff --git a/sdk/cosmos/azure-cosmos/src/test/java/com/azure/cosmos/BatchTestBase.java b/sdk/cosmos/azure-cosmos/src/test/java/com/azure/cosmos/BatchTestBase.java index e07ed3fa01f5..78af0230314f 100644 --- a/sdk/cosmos/azure-cosmos/src/test/java/com/azure/cosmos/BatchTestBase.java +++ b/sdk/cosmos/azure-cosmos/src/test/java/com/azure/cosmos/BatchTestBase.java @@ -5,6 +5,8 @@ import com.azure.cosmos.implementation.ISessionToken; import com.azure.cosmos.implementation.SessionTokenHelper; +import com.azure.cosmos.implementation.VectorSessionToken; +import com.azure.cosmos.implementation.apachecommons.collections.map.UnmodifiableMap; import com.azure.cosmos.implementation.apachecommons.lang.StringUtils; import com.azure.cosmos.models.CosmosItemResponse; import com.azure.cosmos.models.PartitionKey; @@ -12,7 +14,10 @@ import com.fasterxml.jackson.annotation.JsonProperty; import io.netty.handler.codec.http.HttpResponseStatus; import org.assertj.core.api.Assertions; +import org.assertj.core.data.Offset; +import java.lang.reflect.Constructor; +import java.lang.reflect.Field; import java.util.Objects; import java.util.Random; import java.util.UUID; @@ -41,6 +46,13 @@ void createJsonTestDocs(CosmosContainer container) { this.TestDocPk1ExistingD = this.createJsonTestDoc(container, this.partitionKey1); } + void createJsonTestDocs(CosmosAsyncContainer container) { + this.TestDocPk1ExistingA = this.createJsonTestDoc(container, this.partitionKey1); + this.TestDocPk1ExistingB = this.createJsonTestDoc(container, this.partitionKey1); + this.TestDocPk1ExistingC = this.createJsonTestDoc(container, this.partitionKey1); + this.TestDocPk1ExistingD = this.createJsonTestDoc(container, this.partitionKey1); + } + TestDoc populateTestDoc(String partitionKey) { return populateTestDoc(partitionKey, 20); } @@ -98,15 +110,77 @@ TestDoc createJsonTestDoc(CosmosContainer container, String partitionKey, int mi return doc; } + private TestDoc createJsonTestDoc(CosmosAsyncContainer container, String partitionKey) { + return createJsonTestDoc(container, partitionKey, 20); + } + + TestDoc createJsonTestDoc(CosmosAsyncContainer container, String partitionKey, int minDesiredSize) { + TestDoc doc = this.populateTestDoc(partitionKey, minDesiredSize); + CosmosItemResponse createResponse = container.createItem(doc, this.getPartitionKey(partitionKey), null).block(); + assertThat(createResponse.getStatusCode()).isEqualTo(HttpResponseStatus.CREATED.code()); + return doc; + } + public Random getRandom() { return random; } ISessionToken getSessionToken(String sessionToken) { - String[] tokenParts = sessionToken.split(":"); + String[] tokenParts = org.apache.commons.lang3.StringUtils.split(sessionToken, ':'); return SessionTokenHelper.parse(tokenParts[1]); } + String getDifferentLSNToken(String token, long lsnDifferent) throws Exception { + String[] tokenParts = org.apache.commons.lang3.StringUtils.split(token, ':'); + ISessionToken sessionToken = SessionTokenHelper.parse(tokenParts[1]); + ISessionToken differentSessionToken = createSessionToken(sessionToken, sessionToken.getLSN() + lsnDifferent); + return String.format("%s:%s", tokenParts[0], differentSessionToken.convertToString()); + } + + @SuppressWarnings({"unchecked", "rawtypes"}) + private static ISessionToken createSessionToken(ISessionToken from, long globalLSN) throws Exception { + // Creates session token with specified GlobalLSN + if (from instanceof VectorSessionToken) { + VectorSessionToken fromSessionToken = (VectorSessionToken) from; + Field fieldVersion = VectorSessionToken.class.getDeclaredField("version"); + fieldVersion.setAccessible(true); + Long version = (Long) fieldVersion.get(fromSessionToken); + + Field fieldLocalLsnByRegion = VectorSessionToken.class.getDeclaredField("localLsnByRegion"); + fieldLocalLsnByRegion.setAccessible(true); + UnmodifiableMap localLsnByRegion = (UnmodifiableMap) fieldLocalLsnByRegion.get(fromSessionToken); + + Constructor constructor = VectorSessionToken.class.getDeclaredConstructor(long.class, long.class, UnmodifiableMap.class); + constructor.setAccessible(true); + VectorSessionToken vectorSessionToken = constructor.newInstance(version, globalLSN, localLsnByRegion); + return vectorSessionToken; + } else { + throw new IllegalArgumentException(); + } + } + + void verifyBatchProcessed(TransactionalBatchResponse batchResponse, int numberOfOperations) { + this.verifyBatchProcessed(batchResponse, numberOfOperations, HttpResponseStatus.OK); + } + + void verifyBatchProcessed(TransactionalBatchResponse batchResponse, int numberOfOperations, HttpResponseStatus expectedStatusCode) { + assertThat(batchResponse).isNotNull(); + assertThat(batchResponse.getStatusCode()) + .as("Batch server response had StatusCode {0} instead of {1} expected and had ErrorMessage {2}", + batchResponse.getStatusCode(), expectedStatusCode.code()) + .isEqualTo(expectedStatusCode.code()); + + assertThat(batchResponse.size()).isEqualTo(numberOfOperations); + assertThat(batchResponse.getRequestCharge()).isPositive(); + assertThat(batchResponse.getDiagnostics().toString()).isNotEmpty(); + + // Allow a delta since we round both the total charge and the individual operation + // charges to 2 decimal places. + assertThat(batchResponse.getRequestCharge()) + .isCloseTo(batchResponse.getResults().stream().mapToDouble(TransactionalBatchOperationResult::getRequestCharge).sum(), + Offset.offset(0.1)); + } + public static class TestDoc { public String id; public int cost; diff --git a/sdk/cosmos/azure-cosmos/src/test/java/com/azure/cosmos/TransactionalBatchAsyncContainerTest.java b/sdk/cosmos/azure-cosmos/src/test/java/com/azure/cosmos/TransactionalBatchAsyncContainerTest.java new file mode 100644 index 000000000000..4ffdd5975da4 --- /dev/null +++ b/sdk/cosmos/azure-cosmos/src/test/java/com/azure/cosmos/TransactionalBatchAsyncContainerTest.java @@ -0,0 +1,125 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +package com.azure.cosmos; + +import com.azure.cosmos.implementation.HttpConstants; +import com.azure.cosmos.models.CosmosItemResponse; +import io.netty.handler.codec.http.HttpResponseStatus; +import org.assertj.core.api.Assertions; +import org.testng.annotations.AfterClass; +import org.testng.annotations.BeforeClass; +import org.testng.annotations.Factory; +import org.testng.annotations.Test; +import reactor.core.publisher.Mono; + +import static org.assertj.core.api.Assertions.assertThat; + +public class TransactionalBatchAsyncContainerTest extends BatchTestBase { + + private CosmosAsyncClient batchClient; + private CosmosAsyncContainer batchAsyncContainer; + + @Factory(dataProvider = "clientBuildersWithDirectSession") + public TransactionalBatchAsyncContainerTest(CosmosClientBuilder clientBuilder) { + super(clientBuilder); + } + + @BeforeClass(groups = {"simple"}, timeOut = SETUP_TIMEOUT) + public void before_TransactionalBatchAsyncContainerTest() { + assertThat(this.batchClient).isNull(); + this.batchClient = getClientBuilder().buildAsyncClient(); + batchAsyncContainer = getSharedMultiPartitionCosmosContainer(this.batchClient); + } + + @AfterClass(groups = {"simple"}, timeOut = SHUTDOWN_TIMEOUT, alwaysRun = true) + public void afterClass() { + safeCloseAsync(this.batchClient); + } + + @Test(groups = {"simple"}, timeOut = TIMEOUT) + public void batchExecutionRepeat() { + CosmosAsyncContainer container = this.batchAsyncContainer; + this.createJsonTestDocs(container); + + TestDoc firstDoc = this.populateTestDoc(this.partitionKey1); + TestDoc replaceDoc = this.getTestDocCopy(firstDoc); + replaceDoc.setCost(replaceDoc.getCost() + 1); + + Mono batchResponseMono = batchAsyncContainer.executeTransactionalBatch( + TransactionalBatch.createTransactionalBatch(this.getPartitionKey(this.partitionKey1)) + .createItem(firstDoc) + .replaceItem(replaceDoc.getId(), replaceDoc)); + + TransactionalBatchResponse batchResponse1 = batchResponseMono.block(); + this.verifyBatchProcessed(batchResponse1, 2); + + assertThat(batchResponse1.get(0).getStatusCode()).isEqualTo(HttpResponseStatus.CREATED.code()); + assertThat(batchResponse1.get(1).getStatusCode()).isEqualTo(HttpResponseStatus.OK.code()); + + // Block again. + TransactionalBatchResponse batchResponse2 = batchResponseMono.block(); + assertThat(batchResponse2.getStatusCode()).isEqualTo(HttpResponseStatus.CONFLICT.code()); + assertThat(batchResponse2.get(0).getStatusCode()).isEqualTo(HttpResponseStatus.CONFLICT.code()); + assertThat(batchResponse2.get(1).getStatusCode()).isEqualTo(HttpResponseStatus.FAILED_DEPENDENCY.code()); + } + + @Test(groups = {"simple"}, timeOut = TIMEOUT * 100) + public void batchInvalidSessionToken() throws Exception { + CosmosAsyncContainer container = batchAsyncContainer; + this.createJsonTestDocs(container); + + CosmosItemResponse readResponse = container.readItem( + this.TestDocPk1ExistingC.getId(), + this.getPartitionKey(this.partitionKey1), + TestDoc.class).block(); + + assertThat(readResponse.getStatusCode()).isEqualTo(HttpResponseStatus.OK.code()); + String invalidSessionToken = this.getDifferentLSNToken(readResponse.getSessionToken(), 2000); + + { + // Batch without Read operation + TestDoc testDocToCreate = this.populateTestDoc(this.partitionKey1); + TestDoc testDocToReplace = this.getTestDocCopy(this.TestDocPk1ExistingA); + testDocToReplace.setCost(testDocToReplace.getCost() + 1); + TestDoc testDocToUpsert = this.populateTestDoc(this.partitionKey1); + + TransactionalBatchResponse batchResponse = container.executeTransactionalBatch( + TransactionalBatch.createTransactionalBatch(this.getPartitionKey(this.partitionKey1)) + .createItem(testDocToCreate) + .replaceItem(testDocToReplace.getId(), testDocToReplace) + .upsertItem(testDocToUpsert) + .deleteItem(this.TestDocPk1ExistingC.getId()), new TransactionalBatchRequestOptions().setSessionToken(invalidSessionToken)).block(); + + this.verifyBatchProcessed(batchResponse, 4); + + assertThat(batchResponse.get(0).getStatusCode()).isEqualTo(HttpResponseStatus.CREATED.code()); + assertThat(batchResponse.get(1).getStatusCode()).isEqualTo(HttpResponseStatus.OK.code()); + assertThat(batchResponse.get(2).getStatusCode()).isEqualTo(HttpResponseStatus.CREATED.code()); + assertThat(batchResponse.get(3).getStatusCode()).isEqualTo(HttpResponseStatus.NO_CONTENT.code()); + } + + { + // Batch with Read operation + TestDoc testDocToCreate = this.populateTestDoc(this.partitionKey1); + TestDoc testDocToReplace = this.getTestDocCopy(this.TestDocPk1ExistingB); + testDocToReplace.setCost(testDocToReplace.getCost() + 1); + TestDoc testDocToUpsert = this.populateTestDoc(this.partitionKey1); + + try { + container.executeTransactionalBatch( + TransactionalBatch.createTransactionalBatch(this.getPartitionKey(this.partitionKey1)) + .createItem(testDocToCreate) + .replaceItem(testDocToReplace.getId(), testDocToReplace) + .upsertItem(testDocToUpsert) + .deleteItem(this.TestDocPk1ExistingD.getId()) + .readItem(this.TestDocPk1ExistingA.getId()), new TransactionalBatchRequestOptions().setSessionToken(invalidSessionToken)).block(); + + Assertions.fail("Should throw NOT_FOUND/READ_SESSION_NOT_AVAILABLE exception"); + } catch (CosmosException ex) { + assertThat(ex.getStatusCode()).isEqualTo(HttpResponseStatus.NOT_FOUND.code()); + assertThat(ex.getSubStatusCode()).isEqualTo(HttpConstants.SubStatusCodes.READ_SESSION_NOT_AVAILABLE); + } + } + } +} diff --git a/sdk/cosmos/azure-cosmos/src/test/java/com/azure/cosmos/TransactionalBatchTest.java b/sdk/cosmos/azure-cosmos/src/test/java/com/azure/cosmos/TransactionalBatchTest.java index 01a950241a93..ba5afb158c73 100644 --- a/sdk/cosmos/azure-cosmos/src/test/java/com/azure/cosmos/TransactionalBatchTest.java +++ b/sdk/cosmos/azure-cosmos/src/test/java/com/azure/cosmos/TransactionalBatchTest.java @@ -9,7 +9,6 @@ import com.azure.cosmos.models.CosmosItemResponse; import io.netty.handler.codec.http.HttpResponseStatus; import org.assertj.core.api.Assertions; -import org.assertj.core.data.Offset; import org.testng.annotations.AfterClass; import org.testng.annotations.BeforeClass; import org.testng.annotations.Factory; @@ -164,38 +163,222 @@ public void batchItemETagTest() { } @Test(groups = {"simple"}, timeOut = TIMEOUT) - public void batchItemSessionTokenTest() { + public void batchSessionTokenPropertiesTest() throws Exception { CosmosContainer container = batchContainer; this.createJsonTestDocs(container); - TestDoc testDocToCreate = this.populateTestDoc(this.partitionKey1); + TestDoc sampleDoc = this.populateTestDoc(this.partitionKey1); - BatchTestBase.TestDoc testDocToReplace = this.getTestDocCopy(this.TestDocPk1ExistingA); - testDocToReplace.setCost(testDocToReplace.getCost() + 1); + CosmosItemResponse createResponse = container.createItem( + sampleDoc, + this.getPartitionKey(this.partitionKey1), + null); + + String ownerIdCreate = createResponse.getResponseHeaders().get(HttpConstants.HttpHeaders.OWNER_ID); + assertThat(createResponse.getStatusCode()).isEqualTo(HttpResponseStatus.CREATED.code()); + assertThat(ownerIdCreate).isNotEmpty(); CosmosItemResponse readResponse = container.readItem( - this.TestDocPk1ExistingA.getId(), + this.TestDocPk1ExistingC.getId(), this.getPartitionKey(this.partitionKey1), TestDoc.class); assertThat(readResponse.getStatusCode()).isEqualTo(HttpResponseStatus.OK.code()); - ISessionToken beforeRequestSessionToken = this.getSessionToken(readResponse.getResponseHeaders().get(HttpConstants.HttpHeaders.SESSION_TOKEN)); + ISessionToken beforeRequestSessionToken = this.getSessionToken(readResponse.getSessionToken()); - TransactionalBatchResponse batchResponse = container.executeTransactionalBatch( + String readEtagValue = readResponse.getETag(); + TransactionalBatchItemRequestOptions readRequestOption = new TransactionalBatchItemRequestOptions() + .setIfMatchETag(readEtagValue); + + String oldSessionToken = this.getDifferentLSNToken(readResponse.getSessionToken(), -10); + + { + // Batch with only Read operation + TransactionalBatchResponse batchResponse = container.executeTransactionalBatch( + TransactionalBatch.createTransactionalBatch(this.getPartitionKey(this.partitionKey1)) + .readItem(this.TestDocPk1ExistingA.getId()) + .readItem(this.TestDocPk1ExistingC.getId(), readRequestOption), + new TransactionalBatchRequestOptions().setSessionToken(oldSessionToken)); + + this.verifyBatchProcessed(batchResponse, 2, HttpResponseStatus.MULTI_STATUS); + + assertThat(batchResponse.get(0).getStatusCode()).isEqualTo(HttpResponseStatus.OK.code()); + assertThat(batchResponse.get(1).getStatusCode()).isEqualTo(HttpResponseStatus.NOT_MODIFIED.code()); + + ISessionToken afterRequestSessionToken = this.getSessionToken(batchResponse.getSessionToken()); + + assertThat(afterRequestSessionToken.getLSN()) + .as("Response session token should be more than or equal to request session token") + .isGreaterThanOrEqualTo(beforeRequestSessionToken.getLSN()); + + String ownerIdBatch = batchResponse.getResponseHeaders().get(HttpConstants.HttpHeaders.OWNER_ID); + assertThat(ownerIdBatch).isNotEmpty(); + assertThat(ownerIdBatch).isEqualTo(ownerIdCreate); + } + + { + // Batch with write-read operations + TestDoc testDocToCreate = this.populateTestDoc(this.partitionKey1); + + TestDoc testDocToReplace = this.getTestDocCopy(this.TestDocPk1ExistingB); + testDocToReplace.setCost(testDocToReplace.getCost() + 1); + TestDoc testDocToUpsert = this.populateTestDoc(this.partitionKey1); + + TransactionalBatchResponse batchResponse = container.executeTransactionalBatch( + TransactionalBatch.createTransactionalBatch(this.getPartitionKey(this.partitionKey1)) + .createItem(testDocToCreate) + .replaceItem(testDocToReplace.getId(), testDocToReplace) + .upsertItem(testDocToUpsert) + .deleteItem(this.TestDocPk1ExistingD.getId()) + .readItem(this.TestDocPk1ExistingA.getId()) + .readItem(this.TestDocPk1ExistingC.getId(), readRequestOption), + new TransactionalBatchRequestOptions().setSessionToken(oldSessionToken)); + + this.verifyBatchProcessed(batchResponse, 6, HttpResponseStatus.MULTI_STATUS); + + assertThat(batchResponse.get(0).getStatusCode()).isEqualTo(HttpResponseStatus.CREATED.code()); + assertThat(batchResponse.get(1).getStatusCode()).isEqualTo(HttpResponseStatus.OK.code()); + assertThat(batchResponse.get(2).getStatusCode()).isEqualTo(HttpResponseStatus.CREATED.code()); + assertThat(batchResponse.get(3).getStatusCode()).isEqualTo(HttpResponseStatus.NO_CONTENT.code()); + assertThat(batchResponse.get(4).getStatusCode()).isEqualTo(HttpResponseStatus.OK.code()); + assertThat(batchResponse.get(5).getStatusCode()).isEqualTo(HttpResponseStatus.NOT_MODIFIED.code()); + + ISessionToken afterRequestSessionToken = this.getSessionToken(batchResponse.getSessionToken()); + + assertThat(afterRequestSessionToken.getLSN()) + .as("Response session token should be more than request session token") + .isGreaterThan(beforeRequestSessionToken.getLSN()); + + String ownerIdBatch = batchResponse.getResponseHeaders().get(HttpConstants.HttpHeaders.OWNER_ID); + assertThat(ownerIdBatch).isNotEmpty(); + assertThat(ownerIdBatch).isEqualTo(ownerIdCreate); + } + } + + @Test(groups = {"simple"}, timeOut = TIMEOUT) + public void batchErrorSessionToken() { + CosmosContainer container = batchContainer; + this.createJsonTestDocs(container); + + ISessionToken readResponseNotExistsToken = null; + try { + container.readItem( + UUID.randomUUID().toString(), + this.getPartitionKey(this.partitionKey1), + TestDoc.class); + } catch (CosmosException ex) { + readResponseNotExistsToken = this.getSessionToken(ex.getResponseHeaders().get(HttpConstants.HttpHeaders.SESSION_TOKEN)); + + // When this is changed to return non null, batch needs to be modified too. + String ownerIdRead = ex.getResponseHeaders().get(HttpConstants.HttpHeaders.OWNER_ID); + assertThat(ownerIdRead).isNull(); + } + + { + // Only errored read + TransactionalBatchResponse batchResponse = container.executeTransactionalBatch( TransactionalBatch.createTransactionalBatch(this.getPartitionKey(this.partitionKey1)) - .createItem(testDocToCreate) - .replaceItem(testDocToReplace.getId(), testDocToReplace)); + .readItem(UUID.randomUUID().toString())); - this.verifyBatchProcessed(batchResponse, 2); + assertThat(batchResponse.getStatusCode()).isEqualTo(HttpResponseStatus.NOT_FOUND.code()); + assertThat(batchResponse.get(0).getStatusCode()).isEqualTo(HttpResponseStatus.NOT_FOUND.code()); - assertThat(batchResponse.get(0).getStatusCode()).isEqualTo(HttpResponseStatus.CREATED.code()); - assertThat(batchResponse.get(1).getStatusCode()).isEqualTo(HttpResponseStatus.OK.code()); + String ownerIdBatch = batchResponse.getResponseHeaders().get(HttpConstants.HttpHeaders.OWNER_ID); + assertThat(ownerIdBatch).isNull(); + + ISessionToken batchResponseToken = this.getSessionToken(batchResponse.getSessionToken()); + + assertThat(batchResponseToken.getLSN()) + .as("Response session token should be more than or equal to request session token") + .isGreaterThanOrEqualTo(readResponseNotExistsToken.getLSN()); + } + + { + // One valid read one error read + TransactionalBatchResponse batchResponse = container.executeTransactionalBatch( + TransactionalBatch.createTransactionalBatch(this.getPartitionKey(this.partitionKey1)) + .readItem(this.TestDocPk1ExistingA.getId()) + .readItem(UUID.randomUUID().toString())); + + assertThat(batchResponse.getStatusCode()).isEqualTo(HttpResponseStatus.NOT_FOUND.code()); + assertThat(batchResponse.get(0).getStatusCode()).isEqualTo(HttpResponseStatus.FAILED_DEPENDENCY.code()); + assertThat(batchResponse.get(1).getStatusCode()).isEqualTo(HttpResponseStatus.NOT_FOUND.code()); + + String ownerIdBatch = batchResponse.getResponseHeaders().get(HttpConstants.HttpHeaders.OWNER_ID); + assertThat(ownerIdBatch).isNull(); + + ISessionToken batchResponseToken = this.getSessionToken(batchResponse.getSessionToken()); + + assertThat(batchResponseToken.getLSN()) + .as("Response session token should be more than or equal to request session token") + .isGreaterThanOrEqualTo(readResponseNotExistsToken.getLSN()); + } + + { + // One error one valid read + TransactionalBatchResponse batchResponse = container.executeTransactionalBatch( + TransactionalBatch.createTransactionalBatch(this.getPartitionKey(this.partitionKey1)) + .readItem(UUID.randomUUID().toString()) + .readItem(this.TestDocPk1ExistingA.getId())); + + assertThat(batchResponse.getStatusCode()).isEqualTo(HttpResponseStatus.NOT_FOUND.code()); + assertThat(batchResponse.get(0).getStatusCode()).isEqualTo(HttpResponseStatus.NOT_FOUND.code()); + assertThat(batchResponse.get(1).getStatusCode()).isEqualTo(HttpResponseStatus.FAILED_DEPENDENCY.code()); + + String ownerIdBatch = batchResponse.getResponseHeaders().get(HttpConstants.HttpHeaders.OWNER_ID); + assertThat(ownerIdBatch).isNull(); + + ISessionToken batchResponseToken = this.getSessionToken(batchResponse.getSessionToken()); - ISessionToken afterRequestSessionToken = this.getSessionToken(batchResponse.getResponseHeaders().get(HttpConstants.HttpHeaders.SESSION_TOKEN)); - assertThat(afterRequestSessionToken.getLSN()) - .as("Response session token should be more than request session token") - .isGreaterThan(beforeRequestSessionToken.getLSN()); + assertThat(batchResponseToken.getLSN()) + .as("Response session token should be more than or equal to request session token") + .isGreaterThanOrEqualTo(readResponseNotExistsToken.getLSN()); + } + + { + // One valid write and one error + TestDoc testDocToCreate = this.populateTestDoc(this.partitionKey1); + TransactionalBatchResponse batchResponse = container.executeTransactionalBatch( + TransactionalBatch.createTransactionalBatch(this.getPartitionKey(this.partitionKey1)) + .createItem(testDocToCreate) + .readItem(UUID.randomUUID().toString())); + + assertThat(batchResponse.getStatusCode()).isEqualTo(HttpResponseStatus.NOT_FOUND.code()); + assertThat(batchResponse.get(0).getStatusCode()).isEqualTo(HttpResponseStatus.FAILED_DEPENDENCY.code()); + assertThat(batchResponse.get(1).getStatusCode()).isEqualTo(HttpResponseStatus.NOT_FOUND.code()); + + String ownerIdBatch = batchResponse.getResponseHeaders().get(HttpConstants.HttpHeaders.OWNER_ID); + assertThat(ownerIdBatch).isNull(); + + ISessionToken batchResponseToken = this.getSessionToken(batchResponse.getSessionToken()); + + assertThat(batchResponseToken.getLSN()) + .as("Response session token should be more than or equal to request session token") + .isGreaterThanOrEqualTo(readResponseNotExistsToken.getLSN()); + } + + { + // One error one valid write + TestDoc testDocToCreate = this.populateTestDoc(this.partitionKey1); + TransactionalBatchResponse batchResponse = container.executeTransactionalBatch( + TransactionalBatch.createTransactionalBatch(this.getPartitionKey(this.partitionKey1)) + .readItem(UUID.randomUUID().toString()) + .createItem(testDocToCreate)); + + assertThat(batchResponse.getStatusCode()).isEqualTo(HttpResponseStatus.NOT_FOUND.code()); + assertThat(batchResponse.get(0).getStatusCode()).isEqualTo(HttpResponseStatus.NOT_FOUND.code()); + assertThat(batchResponse.get(1).getStatusCode()).isEqualTo(HttpResponseStatus.FAILED_DEPENDENCY.code()); + + String ownerIdBatch = batchResponse.getResponseHeaders().get(HttpConstants.HttpHeaders.OWNER_ID); + assertThat(ownerIdBatch).isNull(); + + ISessionToken batchResponseToken = this.getSessionToken(batchResponse.getSessionToken()); + + assertThat(batchResponseToken.getLSN()) + .as("Response session token should be more than or equal to request session token") + .isGreaterThanOrEqualTo(readResponseNotExistsToken.getLSN()); + } } @Test(groups = {"simple"}, timeOut = TIMEOUT) @@ -407,26 +590,4 @@ private void runWithError( this.verifyNotFound(container, testDocToCreate); this.verifyNotFound(container, anotherTestDocToCreate); } - - private void verifyBatchProcessed(TransactionalBatchResponse batchResponse, int numberOfOperations) { - this.verifyBatchProcessed(batchResponse, numberOfOperations, HttpResponseStatus.OK); - } - - private void verifyBatchProcessed(TransactionalBatchResponse batchResponse, int numberOfOperations, HttpResponseStatus expectedStatusCode) { - assertThat(batchResponse).isNotNull(); - assertThat(batchResponse.getStatusCode()) - .as("Batch server response had StatusCode {0} instead of {1} expected and had ErrorMessage {2}", - batchResponse.getStatusCode(), expectedStatusCode.code()) - .isEqualTo(expectedStatusCode.code()); - - assertThat(batchResponse.size()).isEqualTo(numberOfOperations); - assertThat(batchResponse.getRequestCharge()).isPositive(); - assertThat(batchResponse.getDiagnostics().toString()).isNotEmpty(); - - // Allow a delta since we round both the total charge and the individual operation - // charges to 2 decimal places. - assertThat(batchResponse.getRequestCharge()) - .isCloseTo(batchResponse.getResults().stream().mapToDouble(TransactionalBatchOperationResult::getRequestCharge).sum(), - Offset.offset(0.1)); - } } From 6d8ea0d24ef2c8a8960aa87d7526b639ecf00830 Mon Sep 17 00:00:00 2001 From: Rakesh Kumar Date: Wed, 7 Oct 2020 02:25:57 +0530 Subject: [PATCH 14/22] code review changes Signed-off-by: Rakesh Kumar --- .../java/com/azure/cosmos/BridgeInternal.java | 6 +-- .../azure/cosmos/CosmosAsyncContainer.java | 6 +++ .../com/azure/cosmos/CosmosContainer.java | 6 +++ .../com/azure/cosmos/TransactionalBatch.java | 8 +-- .../TransactionalBatchOperationResult.java | 50 ++++++------------- .../cosmos/TransactionalBatchResponse.java | 41 +++------------ .../implementation/batch/BatchExecUtils.java | 2 +- .../batch/BatchResponseParser.java | 12 ++--- .../cosmos/BatchOperationResultTests.java | 38 ++------------ .../TransactionalBatchAsyncContainerTest.java | 3 -- .../azure/cosmos/TransactionalBatchTest.java | 19 +++---- .../batch/BatchResponsePayloadWriter.java | 12 ++--- .../TransactionalBatchResponseTests.java | 16 +++--- 13 files changed, 74 insertions(+), 145 deletions(-) diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/BridgeInternal.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/BridgeInternal.java index f07f8af91fd8..f63321c12cc0 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/BridgeInternal.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/BridgeInternal.java @@ -626,7 +626,7 @@ public static RequestOptions toRequestOptions(TransactionalBatchRequestOptions t } @Warning(value = INTERNAL_USE_ONLY_WARNING) - public static TransactionalBatchOperationResult createTransactionBatchResult( + public static TransactionalBatchOperationResult createTransactionBatchResult( String eTag, double requestCharge, ObjectNode resourceObject, @@ -634,7 +634,7 @@ public static TransactionalBatchOperationResult createTransactionBatchResult( Duration retryAfter, int subStatusCode) { - return new TransactionalBatchOperationResult<>( + return new TransactionalBatchOperationResult( eTag, requestCharge, resourceObject, @@ -662,7 +662,7 @@ public static TransactionalBatchResponse createTransactionBatchResponse( @Warning(value = INTERNAL_USE_ONLY_WARNING) public static void addTransactionBatchResultInResponse( TransactionalBatchResponse transactionalBatchResponse, - List> transactionalBatchOperationResult) { + List transactionalBatchOperationResult) { transactionalBatchResponse.addAll(transactionalBatchOperationResult); } diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncContainer.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncContainer.java index 8a2667e70853..8d0a3a151c69 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncContainer.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncContainer.java @@ -519,6 +519,9 @@ private T transform(Object object, Class classType) { * for the operation that caused the batch to abort, the value of this field * will indicate the cause of failure. *

+ * If there are issues such as request timeouts, Gone, session not available, network failure + * or if the service somehow returns 5xx then the Mono will return error instead of TransactionalBatchResponse. + *

* Use {@link TransactionalBatchResponse#isSuccessStatusCode} on the response returned to ensure that the * transactional batch succeeded. */ @@ -551,6 +554,9 @@ public Mono executeTransactionalBatch(TransactionalB * for the operation that caused the batch to abort, the value of this field * will indicate the cause of failure. *

+ * If there are issues such as request timeouts, Gone, session not available, network failure + * or if the service somehow returns 5xx then the Mono will return error instead of TransactionalBatchResponse. + *

* Use {@link TransactionalBatchResponse#isSuccessStatusCode} on the response returned to ensure that the * transactional batch succeeded. */ diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosContainer.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosContainer.java index fa991d5f2758..2b7379f1cc53 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosContainer.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosContainer.java @@ -489,6 +489,9 @@ public CosmosItemResponse deleteItem(T item, CosmosItemRequestOption * for the operation that caused the batch to abort, the value of this field * will indicate the cause of failure. *

+ * If there are issues such as request timeouts, Gone, session not available, network failure + * or if the service somehow returns 5xx then this will throw an exception instead of returning a TransactionalBatchResponse. + *

* Use {@link TransactionalBatchResponse#isSuccessStatusCode} on the response returned to ensure that the * transactional batch succeeded. */ @@ -521,6 +524,9 @@ public TransactionalBatchResponse executeTransactionalBatch(TransactionalBatch t * for the operation that caused the batch to abort, the value of this field * will indicate the cause of failure. *

+ * If there are issues such as request timeouts, Gone, session not available, network failure + * or if the service somehow returns 5xx then this will throw an exception instead of returning a TransactionalBatchResponse. + *

* Use {@link TransactionalBatchResponse#isSuccessStatusCode} on the response returned to ensure that the * transactional batch succeeded. */ diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/TransactionalBatch.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/TransactionalBatch.java index 059fdd5eb856..69e45055164d 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/TransactionalBatch.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/TransactionalBatch.java @@ -54,8 +54,8 @@ * * // Look up interested results - e.g., via typed access on operation results * - * TransactionalBatchOperationResult result = response.getOperationResultAtIndex(0, ToDoActivity.class); - * ToDoActivity readActivity = result.getItem(); + * TransactionalBatchOperationResult result = response.get(0); + * ToDoActivity readActivity = result.getItem(ToDoActivity.class); * * } * @@ -74,8 +74,8 @@ * List resultItems = new ArrayList(); * * for (int i = 0; i < response.size(); i++) { - * TransactionalBatchOperationResult result = response.getOperationResultAtIndex(0, ToDoActivity.class); - * resultItems.add(result.getItem()); + * TransactionalBatchOperationResult result = response.get(0); + * resultItems.add(result.getItem(ToDoActivity.class)); * } * * } diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/TransactionalBatchOperationResult.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/TransactionalBatchOperationResult.java index 13eae1b26758..8fc02acc73de 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/TransactionalBatchOperationResult.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/TransactionalBatchOperationResult.java @@ -3,6 +3,7 @@ package com.azure.cosmos; +import com.azure.cosmos.implementation.JsonSerializable; import com.azure.cosmos.util.Beta; import com.fasterxml.jackson.databind.node.ObjectNode; import java.time.Duration; @@ -11,48 +12,17 @@ /** * Represents a result for a specific operation that was part of a {@link TransactionalBatch} request. - * - * @param the type parameter */ @Beta(Beta.SinceVersion.V4_7_0) -public final class TransactionalBatchOperationResult { +public final class TransactionalBatchOperationResult { private final String eTag; private final double requestCharge; - private TResource item; private final ObjectNode resourceObject; private final int statusCode; private final Duration retryAfter; private final int subStatusCode; - /** - * Instantiates a new Transactional batch operation result using a TransactionalBatchOperationResult. - * - * @param other the other TransactionalBatchOperationResult. - */ - TransactionalBatchOperationResult(final TransactionalBatchOperationResult other) { - - checkNotNull(other, "expected non-null other"); - - this.statusCode = other.statusCode; - this.subStatusCode = other.subStatusCode; - this.eTag = other.eTag; - this.requestCharge = other.requestCharge; - this.retryAfter = other.retryAfter; - this.resourceObject = other.resourceObject; - } - - /** - * Instantiates a new Transactional batch operation result using a TransactionalBatchOperationResult and item. - * - * @param result the TransactionalBatchOperationResult. - * @param item the item. - */ - TransactionalBatchOperationResult(TransactionalBatchOperationResult result, TResource item) { - this(result); - this.item = item; - } - /** * Initializes a new instance of the {@link TransactionalBatchOperationResult} class. */ @@ -98,10 +68,20 @@ public double getRequestCharge() { /** * Gets the item associated with the current result. * + * @param the type parameter + * + * @param type class type for which deserialization is needed. + * * @return item associated with the current result. */ - public TResource getItem() { - return this.item; + public T getItem(final Class type) { + T item = null; + + if (this.getResourceObject() != null) { + item = new JsonSerializable(this.getResourceObject()).toObject(type); + } + + return item; } /** @@ -109,7 +89,7 @@ public TResource getItem() { * * @return the retry after */ - public Duration getRetryAfter() { + public Duration getRetryAfterDuration() { return this.retryAfter; } diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/TransactionalBatchResponse.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/TransactionalBatchResponse.java index 22d5391b1147..006b488336b3 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/TransactionalBatchResponse.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/TransactionalBatchResponse.java @@ -3,7 +3,6 @@ package com.azure.cosmos; -import com.azure.cosmos.implementation.JsonSerializable; import com.azure.cosmos.implementation.batch.BatchExecUtils; import com.azure.cosmos.util.Beta; @@ -13,7 +12,6 @@ import java.util.List; import java.util.Map; -import static com.azure.cosmos.implementation.guava25.base.Preconditions.checkArgument; import static com.azure.cosmos.implementation.guava25.base.Preconditions.checkNotNull; /** @@ -25,7 +23,7 @@ public class TransactionalBatchResponse { private final Map responseHeaders; private final int statusCode; private final String errorMessage; - private final List> results; + private final List results; private final int subStatusCode; private final CosmosDiagnostics cosmosDiagnostics; @@ -56,33 +54,6 @@ public class TransactionalBatchResponse { this.results = new ArrayList<>(); } - /** - * Gets the result of the operation at the provided index in the current {@link TransactionalBatchResponse batch}. - *

- * @param the type parameter. - * @param index 0-based index of the operation in the batch whose result needs to be returned. - * de-serialized, when present. - * @param type class type for which deserialization is needed. - * - * @return TransactionalBatchOperationResult containing the individual result of operation. - */ - public TransactionalBatchOperationResult getOperationResultAtIndex( - final int index, - final Class type) { - - checkArgument(index >= 0, "expected non-negative index"); - checkNotNull(type, "expected non-null type"); - - final TransactionalBatchOperationResult result = this.results.get(index); - T item = null; - - if (result.getResourceObject() != null) { - item = new JsonSerializable(result.getResourceObject()).toObject(type); - } - - return new TransactionalBatchOperationResult(result, item); - } - /** * Gets the diagnostics information for the current request to Azure Cosmos DB service. * @@ -172,8 +143,8 @@ public Map getResponseHeaders() { * * @return the amount of time to wait before retrying this or any other request due to throttling. */ - public Duration getRetryAfter() { - return BatchExecUtils.getRetryAfter(this.responseHeaders); + public Duration getRetryAfterDuration() { + return BatchExecUtils.getRetryAfterDuration(this.responseHeaders); } /** @@ -190,7 +161,7 @@ public int getSubStatusCode() { * * @return Results of operation in batch. */ - public List> getResults() { + public List getResults() { return this.results; } @@ -201,7 +172,7 @@ public List> getResults() { * * @return Result of operation at the provided index in the batch. */ - public TransactionalBatchOperationResult get(int index) { + public TransactionalBatchOperationResult get(int index) { return this.results.get(index); } @@ -218,7 +189,7 @@ public Duration getDuration() { return this.cosmosDiagnostics.getDuration(); } - boolean addAll(Collection> collection) { + boolean addAll(Collection collection) { return this.results.addAll(collection); } } diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/batch/BatchExecUtils.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/batch/BatchExecUtils.java index 7767325f4ffe..b2f7cb5a6852 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/batch/BatchExecUtils.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/batch/BatchExecUtils.java @@ -42,7 +42,7 @@ static String getStringOperationType(OperationType operationType) { return null; } - public static Duration getRetryAfter(Map responseHeaders) { + public static Duration getRetryAfterDuration(Map responseHeaders) { long retryIntervalInMilliseconds = 0; if (responseHeaders != null) { diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/batch/BatchResponseParser.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/batch/BatchResponseParser.java index 5162fb264e05..678185410996 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/batch/BatchResponseParser.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/batch/BatchResponseParser.java @@ -86,7 +86,7 @@ public static TransactionalBatchResponse fromDocumentServiceResponse( // When the overall response status code is TooManyRequests, propagate the RetryAfter into the individual operations. Duration retryAfterDuration = Duration.ZERO; if (responseStatusCode == HttpResponseStatus.TOO_MANY_REQUESTS.code()) { - retryAfterDuration = BatchExecUtils.getRetryAfter(documentServiceResponse.getResponseHeaders()); + retryAfterDuration = BatchExecUtils.getRetryAfterDuration(documentServiceResponse.getResponseHeaders()); } BatchResponseParser.createAndPopulateResults(response, request.getOperations(), retryAfterDuration); @@ -103,7 +103,7 @@ private static TransactionalBatchResponse populateFromResponseContent( final ServerBatchRequest request, final boolean shouldPromoteOperationStatus) { - final ArrayList> results = new ArrayList<>(request.getOperations().size()); + final ArrayList results = new ArrayList<>(request.getOperations().size()); final byte[] responseContent = documentServiceResponse.getResponseBodyAsByteArray(); if (responseContent[0] != (byte)HYBRID_V1) { @@ -113,7 +113,7 @@ private static TransactionalBatchResponse populateFromResponseContent( try { final ObjectNode[] objectNodes = mapper.readValue(responseContent, ObjectNode[].class); for (ObjectNode objectInArray : objectNodes) { - final TransactionalBatchOperationResult batchOperationResult = BatchResponseParser.createBatchOperationResultFromJson(objectInArray); + final TransactionalBatchOperationResult batchOperationResult = BatchResponseParser.createBatchOperationResultFromJson(objectInArray); results.add(batchOperationResult); } } catch (IOException ex) { @@ -132,7 +132,7 @@ private static TransactionalBatchResponse populateFromResponseContent( // Status code of the exact operation which failed. if (responseStatusCode == HttpResponseStatus.MULTI_STATUS.code() && shouldPromoteOperationStatus) { - for (TransactionalBatchOperationResult result : results) { + for (TransactionalBatchOperationResult result : results) { if (result.getStatusCode() != HttpResponseStatus.FAILED_DEPENDENCY.code() && result.getStatusCode() >= 400) { responseStatusCode = result.getStatusCode(); @@ -163,7 +163,7 @@ private static TransactionalBatchResponse populateFromResponseContent( * * @return the result */ - private static TransactionalBatchOperationResult createBatchOperationResultFromJson(ObjectNode objectNode) { + private static TransactionalBatchOperationResult createBatchOperationResultFromJson(ObjectNode objectNode) { final JsonSerializable jsonSerializable = new JsonSerializable(objectNode); final int statusCode = jsonSerializable.getInt(BatchRequestResponseConstant.FIELD_STATUS_CODE); @@ -200,7 +200,7 @@ private static TransactionalBatchOperationResult createBatchOperationResultFr private static void createAndPopulateResults(final TransactionalBatchResponse response, final List> operations, final Duration retryAfterDuration) { - final List> results = new ArrayList<>(operations.size()); + final List results = new ArrayList<>(operations.size()); for (int i = 0; i < operations.size(); i++) { results.add( BridgeInternal.createTransactionBatchResult( diff --git a/sdk/cosmos/azure-cosmos/src/test/java/com/azure/cosmos/BatchOperationResultTests.java b/sdk/cosmos/azure-cosmos/src/test/java/com/azure/cosmos/BatchOperationResultTests.java index f5d0e143f4e0..e19ce110d79a 100644 --- a/sdk/cosmos/azure-cosmos/src/test/java/com/azure/cosmos/BatchOperationResultTests.java +++ b/sdk/cosmos/azure-cosmos/src/test/java/com/azure/cosmos/BatchOperationResultTests.java @@ -17,8 +17,8 @@ public class BatchOperationResultTests { private static final int TIMEOUT = 40000; private ObjectNode objectNode = Utils.getSimpleObjectMapper().createObjectNode(); - private TransactionalBatchOperationResult createTestResult() { - TransactionalBatchOperationResult result = BridgeInternal.createTransactionBatchResult( + private TransactionalBatchOperationResult createTestResult() { + TransactionalBatchOperationResult result = BridgeInternal.createTransactionBatchResult( "TestETag", 1.4, objectNode, @@ -32,48 +32,20 @@ private TransactionalBatchOperationResult createTestResult() { @Test(groups = {"unit"}, timeOut = TIMEOUT) public void propertiesAreSetThroughCtor() { - TransactionalBatchOperationResult result = createTestResult(); + TransactionalBatchOperationResult result = createTestResult(); assertThat(result.getStatusCode()).isEqualTo(HttpResponseStatus.OK.code()); assertThat(result.getSubStatusCode()).isEqualTo(HttpConstants.SubStatusCodes.NAME_CACHE_IS_STALE); assertThat(result.getETag()).isEqualTo("TestETag"); assertThat(result.getRequestCharge()).isEqualTo(1.4); - assertThat(result.getRetryAfter()).isEqualTo(Duration.ofMillis(1234)); + assertThat(result.getRetryAfterDuration()).isEqualTo(Duration.ofMillis(1234)); assertThat(result.getResourceObject()).isSameAs(objectNode); } - @Test(groups = {"unit"}, timeOut = TIMEOUT) - public void propertiesAreSetThroughCopyCtor() { - TransactionalBatchOperationResult other = createTestResult(); - TransactionalBatchOperationResult result = new TransactionalBatchOperationResult(other); - - assertThat(other.getStatusCode()).isEqualTo(result.getStatusCode()); - assertThat(other.getSubStatusCode()).isEqualTo(result.getSubStatusCode()); - assertThat(other.getETag()).isEqualTo(result.getETag()); - assertThat(other.getRequestCharge()).isEqualTo(result.getRequestCharge()); - assertThat(other.getRetryAfter()).isEqualTo(result.getRetryAfter()); - assertThat(other.getResourceObject()).isSameAs(result.getResourceObject()); - } - - @Test(groups = {"unit"}, timeOut = TIMEOUT) - public void propertiesAreSetThroughGenericCtor() { - TransactionalBatchOperationResult other = createTestResult(); - Object testObject = new Object(); - TransactionalBatchOperationResult result = new TransactionalBatchOperationResult(other, testObject); - - assertThat(other.getStatusCode()).isEqualTo(result.getStatusCode()); - assertThat(other.getSubStatusCode()).isEqualTo(result.getSubStatusCode()); - assertThat(other.getETag()).isEqualTo(result.getETag()); - assertThat(other.getRequestCharge()).isEqualTo(result.getRequestCharge()); - assertThat(other.getRetryAfter()).isEqualTo(result.getRetryAfter()); - assertThat(other.getResourceObject()).isSameAs(result.getResourceObject()); - assertThat(testObject).isSameAs(result.getItem()); - } - @Test(groups = {"unit"}, timeOut = TIMEOUT) public void isSuccessStatusCodeTrueFor200To299() { for (int x = 100; x < 999; ++x) { - TransactionalBatchOperationResult result = BridgeInternal.createTransactionBatchResult( + TransactionalBatchOperationResult result = BridgeInternal.createTransactionBatchResult( null, 0.0, null, diff --git a/sdk/cosmos/azure-cosmos/src/test/java/com/azure/cosmos/TransactionalBatchAsyncContainerTest.java b/sdk/cosmos/azure-cosmos/src/test/java/com/azure/cosmos/TransactionalBatchAsyncContainerTest.java index 4ffdd5975da4..b27ba8679936 100644 --- a/sdk/cosmos/azure-cosmos/src/test/java/com/azure/cosmos/TransactionalBatchAsyncContainerTest.java +++ b/sdk/cosmos/azure-cosmos/src/test/java/com/azure/cosmos/TransactionalBatchAsyncContainerTest.java @@ -39,9 +39,6 @@ public void afterClass() { @Test(groups = {"simple"}, timeOut = TIMEOUT) public void batchExecutionRepeat() { - CosmosAsyncContainer container = this.batchAsyncContainer; - this.createJsonTestDocs(container); - TestDoc firstDoc = this.populateTestDoc(this.partitionKey1); TestDoc replaceDoc = this.getTestDocCopy(firstDoc); replaceDoc.setCost(replaceDoc.getCost() + 1); diff --git a/sdk/cosmos/azure-cosmos/src/test/java/com/azure/cosmos/TransactionalBatchTest.java b/sdk/cosmos/azure-cosmos/src/test/java/com/azure/cosmos/TransactionalBatchTest.java index ba5afb158c73..6253b5a5cb3b 100644 --- a/sdk/cosmos/azure-cosmos/src/test/java/com/azure/cosmos/TransactionalBatchTest.java +++ b/sdk/cosmos/azure-cosmos/src/test/java/com/azure/cosmos/TransactionalBatchTest.java @@ -46,10 +46,8 @@ public void afterClass() { @Test(groups = {"simple"}, timeOut = TIMEOUT) public void batchOrdered() { CosmosContainer container = this.batchContainer; - this.createJsonTestDocs(container); TestDoc firstDoc = this.populateTestDoc(this.partitionKey1); - TestDoc replaceDoc = this.getTestDocCopy(firstDoc); replaceDoc.setCost(replaceDoc.getCost() + 1); @@ -70,7 +68,6 @@ public void batchOrdered() { @Test(groups = {"simple"}, timeOut = TIMEOUT) public void batchMultipleItemExecution() { CosmosContainer container = this.batchContainer; - this.createJsonTestDocs(container); TestDoc firstDoc = this.populateTestDoc(this.partitionKey1); TestDoc replaceDoc = this.getTestDocCopy(firstDoc); @@ -91,16 +88,16 @@ public void batchMultipleItemExecution() { this.verifyBatchProcessed(batchResponse, 4); assertThat(batchResponse.get(0).getStatusCode()).isEqualTo(HttpResponseStatus.CREATED.code()); - assertThat(batchResponse.getOperationResultAtIndex(0, TestDoc.class).getItem()).isEqualTo(firstDoc); + assertThat(batchResponse.get(0).getItem(TestDoc.class)).isEqualTo(firstDoc); assertThat(batchResponse.get(1).getStatusCode()).isEqualTo(HttpResponseStatus.CREATED.code()); - assertThat(batchResponse.getOperationResultAtIndex(1, EventDoc.class).getItem()).isEqualTo(eventDoc1); + assertThat(batchResponse.get(1).getItem(EventDoc.class)).isEqualTo(eventDoc1); assertThat(batchResponse.get(2).getStatusCode()).isEqualTo(HttpResponseStatus.OK.code()); - assertThat(batchResponse.getOperationResultAtIndex(2, TestDoc.class).getItem()).isEqualTo(replaceDoc); + assertThat(batchResponse.get(2).getItem(TestDoc.class)).isEqualTo(replaceDoc); assertThat(batchResponse.get(3).getStatusCode()).isEqualTo(HttpResponseStatus.OK.code()); - assertThat(batchResponse.getOperationResultAtIndex(3, EventDoc.class).getItem()).isEqualTo(readEventDoc); + assertThat(batchResponse.get(3).getItem(EventDoc.class)).isEqualTo(readEventDoc); // Ensure that the replace overwrote the doc from the first operation this.verifyByRead(container, replaceDoc); @@ -456,9 +453,9 @@ public void batchReadsOnlyTest() { assertThat(batchResponse.get(1).getStatusCode()).isEqualTo(HttpResponseStatus.OK.code()); assertThat(batchResponse.get(2).getStatusCode()).isEqualTo(HttpResponseStatus.OK.code()); - assertThat(batchResponse.getOperationResultAtIndex(0, TestDoc.class).getItem()).isEqualTo(this.TestDocPk1ExistingA); - assertThat(batchResponse.getOperationResultAtIndex(1, TestDoc.class).getItem()).isEqualTo(this.TestDocPk1ExistingB); - assertThat(batchResponse.getOperationResultAtIndex(2, TestDoc.class).getItem()).isEqualTo(this.TestDocPk1ExistingC); + assertThat(batchResponse.get(0).getItem(TestDoc.class)).isEqualTo(this.TestDocPk1ExistingA); + assertThat(batchResponse.get(1).getItem(TestDoc.class)).isEqualTo(this.TestDocPk1ExistingB); + assertThat(batchResponse.get(2).getItem(TestDoc.class)).isEqualTo(this.TestDocPk1ExistingC); } @Test(groups = {"simple"}, timeOut = TIMEOUT) @@ -494,7 +491,7 @@ public void batchCrud() { assertThat(batchResponse.get(4).getStatusCode()).isEqualTo(HttpResponseStatus.OK.code()); assertThat(batchResponse.get(5).getStatusCode()).isEqualTo(HttpResponseStatus.NO_CONTENT.code()); - assertThat(batchResponse.getOperationResultAtIndex(1, TestDoc.class).getItem()).isEqualTo(this.TestDocPk1ExistingC); + assertThat(batchResponse.get(1).getItem(TestDoc.class)).isEqualTo(this.TestDocPk1ExistingC); this.verifyByRead(container, testDocToCreate); this.verifyByRead(container, testDocToReplace); diff --git a/sdk/cosmos/azure-cosmos/src/test/java/com/azure/cosmos/implementation/batch/BatchResponsePayloadWriter.java b/sdk/cosmos/azure-cosmos/src/test/java/com/azure/cosmos/implementation/batch/BatchResponsePayloadWriter.java index eb29c2b61743..4f850dc84c88 100644 --- a/sdk/cosmos/azure-cosmos/src/test/java/com/azure/cosmos/implementation/batch/BatchResponsePayloadWriter.java +++ b/sdk/cosmos/azure-cosmos/src/test/java/com/azure/cosmos/implementation/batch/BatchResponsePayloadWriter.java @@ -12,9 +12,9 @@ class BatchResponsePayloadWriter { - private List> results; + private List results; - BatchResponsePayloadWriter(List> results) { + BatchResponsePayloadWriter(List results) { this.results = results; } @@ -25,7 +25,7 @@ String generatePayload() { private ArrayNode writeOperationResult() { ArrayNode arrayNode = Utils.getSimpleObjectMapper().createArrayNode(); - for(TransactionalBatchOperationResult result : results) { + for(TransactionalBatchOperationResult result : results) { JsonSerializable operationJsonSerializable = writeResult(result); arrayNode.add(operationJsonSerializable.getPropertyBag()); @@ -33,7 +33,7 @@ private ArrayNode writeOperationResult() { return arrayNode; } - private JsonSerializable writeResult(TransactionalBatchOperationResult result) { + private JsonSerializable writeResult(TransactionalBatchOperationResult result) { JsonSerializable jsonSerializable = new JsonSerializable(); jsonSerializable.set(BatchRequestResponseConstant.FIELD_STATUS_CODE, result.getStatusCode()); @@ -41,8 +41,8 @@ private JsonSerializable writeResult(TransactionalBatchOperationResult result jsonSerializable.set(BatchRequestResponseConstant.FIELD_ETAG, result.getETag()); jsonSerializable.set(BatchRequestResponseConstant.FIELD_REQUEST_CHARGE, result.getRequestCharge()); - if(result.getRetryAfter() != null) { - jsonSerializable.set(BatchRequestResponseConstant.FIELD_RETRY_AFTER_MILLISECONDS, result.getRetryAfter().toMillis()); + if(result.getRetryAfterDuration() != null) { + jsonSerializable.set(BatchRequestResponseConstant.FIELD_RETRY_AFTER_MILLISECONDS, result.getRetryAfterDuration().toMillis()); } return jsonSerializable; diff --git a/sdk/cosmos/azure-cosmos/src/test/java/com/azure/cosmos/implementation/batch/TransactionalBatchResponseTests.java b/sdk/cosmos/azure-cosmos/src/test/java/com/azure/cosmos/implementation/batch/TransactionalBatchResponseTests.java index 5189fb975e17..41f52ee88239 100644 --- a/sdk/cosmos/azure-cosmos/src/test/java/com/azure/cosmos/implementation/batch/TransactionalBatchResponseTests.java +++ b/sdk/cosmos/azure-cosmos/src/test/java/com/azure/cosmos/implementation/batch/TransactionalBatchResponseTests.java @@ -31,7 +31,7 @@ public class TransactionalBatchResponseTests { @Test(groups = {"unit"}, timeOut = TIMEOUT) public void validateAllSetValuesInResponse() { - List> results = new ArrayList<>(); + List results = new ArrayList<>(); ItemBatchOperation[] arrayOperations = new ItemBatchOperation[1]; ItemBatchOperation operation = new ItemBatchOperation.Builder(OperationType.Read,0) @@ -45,7 +45,7 @@ public void validateAllSetValuesInResponse() { Arrays.asList(arrayOperations)); // Create dummy result - TransactionalBatchOperationResult transactionalBatchOperationResult = BridgeInternal.createTransactionBatchResult( + TransactionalBatchOperationResult transactionalBatchOperationResult = BridgeInternal.createTransactionBatchResult( operation.getId(), 5.0, null, @@ -82,20 +82,20 @@ public void validateAllSetValuesInResponse() { assertThat(batchResponse.getStatusCode()).isEqualTo(HttpResponseStatus.OK.code()); assertThat(batchResponse.getSessionToken()).isEqualTo("token123"); assertThat(batchResponse.getResponseHeaders()).isEqualTo(headers); - assertThat(batchResponse.getRetryAfter()).isEqualTo(Duration.ofMillis(1234)); + assertThat(batchResponse.getRetryAfterDuration()).isEqualTo(Duration.ofMillis(1234)); assertThat(batchResponse.getSubStatusCode()).isEqualTo(HttpConstants.SubStatusCodes.PARTITION_KEY_RANGE_GONE); // Validate result fields assertThat(batchResponse.get(0).getETag()).isEqualTo(operation.getId()); assertThat(batchResponse.get(0).getRequestCharge()).isEqualTo(5.0); - assertThat(batchResponse.get(0).getRetryAfter()).isEqualTo(Duration.ofMillis(100)); + assertThat(batchResponse.get(0).getRetryAfterDuration()).isEqualTo(Duration.ofMillis(100)); assertThat(batchResponse.get(0).getSubStatusCode()).isEqualTo(HttpConstants.SubStatusCodes.PARTITION_KEY_MISMATCH); assertThat(batchResponse.get(0).getStatusCode()).isEqualTo(HttpResponseStatus.NOT_MODIFIED.code()); } @Test(groups = {"unit"}, timeOut = TIMEOUT) public void validateEmptyHeaderInResponse() { - List> results = new ArrayList<>(); + List results = new ArrayList<>(); ItemBatchOperation[] arrayOperations = new ItemBatchOperation[1]; ItemBatchOperation operation = new ItemBatchOperation.Builder(OperationType.Read,0) @@ -109,7 +109,7 @@ public void validateEmptyHeaderInResponse() { Arrays.asList(arrayOperations)); // Create dummy result - TransactionalBatchOperationResult transactionalBatchOperationResult = BridgeInternal.createTransactionBatchResult( + TransactionalBatchOperationResult transactionalBatchOperationResult = BridgeInternal.createTransactionBatchResult( null, 5.0, null, @@ -137,13 +137,13 @@ public void validateEmptyHeaderInResponse() { assertThat(batchResponse.getRequestCharge()).isEqualTo(0); assertThat(batchResponse.getSessionToken()).isNull(); assertThat(batchResponse.getResponseHeaders()).isEmpty(); - assertThat(batchResponse.getRetryAfter()).isEqualTo(Duration.ZERO); + assertThat(batchResponse.getRetryAfterDuration()).isEqualTo(Duration.ZERO); assertThat(batchResponse.getSubStatusCode()).isEqualTo(0); // Validate result fields assertThat(batchResponse.get(0).getETag()).isNull(); assertThat(batchResponse.get(0).getRequestCharge()).isEqualTo(5.0); - assertThat(batchResponse.get(0).getRetryAfter()).isEqualTo(Duration.ZERO); + assertThat(batchResponse.get(0).getRetryAfterDuration()).isEqualTo(Duration.ZERO); assertThat(batchResponse.get(0).getSubStatusCode()).isEqualTo(0); assertThat(batchResponse.get(0).getStatusCode()).isEqualTo(HttpResponseStatus.NOT_MODIFIED.code()); } From 04ca2396502f73abe13838fdea5dc54c3584b34e Mon Sep 17 00:00:00 2001 From: Rakesh Kumar Date: Wed, 7 Oct 2020 12:01:07 +0530 Subject: [PATCH 15/22] Removing one test case Signed-off-by: Rakesh Kumar --- .../azure/cosmos/TransactionalBatchTest.java | 94 ------------------- 1 file changed, 94 deletions(-) diff --git a/sdk/cosmos/azure-cosmos/src/test/java/com/azure/cosmos/TransactionalBatchTest.java b/sdk/cosmos/azure-cosmos/src/test/java/com/azure/cosmos/TransactionalBatchTest.java index 6253b5a5cb3b..1ad43fa6da03 100644 --- a/sdk/cosmos/azure-cosmos/src/test/java/com/azure/cosmos/TransactionalBatchTest.java +++ b/sdk/cosmos/azure-cosmos/src/test/java/com/azure/cosmos/TransactionalBatchTest.java @@ -159,100 +159,6 @@ public void batchItemETagTest() { } } - @Test(groups = {"simple"}, timeOut = TIMEOUT) - public void batchSessionTokenPropertiesTest() throws Exception { - CosmosContainer container = batchContainer; - this.createJsonTestDocs(container); - - TestDoc sampleDoc = this.populateTestDoc(this.partitionKey1); - - CosmosItemResponse createResponse = container.createItem( - sampleDoc, - this.getPartitionKey(this.partitionKey1), - null); - - String ownerIdCreate = createResponse.getResponseHeaders().get(HttpConstants.HttpHeaders.OWNER_ID); - assertThat(createResponse.getStatusCode()).isEqualTo(HttpResponseStatus.CREATED.code()); - assertThat(ownerIdCreate).isNotEmpty(); - - CosmosItemResponse readResponse = container.readItem( - this.TestDocPk1ExistingC.getId(), - this.getPartitionKey(this.partitionKey1), - TestDoc.class); - - assertThat(readResponse.getStatusCode()).isEqualTo(HttpResponseStatus.OK.code()); - - ISessionToken beforeRequestSessionToken = this.getSessionToken(readResponse.getSessionToken()); - - String readEtagValue = readResponse.getETag(); - TransactionalBatchItemRequestOptions readRequestOption = new TransactionalBatchItemRequestOptions() - .setIfMatchETag(readEtagValue); - - String oldSessionToken = this.getDifferentLSNToken(readResponse.getSessionToken(), -10); - - { - // Batch with only Read operation - TransactionalBatchResponse batchResponse = container.executeTransactionalBatch( - TransactionalBatch.createTransactionalBatch(this.getPartitionKey(this.partitionKey1)) - .readItem(this.TestDocPk1ExistingA.getId()) - .readItem(this.TestDocPk1ExistingC.getId(), readRequestOption), - new TransactionalBatchRequestOptions().setSessionToken(oldSessionToken)); - - this.verifyBatchProcessed(batchResponse, 2, HttpResponseStatus.MULTI_STATUS); - - assertThat(batchResponse.get(0).getStatusCode()).isEqualTo(HttpResponseStatus.OK.code()); - assertThat(batchResponse.get(1).getStatusCode()).isEqualTo(HttpResponseStatus.NOT_MODIFIED.code()); - - ISessionToken afterRequestSessionToken = this.getSessionToken(batchResponse.getSessionToken()); - - assertThat(afterRequestSessionToken.getLSN()) - .as("Response session token should be more than or equal to request session token") - .isGreaterThanOrEqualTo(beforeRequestSessionToken.getLSN()); - - String ownerIdBatch = batchResponse.getResponseHeaders().get(HttpConstants.HttpHeaders.OWNER_ID); - assertThat(ownerIdBatch).isNotEmpty(); - assertThat(ownerIdBatch).isEqualTo(ownerIdCreate); - } - - { - // Batch with write-read operations - TestDoc testDocToCreate = this.populateTestDoc(this.partitionKey1); - - TestDoc testDocToReplace = this.getTestDocCopy(this.TestDocPk1ExistingB); - testDocToReplace.setCost(testDocToReplace.getCost() + 1); - TestDoc testDocToUpsert = this.populateTestDoc(this.partitionKey1); - - TransactionalBatchResponse batchResponse = container.executeTransactionalBatch( - TransactionalBatch.createTransactionalBatch(this.getPartitionKey(this.partitionKey1)) - .createItem(testDocToCreate) - .replaceItem(testDocToReplace.getId(), testDocToReplace) - .upsertItem(testDocToUpsert) - .deleteItem(this.TestDocPk1ExistingD.getId()) - .readItem(this.TestDocPk1ExistingA.getId()) - .readItem(this.TestDocPk1ExistingC.getId(), readRequestOption), - new TransactionalBatchRequestOptions().setSessionToken(oldSessionToken)); - - this.verifyBatchProcessed(batchResponse, 6, HttpResponseStatus.MULTI_STATUS); - - assertThat(batchResponse.get(0).getStatusCode()).isEqualTo(HttpResponseStatus.CREATED.code()); - assertThat(batchResponse.get(1).getStatusCode()).isEqualTo(HttpResponseStatus.OK.code()); - assertThat(batchResponse.get(2).getStatusCode()).isEqualTo(HttpResponseStatus.CREATED.code()); - assertThat(batchResponse.get(3).getStatusCode()).isEqualTo(HttpResponseStatus.NO_CONTENT.code()); - assertThat(batchResponse.get(4).getStatusCode()).isEqualTo(HttpResponseStatus.OK.code()); - assertThat(batchResponse.get(5).getStatusCode()).isEqualTo(HttpResponseStatus.NOT_MODIFIED.code()); - - ISessionToken afterRequestSessionToken = this.getSessionToken(batchResponse.getSessionToken()); - - assertThat(afterRequestSessionToken.getLSN()) - .as("Response session token should be more than request session token") - .isGreaterThan(beforeRequestSessionToken.getLSN()); - - String ownerIdBatch = batchResponse.getResponseHeaders().get(HttpConstants.HttpHeaders.OWNER_ID); - assertThat(ownerIdBatch).isNotEmpty(); - assertThat(ownerIdBatch).isEqualTo(ownerIdCreate); - } - } - @Test(groups = {"simple"}, timeOut = TIMEOUT) public void batchErrorSessionToken() { CosmosContainer container = batchContainer; From 2735e339d39f4a4796eafb7fe3c00cd20ba93c40 Mon Sep 17 00:00:00 2001 From: Rakesh Kumar Date: Wed, 7 Oct 2020 14:42:36 +0530 Subject: [PATCH 16/22] test Signed-off-by: Rakesh Kumar --- .../TransactionalBatchAsyncContainerTest.java | 95 +++++++++++++++++++ 1 file changed, 95 insertions(+) diff --git a/sdk/cosmos/azure-cosmos/src/test/java/com/azure/cosmos/TransactionalBatchAsyncContainerTest.java b/sdk/cosmos/azure-cosmos/src/test/java/com/azure/cosmos/TransactionalBatchAsyncContainerTest.java index b27ba8679936..a06f4336fa6f 100644 --- a/sdk/cosmos/azure-cosmos/src/test/java/com/azure/cosmos/TransactionalBatchAsyncContainerTest.java +++ b/sdk/cosmos/azure-cosmos/src/test/java/com/azure/cosmos/TransactionalBatchAsyncContainerTest.java @@ -4,6 +4,7 @@ package com.azure.cosmos; import com.azure.cosmos.implementation.HttpConstants; +import com.azure.cosmos.implementation.ISessionToken; import com.azure.cosmos.models.CosmosItemResponse; import io.netty.handler.codec.http.HttpResponseStatus; import org.assertj.core.api.Assertions; @@ -119,4 +120,98 @@ public void batchInvalidSessionToken() throws Exception { } } } + + @Test(groups = {"simple"}, timeOut = TIMEOUT) + public void batchSessionTokenPropertiesTest() throws Exception { + CosmosAsyncContainer container = batchAsyncContainer; + this.createJsonTestDocs(container); + + TestDoc sampleDoc = this.populateTestDoc(this.partitionKey1); + + CosmosItemResponse createResponse = container.createItem( + sampleDoc, + this.getPartitionKey(this.partitionKey1), + null).block(); + + String ownerIdCreate = createResponse.getResponseHeaders().get(HttpConstants.HttpHeaders.OWNER_ID); + assertThat(createResponse.getStatusCode()).isEqualTo(HttpResponseStatus.CREATED.code()); + assertThat(ownerIdCreate).isNotEmpty(); + + CosmosItemResponse readResponse = container.readItem( + this.TestDocPk1ExistingC.getId(), + this.getPartitionKey(this.partitionKey1), + TestDoc.class).block(); + + assertThat(readResponse.getStatusCode()).isEqualTo(HttpResponseStatus.OK.code()); + + ISessionToken beforeRequestSessionToken = this.getSessionToken(readResponse.getSessionToken()); + + String readEtagValue = readResponse.getETag(); + TransactionalBatchItemRequestOptions readRequestOption = new TransactionalBatchItemRequestOptions() + .setIfMatchETag(readEtagValue); + + String oldSessionToken = this.getDifferentLSNToken(readResponse.getSessionToken(), -10); + + { + // Batch with only Read operation + TransactionalBatchResponse batchResponse = container.executeTransactionalBatch( + TransactionalBatch.createTransactionalBatch(this.getPartitionKey(this.partitionKey1)) + .readItem(this.TestDocPk1ExistingA.getId()) + .readItem(this.TestDocPk1ExistingC.getId(), readRequestOption), + new TransactionalBatchRequestOptions().setSessionToken(oldSessionToken)).block(); + + this.verifyBatchProcessed(batchResponse, 2, HttpResponseStatus.MULTI_STATUS); + + assertThat(batchResponse.get(0).getStatusCode()).isEqualTo(HttpResponseStatus.OK.code()); + assertThat(batchResponse.get(1).getStatusCode()).isEqualTo(HttpResponseStatus.NOT_MODIFIED.code()); + + ISessionToken afterRequestSessionToken = this.getSessionToken(batchResponse.getSessionToken()); + + assertThat(afterRequestSessionToken.getLSN()) + .as("Response session token should be more than or equal to request session token") + .isGreaterThanOrEqualTo(beforeRequestSessionToken.getLSN()); + + String ownerIdBatch = batchResponse.getResponseHeaders().get(HttpConstants.HttpHeaders.OWNER_ID); + assertThat(ownerIdBatch).isNotEmpty(); + assertThat(ownerIdBatch).isEqualTo(ownerIdCreate); + } + + { + // Batch with write-read operations + TestDoc testDocToCreate = this.populateTestDoc(this.partitionKey1); + + TestDoc testDocToReplace = this.getTestDocCopy(this.TestDocPk1ExistingB); + testDocToReplace.setCost(testDocToReplace.getCost() + 1); + TestDoc testDocToUpsert = this.populateTestDoc(this.partitionKey1); + + TransactionalBatchResponse batchResponse = container.executeTransactionalBatch( + TransactionalBatch.createTransactionalBatch(this.getPartitionKey(this.partitionKey1)) + .createItem(testDocToCreate) + .replaceItem(testDocToReplace.getId(), testDocToReplace) + .upsertItem(testDocToUpsert) + .deleteItem(this.TestDocPk1ExistingD.getId()) + .readItem(this.TestDocPk1ExistingA.getId()) + .readItem(this.TestDocPk1ExistingC.getId(), readRequestOption), + new TransactionalBatchRequestOptions().setSessionToken(oldSessionToken)).block(); + + this.verifyBatchProcessed(batchResponse, 6, HttpResponseStatus.MULTI_STATUS); + + assertThat(batchResponse.get(0).getStatusCode()).isEqualTo(HttpResponseStatus.CREATED.code()); + assertThat(batchResponse.get(1).getStatusCode()).isEqualTo(HttpResponseStatus.OK.code()); + assertThat(batchResponse.get(2).getStatusCode()).isEqualTo(HttpResponseStatus.CREATED.code()); + assertThat(batchResponse.get(3).getStatusCode()).isEqualTo(HttpResponseStatus.NO_CONTENT.code()); + assertThat(batchResponse.get(4).getStatusCode()).isEqualTo(HttpResponseStatus.OK.code()); + assertThat(batchResponse.get(5).getStatusCode()).isEqualTo(HttpResponseStatus.NOT_MODIFIED.code()); + + ISessionToken afterRequestSessionToken = this.getSessionToken(batchResponse.getSessionToken()); + + assertThat(afterRequestSessionToken.getLSN()) + .as("Response session token should be more than request session token") + .isGreaterThan(beforeRequestSessionToken.getLSN()); + + String ownerIdBatch = batchResponse.getResponseHeaders().get(HttpConstants.HttpHeaders.OWNER_ID); + assertThat(ownerIdBatch).isNotEmpty(); + assertThat(ownerIdBatch).isEqualTo(ownerIdCreate); + } + } } From d605ba10aa3032025ac81ad80103568707a1a171 Mon Sep 17 00:00:00 2001 From: Rakesh Kumar Date: Wed, 7 Oct 2020 22:23:17 +0530 Subject: [PATCH 17/22] Fix Signed-off-by: Rakesh Kumar --- .../TransactionalBatchAsyncContainerTest.java | 94 ------------------- 1 file changed, 94 deletions(-) diff --git a/sdk/cosmos/azure-cosmos/src/test/java/com/azure/cosmos/TransactionalBatchAsyncContainerTest.java b/sdk/cosmos/azure-cosmos/src/test/java/com/azure/cosmos/TransactionalBatchAsyncContainerTest.java index a06f4336fa6f..e40cce07c24e 100644 --- a/sdk/cosmos/azure-cosmos/src/test/java/com/azure/cosmos/TransactionalBatchAsyncContainerTest.java +++ b/sdk/cosmos/azure-cosmos/src/test/java/com/azure/cosmos/TransactionalBatchAsyncContainerTest.java @@ -120,98 +120,4 @@ public void batchInvalidSessionToken() throws Exception { } } } - - @Test(groups = {"simple"}, timeOut = TIMEOUT) - public void batchSessionTokenPropertiesTest() throws Exception { - CosmosAsyncContainer container = batchAsyncContainer; - this.createJsonTestDocs(container); - - TestDoc sampleDoc = this.populateTestDoc(this.partitionKey1); - - CosmosItemResponse createResponse = container.createItem( - sampleDoc, - this.getPartitionKey(this.partitionKey1), - null).block(); - - String ownerIdCreate = createResponse.getResponseHeaders().get(HttpConstants.HttpHeaders.OWNER_ID); - assertThat(createResponse.getStatusCode()).isEqualTo(HttpResponseStatus.CREATED.code()); - assertThat(ownerIdCreate).isNotEmpty(); - - CosmosItemResponse readResponse = container.readItem( - this.TestDocPk1ExistingC.getId(), - this.getPartitionKey(this.partitionKey1), - TestDoc.class).block(); - - assertThat(readResponse.getStatusCode()).isEqualTo(HttpResponseStatus.OK.code()); - - ISessionToken beforeRequestSessionToken = this.getSessionToken(readResponse.getSessionToken()); - - String readEtagValue = readResponse.getETag(); - TransactionalBatchItemRequestOptions readRequestOption = new TransactionalBatchItemRequestOptions() - .setIfMatchETag(readEtagValue); - - String oldSessionToken = this.getDifferentLSNToken(readResponse.getSessionToken(), -10); - - { - // Batch with only Read operation - TransactionalBatchResponse batchResponse = container.executeTransactionalBatch( - TransactionalBatch.createTransactionalBatch(this.getPartitionKey(this.partitionKey1)) - .readItem(this.TestDocPk1ExistingA.getId()) - .readItem(this.TestDocPk1ExistingC.getId(), readRequestOption), - new TransactionalBatchRequestOptions().setSessionToken(oldSessionToken)).block(); - - this.verifyBatchProcessed(batchResponse, 2, HttpResponseStatus.MULTI_STATUS); - - assertThat(batchResponse.get(0).getStatusCode()).isEqualTo(HttpResponseStatus.OK.code()); - assertThat(batchResponse.get(1).getStatusCode()).isEqualTo(HttpResponseStatus.NOT_MODIFIED.code()); - - ISessionToken afterRequestSessionToken = this.getSessionToken(batchResponse.getSessionToken()); - - assertThat(afterRequestSessionToken.getLSN()) - .as("Response session token should be more than or equal to request session token") - .isGreaterThanOrEqualTo(beforeRequestSessionToken.getLSN()); - - String ownerIdBatch = batchResponse.getResponseHeaders().get(HttpConstants.HttpHeaders.OWNER_ID); - assertThat(ownerIdBatch).isNotEmpty(); - assertThat(ownerIdBatch).isEqualTo(ownerIdCreate); - } - - { - // Batch with write-read operations - TestDoc testDocToCreate = this.populateTestDoc(this.partitionKey1); - - TestDoc testDocToReplace = this.getTestDocCopy(this.TestDocPk1ExistingB); - testDocToReplace.setCost(testDocToReplace.getCost() + 1); - TestDoc testDocToUpsert = this.populateTestDoc(this.partitionKey1); - - TransactionalBatchResponse batchResponse = container.executeTransactionalBatch( - TransactionalBatch.createTransactionalBatch(this.getPartitionKey(this.partitionKey1)) - .createItem(testDocToCreate) - .replaceItem(testDocToReplace.getId(), testDocToReplace) - .upsertItem(testDocToUpsert) - .deleteItem(this.TestDocPk1ExistingD.getId()) - .readItem(this.TestDocPk1ExistingA.getId()) - .readItem(this.TestDocPk1ExistingC.getId(), readRequestOption), - new TransactionalBatchRequestOptions().setSessionToken(oldSessionToken)).block(); - - this.verifyBatchProcessed(batchResponse, 6, HttpResponseStatus.MULTI_STATUS); - - assertThat(batchResponse.get(0).getStatusCode()).isEqualTo(HttpResponseStatus.CREATED.code()); - assertThat(batchResponse.get(1).getStatusCode()).isEqualTo(HttpResponseStatus.OK.code()); - assertThat(batchResponse.get(2).getStatusCode()).isEqualTo(HttpResponseStatus.CREATED.code()); - assertThat(batchResponse.get(3).getStatusCode()).isEqualTo(HttpResponseStatus.NO_CONTENT.code()); - assertThat(batchResponse.get(4).getStatusCode()).isEqualTo(HttpResponseStatus.OK.code()); - assertThat(batchResponse.get(5).getStatusCode()).isEqualTo(HttpResponseStatus.NOT_MODIFIED.code()); - - ISessionToken afterRequestSessionToken = this.getSessionToken(batchResponse.getSessionToken()); - - assertThat(afterRequestSessionToken.getLSN()) - .as("Response session token should be more than request session token") - .isGreaterThan(beforeRequestSessionToken.getLSN()); - - String ownerIdBatch = batchResponse.getResponseHeaders().get(HttpConstants.HttpHeaders.OWNER_ID); - assertThat(ownerIdBatch).isNotEmpty(); - assertThat(ownerIdBatch).isEqualTo(ownerIdCreate); - } - } } From 668c0e36300035e8b9f566d83c149badd8ed7053 Mon Sep 17 00:00:00 2001 From: Rakesh Kumar Date: Fri, 9 Oct 2020 14:57:26 +0530 Subject: [PATCH 18/22] API change Signed-off-by: Rakesh Kumar --- .../java/com/azure/cosmos/BridgeInternal.java | 22 +- .../azure/cosmos/CosmosItemOperationType.java | 25 ++ ...ions.java => ItemBatchRequestOptions.java} | 17 +- .../com/azure/cosmos/TransactionalBatch.java | 154 ++++++----- .../TransactionalBatchOperationResult.java | 16 +- .../cosmos/TransactionalBatchResponse.java | 24 +- .../implementation/batch/BatchExecUtils.java | 26 +- .../implementation/batch/BatchExecutor.java | 5 +- .../batch/BatchRequestResponseConstant.java | 12 +- .../batch/BatchResponseParser.java | 26 +- .../batch/ItemBatchOperation.java | 163 ------------ .../batch/ServerBatchRequest.java | 46 +++- .../SinglePartitionKeyServerBatchRequest.java | 1 + .../cosmos/models/ItemBatchOperation.java | 70 +++++ .../cosmos/models/ModelBridgeInternal.java | 20 ++ .../cosmos/BatchOperationResultTests.java | 17 +- .../TransactionalBatchAsyncContainerTest.java | 59 +++-- .../azure/cosmos/TransactionalBatchTest.java | 241 +++++++++++------- .../TransactionalBatchResponseTests.java | 56 ++-- 19 files changed, 539 insertions(+), 461 deletions(-) create mode 100644 sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosItemOperationType.java rename sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/{TransactionalBatchItemRequestOptions.java => ItemBatchRequestOptions.java} (70%) delete mode 100644 sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/batch/ItemBatchOperation.java create mode 100644 sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/models/ItemBatchOperation.java diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/BridgeInternal.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/BridgeInternal.java index f63321c12cc0..2e89d599be3c 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/BridgeInternal.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/BridgeInternal.java @@ -26,7 +26,7 @@ import com.azure.cosmos.implementation.ServiceUnavailableException; import com.azure.cosmos.implementation.StoredProcedureResponse; import com.azure.cosmos.implementation.Warning; -import com.azure.cosmos.implementation.batch.ItemBatchOperation; +import com.azure.cosmos.models.ItemBatchOperation; import com.azure.cosmos.implementation.directconnectivity.StoreResponse; import com.azure.cosmos.implementation.directconnectivity.StoreResult; import com.azure.cosmos.implementation.directconnectivity.Uri; @@ -610,16 +610,6 @@ public static Duration getRequestTimeoutFromGatewayConnectionConfig(GatewayConne return gatewayConnectionConfig.getRequestTimeout(); } - @Warning(value = INTERNAL_USE_ONLY_WARNING) - public static List> getOperationsFromTransactionalBatch(TransactionalBatch transactionalBatch) { - return transactionalBatch.getOperations(); - } - - @Warning(value = INTERNAL_USE_ONLY_WARNING) - public static PartitionKey getPartitionKeyFromTransactionalBatch(TransactionalBatch transactionalBatch) { - return transactionalBatch.getPartitionKey(); - } - @Warning(value = INTERNAL_USE_ONLY_WARNING) public static RequestOptions toRequestOptions(TransactionalBatchRequestOptions transactionalBatchRequestOptions) { return transactionalBatchRequestOptions.toRequestOptions(); @@ -632,7 +622,8 @@ public static TransactionalBatchOperationResult createTransactionBatchResult( ObjectNode resourceObject, int statusCode, Duration retryAfter, - int subStatusCode) { + int subStatusCode, + ItemBatchOperation itemBatchOperation) { return new TransactionalBatchOperationResult( eTag, @@ -640,7 +631,8 @@ public static TransactionalBatchOperationResult createTransactionBatchResult( resourceObject, statusCode, retryAfter, - subStatusCode); + subStatusCode, + itemBatchOperation); } @Warning(value = INTERNAL_USE_ONLY_WARNING) @@ -662,8 +654,8 @@ public static TransactionalBatchResponse createTransactionBatchResponse( @Warning(value = INTERNAL_USE_ONLY_WARNING) public static void addTransactionBatchResultInResponse( TransactionalBatchResponse transactionalBatchResponse, - List transactionalBatchOperationResult) { + List transactionalBatchOperationResults) { - transactionalBatchResponse.addAll(transactionalBatchOperationResult); + transactionalBatchResponse.addAll(transactionalBatchOperationResults); } } diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosItemOperationType.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosItemOperationType.java new file mode 100644 index 000000000000..ff82bd4905d4 --- /dev/null +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosItemOperationType.java @@ -0,0 +1,25 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +package com.azure.cosmos; + +import com.azure.cosmos.implementation.batch.BatchRequestResponseConstant; + +public enum CosmosItemOperationType { + + Create(BatchRequestResponseConstant.OPERATION_CREATE), + Delete(BatchRequestResponseConstant.OPERATION_DELETE), + Read(BatchRequestResponseConstant.OPERATION_READ), + Replace(BatchRequestResponseConstant.OPERATION_REPLACE), + Upsert(BatchRequestResponseConstant.OPERATION_UPSERT); + + CosmosItemOperationType(String operationValue) { + this.operationValue = operationValue; + } + + public String getOperationValue() { + return operationValue; + } + + private final String operationValue; +} diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/TransactionalBatchItemRequestOptions.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/ItemBatchRequestOptions.java similarity index 70% rename from sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/TransactionalBatchItemRequestOptions.java rename to sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/ItemBatchRequestOptions.java index dde32da47d0e..0818cf491b03 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/TransactionalBatchItemRequestOptions.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/ItemBatchRequestOptions.java @@ -3,14 +3,14 @@ package com.azure.cosmos; -import com.azure.cosmos.implementation.RequestOptions; import com.azure.cosmos.util.Beta; /** - * Encapsulates options that can be specified for an operation within a {@link TransactionalBatch}. + * Encapsulates options that can be specified for an operation within a for a batch or bulk request. + * Currently used in {@link TransactionalBatch} to pass option. */ @Beta(Beta.SinceVersion.V4_7_0) -public final class TransactionalBatchItemRequestOptions { +public final class ItemBatchRequestOptions { private String ifMatchETag; private String ifNoneMatchETag; @@ -29,7 +29,7 @@ public String getIfMatchETag() { * @param ifMatchETag the ifMatchETag associated with the request. * @return the current request options */ - public TransactionalBatchItemRequestOptions setIfMatchETag(final String ifMatchETag) { + public ItemBatchRequestOptions setIfMatchETag(final String ifMatchETag) { this.ifMatchETag = ifMatchETag; return this; } @@ -49,15 +49,8 @@ public String getIfNoneMatchETag() { * @param ifNoneMatchEtag the ifNoneMatchETag associated with the request. * @return the current request options */ - public TransactionalBatchItemRequestOptions setIfNoneMatchETag(final String ifNoneMatchEtag) { + public ItemBatchRequestOptions setIfNoneMatchETag(final String ifNoneMatchEtag) { this.ifNoneMatchETag = ifNoneMatchEtag; return this; } - - RequestOptions toRequestOptions() { - final RequestOptions requestOptions = new RequestOptions(); - requestOptions.setIfMatchETag(getIfMatchETag()); - requestOptions.setIfNoneMatchETag(getIfNoneMatchETag()); - return requestOptions; - } } diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/TransactionalBatch.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/TransactionalBatch.java index 69e45055164d..172cc645b378 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/TransactionalBatch.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/TransactionalBatch.java @@ -3,9 +3,9 @@ package com.azure.cosmos; -import com.azure.cosmos.implementation.OperationType; import com.azure.cosmos.implementation.apachecommons.collections.list.UnmodifiableList; -import com.azure.cosmos.implementation.batch.ItemBatchOperation; +import com.azure.cosmos.models.ItemBatchOperation; +import com.azure.cosmos.models.ModelBridgeInternal; import com.azure.cosmos.models.PartitionKey; import com.azure.cosmos.util.Beta; @@ -39,11 +39,11 @@ * ToDoActivity test2 = new ToDoActivity(activityType, "shopping", "Done"); * ToDoActivity test3 = new ToDoActivity(activityType, "swimming", "ToBeDone"); * - * TransactionalBatch batch = TransactionalBatch.createTransactionalBatch(new Cosmos.PartitionKey(activityType)) - * .createItem(test1) - * .replaceItem(test2.id, test2) - * .upsertItem(test3) - * .deleteItem("reading"); + * TransactionalBatch batch = TransactionalBatch.createTransactionalBatch(new Cosmos.PartitionKey(activityType)); + * batch.createItem(test1); + * batch.replaceItem(test2.id, test2); + * batch.upsertItem(test3); + * batch.deleteItem("reading"); * * TransactionalBatchResponse response = container.executeTransactionalBatch(batch); * @@ -64,11 +64,11 @@ * 
{@code
  * String activityType = "personal";
  *
- * TransactionalBatch batch = TransactionalBatch.createTransactionalBatch(new Cosmos.PartitionKey(activityType))
- *     .readItem("playing")
- *     .readItem("walking")
- *     .readItem("jogging")
- *     .readItem("running")
+ * TransactionalBatch batch = TransactionalBatch.createTransactionalBatch(new Cosmos.PartitionKey(activityType));
+ * batch.readItem("playing");
+ * batch.readItem("walking");
+ * batch.readItem("jogging");
+ * batch.readItem("running")v
  *
  * TransactionalBatchResponse response = container.executeTransactionalBatch(batch);
  * List resultItems = new ArrayList();
@@ -117,9 +117,9 @@ public static TransactionalBatch createTransactionalBatch(PartitionKey partition
      *
      * @return The transactional batch instance with the operation added.
      */
-    public  TransactionalBatch createItem(T item) {
+    public  ItemBatchOperation createItem(T item) {
         checkNotNull(item, "expected non-null item");
-        return this.createItem(item, new TransactionalBatchItemRequestOptions());
+        return this.createItem(item, new ItemBatchRequestOptions());
     }
 
     /**
@@ -132,22 +132,25 @@ public  TransactionalBatch createItem(T item) {
      *
      * @return The transactional batch instance with the operation added.
      */
-    public  TransactionalBatch createItem(T item, TransactionalBatchItemRequestOptions requestOptions) {
+    public  ItemBatchOperation createItem(T item, ItemBatchRequestOptions requestOptions) {
 
         checkNotNull(item, "expected non-null item");
         if (requestOptions == null) {
-            requestOptions = new TransactionalBatchItemRequestOptions();
+            requestOptions = new ItemBatchRequestOptions();
         }
 
-        this.operations.add(
-            new ItemBatchOperation.Builder(
-                OperationType.Create,
-                this.operations.size())
-                .requestOptions(requestOptions.toRequestOptions())
-                .resource(item)
-                .build());
+        ItemBatchOperation operation = ModelBridgeInternal.createItemBatchOperation(
+            CosmosItemOperationType.Create,
+            this.operations.size(),
+            null,
+            this.getPartitionKey(),
+            requestOptions,
+            item
+        );
 
-        return this;
+        this.operations.add(operation);
+
+        return operation;
     }
 
     /**
@@ -157,9 +160,9 @@ public  TransactionalBatch createItem(T item, TransactionalBatchItemRequestOp
      *
      * @return The transactional batch instance with the operation added.
      */
-    public TransactionalBatch deleteItem(String id) {
+    public ItemBatchOperation deleteItem(String id) {
         checkNotNull(id, "expected non-null id");
-        return this.deleteItem(id, new TransactionalBatchItemRequestOptions());
+        return this.deleteItem(id, new ItemBatchRequestOptions());
     }
 
     /**
@@ -170,19 +173,25 @@ public TransactionalBatch deleteItem(String id) {
      *
      * @return The transactional batch instance with the operation added.
      */
-    public TransactionalBatch deleteItem(String id, TransactionalBatchItemRequestOptions requestOptions) {
+    public ItemBatchOperation deleteItem(String id, ItemBatchRequestOptions requestOptions) {
 
         checkNotNull(id, "expected non-null id");
         if (requestOptions == null) {
-            requestOptions = new TransactionalBatchItemRequestOptions();
+            requestOptions = new ItemBatchRequestOptions();
         }
 
-        this.operations.add(new ItemBatchOperation.Builder(OperationType.Delete, this.operations.size())
-            .requestOptions(requestOptions.toRequestOptions())
-            .id(id)
-            .build());
+        ItemBatchOperation operation = ModelBridgeInternal.createItemBatchOperation(
+            CosmosItemOperationType.Delete,
+            this.operations.size(),
+            id,
+            this.getPartitionKey(),
+            requestOptions,
+            null
+        );
+
+        this.operations.add(operation);
 
-        return this;
+        return operation;
     }
 
     /**
@@ -192,9 +201,9 @@ public TransactionalBatch deleteItem(String id, TransactionalBatchItemRequestOpt
      *
      * @return The transactional batch instance with the operation added.
      */
-    public TransactionalBatch readItem(String id) {
+    public ItemBatchOperation readItem(String id) {
         checkNotNull(id, "expected non-null id");
-        return this.readItem(id, new TransactionalBatchItemRequestOptions());
+        return this.readItem(id, new ItemBatchRequestOptions());
     }
 
     /**
@@ -205,19 +214,25 @@ public TransactionalBatch readItem(String id) {
      *
      * @return The transactional batch instance with the operation added.
      */
-    public TransactionalBatch readItem(String id, TransactionalBatchItemRequestOptions requestOptions) {
+    public ItemBatchOperation readItem(String id, ItemBatchRequestOptions requestOptions) {
 
         checkNotNull(id, "expected non-null id");
         if (requestOptions == null) {
-            requestOptions = new TransactionalBatchItemRequestOptions();
+            requestOptions = new ItemBatchRequestOptions();
         }
 
-        this.operations.add(new ItemBatchOperation.Builder(OperationType.Read, this.operations.size())
-            .requestOptions(requestOptions.toRequestOptions())
-            .id(id)
-            .build());
+        ItemBatchOperation operation = ModelBridgeInternal.createItemBatchOperation(
+            CosmosItemOperationType.Read,
+            this.operations.size(),
+            id,
+            this.getPartitionKey(),
+            requestOptions,
+            null
+        );
+
+        this.operations.add(operation);
 
-        return this;
+        return operation;
     }
 
     /**
@@ -229,10 +244,10 @@ public TransactionalBatch readItem(String id, TransactionalBatchItemRequestOptio
      *
      * @return The transactional batch instance with the operation added.
      */
-    public  TransactionalBatch replaceItem(String id, T item) {
+    public  ItemBatchOperation replaceItem(String id, T item) {
         checkNotNull(id, "expected non-null id");
         checkNotNull(item, "expected non-null item");
-        return this.replaceItem(id, item, new TransactionalBatchItemRequestOptions());
+        return this.replaceItem(id, item, new ItemBatchRequestOptions());
     }
 
     /**
@@ -246,22 +261,27 @@ public  TransactionalBatch replaceItem(String id, T item) {
      *
      * @return The transactional batch instance with the operation added.
      */
-    public  TransactionalBatch replaceItem(
-        String id, T item, TransactionalBatchItemRequestOptions requestOptions) {
+    public  ItemBatchOperation replaceItem(
+        String id, T item, ItemBatchRequestOptions requestOptions) {
 
         checkNotNull(id, "expected non-null id");
         checkNotNull(item, "expected non-null item");
         if (requestOptions == null) {
-            requestOptions = new TransactionalBatchItemRequestOptions();
+            requestOptions = new ItemBatchRequestOptions();
         }
 
-        this.operations.add(new ItemBatchOperation.Builder(OperationType.Replace, this.operations.size())
-            .requestOptions(requestOptions.toRequestOptions())
-            .resource(item)
-            .id(id)
-            .build());
+        ItemBatchOperation operation = ModelBridgeInternal.createItemBatchOperation(
+            CosmosItemOperationType.Replace,
+            this.operations.size(),
+            id,
+            this.getPartitionKey(),
+            requestOptions,
+            item
+        );
 
-        return this;
+        this.operations.add(operation);
+
+        return operation;
     }
 
     /**
@@ -272,9 +292,9 @@ public  TransactionalBatch replaceItem(
      *
      * @return The transactional batch instance with the operation added.
      */
-    public  TransactionalBatch upsertItem(T item) {
+    public  ItemBatchOperation upsertItem(T item) {
         checkNotNull(item, "expected non-null item");
-        return this.upsertItem(item, new TransactionalBatchItemRequestOptions());
+        return this.upsertItem(item, new ItemBatchRequestOptions());
     }
 
     /**
@@ -287,19 +307,25 @@ public  TransactionalBatch upsertItem(T item) {
      *
      * @return The transactional batch instance with the operation added.
      */
-    public  TransactionalBatch upsertItem(T item, TransactionalBatchItemRequestOptions requestOptions) {
+    public  ItemBatchOperation upsertItem(T item, ItemBatchRequestOptions requestOptions) {
 
         checkNotNull(item, "expected non-null item");
         if (requestOptions == null) {
-            requestOptions = new TransactionalBatchItemRequestOptions();
+            requestOptions = new ItemBatchRequestOptions();
         }
 
-        this.operations.add(new ItemBatchOperation.Builder(OperationType.Upsert, this.operations.size())
-            .requestOptions(requestOptions.toRequestOptions())
-            .resource(item)
-            .build());
+        ItemBatchOperation operation = ModelBridgeInternal.createItemBatchOperation(
+            CosmosItemOperationType.Upsert,
+            this.operations.size(),
+            null,
+            this.getPartitionKey(),
+            requestOptions,
+            item
+        );
+
+        this.operations.add(operation);
 
-        return this;
+        return operation;
     }
 
     /**
@@ -307,7 +333,7 @@ public  TransactionalBatch upsertItem(T item, TransactionalBatchItemRequestOp
      *
      * @return The list of operations which are to be executed.
      */
-    List> getOperations() {
+    public List> getOperations() {
         return UnmodifiableList.unmodifiableList(operations);
     }
 
@@ -316,7 +342,7 @@ List> getOperations() {
      *
      * @return The partition key for this batch.
      */
-    PartitionKey getPartitionKey() {
+    public PartitionKey getPartitionKey() {
         return partitionKey;
     }
 }
diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/TransactionalBatchOperationResult.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/TransactionalBatchOperationResult.java
index 8fc02acc73de..46726b905771 100644
--- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/TransactionalBatchOperationResult.java
+++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/TransactionalBatchOperationResult.java
@@ -4,6 +4,7 @@
 package com.azure.cosmos;
 
 import com.azure.cosmos.implementation.JsonSerializable;
+import com.azure.cosmos.models.ItemBatchOperation;
 import com.azure.cosmos.util.Beta;
 import com.fasterxml.jackson.databind.node.ObjectNode;
 import java.time.Duration;
@@ -22,6 +23,7 @@ public final class TransactionalBatchOperationResult {
     private final int statusCode;
     private final Duration retryAfter;
     private final int subStatusCode;
+    private final ItemBatchOperation itemBatchOperation;
 
     /**
      * Initializes a new instance of the {@link TransactionalBatchOperationResult} class.
@@ -31,8 +33,10 @@ public final class TransactionalBatchOperationResult {
                                       ObjectNode resourceObject,
                                       int statusCode,
                                       Duration retryAfter,
-                                      int subStatusCode) {
+                                      int subStatusCode,
+                                      ItemBatchOperation itemBatchOperation) {
         checkNotNull(statusCode, "expected non-null statusCode");
+        checkNotNull(itemBatchOperation, "expected non-null itemBatchOperation");
 
         this.eTag = eTag;
         this.requestCharge = requestCharge;
@@ -40,6 +44,7 @@ public final class TransactionalBatchOperationResult {
         this.statusCode = statusCode;
         this.retryAfter = retryAfter;
         this.subStatusCode = subStatusCode;
+        this.itemBatchOperation = itemBatchOperation;
     }
 
     /**
@@ -123,4 +128,13 @@ public int getStatusCode() {
     ObjectNode getResourceObject() {
         return resourceObject;
     }
+
+    /**
+     * Gets the original ItemBatchOperation for this result.
+     *
+     * @return the ItemBatchOperation.
+     */
+    public ItemBatchOperation getItemBatchOperation() {
+        return itemBatchOperation;
+    }
 }
diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/TransactionalBatchResponse.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/TransactionalBatchResponse.java
index 006b488336b3..31176f241b6a 100644
--- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/TransactionalBatchResponse.java
+++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/TransactionalBatchResponse.java
@@ -8,7 +8,7 @@
 
 import java.time.Duration;
 import java.util.ArrayList;
-import java.util.Collection;
+import java.util.Collections;
 import java.util.List;
 import java.util.Map;
 
@@ -157,23 +157,13 @@ public int getSubStatusCode() {
     }
 
     /**
-     * Get all the results of the operations in batch.
+     * Get all the results of the operations in a batch in an unmodifiable instance so no one can
+     * change it in the down path.
      *
-     * @return Results of operation in batch.
+     * @return Results of operations in a batch.
      */
     public List getResults() {
-        return this.results;
-    }
-
-    /**
-     * Gets the result of the operation at the provided index in the batch.
-     *
-     * @param index 0-based index of the operation in the batch whose result needs to be returned.
-     *
-     * @return Result of operation at the provided index in the batch.
-     */
-    public TransactionalBatchOperationResult get(int index) {
-        return this.results.get(index);
+        return Collections.unmodifiableList(this.results);
     }
 
     /**
@@ -189,7 +179,7 @@ public Duration getDuration() {
         return this.cosmosDiagnostics.getDuration();
     }
 
-    boolean addAll(Collection collection) {
-        return this.results.addAll(collection);
+    void addAll(List collection) {
+        this.results.addAll(collection);
     }
 }
diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/batch/BatchExecUtils.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/batch/BatchExecUtils.java
index b2f7cb5a6852..520794fb561e 100644
--- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/batch/BatchExecUtils.java
+++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/batch/BatchExecUtils.java
@@ -4,7 +4,6 @@
 package com.azure.cosmos.implementation.batch;
 
 import com.azure.cosmos.implementation.HttpConstants;
-import com.azure.cosmos.implementation.OperationType;
 import com.azure.cosmos.implementation.apachecommons.lang.StringUtils;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
@@ -12,12 +11,6 @@
 import java.time.Duration;
 import java.util.Map;
 
-import static com.azure.cosmos.implementation.batch.BatchRequestResponseConstant.OPERATION_CREATE;
-import static com.azure.cosmos.implementation.batch.BatchRequestResponseConstant.OPERATION_DELETE;
-import static com.azure.cosmos.implementation.batch.BatchRequestResponseConstant.OPERATION_READ;
-import static com.azure.cosmos.implementation.batch.BatchRequestResponseConstant.OPERATION_REPLACE;
-import static com.azure.cosmos.implementation.batch.BatchRequestResponseConstant.OPERATION_UPSERT;
-
 /**
  * Util methods for batch requests/response.
  */
@@ -25,23 +18,6 @@ public final class BatchExecUtils {
 
     private final static Logger logger = LoggerFactory.getLogger(BatchExecUtils.class);
 
-    static String getStringOperationType(OperationType operationType) {
-        switch (operationType) {
-            case Create:
-                return OPERATION_CREATE;
-            case Delete:
-                return OPERATION_DELETE;
-            case Read:
-                return OPERATION_READ;
-            case Replace:
-                return OPERATION_REPLACE;
-            case Upsert:
-                return OPERATION_UPSERT;
-        }
-
-        return null;
-    }
-
     public static Duration getRetryAfterDuration(Map responseHeaders) {
         long retryIntervalInMilliseconds = 0;
 
@@ -87,7 +63,7 @@ public static double getRequestCharge(Map responseHeaders) {
         }
 
         try {
-            return Double.valueOf(value);
+            return Double.parseDouble(value);
         } catch (NumberFormatException e) {
             logger.warn("INVALID x-ms-request-charge value {}.", value);
             return 0;
diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/batch/BatchExecutor.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/batch/BatchExecutor.java
index b3160e4ac317..088f3790a1d1 100644
--- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/batch/BatchExecutor.java
+++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/batch/BatchExecutor.java
@@ -9,6 +9,7 @@
 import com.azure.cosmos.TransactionalBatch;
 import com.azure.cosmos.TransactionalBatchRequestOptions;
 import com.azure.cosmos.TransactionalBatchResponse;
+import com.azure.cosmos.models.ItemBatchOperation;
 import reactor.core.publisher.Mono;
 
 import java.util.List;
@@ -38,11 +39,11 @@ public BatchExecutor(
      */
     public final Mono executeAsync() {
 
-        List> operations = BridgeInternal.getOperationsFromTransactionalBatch(this.transactionalBatch);
+        List> operations = this.transactionalBatch.getOperations();
         checkArgument(operations.size() > 0, "Number of operations should be more than 0.");
 
         final SinglePartitionKeyServerBatchRequest request = SinglePartitionKeyServerBatchRequest.createBatchRequest(
-            BridgeInternal.getPartitionKeyFromTransactionalBatch(this.transactionalBatch),
+            this.transactionalBatch.getPartitionKey(),
             operations);
         request.setAtomicBatch(true);
         request.setShouldContinueOnError(false);
diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/batch/BatchRequestResponseConstant.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/batch/BatchRequestResponseConstant.java
index d5fcc9588ac1..84429b20e295 100644
--- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/batch/BatchRequestResponseConstant.java
+++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/batch/BatchRequestResponseConstant.java
@@ -33,10 +33,10 @@ public class BatchRequestResponseConstant {
     static final String FIELD_IS_CLIENTENCRYPTED = "isClientEncrypted";
 
     // Batch supported operation type for json
-    static final String OPERATION_CREATE = "Create";
-    static final String OPERATION_PATCH = "Patch";
-    static final String OPERATION_READ = "Read";
-    static final String OPERATION_UPSERT = "Upsert";
-    static final String OPERATION_DELETE = "Delete";
-    static final String OPERATION_REPLACE = "Replace";
+    public static final String OPERATION_CREATE = "Create";
+    public static final String OPERATION_PATCH = "Patch";
+    public static final String OPERATION_READ = "Read";
+    public static final String OPERATION_UPSERT = "Upsert";
+    public static final String OPERATION_DELETE = "Delete";
+    public static final String OPERATION_REPLACE = "Replace";
 }
diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/batch/BatchResponseParser.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/batch/BatchResponseParser.java
index 678185410996..b18047d89813 100644
--- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/batch/BatchResponseParser.java
+++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/batch/BatchResponseParser.java
@@ -10,6 +10,7 @@
 import com.azure.cosmos.implementation.JsonSerializable;
 import com.azure.cosmos.implementation.RxDocumentServiceResponse;
 import com.azure.cosmos.implementation.Utils;
+import com.azure.cosmos.models.ItemBatchOperation;
 import com.fasterxml.jackson.databind.ObjectMapper;
 import com.fasterxml.jackson.databind.node.ObjectNode;
 import io.netty.handler.codec.http.HttpResponseStatus;
@@ -103,7 +104,7 @@ private static TransactionalBatchResponse populateFromResponseContent(
         final ServerBatchRequest request,
         final boolean shouldPromoteOperationStatus) {
 
-        final ArrayList results = new ArrayList<>(request.getOperations().size());
+        final List results = new ArrayList<>(request.getOperations().size());
         final byte[] responseContent = documentServiceResponse.getResponseBodyAsByteArray();
 
         if (responseContent[0] != (byte)HYBRID_V1) {
@@ -111,10 +112,14 @@ private static TransactionalBatchResponse populateFromResponseContent(
             final ObjectMapper mapper = Utils.getSimpleObjectMapper();
 
             try {
+                final List> itemBatchOperations = request.getOperations();
                 final ObjectNode[] objectNodes = mapper.readValue(responseContent, ObjectNode[].class);
-                for (ObjectNode objectInArray : objectNodes) {
-                    final TransactionalBatchOperationResult batchOperationResult = BatchResponseParser.createBatchOperationResultFromJson(objectInArray);
-                    results.add(batchOperationResult);
+
+                for (int index = 0; index < objectNodes.length; index++) {
+                    ObjectNode objectInArray = objectNodes[index];
+
+                    results.add(
+                        BatchResponseParser.createBatchOperationResultFromJson(objectInArray, itemBatchOperations.get(index)));
                 }
             } catch (IOException ex) {
                 logger.error("Exception in parsing response", ex);
@@ -163,7 +168,10 @@ private static TransactionalBatchResponse populateFromResponseContent(
      *
      * @return the result
      */
-    private static TransactionalBatchOperationResult createBatchOperationResultFromJson(ObjectNode objectNode) {
+    private static TransactionalBatchOperationResult createBatchOperationResultFromJson(
+        ObjectNode objectNode,
+        ItemBatchOperation itemBatchOperation) {
+
         final JsonSerializable jsonSerializable = new JsonSerializable(objectNode);
 
         final int statusCode = jsonSerializable.getInt(BatchRequestResponseConstant.FIELD_STATUS_CODE);
@@ -187,7 +195,8 @@ private static TransactionalBatchOperationResult createBatchOperationResultFromJ
             resourceBody,
             statusCode,
             retryAfterMilliseconds != null ? Duration.ofMillis(retryAfterMilliseconds) : Duration.ZERO,
-            subStatusCode);
+            subStatusCode,
+            itemBatchOperation);
     }
 
     /**
@@ -201,7 +210,7 @@ private static void createAndPopulateResults(final TransactionalBatchResponse re
                                                  final List> operations,
                                                  final Duration retryAfterDuration) {
         final List results = new ArrayList<>(operations.size());
-        for (int i = 0; i < operations.size(); i++) {
+        for (ItemBatchOperation itemBatchOperation : operations) {
             results.add(
                 BridgeInternal.createTransactionBatchResult(
                     null,
@@ -209,7 +218,8 @@ private static void createAndPopulateResults(final TransactionalBatchResponse re
                     null,
                     response.getStatusCode(),
                     retryAfterDuration,
-                    response.getSubStatusCode()
+                    response.getSubStatusCode(),
+                    itemBatchOperation
                 ));
         }
 
diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/batch/ItemBatchOperation.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/batch/ItemBatchOperation.java
deleted file mode 100644
index 0db9ba62c406..000000000000
--- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/batch/ItemBatchOperation.java
+++ /dev/null
@@ -1,163 +0,0 @@
-// Copyright (c) Microsoft Corporation. All rights reserved.
-// Licensed under the MIT License.
-
-package com.azure.cosmos.implementation.batch;
-
-import com.azure.cosmos.implementation.JsonSerializable;
-import com.azure.cosmos.implementation.OperationType;
-import com.azure.cosmos.implementation.RequestOptions;
-import com.azure.cosmos.implementation.apachecommons.lang.StringUtils;
-import com.azure.cosmos.models.PartitionKey;
-
-import static com.azure.cosmos.implementation.guava25.base.Preconditions.checkArgument;
-import static com.azure.cosmos.implementation.guava25.base.Preconditions.checkNotNull;
-
-/**
- * Represents an operation on an item which will be executed as part of a batch request on a container.
- */
-public final class ItemBatchOperation {
-
-    private TResource resource;
-
-    private final String id;
-    private final int operationIndex;
-    private final PartitionKey partitionKey;
-    private String partitionKeyJson;
-    private final OperationType operationType;
-    private final RequestOptions requestOptions;
-
-    private ItemBatchOperation(
-        final OperationType operationType,
-        final int operationIndex,
-        final PartitionKey partitionKey,
-        final String id,
-        final TResource resource,
-        final RequestOptions requestOptions) {
-
-        checkArgument(operationIndex >= 0, "expected operationIndex >= 0, not %s", operationIndex);
-        checkNotNull(operationType, "expected non-null operationType");
-
-        this.operationType = operationType;
-        this.operationIndex = operationIndex;
-        this.partitionKey = partitionKey;
-        this.id = id;
-        this.resource = resource;
-        this.requestOptions = requestOptions;
-    }
-
-    // TODO(rakkuma): Similarly for hybrid row, operation needs to be written in Hybrid row.
-    // Issue: https://github.com/Azure/azure-sdk-for-java/issues/15856
-    static JsonSerializable writeOperation(final ItemBatchOperation operation) {
-        final JsonSerializable jsonSerializable = new JsonSerializable();
-
-        jsonSerializable.set(BatchRequestResponseConstant.FIELD_OPERATION_TYPE, BatchExecUtils.getStringOperationType(operation.getOperationType()));
-
-        if (StringUtils.isNotEmpty(operation.getPartitionKeyJson())) {
-            // Used for non transactional batch.
-            jsonSerializable.set(BatchRequestResponseConstant.FIELD_PARTITION_KEY, operation.getPartitionKeyJson());
-        }
-
-        if (StringUtils.isNotEmpty(operation.getId())) {
-            jsonSerializable.set(BatchRequestResponseConstant.FIELD_ID, operation.getId());
-        }
-
-        if (operation.getResource() != null) {
-            jsonSerializable.set(BatchRequestResponseConstant.FIELD_RESOURCE_BODY, operation.getResource());
-        }
-
-        if (operation.getRequestOptions() != null) {
-            RequestOptions requestOptions = operation.getRequestOptions();
-
-            if (StringUtils.isNotEmpty(requestOptions.getIfMatchETag())) {
-                jsonSerializable.set(BatchRequestResponseConstant.FIELD_IF_MATCH, requestOptions.getIfMatchETag());
-            }
-
-            if (StringUtils.isNotEmpty(requestOptions.getIfNoneMatchETag())) {
-                jsonSerializable.set(BatchRequestResponseConstant.FIELD_IF_NONE_MATCH, requestOptions.getIfNoneMatchETag());
-            }
-        }
-
-        return jsonSerializable;
-    }
-
-    String getId() {
-        return this.id;
-    }
-
-    OperationType getOperationType() {
-        return this.operationType;
-    }
-
-    String getPartitionKeyJson() {
-        return this.partitionKeyJson;
-    }
-
-    int getOperationIndex() {
-        return operationIndex;
-    }
-
-    PartitionKey getPartitionKey() {
-        return partitionKey;
-    }
-
-    void setPartitionKeyJson(String partitionKeyJson) {
-        this.partitionKeyJson = partitionKeyJson;
-    }
-
-    RequestOptions getRequestOptions() {
-        return this.requestOptions;
-    }
-
-    TResource getResource() {
-        return this.resource;
-    }
-
-    public static final class Builder {
-
-        private final OperationType operationType;
-        private final int operationIndex;
-        private String id;
-        private PartitionKey partitionKey;
-        private RequestOptions requestOptions;
-        private TResource resource;
-
-        public Builder(final OperationType type, final int index) {
-
-            checkNotNull(type, "expected non-null type");
-            checkArgument(index >= 0, "expected index >= 0, not %s", index);
-
-            this.operationType = type;
-            this.operationIndex = index;
-        }
-
-        public Builder id(String value) {
-            this.id = value;
-            return this;
-        }
-
-        public Builder partitionKey(PartitionKey value) {
-            this.partitionKey = value;
-            return this;
-        }
-
-        public Builder requestOptions(RequestOptions value) {
-            this.requestOptions = value;
-            return this;
-        }
-
-        public Builder resource(TResource value) {
-            this.resource = value;
-            return this;
-        }
-
-        public ItemBatchOperation build() {
-            return new ItemBatchOperation<>(
-                this.operationType,
-                this.operationIndex,
-                this.partitionKey,
-                this.id,
-                this.resource,
-                this.requestOptions);
-        }
-    }
-}
diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/batch/ServerBatchRequest.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/batch/ServerBatchRequest.java
index 41ac86a20363..f4db876d7369 100644
--- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/batch/ServerBatchRequest.java
+++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/batch/ServerBatchRequest.java
@@ -3,9 +3,12 @@
 
 package com.azure.cosmos.implementation.batch;
 
+import com.azure.cosmos.ItemBatchRequestOptions;
 import com.azure.cosmos.implementation.JsonSerializable;
 import com.azure.cosmos.implementation.Utils;
 import com.azure.cosmos.implementation.apachecommons.collections.list.UnmodifiableList;
+import com.azure.cosmos.implementation.apachecommons.lang.StringUtils;
+import com.azure.cosmos.models.ItemBatchOperation;
 import com.fasterxml.jackson.databind.node.ArrayNode;
 
 import java.util.List;
@@ -59,7 +62,7 @@ final List> createBodyOfBatchRequest(final List operation : operations) {
 
-            final JsonSerializable operationJsonSerializable = ItemBatchOperation.writeOperation(operation);
+            final JsonSerializable operationJsonSerializable = this.writeOperation(operation);
             final int operationSerializedLength = operationJsonSerializable.toString().length();
 
             if (totalOperationCount != 0 &&
@@ -76,8 +79,45 @@ final List> createBodyOfBatchRequest(final List> pendingOperations = operations.subList(totalOperationCount, operations.size());
-        return pendingOperations;
+
+        return operations.subList(totalOperationCount, operations.size());
+    }
+
+    /**
+     * Writes a single operation to JsonSerializable.
+     * TODO(rakkuma): Similarly for hybrid row, operation needs to be written in Hybrid row.
+     * Issue: https://github.com/Azure/azure-sdk-for-java/issues/15856
+     *
+     * @param operation a single operation which needs to be serialized.
+     *
+     * @return instance of JsonSerializable containing values for a operation.
+     */
+    final JsonSerializable writeOperation(final ItemBatchOperation operation) {
+        final JsonSerializable jsonSerializable = new JsonSerializable();
+
+        jsonSerializable.set(BatchRequestResponseConstant.FIELD_OPERATION_TYPE, operation.getOperationType().getOperationValue());
+
+        if (StringUtils.isNotEmpty(operation.getId())) {
+            jsonSerializable.set(BatchRequestResponseConstant.FIELD_ID, operation.getId());
+        }
+
+        if (operation.getItem() != null) {
+            jsonSerializable.set(BatchRequestResponseConstant.FIELD_RESOURCE_BODY, operation.getItem());
+        }
+
+        if (operation.getItemBatchRequestOptions() != null) {
+            ItemBatchRequestOptions requestOptions = operation.getItemBatchRequestOptions();
+
+            if (StringUtils.isNotEmpty(requestOptions.getIfMatchETag())) {
+                jsonSerializable.set(BatchRequestResponseConstant.FIELD_IF_MATCH, requestOptions.getIfMatchETag());
+            }
+
+            if (StringUtils.isNotEmpty(requestOptions.getIfNoneMatchETag())) {
+                jsonSerializable.set(BatchRequestResponseConstant.FIELD_IF_NONE_MATCH, requestOptions.getIfNoneMatchETag());
+            }
+        }
+
+        return jsonSerializable;
     }
 
     public final String getRequestBody() {
diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/batch/SinglePartitionKeyServerBatchRequest.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/batch/SinglePartitionKeyServerBatchRequest.java
index 3d6e11ad5c41..78f6ae216016 100644
--- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/batch/SinglePartitionKeyServerBatchRequest.java
+++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/batch/SinglePartitionKeyServerBatchRequest.java
@@ -3,6 +3,7 @@
 
 package com.azure.cosmos.implementation.batch;
 
+import com.azure.cosmos.models.ItemBatchOperation;
 import com.azure.cosmos.models.PartitionKey;
 import java.util.List;
 
diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/models/ItemBatchOperation.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/models/ItemBatchOperation.java
new file mode 100644
index 000000000000..8f864fed40b6
--- /dev/null
+++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/models/ItemBatchOperation.java
@@ -0,0 +1,70 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+package com.azure.cosmos.models;
+
+import com.azure.cosmos.ItemBatchRequestOptions;
+import com.azure.cosmos.CosmosItemOperationType;
+
+import static com.azure.cosmos.implementation.guava25.base.Preconditions.checkArgument;
+import static com.azure.cosmos.implementation.guava25.base.Preconditions.checkNotNull;
+
+/**
+ * Represents an operation on an item which will be executed as part of a batch request on a container. This will be
+ * serialized and sent in the request.
+ *
+ * @param  The type of item.
+ */
+public final class ItemBatchOperation {
+
+    private T item;
+
+    private final String id;
+    private final int operationIndex;
+    private final PartitionKey partitionKey;
+    private final CosmosItemOperationType operationType;
+    private final ItemBatchRequestOptions requestOptions;
+
+    ItemBatchOperation(
+        final CosmosItemOperationType operationType,
+        final int operationIndex,
+        final PartitionKey partitionKey,
+        final String id,
+        final T item,
+        final ItemBatchRequestOptions requestOptions) {
+
+        checkNotNull(operationType, "expected non-null operationType");
+        checkArgument(operationIndex >= 0, "expected operationIndex >= 0, not %s", operationIndex);
+
+        this.operationType = operationType;
+        this.operationIndex = operationIndex;
+        this.partitionKey = partitionKey;
+        this.id = id;
+        this.item = item;
+        this.requestOptions = requestOptions;
+    }
+
+    public T getItem() {
+        return this.item;
+    }
+
+    public String getId() {
+        return this.id;
+    }
+
+    public int getOperationIndex() {
+        return operationIndex;
+    }
+
+    public PartitionKey getPartitionKey() {
+        return partitionKey;
+    }
+
+    public CosmosItemOperationType getOperationType() {
+        return this.operationType;
+    }
+
+    public ItemBatchRequestOptions getItemBatchRequestOptions() {
+        return this.requestOptions;
+    }
+}
diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/models/ModelBridgeInternal.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/models/ModelBridgeInternal.java
index dcde8b83e57a..426a530cb3a0 100644
--- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/models/ModelBridgeInternal.java
+++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/models/ModelBridgeInternal.java
@@ -3,6 +3,8 @@
 
 package com.azure.cosmos.models;
 
+import com.azure.cosmos.ItemBatchRequestOptions;
+import com.azure.cosmos.CosmosItemOperationType;
 import com.azure.cosmos.implementation.Conflict;
 import com.azure.cosmos.implementation.ConsistencyPolicy;
 import com.azure.cosmos.implementation.CosmosResourceType;
@@ -724,4 +726,22 @@ public static CosmosItemRequestOptions clone(CosmosItemRequestOptions options) {
     public static  int getPayloadLength(CosmosItemResponse cosmosItemResponse) {
         return cosmosItemResponse.responseBodyAsByteArray != null ? cosmosItemResponse.responseBodyAsByteArray.length : 0;
     }
+
+    @Warning(value = INTERNAL_USE_ONLY_WARNING)
+    public static  ItemBatchOperation createItemBatchOperation(
+        final CosmosItemOperationType operationType,
+        final int operationIndex,
+        String id,
+        PartitionKey partitionKey,
+        ItemBatchRequestOptions itemBatchRequestOptions,
+        T resource) {
+
+        return new ItemBatchOperation<>(
+            operationType,
+            operationIndex,
+            partitionKey,
+            id,
+            resource,
+            itemBatchRequestOptions);
+    }
 }
diff --git a/sdk/cosmos/azure-cosmos/src/test/java/com/azure/cosmos/BatchOperationResultTests.java b/sdk/cosmos/azure-cosmos/src/test/java/com/azure/cosmos/BatchOperationResultTests.java
index e19ce110d79a..2394dd51ee4d 100644
--- a/sdk/cosmos/azure-cosmos/src/test/java/com/azure/cosmos/BatchOperationResultTests.java
+++ b/sdk/cosmos/azure-cosmos/src/test/java/com/azure/cosmos/BatchOperationResultTests.java
@@ -5,6 +5,8 @@
 
 import com.azure.cosmos.implementation.HttpConstants;
 import com.azure.cosmos.implementation.Utils;
+import com.azure.cosmos.models.ItemBatchOperation;
+import com.azure.cosmos.models.ModelBridgeInternal;
 import com.fasterxml.jackson.databind.node.ObjectNode;
 import io.netty.handler.codec.http.HttpResponseStatus;
 import org.testng.annotations.Test;
@@ -16,6 +18,14 @@ public class BatchOperationResultTests {
 
     private static final int TIMEOUT = 40000;
     private ObjectNode objectNode = Utils.getSimpleObjectMapper().createObjectNode();
+    private ItemBatchOperation operation = ModelBridgeInternal.createItemBatchOperation(
+        CosmosItemOperationType.Read,
+        1,
+        null,
+        null,
+        null,
+        null
+        );
 
     private TransactionalBatchOperationResult createTestResult() {
         TransactionalBatchOperationResult result = BridgeInternal.createTransactionBatchResult(
@@ -24,7 +34,8 @@ private TransactionalBatchOperationResult createTestResult() {
             objectNode,
             HttpResponseStatus.OK.code(),
             Duration.ofMillis(1234),
-            HttpConstants.SubStatusCodes.NAME_CACHE_IS_STALE
+            HttpConstants.SubStatusCodes.NAME_CACHE_IS_STALE,
+            operation
         );
 
         return result;
@@ -40,6 +51,7 @@ public void propertiesAreSetThroughCtor() {
         assertThat(result.getRequestCharge()).isEqualTo(1.4);
         assertThat(result.getRetryAfterDuration()).isEqualTo(Duration.ofMillis(1234));
         assertThat(result.getResourceObject()).isSameAs(objectNode);
+        assertThat(result.getItemBatchOperation()).isSameAs(operation);
     }
 
     @Test(groups = {"unit"}, timeOut = TIMEOUT)
@@ -51,7 +63,8 @@ public void isSuccessStatusCodeTrueFor200To299() {
                 null,
                 x,
                 null,
-                0
+                0,
+                operation
             );
 
             boolean success = x >= 200 && x <= 299;
diff --git a/sdk/cosmos/azure-cosmos/src/test/java/com/azure/cosmos/TransactionalBatchAsyncContainerTest.java b/sdk/cosmos/azure-cosmos/src/test/java/com/azure/cosmos/TransactionalBatchAsyncContainerTest.java
index e40cce07c24e..4f4a779c8990 100644
--- a/sdk/cosmos/azure-cosmos/src/test/java/com/azure/cosmos/TransactionalBatchAsyncContainerTest.java
+++ b/sdk/cosmos/azure-cosmos/src/test/java/com/azure/cosmos/TransactionalBatchAsyncContainerTest.java
@@ -4,8 +4,8 @@
 package com.azure.cosmos;
 
 import com.azure.cosmos.implementation.HttpConstants;
-import com.azure.cosmos.implementation.ISessionToken;
 import com.azure.cosmos.models.CosmosItemResponse;
+import com.azure.cosmos.models.ItemBatchOperation;
 import io.netty.handler.codec.http.HttpResponseStatus;
 import org.assertj.core.api.Assertions;
 import org.testng.annotations.AfterClass;
@@ -14,6 +14,8 @@
 import org.testng.annotations.Test;
 import reactor.core.publisher.Mono;
 
+import java.util.List;
+
 import static org.assertj.core.api.Assertions.assertThat;
 
 public class TransactionalBatchAsyncContainerTest extends BatchTestBase {
@@ -43,23 +45,23 @@ public void batchExecutionRepeat() {
         TestDoc firstDoc = this.populateTestDoc(this.partitionKey1);
         TestDoc replaceDoc = this.getTestDocCopy(firstDoc);
         replaceDoc.setCost(replaceDoc.getCost() + 1);
+        TransactionalBatch batch = TransactionalBatch.createTransactionalBatch(this.getPartitionKey(this.partitionKey1));
+        batch.createItem(firstDoc);
+        batch.replaceItem(replaceDoc.getId(), replaceDoc);
 
-        Mono batchResponseMono = batchAsyncContainer.executeTransactionalBatch(
-            TransactionalBatch.createTransactionalBatch(this.getPartitionKey(this.partitionKey1))
-                .createItem(firstDoc)
-                .replaceItem(replaceDoc.getId(), replaceDoc));
+        Mono batchResponseMono = batchAsyncContainer.executeTransactionalBatch(batch);
 
         TransactionalBatchResponse batchResponse1 = batchResponseMono.block();
         this.verifyBatchProcessed(batchResponse1, 2);
 
-        assertThat(batchResponse1.get(0).getStatusCode()).isEqualTo(HttpResponseStatus.CREATED.code());
-        assertThat(batchResponse1.get(1).getStatusCode()).isEqualTo(HttpResponseStatus.OK.code());
+        assertThat(batchResponse1.getResults().get(0).getStatusCode()).isEqualTo(HttpResponseStatus.CREATED.code());
+        assertThat(batchResponse1.getResults().get(1).getStatusCode()).isEqualTo(HttpResponseStatus.OK.code());
 
         // Block again.
         TransactionalBatchResponse batchResponse2 = batchResponseMono.block();
         assertThat(batchResponse2.getStatusCode()).isEqualTo(HttpResponseStatus.CONFLICT.code());
-        assertThat(batchResponse2.get(0).getStatusCode()).isEqualTo(HttpResponseStatus.CONFLICT.code());
-        assertThat(batchResponse2.get(1).getStatusCode()).isEqualTo(HttpResponseStatus.FAILED_DEPENDENCY.code());
+        assertThat(batchResponse2.getResults().get(0).getStatusCode()).isEqualTo(HttpResponseStatus.CONFLICT.code());
+        assertThat(batchResponse2.getResults().get(1).getStatusCode()).isEqualTo(HttpResponseStatus.FAILED_DEPENDENCY.code());
     }
 
     @Test(groups = {"simple"}, timeOut = TIMEOUT * 100)
@@ -82,19 +84,26 @@ public void batchInvalidSessionToken() throws Exception {
             testDocToReplace.setCost(testDocToReplace.getCost() + 1);
             TestDoc testDocToUpsert = this.populateTestDoc(this.partitionKey1);
 
+            TransactionalBatch batch = TransactionalBatch.createTransactionalBatch(this.getPartitionKey(this.partitionKey1));
+            batch.createItem(testDocToCreate);
+            batch.replaceItem(testDocToReplace.getId(), testDocToReplace);
+            batch.upsertItem(testDocToUpsert);
+            batch.deleteItem(this.TestDocPk1ExistingC.getId());
+
             TransactionalBatchResponse batchResponse = container.executeTransactionalBatch(
-                TransactionalBatch.createTransactionalBatch(this.getPartitionKey(this.partitionKey1))
-                    .createItem(testDocToCreate)
-                    .replaceItem(testDocToReplace.getId(), testDocToReplace)
-                    .upsertItem(testDocToUpsert)
-                    .deleteItem(this.TestDocPk1ExistingC.getId()), new TransactionalBatchRequestOptions().setSessionToken(invalidSessionToken)).block();
+                batch, new TransactionalBatchRequestOptions().setSessionToken(invalidSessionToken)).block();
 
             this.verifyBatchProcessed(batchResponse, 4);
 
-            assertThat(batchResponse.get(0).getStatusCode()).isEqualTo(HttpResponseStatus.CREATED.code());
-            assertThat(batchResponse.get(1).getStatusCode()).isEqualTo(HttpResponseStatus.OK.code());
-            assertThat(batchResponse.get(2).getStatusCode()).isEqualTo(HttpResponseStatus.CREATED.code());
-            assertThat(batchResponse.get(3).getStatusCode()).isEqualTo(HttpResponseStatus.NO_CONTENT.code());
+            assertThat(batchResponse.getResults().get(0).getStatusCode()).isEqualTo(HttpResponseStatus.CREATED.code());
+            assertThat(batchResponse.getResults().get(1).getStatusCode()).isEqualTo(HttpResponseStatus.OK.code());
+            assertThat(batchResponse.getResults().get(2).getStatusCode()).isEqualTo(HttpResponseStatus.CREATED.code());
+            assertThat(batchResponse.getResults().get(3).getStatusCode()).isEqualTo(HttpResponseStatus.NO_CONTENT.code());
+
+            List> batchOperations = batch.getOperations();
+            for (int index = 0; index < batchOperations.size(); index++) {
+                assertThat(batchResponse.getResults().get(index).getItemBatchOperation()).isEqualTo(batchOperations.get(index));
+            }
         }
 
         {
@@ -104,14 +113,16 @@ public void batchInvalidSessionToken() throws Exception {
             testDocToReplace.setCost(testDocToReplace.getCost() + 1);
             TestDoc testDocToUpsert = this.populateTestDoc(this.partitionKey1);
 
+            TransactionalBatch batch = TransactionalBatch.createTransactionalBatch(this.getPartitionKey(this.partitionKey1));
+            batch.createItem(testDocToCreate);
+            batch.replaceItem(testDocToReplace.getId(), testDocToReplace);
+            batch.upsertItem(testDocToUpsert);
+            batch.deleteItem(this.TestDocPk1ExistingD.getId());
+            batch.readItem(this.TestDocPk1ExistingA.getId());
+
             try {
                 container.executeTransactionalBatch(
-                    TransactionalBatch.createTransactionalBatch(this.getPartitionKey(this.partitionKey1))
-                        .createItem(testDocToCreate)
-                        .replaceItem(testDocToReplace.getId(), testDocToReplace)
-                        .upsertItem(testDocToUpsert)
-                        .deleteItem(this.TestDocPk1ExistingD.getId())
-                        .readItem(this.TestDocPk1ExistingA.getId()), new TransactionalBatchRequestOptions().setSessionToken(invalidSessionToken)).block();
+                    batch, new TransactionalBatchRequestOptions().setSessionToken(invalidSessionToken)).block();
 
                 Assertions.fail("Should throw NOT_FOUND/READ_SESSION_NOT_AVAILABLE exception");
             } catch (CosmosException ex) {
diff --git a/sdk/cosmos/azure-cosmos/src/test/java/com/azure/cosmos/TransactionalBatchTest.java b/sdk/cosmos/azure-cosmos/src/test/java/com/azure/cosmos/TransactionalBatchTest.java
index 1ad43fa6da03..30d5bc03ad28 100644
--- a/sdk/cosmos/azure-cosmos/src/test/java/com/azure/cosmos/TransactionalBatchTest.java
+++ b/sdk/cosmos/azure-cosmos/src/test/java/com/azure/cosmos/TransactionalBatchTest.java
@@ -7,6 +7,7 @@
 import com.azure.cosmos.implementation.ISessionToken;
 import com.azure.cosmos.implementation.guava25.base.Function;
 import com.azure.cosmos.models.CosmosItemResponse;
+import com.azure.cosmos.models.ItemBatchOperation;
 import io.netty.handler.codec.http.HttpResponseStatus;
 import org.assertj.core.api.Assertions;
 import org.testng.annotations.AfterClass;
@@ -14,6 +15,7 @@
 import org.testng.annotations.Factory;
 import org.testng.annotations.Test;
 
+import java.util.List;
 import java.util.UUID;
 
 import static com.azure.cosmos.implementation.batch.BatchRequestResponseConstant.MAX_DIRECT_MODE_BATCH_REQUEST_BODY_SIZE_IN_BYTES;
@@ -51,15 +53,21 @@ public void batchOrdered() {
         TestDoc replaceDoc = this.getTestDocCopy(firstDoc);
         replaceDoc.setCost(replaceDoc.getCost() + 1);
 
-        TransactionalBatchResponse batchResponse = container.executeTransactionalBatch(
-            TransactionalBatch.createTransactionalBatch(this.getPartitionKey(this.partitionKey1))
-                .createItem(firstDoc)
-                .replaceItem(replaceDoc.getId(), replaceDoc));
+        TransactionalBatch batch = TransactionalBatch.createTransactionalBatch(this.getPartitionKey(this.partitionKey1));
+        batch.createItem(firstDoc);
+        batch.replaceItem(replaceDoc.getId(), replaceDoc);
+
+        TransactionalBatchResponse batchResponse = container.executeTransactionalBatch(batch);
 
         this.verifyBatchProcessed(batchResponse, 2);
 
-        assertThat(batchResponse.get(0).getStatusCode()).isEqualTo(HttpResponseStatus.CREATED.code());
-        assertThat(batchResponse.get(1).getStatusCode()).isEqualTo(HttpResponseStatus.OK.code());
+        assertThat(batchResponse.getResults().get(0).getStatusCode()).isEqualTo(HttpResponseStatus.CREATED.code());
+        assertThat(batchResponse.getResults().get(1).getStatusCode()).isEqualTo(HttpResponseStatus.OK.code());
+
+        List> batchOperations = batch.getOperations();
+        for (int index = 0; index < batchOperations.size(); index++) {
+            assertThat(batchResponse.getResults().get(index).getItemBatchOperation()).isEqualTo(batchOperations.get(index));
+        }
 
         // Ensure that the replace overwrote the doc from the first operation
         this.verifyByRead(container, replaceDoc);
@@ -78,29 +86,35 @@ public void batchMultipleItemExecution() {
         CosmosItemResponse createResponse = container.createItem(readEventDoc, this.getPartitionKey(this.partitionKey1), null);
         assertThat(createResponse.getStatusCode()).isEqualTo(HttpResponseStatus.CREATED.code());
 
-        TransactionalBatchResponse batchResponse = container.executeTransactionalBatch(
-            TransactionalBatch.createTransactionalBatch(this.getPartitionKey(this.partitionKey1))
-                .createItem(firstDoc)
-                .createItem(eventDoc1)
-                .replaceItem(replaceDoc.getId(), replaceDoc)
-                .readItem(readEventDoc.getId()));
+        TransactionalBatch batch = TransactionalBatch.createTransactionalBatch(this.getPartitionKey(this.partitionKey1));
+        batch.createItem(firstDoc);
+        batch.createItem(eventDoc1);
+        batch.replaceItem(replaceDoc.getId(), replaceDoc);
+        batch.readItem(readEventDoc.getId());
+
+        TransactionalBatchResponse batchResponse = container.executeTransactionalBatch(batch);
 
         this.verifyBatchProcessed(batchResponse, 4);
 
-        assertThat(batchResponse.get(0).getStatusCode()).isEqualTo(HttpResponseStatus.CREATED.code());
-        assertThat(batchResponse.get(0).getItem(TestDoc.class)).isEqualTo(firstDoc);
+        assertThat(batchResponse.getResults().get(0).getStatusCode()).isEqualTo(HttpResponseStatus.CREATED.code());
+        assertThat(batchResponse.getResults().get(0).getItem(TestDoc.class)).isEqualTo(firstDoc);
 
-        assertThat(batchResponse.get(1).getStatusCode()).isEqualTo(HttpResponseStatus.CREATED.code());
-        assertThat(batchResponse.get(1).getItem(EventDoc.class)).isEqualTo(eventDoc1);
+        assertThat(batchResponse.getResults().get(1).getStatusCode()).isEqualTo(HttpResponseStatus.CREATED.code());
+        assertThat(batchResponse.getResults().get(1).getItem(EventDoc.class)).isEqualTo(eventDoc1);
 
-        assertThat(batchResponse.get(2).getStatusCode()).isEqualTo(HttpResponseStatus.OK.code());
-        assertThat(batchResponse.get(2).getItem(TestDoc.class)).isEqualTo(replaceDoc);
+        assertThat(batchResponse.getResults().get(2).getStatusCode()).isEqualTo(HttpResponseStatus.OK.code());
+        assertThat(batchResponse.getResults().get(2).getItem(TestDoc.class)).isEqualTo(replaceDoc);
 
-        assertThat(batchResponse.get(3).getStatusCode()).isEqualTo(HttpResponseStatus.OK.code());
-        assertThat(batchResponse.get(3).getItem(EventDoc.class)).isEqualTo(readEventDoc);
+        assertThat(batchResponse.getResults().get(3).getStatusCode()).isEqualTo(HttpResponseStatus.OK.code());
+        assertThat(batchResponse.getResults().get(3).getItem(EventDoc.class)).isEqualTo(readEventDoc);
 
         // Ensure that the replace overwrote the doc from the first operation
         this.verifyByRead(container, replaceDoc);
+
+        List> batchOperations = batch.getOperations();
+        for (int index = 0; index < batchOperations.size(); index++) {
+            assertThat(batchResponse.getResults().get(index).getItemBatchOperation()).isEqualTo(batchOperations.get(index));
+        }
     }
 
     @Test(groups = {"simple"}, timeOut = TIMEOUT)
@@ -121,38 +135,40 @@ public void batchItemETagTest() {
 
             assertThat(response.getStatusCode()).isEqualTo(HttpResponseStatus.OK.code());
 
-            TransactionalBatchItemRequestOptions firstReplaceOptions = new TransactionalBatchItemRequestOptions();
+            ItemBatchRequestOptions firstReplaceOptions = new ItemBatchRequestOptions();
             firstReplaceOptions.setIfMatchETag(response.getETag());
 
-            TransactionalBatchResponse batchResponse = container.executeTransactionalBatch(
-                TransactionalBatch.createTransactionalBatch(this.getPartitionKey(this.partitionKey1))
-                    .createItem(testDocToCreate)
-                    .replaceItem(testDocToReplace.getId(), testDocToReplace, firstReplaceOptions));
+            TransactionalBatch batch = TransactionalBatch.createTransactionalBatch(this.getPartitionKey(this.partitionKey1));
+            batch.createItem(testDocToCreate);
+            batch.replaceItem(testDocToReplace.getId(), testDocToReplace, firstReplaceOptions);
+
+            TransactionalBatchResponse batchResponse = container.executeTransactionalBatch(batch);
 
             this.verifyBatchProcessed(batchResponse, 2);
 
-            assertThat(batchResponse.get(0).getStatusCode()).isEqualTo(HttpResponseStatus.CREATED.code());
-            assertThat(batchResponse.get(1).getStatusCode()).isEqualTo(HttpResponseStatus.OK.code());
+            assertThat(batchResponse.getResults().get(0).getStatusCode()).isEqualTo(HttpResponseStatus.CREATED.code());
+            assertThat(batchResponse.getResults().get(1).getStatusCode()).isEqualTo(HttpResponseStatus.OK.code());
 
             // Ensure that the replace overwrote the doc from the first operation
-            this.verifyByRead(container, testDocToCreate, batchResponse.get(0).getETag());
-            this.verifyByRead(container, testDocToReplace, batchResponse.get(1).getETag());
+            this.verifyByRead(container, testDocToCreate, batchResponse.getResults().get(0).getETag());
+            this.verifyByRead(container, testDocToReplace, batchResponse.getResults().get(1).getETag());
         }
 
         {
             TestDoc testDocToReplace = this.getTestDocCopy(this.TestDocPk1ExistingB);
             testDocToReplace.setCost(testDocToReplace.getCost() + 1);
 
-            TransactionalBatchItemRequestOptions replaceOptions = new TransactionalBatchItemRequestOptions();
+            ItemBatchRequestOptions replaceOptions = new ItemBatchRequestOptions();
             replaceOptions.setIfMatchETag(String.valueOf(this.getRandom().nextInt()));
 
-            TransactionalBatchResponse batchResponse = container.executeTransactionalBatch(
-                TransactionalBatch.createTransactionalBatch(this.getPartitionKey(this.partitionKey1))
-                    .replaceItem(testDocToReplace.getId(), testDocToReplace, replaceOptions));
+            TransactionalBatch batch = TransactionalBatch.createTransactionalBatch(this.getPartitionKey(this.partitionKey1));
+            batch.replaceItem(testDocToReplace.getId(), testDocToReplace, replaceOptions);
+
+            TransactionalBatchResponse batchResponse = container.executeTransactionalBatch(batch);
 
             this.verifyBatchProcessed(batchResponse, 1, HttpResponseStatus.PRECONDITION_FAILED);
 
-            assertThat(batchResponse.get(0).getStatusCode()).isEqualTo(HttpResponseStatus.PRECONDITION_FAILED.code());
+            assertThat(batchResponse.getResults().get(0).getStatusCode()).isEqualTo(HttpResponseStatus.PRECONDITION_FAILED.code());
 
             // ensure the item was not updated
             this.verifyByRead(container, this.TestDocPk1ExistingB);
@@ -180,12 +196,13 @@ public void batchErrorSessionToken() {
 
         {
             // Only errored read
-            TransactionalBatchResponse batchResponse = container.executeTransactionalBatch(
-            TransactionalBatch.createTransactionalBatch(this.getPartitionKey(this.partitionKey1))
-                .readItem(UUID.randomUUID().toString()));
+            TransactionalBatch batch = TransactionalBatch.createTransactionalBatch(this.getPartitionKey(this.partitionKey1));
+            batch.readItem(UUID.randomUUID().toString());
+
+            TransactionalBatchResponse batchResponse = container.executeTransactionalBatch(batch);
 
             assertThat(batchResponse.getStatusCode()).isEqualTo(HttpResponseStatus.NOT_FOUND.code());
-            assertThat(batchResponse.get(0).getStatusCode()).isEqualTo(HttpResponseStatus.NOT_FOUND.code());
+            assertThat(batchResponse.getResults().get(0).getStatusCode()).isEqualTo(HttpResponseStatus.NOT_FOUND.code());
 
             String ownerIdBatch = batchResponse.getResponseHeaders().get(HttpConstants.HttpHeaders.OWNER_ID);
             assertThat(ownerIdBatch).isNull();
@@ -199,14 +216,15 @@ public void batchErrorSessionToken() {
 
         {
             // One valid read one error read
-            TransactionalBatchResponse batchResponse = container.executeTransactionalBatch(
-                TransactionalBatch.createTransactionalBatch(this.getPartitionKey(this.partitionKey1))
-                    .readItem(this.TestDocPk1ExistingA.getId())
-                    .readItem(UUID.randomUUID().toString()));
+            TransactionalBatch batch = TransactionalBatch.createTransactionalBatch(this.getPartitionKey(this.partitionKey1));
+            batch.readItem(this.TestDocPk1ExistingA.getId());
+            batch.readItem(UUID.randomUUID().toString());
+
+            TransactionalBatchResponse batchResponse = container.executeTransactionalBatch(batch);
 
             assertThat(batchResponse.getStatusCode()).isEqualTo(HttpResponseStatus.NOT_FOUND.code());
-            assertThat(batchResponse.get(0).getStatusCode()).isEqualTo(HttpResponseStatus.FAILED_DEPENDENCY.code());
-            assertThat(batchResponse.get(1).getStatusCode()).isEqualTo(HttpResponseStatus.NOT_FOUND.code());
+            assertThat(batchResponse.getResults().get(0).getStatusCode()).isEqualTo(HttpResponseStatus.FAILED_DEPENDENCY.code());
+            assertThat(batchResponse.getResults().get(1).getStatusCode()).isEqualTo(HttpResponseStatus.NOT_FOUND.code());
 
             String ownerIdBatch = batchResponse.getResponseHeaders().get(HttpConstants.HttpHeaders.OWNER_ID);
             assertThat(ownerIdBatch).isNull();
@@ -220,14 +238,15 @@ public void batchErrorSessionToken() {
 
         {
             // One error one valid read
-            TransactionalBatchResponse batchResponse = container.executeTransactionalBatch(
-                TransactionalBatch.createTransactionalBatch(this.getPartitionKey(this.partitionKey1))
-                    .readItem(UUID.randomUUID().toString())
-                    .readItem(this.TestDocPk1ExistingA.getId()));
+            TransactionalBatch batch = TransactionalBatch.createTransactionalBatch(this.getPartitionKey(this.partitionKey1));
+            batch.readItem(UUID.randomUUID().toString());
+            batch.readItem(this.TestDocPk1ExistingA.getId());
+
+            TransactionalBatchResponse batchResponse = container.executeTransactionalBatch(batch);
 
             assertThat(batchResponse.getStatusCode()).isEqualTo(HttpResponseStatus.NOT_FOUND.code());
-            assertThat(batchResponse.get(0).getStatusCode()).isEqualTo(HttpResponseStatus.NOT_FOUND.code());
-            assertThat(batchResponse.get(1).getStatusCode()).isEqualTo(HttpResponseStatus.FAILED_DEPENDENCY.code());
+            assertThat(batchResponse.getResults().get(0).getStatusCode()).isEqualTo(HttpResponseStatus.NOT_FOUND.code());
+            assertThat(batchResponse.getResults().get(1).getStatusCode()).isEqualTo(HttpResponseStatus.FAILED_DEPENDENCY.code());
 
             String ownerIdBatch = batchResponse.getResponseHeaders().get(HttpConstants.HttpHeaders.OWNER_ID);
             assertThat(ownerIdBatch).isNull();
@@ -242,14 +261,16 @@ public void batchErrorSessionToken() {
         {
             // One valid write and one error
             TestDoc testDocToCreate = this.populateTestDoc(this.partitionKey1);
-            TransactionalBatchResponse batchResponse = container.executeTransactionalBatch(
-                TransactionalBatch.createTransactionalBatch(this.getPartitionKey(this.partitionKey1))
-                    .createItem(testDocToCreate)
-                    .readItem(UUID.randomUUID().toString()));
+
+            TransactionalBatch batch = TransactionalBatch.createTransactionalBatch(this.getPartitionKey(this.partitionKey1));
+            batch.createItem(testDocToCreate);
+            batch.readItem(UUID.randomUUID().toString());
+
+            TransactionalBatchResponse batchResponse = container.executeTransactionalBatch(batch);
 
             assertThat(batchResponse.getStatusCode()).isEqualTo(HttpResponseStatus.NOT_FOUND.code());
-            assertThat(batchResponse.get(0).getStatusCode()).isEqualTo(HttpResponseStatus.FAILED_DEPENDENCY.code());
-            assertThat(batchResponse.get(1).getStatusCode()).isEqualTo(HttpResponseStatus.NOT_FOUND.code());
+            assertThat(batchResponse.getResults().get(0).getStatusCode()).isEqualTo(HttpResponseStatus.FAILED_DEPENDENCY.code());
+            assertThat(batchResponse.getResults().get(1).getStatusCode()).isEqualTo(HttpResponseStatus.NOT_FOUND.code());
 
             String ownerIdBatch = batchResponse.getResponseHeaders().get(HttpConstants.HttpHeaders.OWNER_ID);
             assertThat(ownerIdBatch).isNull();
@@ -264,14 +285,15 @@ public void batchErrorSessionToken() {
         {
             // One error one valid write
             TestDoc testDocToCreate = this.populateTestDoc(this.partitionKey1);
-            TransactionalBatchResponse batchResponse = container.executeTransactionalBatch(
-                TransactionalBatch.createTransactionalBatch(this.getPartitionKey(this.partitionKey1))
-                    .readItem(UUID.randomUUID().toString())
-                    .createItem(testDocToCreate));
+            TransactionalBatch batch = TransactionalBatch.createTransactionalBatch(this.getPartitionKey(this.partitionKey1));
+            batch.readItem(UUID.randomUUID().toString());
+            batch.createItem(testDocToCreate);
+
+            TransactionalBatchResponse batchResponse = container.executeTransactionalBatch(batch);
 
             assertThat(batchResponse.getStatusCode()).isEqualTo(HttpResponseStatus.NOT_FOUND.code());
-            assertThat(batchResponse.get(0).getStatusCode()).isEqualTo(HttpResponseStatus.NOT_FOUND.code());
-            assertThat(batchResponse.get(1).getStatusCode()).isEqualTo(HttpResponseStatus.FAILED_DEPENDENCY.code());
+            assertThat(batchResponse.getResults().get(0).getStatusCode()).isEqualTo(HttpResponseStatus.NOT_FOUND.code());
+            assertThat(batchResponse.getResults().get(1).getStatusCode()).isEqualTo(HttpResponseStatus.FAILED_DEPENDENCY.code());
 
             String ownerIdBatch = batchResponse.getResponseHeaders().get(HttpConstants.HttpHeaders.OWNER_ID);
             assertThat(ownerIdBatch).isNull();
@@ -339,7 +361,12 @@ public void batchServerResponseTooLarge() {
 
         TransactionalBatchResponse batchResponse = batchContainer.executeTransactionalBatch(batch);
         assertThat(batchResponse.getStatusCode()).isEqualTo(HttpResponseStatus.REQUEST_ENTITY_TOO_LARGE.code());
-        assertThat(batchResponse.get(1).getStatusCode()).isEqualTo(HttpResponseStatus.FAILED_DEPENDENCY.code());
+        assertThat(batchResponse.getResults().get(1).getStatusCode()).isEqualTo(HttpResponseStatus.FAILED_DEPENDENCY.code());
+
+        List> batchOperations = batch.getOperations();
+        for (int index = 0; index < batchOperations.size(); index++) {
+            assertThat(batchResponse.getResults().get(index).getItemBatchOperation()).isEqualTo(batchOperations.get(index));
+        }
     }
 
     @Test(groups = {"simple"}, timeOut = TIMEOUT)
@@ -347,21 +374,27 @@ public void batchReadsOnlyTest() {
         CosmosContainer container = batchContainer;
         this.createJsonTestDocs(container);
 
-        TransactionalBatchResponse batchResponse = container.executeTransactionalBatch(
-            TransactionalBatch.createTransactionalBatch(this.getPartitionKey(this.partitionKey1))
-                .readItem(this.TestDocPk1ExistingA.getId())
-                .readItem(this.TestDocPk1ExistingB.getId())
-                .readItem(this.TestDocPk1ExistingC.getId()));
+        TransactionalBatch batch = TransactionalBatch.createTransactionalBatch(this.getPartitionKey(this.partitionKey1));
+        batch.readItem(this.TestDocPk1ExistingA.getId());
+        batch.readItem(this.TestDocPk1ExistingB.getId());
+        batch.readItem(this.TestDocPk1ExistingC.getId());
+
+        TransactionalBatchResponse batchResponse = container.executeTransactionalBatch(batch);
 
         this.verifyBatchProcessed(batchResponse, 3);
 
-        assertThat(batchResponse.get(0).getStatusCode()).isEqualTo(HttpResponseStatus.OK.code());
-        assertThat(batchResponse.get(1).getStatusCode()).isEqualTo(HttpResponseStatus.OK.code());
-        assertThat(batchResponse.get(2).getStatusCode()).isEqualTo(HttpResponseStatus.OK.code());
+        assertThat(batchResponse.getResults().get(0).getStatusCode()).isEqualTo(HttpResponseStatus.OK.code());
+        assertThat(batchResponse.getResults().get(1).getStatusCode()).isEqualTo(HttpResponseStatus.OK.code());
+        assertThat(batchResponse.getResults().get(2).getStatusCode()).isEqualTo(HttpResponseStatus.OK.code());
+
+        assertThat(batchResponse.getResults().get(0).getItem(TestDoc.class)).isEqualTo(this.TestDocPk1ExistingA);
+        assertThat(batchResponse.getResults().get(1).getItem(TestDoc.class)).isEqualTo(this.TestDocPk1ExistingB);
+        assertThat(batchResponse.getResults().get(2).getItem(TestDoc.class)).isEqualTo(this.TestDocPk1ExistingC);
 
-        assertThat(batchResponse.get(0).getItem(TestDoc.class)).isEqualTo(this.TestDocPk1ExistingA);
-        assertThat(batchResponse.get(1).getItem(TestDoc.class)).isEqualTo(this.TestDocPk1ExistingB);
-        assertThat(batchResponse.get(2).getItem(TestDoc.class)).isEqualTo(this.TestDocPk1ExistingC);
+        List> batchOperations = batch.getOperations();
+        for (int index = 0; index < batchOperations.size(); index++) {
+            assertThat(batchResponse.getResults().get(index).getItemBatchOperation()).isEqualTo(batchOperations.get(index));
+        }
     }
 
     @Test(groups = {"simple"}, timeOut = TIMEOUT)
@@ -378,26 +411,32 @@ public void batchCrud() {
         BatchTestBase.TestDoc testDocToReplace = this.getTestDocCopy(this.TestDocPk1ExistingB);
         testDocToReplace.setCost(testDocToReplace.getCost() + 1);
 
+        TransactionalBatch batch = TransactionalBatch.createTransactionalBatch(this.getPartitionKey(this.partitionKey1));
+        batch.createItem(testDocToCreate);
+        batch.readItem(this.TestDocPk1ExistingC.getId());
+        batch.replaceItem(testDocToReplace.getId(), testDocToReplace);
+        batch.upsertItem(testDocToUpsert);
+        batch.upsertItem(anotherTestDocToUpsert);
+        batch.deleteItem(this.TestDocPk1ExistingD.getId());
+
         // We run CRUD operations where all are expected to return HTTP 2xx.
-        TransactionalBatchResponse batchResponse = container.executeTransactionalBatch(
-            TransactionalBatch.createTransactionalBatch(this.getPartitionKey(this.partitionKey1))
-                .createItem(testDocToCreate)
-                .readItem(this.TestDocPk1ExistingC.getId())
-                .replaceItem(testDocToReplace.getId(), testDocToReplace)
-                .upsertItem(testDocToUpsert)
-                .upsertItem(anotherTestDocToUpsert)
-                .deleteItem(this.TestDocPk1ExistingD.getId()));
+        TransactionalBatchResponse batchResponse = container.executeTransactionalBatch(batch);
 
         this.verifyBatchProcessed(batchResponse, 6);
 
-        assertThat(batchResponse.get(0).getStatusCode()).isEqualTo(HttpResponseStatus.CREATED.code());
-        assertThat(batchResponse.get(1).getStatusCode()).isEqualTo(HttpResponseStatus.OK.code());
-        assertThat(batchResponse.get(2).getStatusCode()).isEqualTo(HttpResponseStatus.OK.code());
-        assertThat(batchResponse.get(3).getStatusCode()).isEqualTo(HttpResponseStatus.CREATED.code());
-        assertThat(batchResponse.get(4).getStatusCode()).isEqualTo(HttpResponseStatus.OK.code());
-        assertThat(batchResponse.get(5).getStatusCode()).isEqualTo(HttpResponseStatus.NO_CONTENT.code());
+        assertThat(batchResponse.getResults().get(0).getStatusCode()).isEqualTo(HttpResponseStatus.CREATED.code());
+        assertThat(batchResponse.getResults().get(1).getStatusCode()).isEqualTo(HttpResponseStatus.OK.code());
+        assertThat(batchResponse.getResults().get(2).getStatusCode()).isEqualTo(HttpResponseStatus.OK.code());
+        assertThat(batchResponse.getResults().get(3).getStatusCode()).isEqualTo(HttpResponseStatus.CREATED.code());
+        assertThat(batchResponse.getResults().get(4).getStatusCode()).isEqualTo(HttpResponseStatus.OK.code());
+        assertThat(batchResponse.getResults().get(5).getStatusCode()).isEqualTo(HttpResponseStatus.NO_CONTENT.code());
+
+        List> batchOperations = batch.getOperations();
+        for (int index = 0; index < batchOperations.size(); index++) {
+            assertThat(batchResponse.getResults().get(index).getItemBatchOperation()).isEqualTo(batchOperations.get(index));
+        }
 
-        assertThat(batchResponse.get(1).getItem(TestDoc.class)).isEqualTo(this.TestDocPk1ExistingC);
+        assertThat(batchResponse.getResults().get(1).getItem(TestDoc.class)).isEqualTo(this.TestDocPk1ExistingC);
 
         this.verifyByRead(container, testDocToCreate);
         this.verifyByRead(container, testDocToReplace);
@@ -430,7 +469,7 @@ public void batchWithReplaceOfStaleEntity() {
         TestDoc staleTestDocToReplace = this.getTestDocCopy(this.TestDocPk1ExistingA);
         staleTestDocToReplace.setCost(staleTestDocToReplace.getCost() + 1);
 
-        TransactionalBatchItemRequestOptions staleReplaceOptions = new TransactionalBatchItemRequestOptions();
+        ItemBatchRequestOptions staleReplaceOptions = new ItemBatchRequestOptions();
         staleReplaceOptions.setIfMatchETag(UUID.randomUUID().toString());
 
         this.runWithError(
@@ -470,25 +509,31 @@ public void batchWithCreateConflict() {
 
     private void runWithError(
         CosmosContainer container,
-        Function appendOperation,
+        Function> appendOperation,
         HttpResponseStatus expectedFailedOperationStatusCode) {
 
         TestDoc testDocToCreate = this.populateTestDoc(this.partitionKey1);
         TestDoc anotherTestDocToCreate = this.populateTestDoc(this.partitionKey1);
 
-        TransactionalBatch batch = TransactionalBatch.createTransactionalBatch(this.getPartitionKey(this.partitionKey1))
-            .createItem(testDocToCreate);
+        TransactionalBatch batch = TransactionalBatch.createTransactionalBatch(this.getPartitionKey(this.partitionKey1));
+        batch.createItem(testDocToCreate);
 
         appendOperation.apply(batch);
 
-        TransactionalBatchResponse batchResponse = container.executeTransactionalBatch(
-            batch.createItem(anotherTestDocToCreate));
+        batch.createItem(anotherTestDocToCreate);
+
+        TransactionalBatchResponse batchResponse = container.executeTransactionalBatch(batch);
 
         this.verifyBatchProcessed(batchResponse, 3, expectedFailedOperationStatusCode);
 
-        assertThat(batchResponse.get(0).getStatusCode()).isEqualTo(HttpResponseStatus.FAILED_DEPENDENCY.code());
-        assertThat(batchResponse.get(1).getStatusCode()).isEqualTo(expectedFailedOperationStatusCode.code());
-        assertThat(batchResponse.get(2).getStatusCode()).isEqualTo(HttpResponseStatus.FAILED_DEPENDENCY.code());
+        assertThat(batchResponse.getResults().get(0).getStatusCode()).isEqualTo(HttpResponseStatus.FAILED_DEPENDENCY.code());
+        assertThat(batchResponse.getResults().get(1).getStatusCode()).isEqualTo(expectedFailedOperationStatusCode.code());
+        assertThat(batchResponse.getResults().get(2).getStatusCode()).isEqualTo(HttpResponseStatus.FAILED_DEPENDENCY.code());
+
+        List> batchOperations = batch.getOperations();
+        for (int index = 0; index < batchOperations.size(); index++) {
+            assertThat(batchResponse.getResults().get(index).getItemBatchOperation()).isEqualTo(batchOperations.get(index));
+        }
 
         this.verifyNotFound(container, testDocToCreate);
         this.verifyNotFound(container, anotherTestDocToCreate);
diff --git a/sdk/cosmos/azure-cosmos/src/test/java/com/azure/cosmos/implementation/batch/TransactionalBatchResponseTests.java b/sdk/cosmos/azure-cosmos/src/test/java/com/azure/cosmos/implementation/batch/TransactionalBatchResponseTests.java
index 41f52ee88239..fc247993f3b9 100644
--- a/sdk/cosmos/azure-cosmos/src/test/java/com/azure/cosmos/implementation/batch/TransactionalBatchResponseTests.java
+++ b/sdk/cosmos/azure-cosmos/src/test/java/com/azure/cosmos/implementation/batch/TransactionalBatchResponseTests.java
@@ -4,12 +4,14 @@
 package com.azure.cosmos.implementation.batch;
 
 import com.azure.cosmos.BridgeInternal;
+import com.azure.cosmos.CosmosItemOperationType;
 import com.azure.cosmos.TransactionalBatchOperationResult;
 import com.azure.cosmos.TransactionalBatchResponse;
 import com.azure.cosmos.implementation.HttpConstants;
-import com.azure.cosmos.implementation.OperationType;
 import com.azure.cosmos.implementation.RxDocumentServiceResponse;
 import com.azure.cosmos.implementation.directconnectivity.StoreResponse;
+import com.azure.cosmos.models.ItemBatchOperation;
+import com.azure.cosmos.models.ModelBridgeInternal;
 import com.azure.cosmos.models.PartitionKey;
 import io.netty.handler.codec.http.HttpResponseStatus;
 import org.testng.annotations.Test;
@@ -34,10 +36,14 @@ public void validateAllSetValuesInResponse() {
         List results = new ArrayList<>();
         ItemBatchOperation[] arrayOperations = new ItemBatchOperation[1];
 
-        ItemBatchOperation operation = new ItemBatchOperation.Builder(OperationType.Read,0)
-            .partitionKey(PartitionKey.NONE)
-            .id("0")
-            .build();
+        ItemBatchOperation operation = ModelBridgeInternal.createItemBatchOperation(
+            CosmosItemOperationType.Read,
+            0,
+            "0",
+            PartitionKey.NONE,
+            null,
+            null
+        );
 
         arrayOperations[0] = operation;
         SinglePartitionKeyServerBatchRequest batchRequest = SinglePartitionKeyServerBatchRequest.createBatchRequest(
@@ -51,7 +57,8 @@ public void validateAllSetValuesInResponse() {
             null,
             HttpResponseStatus.NOT_MODIFIED.code(),
             Duration.ofMillis(100),
-            HttpConstants.SubStatusCodes.PARTITION_KEY_MISMATCH
+            HttpConstants.SubStatusCodes.PARTITION_KEY_MISMATCH,
+            operation
         );
 
         results.add(transactionalBatchOperationResult);
@@ -86,11 +93,12 @@ public void validateAllSetValuesInResponse() {
         assertThat(batchResponse.getSubStatusCode()).isEqualTo(HttpConstants.SubStatusCodes.PARTITION_KEY_RANGE_GONE);
 
         // Validate result fields
-        assertThat(batchResponse.get(0).getETag()).isEqualTo(operation.getId());
-        assertThat(batchResponse.get(0).getRequestCharge()).isEqualTo(5.0);
-        assertThat(batchResponse.get(0).getRetryAfterDuration()).isEqualTo(Duration.ofMillis(100));
-        assertThat(batchResponse.get(0).getSubStatusCode()).isEqualTo(HttpConstants.SubStatusCodes.PARTITION_KEY_MISMATCH);
-        assertThat(batchResponse.get(0).getStatusCode()).isEqualTo(HttpResponseStatus.NOT_MODIFIED.code());
+        assertThat(batchResponse.getResults().get(0).getETag()).isEqualTo(operation.getId());
+        assertThat(batchResponse.getResults().get(0).getRequestCharge()).isEqualTo(5.0);
+        assertThat(batchResponse.getResults().get(0).getRetryAfterDuration()).isEqualTo(Duration.ofMillis(100));
+        assertThat(batchResponse.getResults().get(0).getSubStatusCode()).isEqualTo(HttpConstants.SubStatusCodes.PARTITION_KEY_MISMATCH);
+        assertThat(batchResponse.getResults().get(0).getStatusCode()).isEqualTo(HttpResponseStatus.NOT_MODIFIED.code());
+        assertThat(batchResponse.getResults().get(0).getItemBatchOperation()).isEqualTo(operation);
     }
 
     @Test(groups = {"unit"}, timeOut = TIMEOUT)
@@ -98,10 +106,14 @@ public void validateEmptyHeaderInResponse() {
         List results = new ArrayList<>();
         ItemBatchOperation[] arrayOperations = new ItemBatchOperation[1];
 
-        ItemBatchOperation operation = new ItemBatchOperation.Builder(OperationType.Read,0)
-            .partitionKey(PartitionKey.NONE)
-            .id("0")
-            .build();
+        ItemBatchOperation operation = ModelBridgeInternal.createItemBatchOperation(
+            CosmosItemOperationType.Read,
+            0,
+            "0",
+            PartitionKey.NONE,
+            null,
+            null
+        );
 
         arrayOperations[0] = operation;
         SinglePartitionKeyServerBatchRequest batchRequest = SinglePartitionKeyServerBatchRequest.createBatchRequest(
@@ -115,7 +127,8 @@ public void validateEmptyHeaderInResponse() {
             null,
             HttpResponseStatus.NOT_MODIFIED.code(),
             null,
-            0
+            0,
+            operation
         );
 
         results.add(transactionalBatchOperationResult);
@@ -141,10 +154,11 @@ public void validateEmptyHeaderInResponse() {
         assertThat(batchResponse.getSubStatusCode()).isEqualTo(0);
 
         // Validate result fields
-        assertThat(batchResponse.get(0).getETag()).isNull();
-        assertThat(batchResponse.get(0).getRequestCharge()).isEqualTo(5.0);
-        assertThat(batchResponse.get(0).getRetryAfterDuration()).isEqualTo(Duration.ZERO);
-        assertThat(batchResponse.get(0).getSubStatusCode()).isEqualTo(0);
-        assertThat(batchResponse.get(0).getStatusCode()).isEqualTo(HttpResponseStatus.NOT_MODIFIED.code());
+        assertThat(batchResponse.getResults().get(0).getETag()).isNull();
+        assertThat(batchResponse.getResults().get(0).getRequestCharge()).isEqualTo(5.0);
+        assertThat(batchResponse.getResults().get(0).getRetryAfterDuration()).isEqualTo(Duration.ZERO);
+        assertThat(batchResponse.getResults().get(0).getSubStatusCode()).isEqualTo(0);
+        assertThat(batchResponse.getResults().get(0).getStatusCode()).isEqualTo(HttpResponseStatus.NOT_MODIFIED.code());
+        assertThat(batchResponse.getResults().get(0).getItemBatchOperation()).isEqualTo(operation);
     }
 }

From e67348a64a1ec87361d0599c3a9a6f93505a9e44 Mon Sep 17 00:00:00 2001
From: Rakesh Kumar 
Date: Fri, 9 Oct 2020 17:28:17 +0530
Subject: [PATCH 19/22] Adding CosmosItemoperation interface

Signed-off-by: Rakesh Kumar 
---
 .../java/com/azure/cosmos/BridgeInternal.java |  14 ++-
 .../com/azure/cosmos/CosmosItemOperation.java |  16 +++
 .../azure/cosmos/CosmosItemOperationType.java |   2 +-
 .../com/azure/cosmos/TransactionalBatch.java  |  95 +++++++--------
 ...TransactionalBatchItemRequestOptions.java} |  14 ++-
 .../TransactionalBatchOperationResult.java    |  16 +--
 .../implementation/batch/BatchExecutor.java   |   5 +-
 .../batch/BatchResponseParser.java            |   1 -
 .../batch/ItemBatchOperation.java             | 110 ++++++++++++++++++
 .../batch/ServerBatchRequest.java             |  42 +------
 .../SinglePartitionKeyServerBatchRequest.java |   1 -
 .../cosmos/models/ItemBatchOperation.java     |  70 -----------
 .../cosmos/models/ModelBridgeInternal.java    |  20 ----
 .../cosmos/BatchOperationResultTests.java     |   8 +-
 .../TransactionalBatchAsyncContainerTest.java |   6 +-
 .../azure/cosmos/TransactionalBatchTest.java  |  34 +++---
 .../TransactionalBatchResponseTests.java      |  12 +-
 17 files changed, 238 insertions(+), 228 deletions(-)
 create mode 100644 sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosItemOperation.java
 rename sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/{ItemBatchRequestOptions.java => TransactionalBatchItemRequestOptions.java} (73%)
 create mode 100644 sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/batch/ItemBatchOperation.java
 delete mode 100644 sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/models/ItemBatchOperation.java

diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/BridgeInternal.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/BridgeInternal.java
index 2e89d599be3c..71117a9d8353 100644
--- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/BridgeInternal.java
+++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/BridgeInternal.java
@@ -26,7 +26,7 @@
 import com.azure.cosmos.implementation.ServiceUnavailableException;
 import com.azure.cosmos.implementation.StoredProcedureResponse;
 import com.azure.cosmos.implementation.Warning;
-import com.azure.cosmos.models.ItemBatchOperation;
+import com.azure.cosmos.implementation.batch.ItemBatchOperation;
 import com.azure.cosmos.implementation.directconnectivity.StoreResponse;
 import com.azure.cosmos.implementation.directconnectivity.StoreResult;
 import com.azure.cosmos.implementation.directconnectivity.Uri;
@@ -610,6 +610,16 @@ public static Duration getRequestTimeoutFromGatewayConnectionConfig(GatewayConne
         return gatewayConnectionConfig.getRequestTimeout();
     }
 
+    @Warning(value = INTERNAL_USE_ONLY_WARNING)
+    public static List> getOperationsFromTransactionalBatch(TransactionalBatch transactionalBatch) {
+        return transactionalBatch.getOperationsInternal();
+    }
+
+    @Warning(value = INTERNAL_USE_ONLY_WARNING)
+    public static String getOperationValueForCosmosItemOperationType(CosmosItemOperationType cosmosItemOperationType) {
+        return cosmosItemOperationType.getOperationValue();
+    }
+
     @Warning(value = INTERNAL_USE_ONLY_WARNING)
     public static RequestOptions toRequestOptions(TransactionalBatchRequestOptions transactionalBatchRequestOptions) {
         return transactionalBatchRequestOptions.toRequestOptions();
@@ -623,7 +633,7 @@ public static TransactionalBatchOperationResult createTransactionBatchResult(
         int statusCode,
         Duration retryAfter,
         int subStatusCode,
-        ItemBatchOperation itemBatchOperation) {
+        CosmosItemOperation itemBatchOperation) {
 
         return new TransactionalBatchOperationResult(
             eTag,
diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosItemOperation.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosItemOperation.java
new file mode 100644
index 000000000000..fc8b57885081
--- /dev/null
+++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosItemOperation.java
@@ -0,0 +1,16 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+package com.azure.cosmos;
+
+import com.azure.cosmos.models.PartitionKey;
+
+public interface CosmosItemOperation {
+    String getId();
+
+    PartitionKey getPartitionKeyValue();
+
+    CosmosItemOperationType getOperationType();
+
+     T getItem();
+}
diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosItemOperationType.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosItemOperationType.java
index ff82bd4905d4..e16c38754915 100644
--- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosItemOperationType.java
+++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosItemOperationType.java
@@ -17,7 +17,7 @@ public enum CosmosItemOperationType {
         this.operationValue = operationValue;
     }
 
-    public String getOperationValue() {
+    String getOperationValue() {
         return operationValue;
     }
 
diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/TransactionalBatch.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/TransactionalBatch.java
index 172cc645b378..aee0b13f2a3c 100644
--- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/TransactionalBatch.java
+++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/TransactionalBatch.java
@@ -4,8 +4,7 @@
 package com.azure.cosmos;
 
 import com.azure.cosmos.implementation.apachecommons.collections.list.UnmodifiableList;
-import com.azure.cosmos.models.ItemBatchOperation;
-import com.azure.cosmos.models.ModelBridgeInternal;
+import com.azure.cosmos.implementation.batch.ItemBatchOperation;
 import com.azure.cosmos.models.PartitionKey;
 import com.azure.cosmos.util.Beta;
 
@@ -117,9 +116,9 @@ public static TransactionalBatch createTransactionalBatch(PartitionKey partition
      *
      * @return The transactional batch instance with the operation added.
      */
-    public  ItemBatchOperation createItem(T item) {
+    public  CosmosItemOperation createItem(T item) {
         checkNotNull(item, "expected non-null item");
-        return this.createItem(item, new ItemBatchRequestOptions());
+        return this.createItem(item, new TransactionalBatchItemRequestOptions());
     }
 
     /**
@@ -132,19 +131,18 @@ public  ItemBatchOperation createItem(T item) {
      *
      * @return The transactional batch instance with the operation added.
      */
-    public  ItemBatchOperation createItem(T item, ItemBatchRequestOptions requestOptions) {
+    public  CosmosItemOperation createItem(T item, TransactionalBatchItemRequestOptions requestOptions) {
 
         checkNotNull(item, "expected non-null item");
         if (requestOptions == null) {
-            requestOptions = new ItemBatchRequestOptions();
+            requestOptions = new TransactionalBatchItemRequestOptions();
         }
 
-        ItemBatchOperation operation = ModelBridgeInternal.createItemBatchOperation(
+        ItemBatchOperation operation = new ItemBatchOperation(
             CosmosItemOperationType.Create,
-            this.operations.size(),
             null,
-            this.getPartitionKey(),
-            requestOptions,
+            this.getPartitionKeyValue(),
+            requestOptions.toRequestOptions(),
             item
         );
 
@@ -160,9 +158,9 @@ public  ItemBatchOperation createItem(T item, ItemBatchRequestOptions requ
      *
      * @return The transactional batch instance with the operation added.
      */
-    public ItemBatchOperation deleteItem(String id) {
+    public CosmosItemOperation deleteItem(String id) {
         checkNotNull(id, "expected non-null id");
-        return this.deleteItem(id, new ItemBatchRequestOptions());
+        return this.deleteItem(id, new TransactionalBatchItemRequestOptions());
     }
 
     /**
@@ -173,19 +171,18 @@ public ItemBatchOperation deleteItem(String id) {
      *
      * @return The transactional batch instance with the operation added.
      */
-    public ItemBatchOperation deleteItem(String id, ItemBatchRequestOptions requestOptions) {
+    public CosmosItemOperation deleteItem(String id, TransactionalBatchItemRequestOptions requestOptions) {
 
         checkNotNull(id, "expected non-null id");
         if (requestOptions == null) {
-            requestOptions = new ItemBatchRequestOptions();
+            requestOptions = new TransactionalBatchItemRequestOptions();
         }
 
-        ItemBatchOperation operation = ModelBridgeInternal.createItemBatchOperation(
+        ItemBatchOperation operation = new ItemBatchOperation<>(
             CosmosItemOperationType.Delete,
-            this.operations.size(),
             id,
-            this.getPartitionKey(),
-            requestOptions,
+            this.getPartitionKeyValue(),
+            requestOptions.toRequestOptions(),
             null
         );
 
@@ -201,9 +198,9 @@ public ItemBatchOperation deleteItem(String id, ItemBatchRequestOptions reque
      *
      * @return The transactional batch instance with the operation added.
      */
-    public ItemBatchOperation readItem(String id) {
+    public CosmosItemOperation readItem(String id) {
         checkNotNull(id, "expected non-null id");
-        return this.readItem(id, new ItemBatchRequestOptions());
+        return this.readItem(id, new TransactionalBatchItemRequestOptions());
     }
 
     /**
@@ -214,19 +211,18 @@ public ItemBatchOperation readItem(String id) {
      *
      * @return The transactional batch instance with the operation added.
      */
-    public ItemBatchOperation readItem(String id, ItemBatchRequestOptions requestOptions) {
+    public CosmosItemOperation readItem(String id, TransactionalBatchItemRequestOptions requestOptions) {
 
         checkNotNull(id, "expected non-null id");
         if (requestOptions == null) {
-            requestOptions = new ItemBatchRequestOptions();
+            requestOptions = new TransactionalBatchItemRequestOptions();
         }
 
-        ItemBatchOperation operation = ModelBridgeInternal.createItemBatchOperation(
+        ItemBatchOperation operation = new ItemBatchOperation<>(
             CosmosItemOperationType.Read,
-            this.operations.size(),
             id,
-            this.getPartitionKey(),
-            requestOptions,
+            this.getPartitionKeyValue(),
+            requestOptions.toRequestOptions(),
             null
         );
 
@@ -244,10 +240,10 @@ public ItemBatchOperation readItem(String id, ItemBatchRequestOptions request
      *
      * @return The transactional batch instance with the operation added.
      */
-    public  ItemBatchOperation replaceItem(String id, T item) {
+    public  CosmosItemOperation replaceItem(String id, T item) {
         checkNotNull(id, "expected non-null id");
         checkNotNull(item, "expected non-null item");
-        return this.replaceItem(id, item, new ItemBatchRequestOptions());
+        return this.replaceItem(id, item, new TransactionalBatchItemRequestOptions());
     }
 
     /**
@@ -261,21 +257,20 @@ public  ItemBatchOperation replaceItem(String id, T item) {
      *
      * @return The transactional batch instance with the operation added.
      */
-    public  ItemBatchOperation replaceItem(
-        String id, T item, ItemBatchRequestOptions requestOptions) {
+    public  CosmosItemOperation replaceItem(
+        String id, T item, TransactionalBatchItemRequestOptions requestOptions) {
 
         checkNotNull(id, "expected non-null id");
         checkNotNull(item, "expected non-null item");
         if (requestOptions == null) {
-            requestOptions = new ItemBatchRequestOptions();
+            requestOptions = new TransactionalBatchItemRequestOptions();
         }
 
-        ItemBatchOperation operation = ModelBridgeInternal.createItemBatchOperation(
+        ItemBatchOperation operation = new ItemBatchOperation(
             CosmosItemOperationType.Replace,
-            this.operations.size(),
             id,
-            this.getPartitionKey(),
-            requestOptions,
+            this.getPartitionKeyValue(),
+            requestOptions.toRequestOptions(),
             item
         );
 
@@ -292,9 +287,9 @@ public  ItemBatchOperation replaceItem(
      *
      * @return The transactional batch instance with the operation added.
      */
-    public  ItemBatchOperation upsertItem(T item) {
+    public  CosmosItemOperation upsertItem(T item) {
         checkNotNull(item, "expected non-null item");
-        return this.upsertItem(item, new ItemBatchRequestOptions());
+        return this.upsertItem(item, new TransactionalBatchItemRequestOptions());
     }
 
     /**
@@ -307,19 +302,18 @@ public  ItemBatchOperation upsertItem(T item) {
      *
      * @return The transactional batch instance with the operation added.
      */
-    public  ItemBatchOperation upsertItem(T item, ItemBatchRequestOptions requestOptions) {
+    public  CosmosItemOperation upsertItem(T item, TransactionalBatchItemRequestOptions requestOptions) {
 
         checkNotNull(item, "expected non-null item");
         if (requestOptions == null) {
-            requestOptions = new ItemBatchRequestOptions();
+            requestOptions = new TransactionalBatchItemRequestOptions();
         }
 
-        ItemBatchOperation operation = ModelBridgeInternal.createItemBatchOperation(
+        ItemBatchOperation operation = new ItemBatchOperation(
             CosmosItemOperationType.Upsert,
-            this.operations.size(),
             null,
-            this.getPartitionKey(),
-            requestOptions,
+            this.getPartitionKeyValue(),
+            requestOptions.toRequestOptions(),
             item
         );
 
@@ -329,11 +323,20 @@ public  ItemBatchOperation upsertItem(T item, ItemBatchRequestOptions requ
     }
 
     /**
-     * Return the list of operation in an unmodifiable instace so no one can change it in the down path.
+     * Return the list of operation in an unmodifiable instance  so no one can change it in the down path.
      *
      * @return The list of operations which are to be executed.
      */
-    public List> getOperations() {
+    List> getOperationsInternal() {
+        return UnmodifiableList.unmodifiableList(operations);
+    }
+
+    /**
+     * Return the list of operation in an unmodifiable instance  so no one can change it in the down path.
+     *
+     * @return The list of operations which are to be executed.
+     */
+    public List getOperations() {
         return UnmodifiableList.unmodifiableList(operations);
     }
 
@@ -342,7 +345,7 @@ public List> getOperations() {
      *
      * @return The partition key for this batch.
      */
-    public PartitionKey getPartitionKey() {
+    public PartitionKey getPartitionKeyValue() {
         return partitionKey;
     }
 }
diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/ItemBatchRequestOptions.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/TransactionalBatchItemRequestOptions.java
similarity index 73%
rename from sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/ItemBatchRequestOptions.java
rename to sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/TransactionalBatchItemRequestOptions.java
index 0818cf491b03..3776eb53cd2d 100644
--- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/ItemBatchRequestOptions.java
+++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/TransactionalBatchItemRequestOptions.java
@@ -3,6 +3,7 @@
 
 package com.azure.cosmos;
 
+import com.azure.cosmos.implementation.RequestOptions;
 import com.azure.cosmos.util.Beta;
 
 /**
@@ -10,7 +11,7 @@
  * Currently used in {@link TransactionalBatch} to pass option.
  */
 @Beta(Beta.SinceVersion.V4_7_0)
-public final class ItemBatchRequestOptions {
+public final class TransactionalBatchItemRequestOptions {
     private String ifMatchETag;
     private String ifNoneMatchETag;
 
@@ -29,7 +30,7 @@ public String getIfMatchETag() {
      * @param ifMatchETag the ifMatchETag associated with the request.
      * @return the current request options
      */
-    public ItemBatchRequestOptions setIfMatchETag(final String ifMatchETag) {
+    public TransactionalBatchItemRequestOptions setIfMatchETag(final String ifMatchETag) {
         this.ifMatchETag = ifMatchETag;
         return this;
     }
@@ -49,8 +50,15 @@ public String getIfNoneMatchETag() {
      * @param ifNoneMatchEtag the ifNoneMatchETag associated with the request.
      * @return the current request options
      */
-    public ItemBatchRequestOptions setIfNoneMatchETag(final String ifNoneMatchEtag) {
+    public TransactionalBatchItemRequestOptions setIfNoneMatchETag(final String ifNoneMatchEtag) {
         this.ifNoneMatchETag = ifNoneMatchEtag;
         return this;
     }
+
+    RequestOptions toRequestOptions() {
+        final RequestOptions requestOptions = new RequestOptions();
+        requestOptions.setIfMatchETag(getIfMatchETag());
+        requestOptions.setIfNoneMatchETag(getIfNoneMatchETag());
+        return requestOptions;
+    }
 }
diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/TransactionalBatchOperationResult.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/TransactionalBatchOperationResult.java
index 46726b905771..a2fe6c2531ce 100644
--- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/TransactionalBatchOperationResult.java
+++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/TransactionalBatchOperationResult.java
@@ -4,7 +4,7 @@
 package com.azure.cosmos;
 
 import com.azure.cosmos.implementation.JsonSerializable;
-import com.azure.cosmos.models.ItemBatchOperation;
+import com.azure.cosmos.implementation.batch.ItemBatchOperation;
 import com.azure.cosmos.util.Beta;
 import com.fasterxml.jackson.databind.node.ObjectNode;
 import java.time.Duration;
@@ -23,7 +23,7 @@ public final class TransactionalBatchOperationResult {
     private final int statusCode;
     private final Duration retryAfter;
     private final int subStatusCode;
-    private final ItemBatchOperation itemBatchOperation;
+    private final CosmosItemOperation cosmosItemOperation;
 
     /**
      * Initializes a new instance of the {@link TransactionalBatchOperationResult} class.
@@ -34,9 +34,9 @@ public final class TransactionalBatchOperationResult {
                                       int statusCode,
                                       Duration retryAfter,
                                       int subStatusCode,
-                                      ItemBatchOperation itemBatchOperation) {
+                                      CosmosItemOperation cosmosItemOperation) {
         checkNotNull(statusCode, "expected non-null statusCode");
-        checkNotNull(itemBatchOperation, "expected non-null itemBatchOperation");
+        checkNotNull(cosmosItemOperation, "expected non-null cosmosItemOperation");
 
         this.eTag = eTag;
         this.requestCharge = requestCharge;
@@ -44,7 +44,7 @@ public final class TransactionalBatchOperationResult {
         this.statusCode = statusCode;
         this.retryAfter = retryAfter;
         this.subStatusCode = subStatusCode;
-        this.itemBatchOperation = itemBatchOperation;
+        this.cosmosItemOperation = cosmosItemOperation;
     }
 
     /**
@@ -130,11 +130,11 @@ ObjectNode getResourceObject() {
     }
 
     /**
-     * Gets the original ItemBatchOperation for this result.
+     * Gets the original operation for this result.
      *
      * @return the ItemBatchOperation.
      */
-    public ItemBatchOperation getItemBatchOperation() {
-        return itemBatchOperation;
+    public CosmosItemOperation getOperation() {
+        return cosmosItemOperation;
     }
 }
diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/batch/BatchExecutor.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/batch/BatchExecutor.java
index 088f3790a1d1..8cc034b75e45 100644
--- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/batch/BatchExecutor.java
+++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/batch/BatchExecutor.java
@@ -9,7 +9,6 @@
 import com.azure.cosmos.TransactionalBatch;
 import com.azure.cosmos.TransactionalBatchRequestOptions;
 import com.azure.cosmos.TransactionalBatchResponse;
-import com.azure.cosmos.models.ItemBatchOperation;
 import reactor.core.publisher.Mono;
 
 import java.util.List;
@@ -39,11 +38,11 @@ public BatchExecutor(
      */
     public final Mono executeAsync() {
 
-        List> operations = this.transactionalBatch.getOperations();
+        List> operations = BridgeInternal.getOperationsFromTransactionalBatch(this.transactionalBatch);
         checkArgument(operations.size() > 0, "Number of operations should be more than 0.");
 
         final SinglePartitionKeyServerBatchRequest request = SinglePartitionKeyServerBatchRequest.createBatchRequest(
-            this.transactionalBatch.getPartitionKey(),
+            this.transactionalBatch.getPartitionKeyValue(),
             operations);
         request.setAtomicBatch(true);
         request.setShouldContinueOnError(false);
diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/batch/BatchResponseParser.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/batch/BatchResponseParser.java
index b18047d89813..dc84c72eff22 100644
--- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/batch/BatchResponseParser.java
+++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/batch/BatchResponseParser.java
@@ -10,7 +10,6 @@
 import com.azure.cosmos.implementation.JsonSerializable;
 import com.azure.cosmos.implementation.RxDocumentServiceResponse;
 import com.azure.cosmos.implementation.Utils;
-import com.azure.cosmos.models.ItemBatchOperation;
 import com.fasterxml.jackson.databind.ObjectMapper;
 import com.fasterxml.jackson.databind.node.ObjectNode;
 import io.netty.handler.codec.http.HttpResponseStatus;
diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/batch/ItemBatchOperation.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/batch/ItemBatchOperation.java
new file mode 100644
index 000000000000..af89cfa819c3
--- /dev/null
+++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/batch/ItemBatchOperation.java
@@ -0,0 +1,110 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+package com.azure.cosmos.implementation.batch;
+
+import com.azure.cosmos.BridgeInternal;
+import com.azure.cosmos.CosmosItemOperation;
+import com.azure.cosmos.CosmosItemOperationType;
+import com.azure.cosmos.implementation.JsonSerializable;
+import com.azure.cosmos.implementation.RequestOptions;
+import com.azure.cosmos.implementation.apachecommons.lang.StringUtils;
+import com.azure.cosmos.models.PartitionKey;
+
+import static com.azure.cosmos.implementation.guava25.base.Preconditions.checkNotNull;
+
+/**
+ * Represents an operation on an item which will be executed as part of a batch request on a container. This will be
+ * serialized and sent in the request.
+ *
+ * @param  The type of item.
+ */
+public final class ItemBatchOperation implements CosmosItemOperation {
+
+    private TInternal item;
+
+    private final String id;
+    private final PartitionKey partitionKey;
+    private final CosmosItemOperationType operationType;
+    private final RequestOptions requestOptions;
+
+    public ItemBatchOperation(
+        final CosmosItemOperationType operationType,
+        final String id,
+        final PartitionKey partitionKey,
+        final RequestOptions requestOptions,
+        final TInternal item) {
+
+        checkNotNull(operationType, "expected non-null operationType");
+
+        this.operationType = operationType;
+        this.partitionKey = partitionKey;
+        this.id = id;
+        this.item = item;
+        this.requestOptions = requestOptions;
+    }
+
+    /**
+     * Writes a single operation to JsonSerializable.
+     * TODO(rakkuma): Similarly for hybrid row, operation needs to be written in Hybrid row.
+     * Issue: https://github.com/Azure/azure-sdk-for-java/issues/15856
+     *
+     * @param operation a single operation which needs to be serialized.
+     *
+     * @return instance of JsonSerializable containing values for a operation.
+     */
+    static JsonSerializable writeOperation(final ItemBatchOperation operation) {
+        final JsonSerializable jsonSerializable = new JsonSerializable();
+
+        jsonSerializable.set(
+            BatchRequestResponseConstant.FIELD_OPERATION_TYPE,
+            BridgeInternal.getOperationValueForCosmosItemOperationType(operation.getOperationType()));
+
+        if (StringUtils.isNotEmpty(operation.getId())) {
+            jsonSerializable.set(BatchRequestResponseConstant.FIELD_ID, operation.getId());
+        }
+
+        if (operation.getItemInternal() != null) {
+            jsonSerializable.set(BatchRequestResponseConstant.FIELD_RESOURCE_BODY, operation.getItemInternal());
+        }
+
+        if (operation.getRequestOptions() != null) {
+            RequestOptions requestOptions = operation.getRequestOptions();
+
+            if (StringUtils.isNotEmpty(requestOptions.getIfMatchETag())) {
+                jsonSerializable.set(BatchRequestResponseConstant.FIELD_IF_MATCH, requestOptions.getIfMatchETag());
+            }
+
+            if (StringUtils.isNotEmpty(requestOptions.getIfNoneMatchETag())) {
+                jsonSerializable.set(BatchRequestResponseConstant.FIELD_IF_NONE_MATCH, requestOptions.getIfNoneMatchETag());
+            }
+        }
+
+        return jsonSerializable;
+    }
+
+    TInternal getItemInternal() {
+        return this.item;
+    }
+
+    @SuppressWarnings("unchecked")
+    public  T getItem() {
+        return (T)this.item;
+    }
+
+    public String getId() {
+        return this.id;
+    }
+
+    public PartitionKey getPartitionKeyValue() {
+        return partitionKey;
+    }
+
+    public CosmosItemOperationType getOperationType() {
+        return this.operationType;
+    }
+
+    public RequestOptions getRequestOptions() {
+        return this.requestOptions;
+    }
+}
diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/batch/ServerBatchRequest.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/batch/ServerBatchRequest.java
index f4db876d7369..e942e1abbbb9 100644
--- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/batch/ServerBatchRequest.java
+++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/batch/ServerBatchRequest.java
@@ -3,12 +3,11 @@
 
 package com.azure.cosmos.implementation.batch;
 
-import com.azure.cosmos.ItemBatchRequestOptions;
+import com.azure.cosmos.TransactionalBatchItemRequestOptions;
 import com.azure.cosmos.implementation.JsonSerializable;
 import com.azure.cosmos.implementation.Utils;
 import com.azure.cosmos.implementation.apachecommons.collections.list.UnmodifiableList;
 import com.azure.cosmos.implementation.apachecommons.lang.StringUtils;
-import com.azure.cosmos.models.ItemBatchOperation;
 import com.fasterxml.jackson.databind.node.ArrayNode;
 
 import java.util.List;
@@ -62,7 +61,7 @@ final List> createBodyOfBatchRequest(final List operation : operations) {
 
-            final JsonSerializable operationJsonSerializable = this.writeOperation(operation);
+            final JsonSerializable operationJsonSerializable = ItemBatchOperation.writeOperation(operation);
             final int operationSerializedLength = operationJsonSerializable.toString().length();
 
             if (totalOperationCount != 0 &&
@@ -83,43 +82,6 @@ final List> createBodyOfBatchRequest(final List operation) {
-        final JsonSerializable jsonSerializable = new JsonSerializable();
-
-        jsonSerializable.set(BatchRequestResponseConstant.FIELD_OPERATION_TYPE, operation.getOperationType().getOperationValue());
-
-        if (StringUtils.isNotEmpty(operation.getId())) {
-            jsonSerializable.set(BatchRequestResponseConstant.FIELD_ID, operation.getId());
-        }
-
-        if (operation.getItem() != null) {
-            jsonSerializable.set(BatchRequestResponseConstant.FIELD_RESOURCE_BODY, operation.getItem());
-        }
-
-        if (operation.getItemBatchRequestOptions() != null) {
-            ItemBatchRequestOptions requestOptions = operation.getItemBatchRequestOptions();
-
-            if (StringUtils.isNotEmpty(requestOptions.getIfMatchETag())) {
-                jsonSerializable.set(BatchRequestResponseConstant.FIELD_IF_MATCH, requestOptions.getIfMatchETag());
-            }
-
-            if (StringUtils.isNotEmpty(requestOptions.getIfNoneMatchETag())) {
-                jsonSerializable.set(BatchRequestResponseConstant.FIELD_IF_NONE_MATCH, requestOptions.getIfNoneMatchETag());
-            }
-        }
-
-        return jsonSerializable;
-    }
-
     public final String getRequestBody() {
         checkState(this.requestBody != null, "expected non-null body");
 
diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/batch/SinglePartitionKeyServerBatchRequest.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/batch/SinglePartitionKeyServerBatchRequest.java
index 78f6ae216016..3d6e11ad5c41 100644
--- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/batch/SinglePartitionKeyServerBatchRequest.java
+++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/batch/SinglePartitionKeyServerBatchRequest.java
@@ -3,7 +3,6 @@
 
 package com.azure.cosmos.implementation.batch;
 
-import com.azure.cosmos.models.ItemBatchOperation;
 import com.azure.cosmos.models.PartitionKey;
 import java.util.List;
 
diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/models/ItemBatchOperation.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/models/ItemBatchOperation.java
deleted file mode 100644
index 8f864fed40b6..000000000000
--- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/models/ItemBatchOperation.java
+++ /dev/null
@@ -1,70 +0,0 @@
-// Copyright (c) Microsoft Corporation. All rights reserved.
-// Licensed under the MIT License.
-
-package com.azure.cosmos.models;
-
-import com.azure.cosmos.ItemBatchRequestOptions;
-import com.azure.cosmos.CosmosItemOperationType;
-
-import static com.azure.cosmos.implementation.guava25.base.Preconditions.checkArgument;
-import static com.azure.cosmos.implementation.guava25.base.Preconditions.checkNotNull;
-
-/**
- * Represents an operation on an item which will be executed as part of a batch request on a container. This will be
- * serialized and sent in the request.
- *
- * @param  The type of item.
- */
-public final class ItemBatchOperation {
-
-    private T item;
-
-    private final String id;
-    private final int operationIndex;
-    private final PartitionKey partitionKey;
-    private final CosmosItemOperationType operationType;
-    private final ItemBatchRequestOptions requestOptions;
-
-    ItemBatchOperation(
-        final CosmosItemOperationType operationType,
-        final int operationIndex,
-        final PartitionKey partitionKey,
-        final String id,
-        final T item,
-        final ItemBatchRequestOptions requestOptions) {
-
-        checkNotNull(operationType, "expected non-null operationType");
-        checkArgument(operationIndex >= 0, "expected operationIndex >= 0, not %s", operationIndex);
-
-        this.operationType = operationType;
-        this.operationIndex = operationIndex;
-        this.partitionKey = partitionKey;
-        this.id = id;
-        this.item = item;
-        this.requestOptions = requestOptions;
-    }
-
-    public T getItem() {
-        return this.item;
-    }
-
-    public String getId() {
-        return this.id;
-    }
-
-    public int getOperationIndex() {
-        return operationIndex;
-    }
-
-    public PartitionKey getPartitionKey() {
-        return partitionKey;
-    }
-
-    public CosmosItemOperationType getOperationType() {
-        return this.operationType;
-    }
-
-    public ItemBatchRequestOptions getItemBatchRequestOptions() {
-        return this.requestOptions;
-    }
-}
diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/models/ModelBridgeInternal.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/models/ModelBridgeInternal.java
index 426a530cb3a0..dcde8b83e57a 100644
--- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/models/ModelBridgeInternal.java
+++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/models/ModelBridgeInternal.java
@@ -3,8 +3,6 @@
 
 package com.azure.cosmos.models;
 
-import com.azure.cosmos.ItemBatchRequestOptions;
-import com.azure.cosmos.CosmosItemOperationType;
 import com.azure.cosmos.implementation.Conflict;
 import com.azure.cosmos.implementation.ConsistencyPolicy;
 import com.azure.cosmos.implementation.CosmosResourceType;
@@ -726,22 +724,4 @@ public static CosmosItemRequestOptions clone(CosmosItemRequestOptions options) {
     public static  int getPayloadLength(CosmosItemResponse cosmosItemResponse) {
         return cosmosItemResponse.responseBodyAsByteArray != null ? cosmosItemResponse.responseBodyAsByteArray.length : 0;
     }
-
-    @Warning(value = INTERNAL_USE_ONLY_WARNING)
-    public static  ItemBatchOperation createItemBatchOperation(
-        final CosmosItemOperationType operationType,
-        final int operationIndex,
-        String id,
-        PartitionKey partitionKey,
-        ItemBatchRequestOptions itemBatchRequestOptions,
-        T resource) {
-
-        return new ItemBatchOperation<>(
-            operationType,
-            operationIndex,
-            partitionKey,
-            id,
-            resource,
-            itemBatchRequestOptions);
-    }
 }
diff --git a/sdk/cosmos/azure-cosmos/src/test/java/com/azure/cosmos/BatchOperationResultTests.java b/sdk/cosmos/azure-cosmos/src/test/java/com/azure/cosmos/BatchOperationResultTests.java
index 2394dd51ee4d..75cab913cb76 100644
--- a/sdk/cosmos/azure-cosmos/src/test/java/com/azure/cosmos/BatchOperationResultTests.java
+++ b/sdk/cosmos/azure-cosmos/src/test/java/com/azure/cosmos/BatchOperationResultTests.java
@@ -5,8 +5,7 @@
 
 import com.azure.cosmos.implementation.HttpConstants;
 import com.azure.cosmos.implementation.Utils;
-import com.azure.cosmos.models.ItemBatchOperation;
-import com.azure.cosmos.models.ModelBridgeInternal;
+import com.azure.cosmos.implementation.batch.ItemBatchOperation;
 import com.fasterxml.jackson.databind.node.ObjectNode;
 import io.netty.handler.codec.http.HttpResponseStatus;
 import org.testng.annotations.Test;
@@ -18,9 +17,8 @@ public class BatchOperationResultTests {
 
     private static final int TIMEOUT = 40000;
     private ObjectNode objectNode = Utils.getSimpleObjectMapper().createObjectNode();
-    private ItemBatchOperation operation = ModelBridgeInternal.createItemBatchOperation(
+    private ItemBatchOperation operation = new ItemBatchOperation<>(
         CosmosItemOperationType.Read,
-        1,
         null,
         null,
         null,
@@ -51,7 +49,7 @@ public void propertiesAreSetThroughCtor() {
         assertThat(result.getRequestCharge()).isEqualTo(1.4);
         assertThat(result.getRetryAfterDuration()).isEqualTo(Duration.ofMillis(1234));
         assertThat(result.getResourceObject()).isSameAs(objectNode);
-        assertThat(result.getItemBatchOperation()).isSameAs(operation);
+        assertThat(result.getOperation()).isSameAs(operation);
     }
 
     @Test(groups = {"unit"}, timeOut = TIMEOUT)
diff --git a/sdk/cosmos/azure-cosmos/src/test/java/com/azure/cosmos/TransactionalBatchAsyncContainerTest.java b/sdk/cosmos/azure-cosmos/src/test/java/com/azure/cosmos/TransactionalBatchAsyncContainerTest.java
index 4f4a779c8990..220370b915db 100644
--- a/sdk/cosmos/azure-cosmos/src/test/java/com/azure/cosmos/TransactionalBatchAsyncContainerTest.java
+++ b/sdk/cosmos/azure-cosmos/src/test/java/com/azure/cosmos/TransactionalBatchAsyncContainerTest.java
@@ -5,7 +5,7 @@
 
 import com.azure.cosmos.implementation.HttpConstants;
 import com.azure.cosmos.models.CosmosItemResponse;
-import com.azure.cosmos.models.ItemBatchOperation;
+import com.azure.cosmos.implementation.batch.ItemBatchOperation;
 import io.netty.handler.codec.http.HttpResponseStatus;
 import org.assertj.core.api.Assertions;
 import org.testng.annotations.AfterClass;
@@ -100,9 +100,9 @@ public void batchInvalidSessionToken() throws Exception {
             assertThat(batchResponse.getResults().get(2).getStatusCode()).isEqualTo(HttpResponseStatus.CREATED.code());
             assertThat(batchResponse.getResults().get(3).getStatusCode()).isEqualTo(HttpResponseStatus.NO_CONTENT.code());
 
-            List> batchOperations = batch.getOperations();
+            List> batchOperations = batch.getOperationsInternal();
             for (int index = 0; index < batchOperations.size(); index++) {
-                assertThat(batchResponse.getResults().get(index).getItemBatchOperation()).isEqualTo(batchOperations.get(index));
+                assertThat(batchResponse.getResults().get(index).getOperation()).isEqualTo(batchOperations.get(index));
             }
         }
 
diff --git a/sdk/cosmos/azure-cosmos/src/test/java/com/azure/cosmos/TransactionalBatchTest.java b/sdk/cosmos/azure-cosmos/src/test/java/com/azure/cosmos/TransactionalBatchTest.java
index 30d5bc03ad28..ee76747317c0 100644
--- a/sdk/cosmos/azure-cosmos/src/test/java/com/azure/cosmos/TransactionalBatchTest.java
+++ b/sdk/cosmos/azure-cosmos/src/test/java/com/azure/cosmos/TransactionalBatchTest.java
@@ -7,7 +7,7 @@
 import com.azure.cosmos.implementation.ISessionToken;
 import com.azure.cosmos.implementation.guava25.base.Function;
 import com.azure.cosmos.models.CosmosItemResponse;
-import com.azure.cosmos.models.ItemBatchOperation;
+import com.azure.cosmos.implementation.batch.ItemBatchOperation;
 import io.netty.handler.codec.http.HttpResponseStatus;
 import org.assertj.core.api.Assertions;
 import org.testng.annotations.AfterClass;
@@ -64,9 +64,9 @@ public void batchOrdered() {
         assertThat(batchResponse.getResults().get(0).getStatusCode()).isEqualTo(HttpResponseStatus.CREATED.code());
         assertThat(batchResponse.getResults().get(1).getStatusCode()).isEqualTo(HttpResponseStatus.OK.code());
 
-        List> batchOperations = batch.getOperations();
+        List> batchOperations = batch.getOperationsInternal();
         for (int index = 0; index < batchOperations.size(); index++) {
-            assertThat(batchResponse.getResults().get(index).getItemBatchOperation()).isEqualTo(batchOperations.get(index));
+            assertThat(batchResponse.getResults().get(index).getOperation()).isEqualTo(batchOperations.get(index));
         }
 
         // Ensure that the replace overwrote the doc from the first operation
@@ -111,9 +111,9 @@ public void batchMultipleItemExecution() {
         // Ensure that the replace overwrote the doc from the first operation
         this.verifyByRead(container, replaceDoc);
 
-        List> batchOperations = batch.getOperations();
+        List> batchOperations = batch.getOperationsInternal();
         for (int index = 0; index < batchOperations.size(); index++) {
-            assertThat(batchResponse.getResults().get(index).getItemBatchOperation()).isEqualTo(batchOperations.get(index));
+            assertThat(batchResponse.getResults().get(index).getOperation()).isEqualTo(batchOperations.get(index));
         }
     }
 
@@ -135,7 +135,7 @@ public void batchItemETagTest() {
 
             assertThat(response.getStatusCode()).isEqualTo(HttpResponseStatus.OK.code());
 
-            ItemBatchRequestOptions firstReplaceOptions = new ItemBatchRequestOptions();
+            TransactionalBatchItemRequestOptions firstReplaceOptions = new TransactionalBatchItemRequestOptions();
             firstReplaceOptions.setIfMatchETag(response.getETag());
 
             TransactionalBatch batch = TransactionalBatch.createTransactionalBatch(this.getPartitionKey(this.partitionKey1));
@@ -158,7 +158,7 @@ public void batchItemETagTest() {
             TestDoc testDocToReplace = this.getTestDocCopy(this.TestDocPk1ExistingB);
             testDocToReplace.setCost(testDocToReplace.getCost() + 1);
 
-            ItemBatchRequestOptions replaceOptions = new ItemBatchRequestOptions();
+            TransactionalBatchItemRequestOptions replaceOptions = new TransactionalBatchItemRequestOptions();
             replaceOptions.setIfMatchETag(String.valueOf(this.getRandom().nextInt()));
 
             TransactionalBatch batch = TransactionalBatch.createTransactionalBatch(this.getPartitionKey(this.partitionKey1));
@@ -363,9 +363,9 @@ public void batchServerResponseTooLarge() {
         assertThat(batchResponse.getStatusCode()).isEqualTo(HttpResponseStatus.REQUEST_ENTITY_TOO_LARGE.code());
         assertThat(batchResponse.getResults().get(1).getStatusCode()).isEqualTo(HttpResponseStatus.FAILED_DEPENDENCY.code());
 
-        List> batchOperations = batch.getOperations();
+        List> batchOperations = batch.getOperationsInternal();
         for (int index = 0; index < batchOperations.size(); index++) {
-            assertThat(batchResponse.getResults().get(index).getItemBatchOperation()).isEqualTo(batchOperations.get(index));
+            assertThat(batchResponse.getResults().get(index).getOperation()).isEqualTo(batchOperations.get(index));
         }
     }
 
@@ -391,9 +391,9 @@ public void batchReadsOnlyTest() {
         assertThat(batchResponse.getResults().get(1).getItem(TestDoc.class)).isEqualTo(this.TestDocPk1ExistingB);
         assertThat(batchResponse.getResults().get(2).getItem(TestDoc.class)).isEqualTo(this.TestDocPk1ExistingC);
 
-        List> batchOperations = batch.getOperations();
+        List> batchOperations = batch.getOperationsInternal();
         for (int index = 0; index < batchOperations.size(); index++) {
-            assertThat(batchResponse.getResults().get(index).getItemBatchOperation()).isEqualTo(batchOperations.get(index));
+            assertThat(batchResponse.getResults().get(index).getOperation()).isEqualTo(batchOperations.get(index));
         }
     }
 
@@ -431,9 +431,9 @@ public void batchCrud() {
         assertThat(batchResponse.getResults().get(4).getStatusCode()).isEqualTo(HttpResponseStatus.OK.code());
         assertThat(batchResponse.getResults().get(5).getStatusCode()).isEqualTo(HttpResponseStatus.NO_CONTENT.code());
 
-        List> batchOperations = batch.getOperations();
+        List> batchOperations = batch.getOperationsInternal();
         for (int index = 0; index < batchOperations.size(); index++) {
-            assertThat(batchResponse.getResults().get(index).getItemBatchOperation()).isEqualTo(batchOperations.get(index));
+            assertThat(batchResponse.getResults().get(index).getOperation()).isEqualTo(batchOperations.get(index));
         }
 
         assertThat(batchResponse.getResults().get(1).getItem(TestDoc.class)).isEqualTo(this.TestDocPk1ExistingC);
@@ -469,7 +469,7 @@ public void batchWithReplaceOfStaleEntity() {
         TestDoc staleTestDocToReplace = this.getTestDocCopy(this.TestDocPk1ExistingA);
         staleTestDocToReplace.setCost(staleTestDocToReplace.getCost() + 1);
 
-        ItemBatchRequestOptions staleReplaceOptions = new ItemBatchRequestOptions();
+        TransactionalBatchItemRequestOptions staleReplaceOptions = new TransactionalBatchItemRequestOptions();
         staleReplaceOptions.setIfMatchETag(UUID.randomUUID().toString());
 
         this.runWithError(
@@ -509,7 +509,7 @@ public void batchWithCreateConflict() {
 
     private void runWithError(
         CosmosContainer container,
-        Function> appendOperation,
+        Function appendOperation,
         HttpResponseStatus expectedFailedOperationStatusCode) {
 
         TestDoc testDocToCreate = this.populateTestDoc(this.partitionKey1);
@@ -530,9 +530,9 @@ private void runWithError(
         assertThat(batchResponse.getResults().get(1).getStatusCode()).isEqualTo(expectedFailedOperationStatusCode.code());
         assertThat(batchResponse.getResults().get(2).getStatusCode()).isEqualTo(HttpResponseStatus.FAILED_DEPENDENCY.code());
 
-        List> batchOperations = batch.getOperations();
+        List> batchOperations = batch.getOperationsInternal();
         for (int index = 0; index < batchOperations.size(); index++) {
-            assertThat(batchResponse.getResults().get(index).getItemBatchOperation()).isEqualTo(batchOperations.get(index));
+            assertThat(batchResponse.getResults().get(index).getOperation()).isEqualTo(batchOperations.get(index));
         }
 
         this.verifyNotFound(container, testDocToCreate);
diff --git a/sdk/cosmos/azure-cosmos/src/test/java/com/azure/cosmos/implementation/batch/TransactionalBatchResponseTests.java b/sdk/cosmos/azure-cosmos/src/test/java/com/azure/cosmos/implementation/batch/TransactionalBatchResponseTests.java
index fc247993f3b9..9109c85726b0 100644
--- a/sdk/cosmos/azure-cosmos/src/test/java/com/azure/cosmos/implementation/batch/TransactionalBatchResponseTests.java
+++ b/sdk/cosmos/azure-cosmos/src/test/java/com/azure/cosmos/implementation/batch/TransactionalBatchResponseTests.java
@@ -10,8 +10,6 @@
 import com.azure.cosmos.implementation.HttpConstants;
 import com.azure.cosmos.implementation.RxDocumentServiceResponse;
 import com.azure.cosmos.implementation.directconnectivity.StoreResponse;
-import com.azure.cosmos.models.ItemBatchOperation;
-import com.azure.cosmos.models.ModelBridgeInternal;
 import com.azure.cosmos.models.PartitionKey;
 import io.netty.handler.codec.http.HttpResponseStatus;
 import org.testng.annotations.Test;
@@ -36,9 +34,8 @@ public void validateAllSetValuesInResponse() {
         List results = new ArrayList<>();
         ItemBatchOperation[] arrayOperations = new ItemBatchOperation[1];
 
-        ItemBatchOperation operation = ModelBridgeInternal.createItemBatchOperation(
+        ItemBatchOperation operation = new ItemBatchOperation<>(
             CosmosItemOperationType.Read,
-            0,
             "0",
             PartitionKey.NONE,
             null,
@@ -98,7 +95,7 @@ public void validateAllSetValuesInResponse() {
         assertThat(batchResponse.getResults().get(0).getRetryAfterDuration()).isEqualTo(Duration.ofMillis(100));
         assertThat(batchResponse.getResults().get(0).getSubStatusCode()).isEqualTo(HttpConstants.SubStatusCodes.PARTITION_KEY_MISMATCH);
         assertThat(batchResponse.getResults().get(0).getStatusCode()).isEqualTo(HttpResponseStatus.NOT_MODIFIED.code());
-        assertThat(batchResponse.getResults().get(0).getItemBatchOperation()).isEqualTo(operation);
+        assertThat(batchResponse.getResults().get(0).getOperation()).isEqualTo(operation);
     }
 
     @Test(groups = {"unit"}, timeOut = TIMEOUT)
@@ -106,9 +103,8 @@ public void validateEmptyHeaderInResponse() {
         List results = new ArrayList<>();
         ItemBatchOperation[] arrayOperations = new ItemBatchOperation[1];
 
-        ItemBatchOperation operation = ModelBridgeInternal.createItemBatchOperation(
+        ItemBatchOperation operation = new ItemBatchOperation<>(
             CosmosItemOperationType.Read,
-            0,
             "0",
             PartitionKey.NONE,
             null,
@@ -159,6 +155,6 @@ public void validateEmptyHeaderInResponse() {
         assertThat(batchResponse.getResults().get(0).getRetryAfterDuration()).isEqualTo(Duration.ZERO);
         assertThat(batchResponse.getResults().get(0).getSubStatusCode()).isEqualTo(0);
         assertThat(batchResponse.getResults().get(0).getStatusCode()).isEqualTo(HttpResponseStatus.NOT_MODIFIED.code());
-        assertThat(batchResponse.getResults().get(0).getItemBatchOperation()).isEqualTo(operation);
+        assertThat(batchResponse.getResults().get(0).getOperation()).isEqualTo(operation);
     }
 }

From 17ea44b3779d14183752e31322b7540bb44292da Mon Sep 17 00:00:00 2001
From: Rakesh Kumar 
Date: Fri, 9 Oct 2020 22:23:12 +0530
Subject: [PATCH 20/22] Fixing comment

Signed-off-by: Rakesh Kumar 
---
 .../com/azure/cosmos/CosmosItemOperation.java |  2 +
 .../azure/cosmos/CosmosItemOperationType.java | 12 +--
 .../com/azure/cosmos/TransactionalBatch.java  | 56 +++++++-------
 .../TransactionalBatchItemRequestOptions.java |  3 +-
 .../TransactionalBatchOperationResult.java    |  1 -
 .../implementation/RxDocumentClientImpl.java  |  5 +-
 .../batch/ServerBatchRequest.java             | 12 +--
 .../SinglePartitionKeyServerBatchRequest.java |  2 +-
 .../cosmos/BatchOperationResultTests.java     |  2 +-
 .../TransactionalBatchAsyncContainerTest.java | 22 +++---
 .../azure/cosmos/TransactionalBatchTest.java  | 74 +++++++++----------
 .../TransactionalBatchResponseTests.java      |  4 +-
 12 files changed, 99 insertions(+), 96 deletions(-)

diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosItemOperation.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosItemOperation.java
index fc8b57885081..d7c78a07c35a 100644
--- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosItemOperation.java
+++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosItemOperation.java
@@ -4,7 +4,9 @@
 package com.azure.cosmos;
 
 import com.azure.cosmos.models.PartitionKey;
+import com.azure.cosmos.util.Beta;
 
+@Beta(Beta.SinceVersion.V4_7_0)
 public interface CosmosItemOperation {
     String getId();
 
diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosItemOperationType.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosItemOperationType.java
index e16c38754915..7054d2839066 100644
--- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosItemOperationType.java
+++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosItemOperationType.java
@@ -4,14 +4,16 @@
 package com.azure.cosmos;
 
 import com.azure.cosmos.implementation.batch.BatchRequestResponseConstant;
+import com.azure.cosmos.util.Beta;
 
+@Beta(Beta.SinceVersion.V4_7_0)
 public enum CosmosItemOperationType {
 
-    Create(BatchRequestResponseConstant.OPERATION_CREATE),
-    Delete(BatchRequestResponseConstant.OPERATION_DELETE),
-    Read(BatchRequestResponseConstant.OPERATION_READ),
-    Replace(BatchRequestResponseConstant.OPERATION_REPLACE),
-    Upsert(BatchRequestResponseConstant.OPERATION_UPSERT);
+    CREATE(BatchRequestResponseConstant.OPERATION_CREATE),
+    DELETE(BatchRequestResponseConstant.OPERATION_DELETE),
+    READ(BatchRequestResponseConstant.OPERATION_READ),
+    REPLACE(BatchRequestResponseConstant.OPERATION_REPLACE),
+    UPSERT(BatchRequestResponseConstant.OPERATION_UPSERT);
 
     CosmosItemOperationType(String operationValue) {
         this.operationValue = operationValue;
diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/TransactionalBatch.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/TransactionalBatch.java
index aee0b13f2a3c..0d852c64883f 100644
--- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/TransactionalBatch.java
+++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/TransactionalBatch.java
@@ -39,10 +39,10 @@
  * ToDoActivity test3 = new ToDoActivity(activityType, "swimming", "ToBeDone");
  *
  * TransactionalBatch batch = TransactionalBatch.createTransactionalBatch(new Cosmos.PartitionKey(activityType));
- * batch.createItem(test1);
- * batch.replaceItem(test2.id, test2);
- * batch.upsertItem(test3);
- * batch.deleteItem("reading");
+ * batch.createItemOperation(test1);
+ * batch.replaceItemOperation(test2.id, test2);
+ * batch.upsertItemOperation(test3);
+ * batch.deleteItemOperation("reading");
  *
  * TransactionalBatchResponse response = container.executeTransactionalBatch(batch);
  *
@@ -64,10 +64,10 @@
  * String activityType = "personal";
  *
  * TransactionalBatch batch = TransactionalBatch.createTransactionalBatch(new Cosmos.PartitionKey(activityType));
- * batch.readItem("playing");
- * batch.readItem("walking");
- * batch.readItem("jogging");
- * batch.readItem("running")v
+ * batch.readItemOperation("playing");
+ * batch.readItemOperation("walking");
+ * batch.readItemOperation("jogging");
+ * batch.readItemOperation("running");
  *
  * TransactionalBatchResponse response = container.executeTransactionalBatch(batch);
  * List resultItems = new ArrayList();
@@ -116,9 +116,9 @@ public static TransactionalBatch createTransactionalBatch(PartitionKey partition
      *
      * @return The transactional batch instance with the operation added.
      */
-    public  CosmosItemOperation createItem(T item) {
+    public  CosmosItemOperation createItemOperation(T item) {
         checkNotNull(item, "expected non-null item");
-        return this.createItem(item, new TransactionalBatchItemRequestOptions());
+        return this.createItemOperation(item, new TransactionalBatchItemRequestOptions());
     }
 
     /**
@@ -131,7 +131,7 @@ public  CosmosItemOperation createItem(T item) {
      *
      * @return The transactional batch instance with the operation added.
      */
-    public  CosmosItemOperation createItem(T item, TransactionalBatchItemRequestOptions requestOptions) {
+    public  CosmosItemOperation createItemOperation(T item, TransactionalBatchItemRequestOptions requestOptions) {
 
         checkNotNull(item, "expected non-null item");
         if (requestOptions == null) {
@@ -139,7 +139,7 @@ public  CosmosItemOperation createItem(T item, TransactionalBatchItemRequestO
         }
 
         ItemBatchOperation operation = new ItemBatchOperation(
-            CosmosItemOperationType.Create,
+            CosmosItemOperationType.CREATE,
             null,
             this.getPartitionKeyValue(),
             requestOptions.toRequestOptions(),
@@ -158,9 +158,9 @@ public  CosmosItemOperation createItem(T item, TransactionalBatchItemRequestO
      *
      * @return The transactional batch instance with the operation added.
      */
-    public CosmosItemOperation deleteItem(String id) {
+    public CosmosItemOperation deleteItemOperation(String id) {
         checkNotNull(id, "expected non-null id");
-        return this.deleteItem(id, new TransactionalBatchItemRequestOptions());
+        return this.deleteItemOperation(id, new TransactionalBatchItemRequestOptions());
     }
 
     /**
@@ -171,7 +171,7 @@ public CosmosItemOperation deleteItem(String id) {
      *
      * @return The transactional batch instance with the operation added.
      */
-    public CosmosItemOperation deleteItem(String id, TransactionalBatchItemRequestOptions requestOptions) {
+    public CosmosItemOperation deleteItemOperation(String id, TransactionalBatchItemRequestOptions requestOptions) {
 
         checkNotNull(id, "expected non-null id");
         if (requestOptions == null) {
@@ -179,7 +179,7 @@ public CosmosItemOperation deleteItem(String id, TransactionalBatchItemRequestOp
         }
 
         ItemBatchOperation operation = new ItemBatchOperation<>(
-            CosmosItemOperationType.Delete,
+            CosmosItemOperationType.DELETE,
             id,
             this.getPartitionKeyValue(),
             requestOptions.toRequestOptions(),
@@ -198,9 +198,9 @@ public CosmosItemOperation deleteItem(String id, TransactionalBatchItemRequestOp
      *
      * @return The transactional batch instance with the operation added.
      */
-    public CosmosItemOperation readItem(String id) {
+    public CosmosItemOperation readItemOperation(String id) {
         checkNotNull(id, "expected non-null id");
-        return this.readItem(id, new TransactionalBatchItemRequestOptions());
+        return this.readItemOperation(id, new TransactionalBatchItemRequestOptions());
     }
 
     /**
@@ -211,7 +211,7 @@ public CosmosItemOperation readItem(String id) {
      *
      * @return The transactional batch instance with the operation added.
      */
-    public CosmosItemOperation readItem(String id, TransactionalBatchItemRequestOptions requestOptions) {
+    public CosmosItemOperation readItemOperation(String id, TransactionalBatchItemRequestOptions requestOptions) {
 
         checkNotNull(id, "expected non-null id");
         if (requestOptions == null) {
@@ -219,7 +219,7 @@ public CosmosItemOperation readItem(String id, TransactionalBatchItemRequestOpti
         }
 
         ItemBatchOperation operation = new ItemBatchOperation<>(
-            CosmosItemOperationType.Read,
+            CosmosItemOperationType.READ,
             id,
             this.getPartitionKeyValue(),
             requestOptions.toRequestOptions(),
@@ -240,10 +240,10 @@ public CosmosItemOperation readItem(String id, TransactionalBatchItemRequestOpti
      *
      * @return The transactional batch instance with the operation added.
      */
-    public  CosmosItemOperation replaceItem(String id, T item) {
+    public  CosmosItemOperation replaceItemOperation(String id, T item) {
         checkNotNull(id, "expected non-null id");
         checkNotNull(item, "expected non-null item");
-        return this.replaceItem(id, item, new TransactionalBatchItemRequestOptions());
+        return this.replaceItemOperation(id, item, new TransactionalBatchItemRequestOptions());
     }
 
     /**
@@ -257,7 +257,7 @@ public  CosmosItemOperation replaceItem(String id, T item) {
      *
      * @return The transactional batch instance with the operation added.
      */
-    public  CosmosItemOperation replaceItem(
+    public  CosmosItemOperation replaceItemOperation(
         String id, T item, TransactionalBatchItemRequestOptions requestOptions) {
 
         checkNotNull(id, "expected non-null id");
@@ -267,7 +267,7 @@ public  CosmosItemOperation replaceItem(
         }
 
         ItemBatchOperation operation = new ItemBatchOperation(
-            CosmosItemOperationType.Replace,
+            CosmosItemOperationType.REPLACE,
             id,
             this.getPartitionKeyValue(),
             requestOptions.toRequestOptions(),
@@ -287,9 +287,9 @@ public  CosmosItemOperation replaceItem(
      *
      * @return The transactional batch instance with the operation added.
      */
-    public  CosmosItemOperation upsertItem(T item) {
+    public  CosmosItemOperation upsertItemOperation(T item) {
         checkNotNull(item, "expected non-null item");
-        return this.upsertItem(item, new TransactionalBatchItemRequestOptions());
+        return this.upsertItemOperation(item, new TransactionalBatchItemRequestOptions());
     }
 
     /**
@@ -302,7 +302,7 @@ public  CosmosItemOperation upsertItem(T item) {
      *
      * @return The transactional batch instance with the operation added.
      */
-    public  CosmosItemOperation upsertItem(T item, TransactionalBatchItemRequestOptions requestOptions) {
+    public  CosmosItemOperation upsertItemOperation(T item, TransactionalBatchItemRequestOptions requestOptions) {
 
         checkNotNull(item, "expected non-null item");
         if (requestOptions == null) {
@@ -310,7 +310,7 @@ public  CosmosItemOperation upsertItem(T item, TransactionalBatchItemRequestO
         }
 
         ItemBatchOperation operation = new ItemBatchOperation(
-            CosmosItemOperationType.Upsert,
+            CosmosItemOperationType.UPSERT,
             null,
             this.getPartitionKeyValue(),
             requestOptions.toRequestOptions(),
diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/TransactionalBatchItemRequestOptions.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/TransactionalBatchItemRequestOptions.java
index 3776eb53cd2d..dde32da47d0e 100644
--- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/TransactionalBatchItemRequestOptions.java
+++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/TransactionalBatchItemRequestOptions.java
@@ -7,8 +7,7 @@
 import com.azure.cosmos.util.Beta;
 
 /**
- * Encapsulates options that can be specified for an operation within a for a batch or bulk request.
- * Currently used in {@link TransactionalBatch} to pass option.
+ * Encapsulates options that can be specified for an operation within a {@link TransactionalBatch}.
  */
 @Beta(Beta.SinceVersion.V4_7_0)
 public final class TransactionalBatchItemRequestOptions {
diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/TransactionalBatchOperationResult.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/TransactionalBatchOperationResult.java
index a2fe6c2531ce..7598ca2ae9c3 100644
--- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/TransactionalBatchOperationResult.java
+++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/TransactionalBatchOperationResult.java
@@ -4,7 +4,6 @@
 package com.azure.cosmos;
 
 import com.azure.cosmos.implementation.JsonSerializable;
-import com.azure.cosmos.implementation.batch.ItemBatchOperation;
 import com.azure.cosmos.util.Beta;
 import com.fasterxml.jackson.databind.node.ObjectNode;
 import java.time.Duration;
diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/RxDocumentClientImpl.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/RxDocumentClientImpl.java
index a03f17c1f240..33e6f3f8334a 100644
--- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/RxDocumentClientImpl.java
+++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/RxDocumentClientImpl.java
@@ -60,7 +60,6 @@
 import java.net.URI;
 import java.net.URLEncoder;
 import java.nio.ByteBuffer;
-import java.nio.charset.StandardCharsets;
 import java.time.Instant;
 import java.util.ArrayList;
 import java.util.Collections;
@@ -1272,7 +1271,7 @@ private Mono getBatchDocumentRequest(DocumentClientRet
         checkNotNull(serverBatchRequest, "expected non null serverBatchRequest");
 
         Instant serializationStartTimeUTC = Instant.now();
-        ByteBuffer content = ByteBuffer.wrap(serverBatchRequest.getRequestBody().getBytes(StandardCharsets.UTF_8));
+        ByteBuffer content = serverBatchRequest.getRequestBodyAsByteBuffer();
         Instant serializationEndTimeUTC = Instant.now();
 
         SerializationDiagnosticsContext.SerializationDiagnostics serializationDiagnostics = new SerializationDiagnosticsContext.SerializationDiagnostics(
@@ -1316,7 +1315,7 @@ private RxDocumentServiceRequest addBatchHeaders(RxDocumentServiceRequest reques
 
         if(serverBatchRequest instanceof SinglePartitionKeyServerBatchRequest) {
 
-            PartitionKey partitionKey = ((SinglePartitionKeyServerBatchRequest) serverBatchRequest).getPartitionKey();
+            PartitionKey partitionKey = ((SinglePartitionKeyServerBatchRequest) serverBatchRequest).getPartitionKeyValue();
             PartitionKeyInternal partitionKeyInternal;
 
             if (partitionKey.equals(PartitionKey.NONE)) {
diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/batch/ServerBatchRequest.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/batch/ServerBatchRequest.java
index e942e1abbbb9..5878aace0d82 100644
--- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/batch/ServerBatchRequest.java
+++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/batch/ServerBatchRequest.java
@@ -3,13 +3,12 @@
 
 package com.azure.cosmos.implementation.batch;
 
-import com.azure.cosmos.TransactionalBatchItemRequestOptions;
 import com.azure.cosmos.implementation.JsonSerializable;
 import com.azure.cosmos.implementation.Utils;
 import com.azure.cosmos.implementation.apachecommons.collections.list.UnmodifiableList;
-import com.azure.cosmos.implementation.apachecommons.lang.StringUtils;
 import com.fasterxml.jackson.databind.node.ArrayNode;
 
+import java.nio.ByteBuffer;
 import java.util.List;
 
 import static com.azure.cosmos.implementation.guava25.base.Preconditions.checkNotNull;
@@ -23,7 +22,7 @@ public abstract class ServerBatchRequest {
     private final int maxBodyLength;
     private final int maxOperationCount;
 
-    private String requestBody;
+    private ByteBuffer requestBody;
     private List> operations;
     private boolean isAtomicBatch = false;
     private boolean shouldContinueOnError = false;
@@ -62,6 +61,9 @@ final List> createBodyOfBatchRequest(final List operation : operations) {
 
             final JsonSerializable operationJsonSerializable = ItemBatchOperation.writeOperation(operation);
+
+            // TODO(rakkuma): If the string contains unicode the byte encoding len will be more. Fix it.
+            // Issue: https://github.com/Azure/azure-sdk-for-java/issues/16112
             final int operationSerializedLength = operationJsonSerializable.toString().length();
 
             if (totalOperationCount != 0 &&
@@ -76,13 +78,13 @@ final List> createBodyOfBatchRequest(final List operation = new ItemBatchOperation<>(
-        CosmosItemOperationType.Read,
+        CosmosItemOperationType.READ,
         null,
         null,
         null,
diff --git a/sdk/cosmos/azure-cosmos/src/test/java/com/azure/cosmos/TransactionalBatchAsyncContainerTest.java b/sdk/cosmos/azure-cosmos/src/test/java/com/azure/cosmos/TransactionalBatchAsyncContainerTest.java
index 220370b915db..816600693ed6 100644
--- a/sdk/cosmos/azure-cosmos/src/test/java/com/azure/cosmos/TransactionalBatchAsyncContainerTest.java
+++ b/sdk/cosmos/azure-cosmos/src/test/java/com/azure/cosmos/TransactionalBatchAsyncContainerTest.java
@@ -46,8 +46,8 @@ public void batchExecutionRepeat() {
         TestDoc replaceDoc = this.getTestDocCopy(firstDoc);
         replaceDoc.setCost(replaceDoc.getCost() + 1);
         TransactionalBatch batch = TransactionalBatch.createTransactionalBatch(this.getPartitionKey(this.partitionKey1));
-        batch.createItem(firstDoc);
-        batch.replaceItem(replaceDoc.getId(), replaceDoc);
+        batch.createItemOperation(firstDoc);
+        batch.replaceItemOperation(replaceDoc.getId(), replaceDoc);
 
         Mono batchResponseMono = batchAsyncContainer.executeTransactionalBatch(batch);
 
@@ -85,10 +85,10 @@ public void batchInvalidSessionToken() throws Exception {
             TestDoc testDocToUpsert = this.populateTestDoc(this.partitionKey1);
 
             TransactionalBatch batch = TransactionalBatch.createTransactionalBatch(this.getPartitionKey(this.partitionKey1));
-            batch.createItem(testDocToCreate);
-            batch.replaceItem(testDocToReplace.getId(), testDocToReplace);
-            batch.upsertItem(testDocToUpsert);
-            batch.deleteItem(this.TestDocPk1ExistingC.getId());
+            batch.createItemOperation(testDocToCreate);
+            batch.replaceItemOperation(testDocToReplace.getId(), testDocToReplace);
+            batch.upsertItemOperation(testDocToUpsert);
+            batch.deleteItemOperation(this.TestDocPk1ExistingC.getId());
 
             TransactionalBatchResponse batchResponse = container.executeTransactionalBatch(
                 batch, new TransactionalBatchRequestOptions().setSessionToken(invalidSessionToken)).block();
@@ -114,11 +114,11 @@ public void batchInvalidSessionToken() throws Exception {
             TestDoc testDocToUpsert = this.populateTestDoc(this.partitionKey1);
 
             TransactionalBatch batch = TransactionalBatch.createTransactionalBatch(this.getPartitionKey(this.partitionKey1));
-            batch.createItem(testDocToCreate);
-            batch.replaceItem(testDocToReplace.getId(), testDocToReplace);
-            batch.upsertItem(testDocToUpsert);
-            batch.deleteItem(this.TestDocPk1ExistingD.getId());
-            batch.readItem(this.TestDocPk1ExistingA.getId());
+            batch.createItemOperation(testDocToCreate);
+            batch.replaceItemOperation(testDocToReplace.getId(), testDocToReplace);
+            batch.upsertItemOperation(testDocToUpsert);
+            batch.deleteItemOperation(this.TestDocPk1ExistingD.getId());
+            batch.readItemOperation(this.TestDocPk1ExistingA.getId());
 
             try {
                 container.executeTransactionalBatch(
diff --git a/sdk/cosmos/azure-cosmos/src/test/java/com/azure/cosmos/TransactionalBatchTest.java b/sdk/cosmos/azure-cosmos/src/test/java/com/azure/cosmos/TransactionalBatchTest.java
index ee76747317c0..562f89e863d8 100644
--- a/sdk/cosmos/azure-cosmos/src/test/java/com/azure/cosmos/TransactionalBatchTest.java
+++ b/sdk/cosmos/azure-cosmos/src/test/java/com/azure/cosmos/TransactionalBatchTest.java
@@ -54,8 +54,8 @@ public void batchOrdered() {
         replaceDoc.setCost(replaceDoc.getCost() + 1);
 
         TransactionalBatch batch = TransactionalBatch.createTransactionalBatch(this.getPartitionKey(this.partitionKey1));
-        batch.createItem(firstDoc);
-        batch.replaceItem(replaceDoc.getId(), replaceDoc);
+        batch.createItemOperation(firstDoc);
+        batch.replaceItemOperation(replaceDoc.getId(), replaceDoc);
 
         TransactionalBatchResponse batchResponse = container.executeTransactionalBatch(batch);
 
@@ -87,10 +87,10 @@ public void batchMultipleItemExecution() {
         assertThat(createResponse.getStatusCode()).isEqualTo(HttpResponseStatus.CREATED.code());
 
         TransactionalBatch batch = TransactionalBatch.createTransactionalBatch(this.getPartitionKey(this.partitionKey1));
-        batch.createItem(firstDoc);
-        batch.createItem(eventDoc1);
-        batch.replaceItem(replaceDoc.getId(), replaceDoc);
-        batch.readItem(readEventDoc.getId());
+        batch.createItemOperation(firstDoc);
+        batch.createItemOperation(eventDoc1);
+        batch.replaceItemOperation(replaceDoc.getId(), replaceDoc);
+        batch.readItemOperation(readEventDoc.getId());
 
         TransactionalBatchResponse batchResponse = container.executeTransactionalBatch(batch);
 
@@ -139,8 +139,8 @@ public void batchItemETagTest() {
             firstReplaceOptions.setIfMatchETag(response.getETag());
 
             TransactionalBatch batch = TransactionalBatch.createTransactionalBatch(this.getPartitionKey(this.partitionKey1));
-            batch.createItem(testDocToCreate);
-            batch.replaceItem(testDocToReplace.getId(), testDocToReplace, firstReplaceOptions);
+            batch.createItemOperation(testDocToCreate);
+            batch.replaceItemOperation(testDocToReplace.getId(), testDocToReplace, firstReplaceOptions);
 
             TransactionalBatchResponse batchResponse = container.executeTransactionalBatch(batch);
 
@@ -162,7 +162,7 @@ public void batchItemETagTest() {
             replaceOptions.setIfMatchETag(String.valueOf(this.getRandom().nextInt()));
 
             TransactionalBatch batch = TransactionalBatch.createTransactionalBatch(this.getPartitionKey(this.partitionKey1));
-            batch.replaceItem(testDocToReplace.getId(), testDocToReplace, replaceOptions);
+            batch.replaceItemOperation(testDocToReplace.getId(), testDocToReplace, replaceOptions);
 
             TransactionalBatchResponse batchResponse = container.executeTransactionalBatch(batch);
 
@@ -197,7 +197,7 @@ public void batchErrorSessionToken() {
         {
             // Only errored read
             TransactionalBatch batch = TransactionalBatch.createTransactionalBatch(this.getPartitionKey(this.partitionKey1));
-            batch.readItem(UUID.randomUUID().toString());
+            batch.readItemOperation(UUID.randomUUID().toString());
 
             TransactionalBatchResponse batchResponse = container.executeTransactionalBatch(batch);
 
@@ -217,8 +217,8 @@ public void batchErrorSessionToken() {
         {
             // One valid read one error read
             TransactionalBatch batch = TransactionalBatch.createTransactionalBatch(this.getPartitionKey(this.partitionKey1));
-            batch.readItem(this.TestDocPk1ExistingA.getId());
-            batch.readItem(UUID.randomUUID().toString());
+            batch.readItemOperation(this.TestDocPk1ExistingA.getId());
+            batch.readItemOperation(UUID.randomUUID().toString());
 
             TransactionalBatchResponse batchResponse = container.executeTransactionalBatch(batch);
 
@@ -239,8 +239,8 @@ public void batchErrorSessionToken() {
         {
             // One error one valid read
             TransactionalBatch batch = TransactionalBatch.createTransactionalBatch(this.getPartitionKey(this.partitionKey1));
-            batch.readItem(UUID.randomUUID().toString());
-            batch.readItem(this.TestDocPk1ExistingA.getId());
+            batch.readItemOperation(UUID.randomUUID().toString());
+            batch.readItemOperation(this.TestDocPk1ExistingA.getId());
 
             TransactionalBatchResponse batchResponse = container.executeTransactionalBatch(batch);
 
@@ -263,8 +263,8 @@ public void batchErrorSessionToken() {
             TestDoc testDocToCreate = this.populateTestDoc(this.partitionKey1);
 
             TransactionalBatch batch = TransactionalBatch.createTransactionalBatch(this.getPartitionKey(this.partitionKey1));
-            batch.createItem(testDocToCreate);
-            batch.readItem(UUID.randomUUID().toString());
+            batch.createItemOperation(testDocToCreate);
+            batch.readItemOperation(UUID.randomUUID().toString());
 
             TransactionalBatchResponse batchResponse = container.executeTransactionalBatch(batch);
 
@@ -286,8 +286,8 @@ public void batchErrorSessionToken() {
             // One error one valid write
             TestDoc testDocToCreate = this.populateTestDoc(this.partitionKey1);
             TransactionalBatch batch = TransactionalBatch.createTransactionalBatch(this.getPartitionKey(this.partitionKey1));
-            batch.readItem(UUID.randomUUID().toString());
-            batch.createItem(testDocToCreate);
+            batch.readItemOperation(UUID.randomUUID().toString());
+            batch.createItemOperation(testDocToCreate);
 
             TransactionalBatchResponse batchResponse = container.executeTransactionalBatch(batch);
 
@@ -314,7 +314,7 @@ public void batchWithTooManyOperationsTest() {
         TransactionalBatch batch = TransactionalBatch.createTransactionalBatch(this.getPartitionKey(this.partitionKey1));
 
         for (int i = 0; i < operationCount; i++) {
-            batch.readItem("someId");
+            batch.readItemOperation("someId");
         }
 
         try {
@@ -336,7 +336,7 @@ public void batchLargerThanServerRequest() {
 
         for (int i = 0; i < operationCount; i++) {
             TestDoc doc = this.populateTestDoc(this.partitionKey1, appxDocSize);
-            batch.createItem(doc);
+            batch.createItemOperation(doc);
         }
 
         try {
@@ -356,7 +356,7 @@ public void batchServerResponseTooLarge() {
 
         TransactionalBatch batch = TransactionalBatch.createTransactionalBatch(this.getPartitionKey(this.partitionKey1));
         for (int i = 0; i < operationCount; i++) {
-            batch.readItem(doc.getId());
+            batch.readItemOperation(doc.getId());
         }
 
         TransactionalBatchResponse batchResponse = batchContainer.executeTransactionalBatch(batch);
@@ -375,9 +375,9 @@ public void batchReadsOnlyTest() {
         this.createJsonTestDocs(container);
 
         TransactionalBatch batch = TransactionalBatch.createTransactionalBatch(this.getPartitionKey(this.partitionKey1));
-        batch.readItem(this.TestDocPk1ExistingA.getId());
-        batch.readItem(this.TestDocPk1ExistingB.getId());
-        batch.readItem(this.TestDocPk1ExistingC.getId());
+        batch.readItemOperation(this.TestDocPk1ExistingA.getId());
+        batch.readItemOperation(this.TestDocPk1ExistingB.getId());
+        batch.readItemOperation(this.TestDocPk1ExistingC.getId());
 
         TransactionalBatchResponse batchResponse = container.executeTransactionalBatch(batch);
 
@@ -412,12 +412,12 @@ public void batchCrud() {
         testDocToReplace.setCost(testDocToReplace.getCost() + 1);
 
         TransactionalBatch batch = TransactionalBatch.createTransactionalBatch(this.getPartitionKey(this.partitionKey1));
-        batch.createItem(testDocToCreate);
-        batch.readItem(this.TestDocPk1ExistingC.getId());
-        batch.replaceItem(testDocToReplace.getId(), testDocToReplace);
-        batch.upsertItem(testDocToUpsert);
-        batch.upsertItem(anotherTestDocToUpsert);
-        batch.deleteItem(this.TestDocPk1ExistingD.getId());
+        batch.createItemOperation(testDocToCreate);
+        batch.readItemOperation(this.TestDocPk1ExistingC.getId());
+        batch.replaceItemOperation(testDocToReplace.getId(), testDocToReplace);
+        batch.upsertItemOperation(testDocToUpsert);
+        batch.upsertItemOperation(anotherTestDocToUpsert);
+        batch.deleteItemOperation(this.TestDocPk1ExistingD.getId());
 
         // We run CRUD operations where all are expected to return HTTP 2xx.
         TransactionalBatchResponse batchResponse = container.executeTransactionalBatch(batch);
@@ -450,7 +450,7 @@ public void batchWithInvalidCreateTest() {
         // partition key mismatch between doc and and value passed in to the operation
         this.runWithError(
             batchContainer,
-            batch -> batch.createItem(this.populateTestDoc(UUID.randomUUID().toString())),
+            batch -> batch.createItemOperation(this.populateTestDoc(UUID.randomUUID().toString())),
             HttpResponseStatus.BAD_REQUEST);
     }
 
@@ -458,7 +458,7 @@ public void batchWithInvalidCreateTest() {
     public void batchWithReadOfNonExistentEntityTest() {
         this.runWithError(
             batchContainer,
-            batch -> batch.readItem(UUID.randomUUID().toString()),
+            batch -> batch.readItemOperation(UUID.randomUUID().toString()),
             HttpResponseStatus.NOT_FOUND);
     }
 
@@ -474,7 +474,7 @@ public void batchWithReplaceOfStaleEntity() {
 
         this.runWithError(
             batchContainer,
-            batch -> batch.replaceItem(staleTestDocToReplace.getId(), staleTestDocToReplace, staleReplaceOptions),
+            batch -> batch.replaceItemOperation(staleTestDocToReplace.getId(), staleTestDocToReplace, staleReplaceOptions),
             HttpResponseStatus.PRECONDITION_FAILED);
 
         // make sure the stale doc hasn't changed
@@ -485,7 +485,7 @@ public void batchWithReplaceOfStaleEntity() {
     public void batchWithDeleteOfNonExistentEntity() {
         this.runWithError(
             batchContainer,
-            batch -> batch.deleteItem(UUID.randomUUID().toString()),
+            batch -> batch.deleteItemOperation(UUID.randomUUID().toString()),
             HttpResponseStatus.NOT_FOUND);
     }
 
@@ -499,7 +499,7 @@ public void batchWithCreateConflict() {
 
         this.runWithError(
             batchContainer,
-            batch -> batch.createItem(conflictingTestDocToCreate),
+            batch -> batch.createItemOperation(conflictingTestDocToCreate),
             HttpResponseStatus.CONFLICT);
 
         // make sure the conflicted doc hasn't changed
@@ -516,11 +516,11 @@ private void runWithError(
         TestDoc anotherTestDocToCreate = this.populateTestDoc(this.partitionKey1);
 
         TransactionalBatch batch = TransactionalBatch.createTransactionalBatch(this.getPartitionKey(this.partitionKey1));
-        batch.createItem(testDocToCreate);
+        batch.createItemOperation(testDocToCreate);
 
         appendOperation.apply(batch);
 
-        batch.createItem(anotherTestDocToCreate);
+        batch.createItemOperation(anotherTestDocToCreate);
 
         TransactionalBatchResponse batchResponse = container.executeTransactionalBatch(batch);
 
diff --git a/sdk/cosmos/azure-cosmos/src/test/java/com/azure/cosmos/implementation/batch/TransactionalBatchResponseTests.java b/sdk/cosmos/azure-cosmos/src/test/java/com/azure/cosmos/implementation/batch/TransactionalBatchResponseTests.java
index 9109c85726b0..a72173db90ce 100644
--- a/sdk/cosmos/azure-cosmos/src/test/java/com/azure/cosmos/implementation/batch/TransactionalBatchResponseTests.java
+++ b/sdk/cosmos/azure-cosmos/src/test/java/com/azure/cosmos/implementation/batch/TransactionalBatchResponseTests.java
@@ -35,7 +35,7 @@ public void validateAllSetValuesInResponse() {
         ItemBatchOperation[] arrayOperations = new ItemBatchOperation[1];
 
         ItemBatchOperation operation = new ItemBatchOperation<>(
-            CosmosItemOperationType.Read,
+            CosmosItemOperationType.READ,
             "0",
             PartitionKey.NONE,
             null,
@@ -104,7 +104,7 @@ public void validateEmptyHeaderInResponse() {
         ItemBatchOperation[] arrayOperations = new ItemBatchOperation[1];
 
         ItemBatchOperation operation = new ItemBatchOperation<>(
-            CosmosItemOperationType.Read,
+            CosmosItemOperationType.READ,
             "0",
             PartitionKey.NONE,
             null,

From 4c4893d8b4ee66d57b9f2b0db9295252809bbd88 Mon Sep 17 00:00:00 2001
From: Rakesh Kumar 
Date: Fri, 9 Oct 2020 23:20:59 +0530
Subject: [PATCH 21/22] Fix

Signed-off-by: Rakesh Kumar 
---
 .../java/com/azure/cosmos/BridgeInternal.java | 10 +----
 .../com/azure/cosmos/TransactionalBatch.java  |  9 ----
 .../TransactionalBatchOperationResult.java    |  2 +-
 .../implementation/batch/BatchExecutor.java   |  3 +-
 .../batch/BatchResponseParser.java            | 15 ++++---
 .../batch/ItemBatchOperation.java             | 18 ++++----
 .../batch/ServerBatchRequest.java             | 43 +++++++++++--------
 .../SinglePartitionKeyServerBatchRequest.java |  3 +-
 .../TransactionalBatchAsyncContainerTest.java |  3 +-
 .../azure/cosmos/TransactionalBatchTest.java  | 13 +++---
 10 files changed, 54 insertions(+), 65 deletions(-)

diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/BridgeInternal.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/BridgeInternal.java
index 71117a9d8353..aa4250955692 100644
--- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/BridgeInternal.java
+++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/BridgeInternal.java
@@ -26,7 +26,6 @@
 import com.azure.cosmos.implementation.ServiceUnavailableException;
 import com.azure.cosmos.implementation.StoredProcedureResponse;
 import com.azure.cosmos.implementation.Warning;
-import com.azure.cosmos.implementation.batch.ItemBatchOperation;
 import com.azure.cosmos.implementation.directconnectivity.StoreResponse;
 import com.azure.cosmos.implementation.directconnectivity.StoreResult;
 import com.azure.cosmos.implementation.directconnectivity.Uri;
@@ -610,11 +609,6 @@ public static Duration getRequestTimeoutFromGatewayConnectionConfig(GatewayConne
         return gatewayConnectionConfig.getRequestTimeout();
     }
 
-    @Warning(value = INTERNAL_USE_ONLY_WARNING)
-    public static List> getOperationsFromTransactionalBatch(TransactionalBatch transactionalBatch) {
-        return transactionalBatch.getOperationsInternal();
-    }
-
     @Warning(value = INTERNAL_USE_ONLY_WARNING)
     public static String getOperationValueForCosmosItemOperationType(CosmosItemOperationType cosmosItemOperationType) {
         return cosmosItemOperationType.getOperationValue();
@@ -633,7 +627,7 @@ public static TransactionalBatchOperationResult createTransactionBatchResult(
         int statusCode,
         Duration retryAfter,
         int subStatusCode,
-        CosmosItemOperation itemBatchOperation) {
+        CosmosItemOperation cosmosItemOperation) {
 
         return new TransactionalBatchOperationResult(
             eTag,
@@ -642,7 +636,7 @@ public static TransactionalBatchOperationResult createTransactionBatchResult(
             statusCode,
             retryAfter,
             subStatusCode,
-            itemBatchOperation);
+            cosmosItemOperation);
     }
 
     @Warning(value = INTERNAL_USE_ONLY_WARNING)
diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/TransactionalBatch.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/TransactionalBatch.java
index 0d852c64883f..697b008f2211 100644
--- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/TransactionalBatch.java
+++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/TransactionalBatch.java
@@ -322,15 +322,6 @@ public  CosmosItemOperation upsertItemOperation(T item, TransactionalBatchIte
         return operation;
     }
 
-    /**
-     * Return the list of operation in an unmodifiable instance  so no one can change it in the down path.
-     *
-     * @return The list of operations which are to be executed.
-     */
-    List> getOperationsInternal() {
-        return UnmodifiableList.unmodifiableList(operations);
-    }
-
     /**
      * Return the list of operation in an unmodifiable instance  so no one can change it in the down path.
      *
diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/TransactionalBatchOperationResult.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/TransactionalBatchOperationResult.java
index 7598ca2ae9c3..63c5f95c062a 100644
--- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/TransactionalBatchOperationResult.java
+++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/TransactionalBatchOperationResult.java
@@ -131,7 +131,7 @@ ObjectNode getResourceObject() {
     /**
      * Gets the original operation for this result.
      *
-     * @return the ItemBatchOperation.
+     * @return the CosmosItemOperation.
      */
     public CosmosItemOperation getOperation() {
         return cosmosItemOperation;
diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/batch/BatchExecutor.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/batch/BatchExecutor.java
index 8cc034b75e45..91c68c39c753 100644
--- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/batch/BatchExecutor.java
+++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/batch/BatchExecutor.java
@@ -6,6 +6,7 @@
 import com.azure.cosmos.BridgeInternal;
 import com.azure.cosmos.CosmosAsyncContainer;
 import com.azure.cosmos.CosmosBridgeInternal;
+import com.azure.cosmos.CosmosItemOperation;
 import com.azure.cosmos.TransactionalBatch;
 import com.azure.cosmos.TransactionalBatchRequestOptions;
 import com.azure.cosmos.TransactionalBatchResponse;
@@ -38,7 +39,7 @@ public BatchExecutor(
      */
     public final Mono executeAsync() {
 
-        List> operations = BridgeInternal.getOperationsFromTransactionalBatch(this.transactionalBatch);
+        List operations = this.transactionalBatch.getOperations();
         checkArgument(operations.size() > 0, "Number of operations should be more than 0.");
 
         final SinglePartitionKeyServerBatchRequest request = SinglePartitionKeyServerBatchRequest.createBatchRequest(
diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/batch/BatchResponseParser.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/batch/BatchResponseParser.java
index dc84c72eff22..378561ae595c 100644
--- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/batch/BatchResponseParser.java
+++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/batch/BatchResponseParser.java
@@ -4,6 +4,7 @@
 package com.azure.cosmos.implementation.batch;
 
 import com.azure.cosmos.BridgeInternal;
+import com.azure.cosmos.CosmosItemOperation;
 import com.azure.cosmos.TransactionalBatchOperationResult;
 import com.azure.cosmos.TransactionalBatchResponse;
 import com.azure.cosmos.implementation.HttpConstants;
@@ -111,14 +112,14 @@ private static TransactionalBatchResponse populateFromResponseContent(
             final ObjectMapper mapper = Utils.getSimpleObjectMapper();
 
             try {
-                final List> itemBatchOperations = request.getOperations();
+                final List cosmosItemOperations = request.getOperations();
                 final ObjectNode[] objectNodes = mapper.readValue(responseContent, ObjectNode[].class);
 
                 for (int index = 0; index < objectNodes.length; index++) {
                     ObjectNode objectInArray = objectNodes[index];
 
                     results.add(
-                        BatchResponseParser.createBatchOperationResultFromJson(objectInArray, itemBatchOperations.get(index)));
+                        BatchResponseParser.createBatchOperationResultFromJson(objectInArray, cosmosItemOperations.get(index)));
                 }
             } catch (IOException ex) {
                 logger.error("Exception in parsing response", ex);
@@ -169,7 +170,7 @@ private static TransactionalBatchResponse populateFromResponseContent(
      */
     private static TransactionalBatchOperationResult createBatchOperationResultFromJson(
         ObjectNode objectNode,
-        ItemBatchOperation itemBatchOperation) {
+        CosmosItemOperation cosmosItemOperation) {
 
         final JsonSerializable jsonSerializable = new JsonSerializable(objectNode);
 
@@ -195,7 +196,7 @@ private static TransactionalBatchOperationResult createBatchOperationResultFromJ
             statusCode,
             retryAfterMilliseconds != null ? Duration.ofMillis(retryAfterMilliseconds) : Duration.ZERO,
             subStatusCode,
-            itemBatchOperation);
+            cosmosItemOperation);
     }
 
     /**
@@ -206,10 +207,10 @@ private static TransactionalBatchOperationResult createBatchOperationResultFromJ
      * @param retryAfterDuration retryAfterDuration.
      * */
     private static void createAndPopulateResults(final TransactionalBatchResponse response,
-                                                 final List> operations,
+                                                 final List operations,
                                                  final Duration retryAfterDuration) {
         final List results = new ArrayList<>(operations.size());
-        for (ItemBatchOperation itemBatchOperation : operations) {
+        for (CosmosItemOperation cosmosItemOperation : operations) {
             results.add(
                 BridgeInternal.createTransactionBatchResult(
                     null,
@@ -218,7 +219,7 @@ private static void createAndPopulateResults(final TransactionalBatchResponse re
                     response.getStatusCode(),
                     retryAfterDuration,
                     response.getSubStatusCode(),
-                    itemBatchOperation
+                    cosmosItemOperation
                 ));
         }
 
diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/batch/ItemBatchOperation.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/batch/ItemBatchOperation.java
index af89cfa819c3..7ff74b406c47 100644
--- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/batch/ItemBatchOperation.java
+++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/batch/ItemBatchOperation.java
@@ -49,27 +49,25 @@ public ItemBatchOperation(
      * TODO(rakkuma): Similarly for hybrid row, operation needs to be written in Hybrid row.
      * Issue: https://github.com/Azure/azure-sdk-for-java/issues/15856
      *
-     * @param operation a single operation which needs to be serialized.
-     *
      * @return instance of JsonSerializable containing values for a operation.
      */
-    static JsonSerializable writeOperation(final ItemBatchOperation operation) {
+    JsonSerializable serializeOperation() {
         final JsonSerializable jsonSerializable = new JsonSerializable();
 
         jsonSerializable.set(
             BatchRequestResponseConstant.FIELD_OPERATION_TYPE,
-            BridgeInternal.getOperationValueForCosmosItemOperationType(operation.getOperationType()));
+            BridgeInternal.getOperationValueForCosmosItemOperationType(this.getOperationType()));
 
-        if (StringUtils.isNotEmpty(operation.getId())) {
-            jsonSerializable.set(BatchRequestResponseConstant.FIELD_ID, operation.getId());
+        if (StringUtils.isNotEmpty(this.getId())) {
+            jsonSerializable.set(BatchRequestResponseConstant.FIELD_ID, this.getId());
         }
 
-        if (operation.getItemInternal() != null) {
-            jsonSerializable.set(BatchRequestResponseConstant.FIELD_RESOURCE_BODY, operation.getItemInternal());
+        if (this.getItemInternal() != null) {
+            jsonSerializable.set(BatchRequestResponseConstant.FIELD_RESOURCE_BODY, this.getItemInternal());
         }
 
-        if (operation.getRequestOptions() != null) {
-            RequestOptions requestOptions = operation.getRequestOptions();
+        if (this.getRequestOptions() != null) {
+            RequestOptions requestOptions = this.getRequestOptions();
 
             if (StringUtils.isNotEmpty(requestOptions.getIfMatchETag())) {
                 jsonSerializable.set(BatchRequestResponseConstant.FIELD_IF_MATCH, requestOptions.getIfMatchETag());
diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/batch/ServerBatchRequest.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/batch/ServerBatchRequest.java
index 5878aace0d82..2c777f71a027 100644
--- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/batch/ServerBatchRequest.java
+++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/batch/ServerBatchRequest.java
@@ -3,6 +3,7 @@
 
 package com.azure.cosmos.implementation.batch;
 
+import com.azure.cosmos.CosmosItemOperation;
 import com.azure.cosmos.implementation.JsonSerializable;
 import com.azure.cosmos.implementation.Utils;
 import com.azure.cosmos.implementation.apachecommons.collections.list.UnmodifiableList;
@@ -23,7 +24,7 @@ public abstract class ServerBatchRequest {
     private final int maxOperationCount;
 
     private ByteBuffer requestBody;
-    private List> operations;
+    private List operations;
     private boolean isAtomicBatch = false;
     private boolean shouldContinueOnError = false;
 
@@ -49,7 +50,7 @@ public abstract class ServerBatchRequest {
      *
      * @return Any pending operations that were not included in the request.
      */
-    final List> createBodyOfBatchRequest(final List> operations) {
+    List createBodyOfBatchRequest(final List operations) {
 
         checkNotNull(operations, "expected non-null operations");
 
@@ -58,24 +59,28 @@ final List> createBodyOfBatchRequest(final List operation : operations) {
+        for(CosmosItemOperation operation : operations) {
+            if (operation instanceof ItemBatchOperation) {
+                final ItemBatchOperation itemBatchOperation = (ItemBatchOperation) operation;
+                final JsonSerializable operationJsonSerializable = itemBatchOperation.serializeOperation();
 
-            final JsonSerializable operationJsonSerializable = ItemBatchOperation.writeOperation(operation);
+                // TODO(rakkuma): If the string contains unicode the byte encoding len will be more. Fix it.
+                // Issue: https://github.com/Azure/azure-sdk-for-java/issues/16112
+                final int operationSerializedLength = operationJsonSerializable.toString().length();
 
-            // TODO(rakkuma): If the string contains unicode the byte encoding len will be more. Fix it.
-            // Issue: https://github.com/Azure/azure-sdk-for-java/issues/16112
-            final int operationSerializedLength = operationJsonSerializable.toString().length();
+                if (totalOperationCount != 0 &&
+                    (totalSerializedLength + operationSerializedLength > this.maxBodyLength || totalOperationCount + 1 > this.maxOperationCount)) {
+                    // Apply the limit only if at least there is one operation in selected operations.
+                    break;
+                }
 
-            if (totalOperationCount != 0 &&
-                (totalSerializedLength + operationSerializedLength > this.maxBodyLength || totalOperationCount + 1 > this.maxOperationCount)) {
-                // Apply the limit only if at least there is one operation in selected operations.
-                break;
-            }
-
-            totalSerializedLength += operationSerializedLength;
-            totalOperationCount++;
+                totalSerializedLength += operationSerializedLength;
+                totalOperationCount++;
 
-            arrayNode.add(operationJsonSerializable.getPropertyBag());
+                arrayNode.add(operationJsonSerializable.getPropertyBag());
+            } else {
+                throw new UnsupportedOperationException("Unknown CosmosItemOperation.");
+            }
         }
 
         this.requestBody = ByteBuffer.wrap(Utils.getUTF8Bytes(arrayNode.toString()));
@@ -91,13 +96,13 @@ public final ByteBuffer getRequestBodyAsByteBuffer() {
     }
 
     /**
-     * Gets the list of {@link ItemBatchOperation operations} in this {@link ServerBatchRequest batch request}.
+     * Gets the list of {@link CosmosItemOperation operations} in this {@link ServerBatchRequest batch request}.
      *
      * The list returned by this method is unmodifiable.
      *
-     * @return the list of {@link ItemBatchOperation operations} in this {@link ServerBatchRequest batch request}.
+     * @return the list of {@link CosmosItemOperation operations} in this {@link ServerBatchRequest batch request}.
      */
-    public final List> getOperations() {
+    public final List getOperations() {
         return UnmodifiableList.unmodifiableList(this.operations);
     }
 
diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/batch/SinglePartitionKeyServerBatchRequest.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/batch/SinglePartitionKeyServerBatchRequest.java
index 1630b94f630b..0a433fd62952 100644
--- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/batch/SinglePartitionKeyServerBatchRequest.java
+++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/batch/SinglePartitionKeyServerBatchRequest.java
@@ -3,6 +3,7 @@
 
 package com.azure.cosmos.implementation.batch;
 
+import com.azure.cosmos.CosmosItemOperation;
 import com.azure.cosmos.models.PartitionKey;
 import java.util.List;
 
@@ -34,7 +35,7 @@ private SinglePartitionKeyServerBatchRequest(final PartitionKey partitionKey) {
      */
     static SinglePartitionKeyServerBatchRequest createBatchRequest(
         final PartitionKey partitionKey,
-        final List> operations) {
+        final List operations) {
 
         checkNotNull(partitionKey, "expected non-null partitionKey");
         checkNotNull(operations, "expected non-null operations");
diff --git a/sdk/cosmos/azure-cosmos/src/test/java/com/azure/cosmos/TransactionalBatchAsyncContainerTest.java b/sdk/cosmos/azure-cosmos/src/test/java/com/azure/cosmos/TransactionalBatchAsyncContainerTest.java
index 816600693ed6..1965a8691f0c 100644
--- a/sdk/cosmos/azure-cosmos/src/test/java/com/azure/cosmos/TransactionalBatchAsyncContainerTest.java
+++ b/sdk/cosmos/azure-cosmos/src/test/java/com/azure/cosmos/TransactionalBatchAsyncContainerTest.java
@@ -5,7 +5,6 @@
 
 import com.azure.cosmos.implementation.HttpConstants;
 import com.azure.cosmos.models.CosmosItemResponse;
-import com.azure.cosmos.implementation.batch.ItemBatchOperation;
 import io.netty.handler.codec.http.HttpResponseStatus;
 import org.assertj.core.api.Assertions;
 import org.testng.annotations.AfterClass;
@@ -100,7 +99,7 @@ public void batchInvalidSessionToken() throws Exception {
             assertThat(batchResponse.getResults().get(2).getStatusCode()).isEqualTo(HttpResponseStatus.CREATED.code());
             assertThat(batchResponse.getResults().get(3).getStatusCode()).isEqualTo(HttpResponseStatus.NO_CONTENT.code());
 
-            List> batchOperations = batch.getOperationsInternal();
+            List batchOperations = batch.getOperations();
             for (int index = 0; index < batchOperations.size(); index++) {
                 assertThat(batchResponse.getResults().get(index).getOperation()).isEqualTo(batchOperations.get(index));
             }
diff --git a/sdk/cosmos/azure-cosmos/src/test/java/com/azure/cosmos/TransactionalBatchTest.java b/sdk/cosmos/azure-cosmos/src/test/java/com/azure/cosmos/TransactionalBatchTest.java
index 562f89e863d8..8dd694fd552a 100644
--- a/sdk/cosmos/azure-cosmos/src/test/java/com/azure/cosmos/TransactionalBatchTest.java
+++ b/sdk/cosmos/azure-cosmos/src/test/java/com/azure/cosmos/TransactionalBatchTest.java
@@ -7,7 +7,6 @@
 import com.azure.cosmos.implementation.ISessionToken;
 import com.azure.cosmos.implementation.guava25.base.Function;
 import com.azure.cosmos.models.CosmosItemResponse;
-import com.azure.cosmos.implementation.batch.ItemBatchOperation;
 import io.netty.handler.codec.http.HttpResponseStatus;
 import org.assertj.core.api.Assertions;
 import org.testng.annotations.AfterClass;
@@ -64,7 +63,7 @@ public void batchOrdered() {
         assertThat(batchResponse.getResults().get(0).getStatusCode()).isEqualTo(HttpResponseStatus.CREATED.code());
         assertThat(batchResponse.getResults().get(1).getStatusCode()).isEqualTo(HttpResponseStatus.OK.code());
 
-        List> batchOperations = batch.getOperationsInternal();
+        List batchOperations = batch.getOperations();
         for (int index = 0; index < batchOperations.size(); index++) {
             assertThat(batchResponse.getResults().get(index).getOperation()).isEqualTo(batchOperations.get(index));
         }
@@ -111,7 +110,7 @@ public void batchMultipleItemExecution() {
         // Ensure that the replace overwrote the doc from the first operation
         this.verifyByRead(container, replaceDoc);
 
-        List> batchOperations = batch.getOperationsInternal();
+        List batchOperations = batch.getOperations();
         for (int index = 0; index < batchOperations.size(); index++) {
             assertThat(batchResponse.getResults().get(index).getOperation()).isEqualTo(batchOperations.get(index));
         }
@@ -363,7 +362,7 @@ public void batchServerResponseTooLarge() {
         assertThat(batchResponse.getStatusCode()).isEqualTo(HttpResponseStatus.REQUEST_ENTITY_TOO_LARGE.code());
         assertThat(batchResponse.getResults().get(1).getStatusCode()).isEqualTo(HttpResponseStatus.FAILED_DEPENDENCY.code());
 
-        List> batchOperations = batch.getOperationsInternal();
+        List batchOperations = batch.getOperations();
         for (int index = 0; index < batchOperations.size(); index++) {
             assertThat(batchResponse.getResults().get(index).getOperation()).isEqualTo(batchOperations.get(index));
         }
@@ -391,7 +390,7 @@ public void batchReadsOnlyTest() {
         assertThat(batchResponse.getResults().get(1).getItem(TestDoc.class)).isEqualTo(this.TestDocPk1ExistingB);
         assertThat(batchResponse.getResults().get(2).getItem(TestDoc.class)).isEqualTo(this.TestDocPk1ExistingC);
 
-        List> batchOperations = batch.getOperationsInternal();
+        List batchOperations = batch.getOperations();
         for (int index = 0; index < batchOperations.size(); index++) {
             assertThat(batchResponse.getResults().get(index).getOperation()).isEqualTo(batchOperations.get(index));
         }
@@ -431,7 +430,7 @@ public void batchCrud() {
         assertThat(batchResponse.getResults().get(4).getStatusCode()).isEqualTo(HttpResponseStatus.OK.code());
         assertThat(batchResponse.getResults().get(5).getStatusCode()).isEqualTo(HttpResponseStatus.NO_CONTENT.code());
 
-        List> batchOperations = batch.getOperationsInternal();
+        List batchOperations = batch.getOperations();
         for (int index = 0; index < batchOperations.size(); index++) {
             assertThat(batchResponse.getResults().get(index).getOperation()).isEqualTo(batchOperations.get(index));
         }
@@ -530,7 +529,7 @@ private void runWithError(
         assertThat(batchResponse.getResults().get(1).getStatusCode()).isEqualTo(expectedFailedOperationStatusCode.code());
         assertThat(batchResponse.getResults().get(2).getStatusCode()).isEqualTo(HttpResponseStatus.FAILED_DEPENDENCY.code());
 
-        List> batchOperations = batch.getOperationsInternal();
+        List batchOperations = batch.getOperations();
         for (int index = 0; index < batchOperations.size(); index++) {
             assertThat(batchResponse.getResults().get(index).getOperation()).isEqualTo(batchOperations.get(index));
         }

From b70749589b0d1c23cf9b832fd75544fb0e9eafdf Mon Sep 17 00:00:00 2001
From: Rakesh Kumar 
Date: Sat, 10 Oct 2020 01:15:56 +0530
Subject: [PATCH 22/22] Reverting byte buffer change

Signed-off-by: Rakesh Kumar 
---
 .../cosmos/implementation/RxDocumentClientImpl.java |  2 +-
 .../implementation/batch/ServerBatchRequest.java    | 13 +++++++------
 2 files changed, 8 insertions(+), 7 deletions(-)

diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/RxDocumentClientImpl.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/RxDocumentClientImpl.java
index 33e6f3f8334a..616ddec441d7 100644
--- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/RxDocumentClientImpl.java
+++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/RxDocumentClientImpl.java
@@ -1271,7 +1271,7 @@ private Mono getBatchDocumentRequest(DocumentClientRet
         checkNotNull(serverBatchRequest, "expected non null serverBatchRequest");
 
         Instant serializationStartTimeUTC = Instant.now();
-        ByteBuffer content = serverBatchRequest.getRequestBodyAsByteBuffer();
+        ByteBuffer content = ByteBuffer.wrap(Utils.getUTF8Bytes(serverBatchRequest.getRequestBody()));
         Instant serializationEndTimeUTC = Instant.now();
 
         SerializationDiagnosticsContext.SerializationDiagnostics serializationDiagnostics = new SerializationDiagnosticsContext.SerializationDiagnostics(
diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/batch/ServerBatchRequest.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/batch/ServerBatchRequest.java
index 2c777f71a027..c55cdef8f763 100644
--- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/batch/ServerBatchRequest.java
+++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/batch/ServerBatchRequest.java
@@ -9,7 +9,6 @@
 import com.azure.cosmos.implementation.apachecommons.collections.list.UnmodifiableList;
 import com.fasterxml.jackson.databind.node.ArrayNode;
 
-import java.nio.ByteBuffer;
 import java.util.List;
 
 import static com.azure.cosmos.implementation.guava25.base.Preconditions.checkNotNull;
@@ -23,7 +22,7 @@ public abstract class ServerBatchRequest {
     private final int maxBodyLength;
     private final int maxOperationCount;
 
-    private ByteBuffer requestBody;
+    private String requestBody;
     private List operations;
     private boolean isAtomicBatch = false;
     private boolean shouldContinueOnError = false;
@@ -50,7 +49,7 @@ public abstract class ServerBatchRequest {
      *
      * @return Any pending operations that were not included in the request.
      */
-    List createBodyOfBatchRequest(final List operations) {
+    final List createBodyOfBatchRequest(final List operations) {
 
         checkNotNull(operations, "expected non-null operations");
 
@@ -83,13 +82,15 @@ List createBodyOfBatchRequest(final List