Adding TransactionalBatchRequestOption and test cases

Signed-off-by: Rakesh Kumar <rakkuma@microsoft.com>

Microsoft.Azure.Cosmos/src/Batch/BatchCore.cs

@@ -214,7 +214,7 @@ public override Task<TransactionalBatchResponse> ExecuteAsync(
  /// <param name="cancellationToken">(Optional) <see cref="CancellationToken"/> representing request cancellation.</param>
  /// <returns>An awaitable <see cref="TransactionalBatchResponse"/> which contains the completion status and results of each operation.</returns>
  public override Task<TransactionalBatchResponse> ExecuteAsync(
- ItemRequestOptions requestOptions,
+ TransactionalBatchRequestOption requestOptions,
  CancellationToken cancellationToken = default(CancellationToken))
  {
  return this.container.ClientContext.OperationHelperAsync(

Microsoft.Azure.Cosmos/src/Batch/PartitionKeyRangeBatchResponse.cs

@@ -50,7 +50,7 @@ internal PartitionKeyRangeBatchResponse(
  }
 
  itemBatchOperations.AddRange(serverResponse.Operations);
- this.RequestCharge += serverResponse.RequestCharge;
+ this.Headers = serverResponse.Headers;
 
  if (!string.IsNullOrEmpty(serverResponse.ErrorMessage))
  {

Microsoft.Azure.Cosmos/src/Batch/TransactionalBatch.cs

@@ -227,7 +227,7 @@ public abstract Task<TransactionalBatchResponse> ExecuteAsync(
  /// <summary>
  /// Executes the transactional batch at the Azure Cosmos service as an asynchronous operation.
  /// </summary>
- /// <param name="requestOptions">Options that apply to the batch.</param>
+ /// <param name="requestOptions">Options that apply specifically to batch request.</param>
  /// <param name="cancellationToken">(Optional) Cancellation token representing request cancellation.</param>
  /// <returns>An awaitable response which contains details of execution of the transactional batch.
  /// <para>
@@ -253,7 +253,7 @@ public abstract Task<TransactionalBatchResponse> ExecuteAsync(
  /// Use <see cref="TransactionalBatchResponse.IsSuccessStatusCode"/> on the response returned to ensure that the transactional batch succeeded.
  /// </remarks>
  public abstract Task<TransactionalBatchResponse> ExecuteAsync(
- ItemRequestOptions requestOptions,
+ TransactionalBatchRequestOption requestOptions,
  CancellationToken cancellationToken = default(CancellationToken));
  }
 }

Microsoft.Azure.Cosmos/src/Batch/TransactionalBatchRequestOption.cs

@@ -0,0 +1,59 @@
+//------------------------------------------------------------
+// Copyright (c) Microsoft Corporation. All rights reserved.
+//------------------------------------------------------------
+
+namespace Microsoft.Azure.Cosmos
+{
+ using Microsoft.Azure.Documents;
+
+ /// <summary>
+ /// Cosmos batch request options
+ /// </summary>
+ public class TransactionalBatchRequestOption : RequestOptions
+ {
+ /// <summary>
+ /// Gets or sets the consistency level required for the request in the Azure Cosmos DB service.
+ /// </summary>
+ /// <value>
+ /// The consistency level required for the request.
+ /// </value>
+ /// <remarks>
+ /// Azure Cosmos DB offers 5 different consistency levels. Strong, Bounded Staleness, Session, Consistent Prefix and Eventual - in order of strongest to weakest consistency. <see cref="ConnectionPolicy"/>
+ /// <para>
+ /// While this is set at a database account level, Azure Cosmos DB allows a developer to override the default consistency level
+ /// for each individual request.
+ /// </para>
+ /// </remarks>
+ public ConsistencyLevel? ConsistencyLevel
+ {
+ get => this.BaseConsistencyLevel;
+ set => this.BaseConsistencyLevel = value;
+ }
+
+ /// <summary>
+ /// Gets or sets the boolean to only return the headers and status code in
+ /// the Cosmos DB response for operations in the transactional batch request.
+ /// This removes the resource from the response. This reduces networking and CPU load by not sending
+ /// the resource back over the network and serializing it on the client.
+ /// </summary>
+ /// <remarks>
+ /// This is optimal for workloads where the returned resource is not used.
+ /// </remarks>
+ public bool? EnableContentResponseOnOperations { get; set; }
+
+ /// <summary>
+ /// Fill the CosmosRequestMessage headers with the set properties
+ /// </summary>
+ /// <param name="request">The <see cref="RequestMessage"/></param>
+ internal override void PopulateRequestOptions(RequestMessage request)
+ {
+ if (this.EnableContentResponseOnOperations.HasValue &&
+ !this.EnableContentResponseOnOperations.Value)
+ {
+ request.Headers.Add(HttpConstants.HttpHeaders.Prefer, HttpConstants.HttpHeaderValues.PreferReturnMinimal);
+ }
+
+ base.PopulateRequestOptions(request);
+ }
+ }
+}

Microsoft.Azure.Cosmos/src/Batch/TransactionalBatchResponse.cs

@@ -44,9 +44,7 @@ internal TransactionalBatchResponse(
  : this(statusCode,
  subStatusCode,
  errorMessage,
- requestCharge: 0,
- retryAfter: null,
- activityId: Guid.Empty.ToString(),
+ new Headers(),
  diagnosticsContext: diagnosticsContext,
  operations: operations,
  serializer: null)
@@ -65,9 +63,7 @@ private TransactionalBatchResponse(
  HttpStatusCode statusCode,
  SubStatusCodes subStatusCode,
  string errorMessage,
- double requestCharge,
- TimeSpan? retryAfter,
- string activityId,
+ Headers headers,
  CosmosDiagnosticsContext diagnosticsContext,
  IReadOnlyList<ItemBatchOperation> operations,
  CosmosSerializerCore serializer)
@@ -77,30 +73,33 @@ private TransactionalBatchResponse(
  this.ErrorMessage = errorMessage;
  this.Operations = operations;
  this.SerializerCore = serializer;
- this.RequestCharge = requestCharge;
- this.RetryAfter = retryAfter;
- this.ActivityId = activityId;
+ this.Headers = headers;
  this.Diagnostics = diagnosticsContext.Diagnostics;
  this.DiagnosticsContext = diagnosticsContext ?? throw new ArgumentNullException(nameof(diagnosticsContext));
  }
 
+ /// <summary>
+ /// Gets the current <see cref="ResponseMessage"/> HTTP headers.
+ /// </summary>
+ public virtual Headers Headers { get; set; }
+
  /// <summary>
  /// Gets the ActivityId that identifies the server request made to execute the batch.
  /// </summary>
- public virtual string ActivityId { get; }
+ public virtual string ActivityId => this.Headers?.ActivityId;
 
  /// <summary>
  /// Gets the request charge for the batch request.
  /// </summary>
  /// <value>
  /// The request charge measured in request units.
  /// </value>
- public virtual double RequestCharge { get; internal set; }
+ public virtual double RequestCharge => this.Headers?.RequestCharge ?? 0;
 
  /// <summary>
  /// Gets the amount of time to wait before retrying this or any other request within Cosmos container or collection due to throttling.
  /// </summary>
- public virtual TimeSpan? RetryAfter { get; }
+ public virtual TimeSpan? RetryAfter => this.Headers?.RetryAfter;
 
  /// <summary>
  /// Gets the completion status code of the batch request.
@@ -252,9 +251,7 @@ internal static async Task<TransactionalBatchResponse> FromResponseMessageAsync(
  HttpStatusCode.InternalServerError,
  SubStatusCodes.Unknown,
  ClientResources.ServerResponseDeserializationFailure,
- responseMessage.Headers.RequestCharge,
- responseMessage.Headers.RetryAfter,
- responseMessage.Headers.ActivityId,
+ responseMessage.Headers,
  responseMessage.DiagnosticsContext,
  serverRequest.Operations,
  serializer);
@@ -268,9 +265,7 @@ internal static async Task<TransactionalBatchResponse> FromResponseMessageAsync(
  responseMessage.StatusCode,
  responseMessage.Headers.SubStatusCode,
  responseMessage.ErrorMessage,
- responseMessage.Headers.RequestCharge,
- responseMessage.Headers.RetryAfter,
- responseMessage.Headers.ActivityId,
+ responseMessage.Headers,
  responseMessage.DiagnosticsContext,
  serverRequest.Operations,
  serializer);
@@ -286,9 +281,7 @@ internal static async Task<TransactionalBatchResponse> FromResponseMessageAsync(
  HttpStatusCode.InternalServerError,
  SubStatusCodes.Unknown,
  ClientResources.InvalidServerResponse,
- responseMessage.Headers.RequestCharge,
- responseMessage.Headers.RetryAfter,
- responseMessage.Headers.ActivityId,
+ responseMessage.Headers,
  responseMessage.DiagnosticsContext,
  serverRequest.Operations,
  serializer);
@@ -382,9 +375,7 @@ record =>
  responseStatusCode,
  responseSubStatusCode,
  responseMessage.ErrorMessage,
- responseMessage.Headers.RequestCharge,
- responseMessage.Headers.RetryAfter,
- responseMessage.Headers.ActivityId,
+ responseMessage.Headers,
  responseMessage.DiagnosticsContext,
  serverRequest.Operations,
  serializer);

Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.EmulatorTests/Batch/BatchSinglePartitionKeyTests.cs

@@ -171,6 +171,38 @@ public async Task BatchItemETagAsync()
  }
  }
 
+ [TestMethod]
+ [Owner("rakkuma")]
+ [Description("Verify session token received from batch operations")]
+ public async Task BatchItemSessionTokenAsync()
+ {
+ Container container = BatchTestBase.JsonContainer;
+ await this.CreateJsonTestDocsAsync(container);
+
+ TestDoc testDocToCreate = BatchTestBase.PopulateTestDoc(this.PartitionKey1);
+
+ TestDoc testDocToReplace = this.GetTestDocCopy(this.TestDocPk1ExistingA);
+ testDocToReplace.Cost++;
+
+ ItemResponse<TestDoc> readResponse = await BatchTestBase.JsonContainer.ReadItemAsync<TestDoc>(
+ this.TestDocPk1ExistingA.Id,
+ BatchTestBase.GetPartitionKey(this.PartitionKey1));
+
+ ISessionToken beforeRequestSessionToken = BatchTestBase.GetSessionToken(readResponse.Headers.Session);
+
+ TransactionalBatchResponse batchResponse = await new BatchCore((ContainerInlineCore)container, BatchTestBase.GetPartitionKey(this.PartitionKey1))
+ .CreateItem(testDocToCreate)
+ .ReplaceItem(testDocToReplace.Id, testDocToReplace)
+ .ExecuteAsync();
+
+ BatchSinglePartitionKeyTests.VerifyBatchProcessed(batchResponse, numberOfOperations: 2);
+ Assert.AreEqual(HttpStatusCode.Created, batchResponse[0].StatusCode);
+ Assert.AreEqual(HttpStatusCode.OK, batchResponse[1].StatusCode);
+
+ ISessionToken afterRequestSessionToken = BatchTestBase.GetSessionToken(batchResponse.Headers.Session);
+ Assert.IsTrue(afterRequestSessionToken.LSN > beforeRequestSessionToken.LSN, "Response session token should be more than request session token");
+ }
+
  [TestMethod]
  [Owner("abpai")]
  [Description("Verify TTL passed to binary passthrough batch operations flow as expected")]
@@ -320,7 +352,7 @@ public async Task BatchReadsOnlyAsync()
 
  private async Task<TransactionalBatchResponse> RunCrudAsync(bool isStream, bool isSchematized, bool useEpk, Container container)
  {
- ItemRequestOptions batchOptions = null;
+ TransactionalBatchRequestOption batchOptions = null;
  if (isSchematized)
  {
  await this.CreateSchematizedTestDocsAsync(container);

Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.EmulatorTests/Batch/BatchTestBase.cs

@@ -292,8 +292,8 @@ protected static async Task VerifyNotFoundAsync(Container container, TestDoc doc
  Assert.AreEqual(HttpStatusCode.NotFound, response.StatusCode);
  }
 
- protected static ItemRequestOptions GetUpdatedBatchRequestOptions(
- ItemRequestOptions batchOptions = null,
+ protected static TransactionalBatchRequestOption GetUpdatedBatchRequestOptions(
+ TransactionalBatchRequestOption batchOptions = null,
  bool isSchematized = false,
  bool useEpk = false,
  object partitionKey = null)
@@ -302,7 +302,7 @@ protected static ItemRequestOptions GetUpdatedBatchRequestOptions(
  {
  if (batchOptions == null)
  {
- batchOptions = new ItemRequestOptions();
+ batchOptions = new TransactionalBatchRequestOption();
  }
 
  Dictionary<string, object> properties = new Dictionary<string, object>()
@@ -412,6 +412,12 @@ protected static byte[] HexStringToBytes(string input)
  return bytes;
  }
 
+ internal static ISessionToken GetSessionToken(string sessionToken)
+ {
+ string[] tokenParts = sessionToken.Split(':');
+ return SessionTokenHelper.Parse(tokenParts[1]);
+ }
+
  private static bool PopulateRequestOptions(RequestOptions requestOptions, TestDoc doc, bool isSchematized, bool useEpk, int? ttlInSeconds)
  {
  Dictionary<string, object> properties = new Dictionary<string, object>();

Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.EmulatorTests/CosmosContentResponseTests.cs

@@ -244,6 +244,52 @@ public async Task ItemBatchNoResponseTest()
  this.ValidateResponse(response, noResponseItemCount);
  }
 
+ [TestMethod]
+ public async Task TransactionalBatchNoResponseTest()
+ {
+ TransactionalBatchRequestOption requestOptions = new TransactionalBatchRequestOption()
+ {
+ EnableContentResponseOnOperations = false
+ };
+
+ string pkId = "TestTransactionalBatchId";
+ TransactionalBatch batch = this.container.CreateTransactionalBatch(new PartitionKey(pkId));
+
+ int noResponseItemCount = 100;
+ for (int i = 0; i < noResponseItemCount; i++)
+ {
+ ToDoActivity item = ToDoActivity.CreateRandomToDoActivity(pk: pkId);
+ batch.CreateItem<ToDoActivity>(item);
+ }
+
+ TransactionalBatchResponse response = await batch.ExecuteAsync(requestOptions);
+ Assert.AreEqual(100, response.Count);
+ this.ValidateResponse(response, noResponseItemCount);
+
+ pkId = "TestTransactionalBatchId2";
+ batch = this.container.CreateTransactionalBatch(new PartitionKey(pkId));
+
+ noResponseItemCount = 0;
+ for (int i = 0; i < 10; i++)
+ {
+ ToDoActivity item = ToDoActivity.CreateRandomToDoActivity(pk: pkId);
+ batch.CreateItem<ToDoActivity>(item);
+ noResponseItemCount++;
+ ToDoActivity item2 = ToDoActivity.CreateRandomToDoActivity(pk: pkId);
+ item2.id = item.id;
+ batch.ReplaceItem<ToDoActivity>(item2.id, item2);
+ noResponseItemCount++;
+
+ // Even Read won't return response
+ batch.ReadItem(item2.id);
+ noResponseItemCount++;
+ }
+
+ response = await batch.ExecuteAsync(requestOptions);
+ Assert.AreEqual(noResponseItemCount, response.Count);
+ this.ValidateResponse(response, noResponseItemCount);
+ }
+
  [TestMethod]
  public async Task ItemBulkNoResponseTest()
  {

Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.EmulatorTests/CosmosDiagnosticsTests.cs

@@ -27,7 +27,7 @@ public class CosmosDiagnosticsTests : BaseCosmosClientHelper
  private Container Container = null;
  private ContainerProperties containerSettings = null;
 
- private static readonly ItemRequestOptions RequestOptionDisableDiagnostic = new ItemRequestOptions()
+ private static readonly TransactionalBatchRequestOption RequestOptionDisableDiagnostic = new TransactionalBatchRequestOption()
  {
  DiagnosticContextFactory = () => EmptyCosmosDiagnosticsContext.Singleton
  };
@@ -398,7 +398,7 @@ public async Task BatchOperationDiagnostic(bool disableDiagnostics)
  batch.ReadItem(createItems[i].id);
  }
 
- ItemRequestOptions requestOptions = disableDiagnostics ? RequestOptionDisableDiagnostic : null;
+ TransactionalBatchRequestOption requestOptions = disableDiagnostics ? RequestOptionDisableDiagnostic : null;
  TransactionalBatchResponse response = await batch.ExecuteAsync(requestOptions);
 
  Assert.IsNotNull(response);

Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Batch/BatchUnitTests.cs

@@ -31,19 +31,19 @@ public class BatchUnitTests
  public async Task BatchInvalidOptionsAsync()
  {
  Container container = BatchUnitTests.GetContainer();
- List<ItemRequestOptions> badBatchOptionsList = new List<ItemRequestOptions>()
+ List<TransactionalBatchRequestOption> badBatchOptionsList = new List<TransactionalBatchRequestOption>()
  {
- new ItemRequestOptions()
+ new TransactionalBatchRequestOption()
  {
  IfMatchEtag = "cond",
  },
- new ItemRequestOptions()
+ new TransactionalBatchRequestOption()
  {
  IfNoneMatchEtag = "cond2",
  }
  };
 
- foreach (ItemRequestOptions batchOptions in badBatchOptionsList)
+ foreach (TransactionalBatchRequestOption batchOptions in badBatchOptionsList)
  {
  BatchCore batch = (BatchCore)
  new BatchCore((ContainerInternal)container, new Cosmos.PartitionKey(BatchUnitTests.PartitionKey1))
@@ -476,8 +476,8 @@ private static void VerifyBatchItemRequestOptionsAreEqual(TransactionalBatchItem
  private static async Task VerifyExceptionThrownOnExecuteAsync(
  TransactionalBatch batch,
  Type expectedTypeOfException,
- string expectedExceptionMessage = null, 
- ItemRequestOptions requestOptions = null)
+ string expectedExceptionMessage = null,
+ TransactionalBatchRequestOption requestOptions = null)
  {
  bool wasExceptionThrown = false;
  try