Transactional batch for Java

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> T transform(Object object, Class<T> 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.
+     * <p>
+     * If the transactional batch executes successfully, the value returned by {@link
+     * TransactionalBatchResponse#getResponseStatus} on the response returned will be set to 200}.
+     * <p>
+     * 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.
+     * <p>
+     * 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.
+     * <p>
+     * Use {@link TransactionalBatchResponse#isSuccessStatusCode} on the response returned to ensure that the
+     * transactional batch succeeded.
+     */
+    @Beta(Beta.SinceVersion.V4_4_0)
+    public Mono<TransactionalBatchResponse> 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.
+     * <p>
+     * If the transactional batch executes successfully, the value returned by {@link
+     * TransactionalBatchResponse#getResponseStatus} on the response returned will be set to 200}.
+     * <p>
+     * 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.
+     * <p>
+     * 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.
+     * <p>
+     * Use {@link TransactionalBatchResponse#isSuccessStatusCode} on the response returned to ensure that the
+     * transactional batch succeeded.
+     */
+    @Beta(Beta.SinceVersion.V4_4_0)
+    public Mono<TransactionalBatchResponse> executeTransactionalBatch(
+        TransactionalBatch transactionalBatch,
+        TransactionalBatchRequestOptions requestOptions) {
+
+        if (requestOptions == null) {
+            requestOptions = new TransactionalBatchRequestOptions();
+        }
+
+        final BatchExecutor executor = new BatchExecutor(this, transactionalBatch, requestOptions);
+        final Mono<TransactionalBatchResponse> responseMono = executor.executeAsync();
+
+        return withContext(context -> database.getClient().getTracerProvider().
+            traceEnabledBatchResponsePublisher(responseMono,
+                context,
+                this.batchSpanName,
+                database.getId(),
+                database.getClient().getServiceEndpoint()));
+    }
+
     /**
      * Reads an item.
      * <p>

sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosContainer.java

@@ -270,6 +270,19 @@ private CosmosItemResponse<Object> blockDeleteItemResponse(Mono<CosmosItemRespon
         }
     }
 
+    private TransactionalBatchResponse blockBatchResponse(Mono<TransactionalBatchResponse> 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 <T> CosmosItemResponse<Object> 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.
+     * <p>
+     * If the transactional batch executes successfully, the value returned by {@link
+     * TransactionalBatchResponse#getResponseStatus} on the response returned will be set to 200}.
+     * <p>
+     * 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.
+     * <p>
+     * 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.
+     * <p>
+     * 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.
+     * <p>
+     * If the transactional batch executes successfully, the value returned by {@link
+     * TransactionalBatchResponse#getResponseStatus} on the response returned will be set to 200}.
+     * <p>
+     * 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.
+     * <p>
+     * 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.
+     * <p>
+     * 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.
      *