CreateItem will only retry for auto-extracted partition key in-case of container re-creation

Microsoft.Azure.Cosmos/src/Resource/Container/ContainerCore.Items.cs

@@ -6,7 +6,6 @@ namespace Microsoft.Azure.Cosmos
 {
  using System;
  using System.Collections.Generic;
- using System.Diagnostics.Contracts;
  using System.Globalization;
  using System.IO;
  using System.Linq;
@@ -48,7 +47,6 @@ public override Task<ResponseMessage> CreateItemStreamAsync(
  streamPayload,
  OperationType.Create,
  requestOptions,
- extractPartitionKeyIfNeeded: false,
  cancellationToken: cancellationToken);
  }
 
@@ -66,7 +64,7 @@ public override Task<ItemResponse<T>> CreateItemAsync<T>(
  Task<ResponseMessage> response = this.ExtractPartitionKeyAndProcessItemStreamAsync(
  partitionKey: partitionKey,
  itemId: null,
- streamPayload: this.ClientContext.CosmosSerializer.ToStream<T>(item),
+ item: item,
  operationType: OperationType.Create,
  requestOptions: requestOptions,
  cancellationToken: cancellationToken);
@@ -86,7 +84,6 @@ public override Task<ResponseMessage> ReadItemStreamAsync(
  null,
  OperationType.Read,
  requestOptions,
- extractPartitionKeyIfNeeded: false,
  cancellationToken: cancellationToken);
  }
 
@@ -117,7 +114,6 @@ public override Task<ResponseMessage> UpsertItemStreamAsync(
  streamPayload,
  OperationType.Upsert,
  requestOptions,
- extractPartitionKeyIfNeeded: false,
  cancellationToken: cancellationToken);
  }
 
@@ -135,7 +131,7 @@ public override Task<ItemResponse<T>> UpsertItemAsync<T>(
  Task<ResponseMessage> response = this.ExtractPartitionKeyAndProcessItemStreamAsync(
  partitionKey: partitionKey,
  itemId: null,
- streamPayload: this.ClientContext.CosmosSerializer.ToStream<T>(item),
+ item: item,
  operationType: OperationType.Upsert,
  requestOptions: requestOptions,
  cancellationToken: cancellationToken);
@@ -156,7 +152,6 @@ public override Task<ResponseMessage> ReplaceItemStreamAsync(
  streamPayload,
  OperationType.Replace,
  requestOptions,
- extractPartitionKeyIfNeeded: false,
  cancellationToken: cancellationToken);
  }
 
@@ -180,7 +175,7 @@ public override Task<ItemResponse<T>> ReplaceItemAsync<T>(
  Task<ResponseMessage> response = this.ExtractPartitionKeyAndProcessItemStreamAsync(
  partitionKey: partitionKey,
  itemId: id,
- streamPayload: this.ClientContext.CosmosSerializer.ToStream<T>(item),
+ item: item,
  operationType: OperationType.Replace,
  requestOptions: requestOptions,
  cancellationToken: cancellationToken);
@@ -200,7 +195,6 @@ public override Task<ResponseMessage> DeleteItemStreamAsync(
  null,
  OperationType.Delete,
  requestOptions,
- extractPartitionKeyIfNeeded: false,
  cancellationToken: cancellationToken);
  }
 
@@ -511,24 +505,39 @@ internal FeedIterator GetItemQueryStreamIteratorInternal(
  // Extracted partition key might be invalid as CollectionCache might be stale.
  // Stale collection cache is refreshed through PartitionKeyMismatchRetryPolicy
  // and partition-key is extracted again. 
- internal async Task<ResponseMessage> ExtractPartitionKeyAndProcessItemStreamAsync(
+ internal async Task<ResponseMessage> ExtractPartitionKeyAndProcessItemStreamAsync<T>(
  PartitionKey? partitionKey,
  string itemId,
- Stream streamPayload,
+ T item,
  OperationType operationType,
  RequestOptions requestOptions,
  CancellationToken cancellationToken)
  {
+ Stream streamPayload = this.ClientContext.CosmosSerializer.ToStream<T>(item);
+
+ // User specified PK value, no need to extract it
+ if (partitionKey.HasValue)
+ {
+ return await this.ProcessItemStreamAsync(
+ partitionKey,
+ itemId,
+ streamPayload,
+ operationType,
+ requestOptions,
+ cancellationToken: cancellationToken);
+ }
+
  PartitionKeyMismatchRetryPolicy requestRetryPolicy = null;
  while (true)
  {
+ partitionKey = await this.GetPartitionKeyValueFromStreamAsync(streamPayload, cancellationToken);
+
  ResponseMessage responseMessage = await this.ProcessItemStreamAsync(
  partitionKey,
  itemId,
  streamPayload,
  operationType,
  requestOptions,
- extractPartitionKeyIfNeeded: true,
  cancellationToken: cancellationToken);
 
  if (responseMessage.IsSuccessStatusCode)
@@ -555,33 +564,27 @@ internal async Task<ResponseMessage> ProcessItemStreamAsync(
  Stream streamPayload,
  OperationType operationType,
  RequestOptions requestOptions,
- bool extractPartitionKeyIfNeeded,
  CancellationToken cancellationToken = default(CancellationToken))
  {
  if (requestOptions != null && requestOptions.IsEffectivePartitionKeyRouting)
  {
  partitionKey = null;
  }
 
- if (extractPartitionKeyIfNeeded && partitionKey == null)
- {
- partitionKey = await this.GetPartitionKeyValueFromStreamAsync(streamPayload, cancellationToken);
- }
-
  ContainerCore.ValidatePartitionKey(partitionKey, requestOptions);
  Uri resourceUri = this.GetResourceUri(requestOptions, operationType, itemId);
 
  return await this.ClientContext.ProcessResourceOperationStreamAsync(
- resourceUri,
- ResourceType.Document,
- operationType,
- requestOptions,
- this,
- partitionKey,
- itemId,
- streamPayload,
- null,
- cancellationToken);
+ resourceUri: resourceUri,
+ resourceType: ResourceType.Document,
+ operationType: operationType,
+ requestOptions: requestOptions,
+ cosmosContainerCore: this,
+ partitionKey: partitionKey,
+ itemId: itemId,
+ streamPayload: streamPayload,
+ requestEnricher: null,
+ cancellationToken: cancellationToken);
  }
 
  internal async Task<PartitionKey> GetPartitionKeyValueFromStreamAsync(

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

@@ -38,7 +38,7 @@ public async Task Cleanup()
  [TestMethod]
  public async Task TestCustomPropertyWithHandler()
  { 
- CustomHandler testHandler = new CustomHandler();
+ RequestHandlerHelper testHandler = new RequestHandlerHelper();
 
  // Add the random guid to the property
  Guid randomGuid = Guid.NewGuid();
@@ -111,20 +111,5 @@ public class ToDoActivity
  public string description { get; set; }
  public string status { get; set; }
  }
-
- public class CustomHandler : RequestHandler
- {
- public Action<RequestMessage> UpdateRequestMessage = null;
-
- public override Task<ResponseMessage> SendAsync(RequestMessage request, CancellationToken cancellationToken)
- {
- if (this.UpdateRequestMessage != null)
- {
- this.UpdateRequestMessage(request);
- }
-
- return base.SendAsync(request, cancellationToken);
- }
- }
  }
 }

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

@@ -1363,6 +1363,52 @@ internal async Task TestPartitionKeyDefinitionOnContainerRecreateFromDifferentPa
  }
  }
 
+ [TestMethod]
+ public async Task TestInvalidPartitionKeyException()
+ {
+ RequestHandlerHelper testHandler = new RequestHandlerHelper();
+ int createItemCount = 0;
+ string partitionKey = null;
+ testHandler.UpdateRequestMessage = (requestMessage) =>
+ {
+ if (requestMessage.ResourceType == ResourceType.Document)
+ {
+ createItemCount++;
+ string pk = requestMessage.Headers.PartitionKey;
+ Assert.AreNotEqual(partitionKey, pk, $"Same PK value should not be sent again. PK value:{pk}");
+ partitionKey = pk;
+ }
+ };
+
+ using (CosmosClient client = TestCommon.CreateCosmosClient((builder) => builder.AddCustomHandlers(testHandler)))
+ {
+ Cosmos.Database database = null;
+ try
+ {
+ database = await client.CreateDatabaseAsync("TestInvalidPartitionKey" + Guid.NewGuid().ToString());
+ Container container = await database.CreateContainerAsync(id: "coll1", partitionKeyPath: "/doesnotexist");
+ ToDoActivity toDoActivity = ToDoActivity.CreateRandomToDoActivity();
+
+ ItemResponse<ToDoActivity> response = await container.CreateItemAsync(toDoActivity, partitionKey: new Cosmos.PartitionKey(toDoActivity.status));
+ Assert.Fail("Create item should fail with wrong partition key value");
+ }
+ catch (CosmosException ce)
+ {
+ Assert.AreEqual(HttpStatusCode.BadRequest, ce.StatusCode);
+ Assert.AreEqual(SubStatusCodes.PartitionKeyMismatch, (SubStatusCodes)ce.SubStatusCode);
+ }
+ finally
+ {
+ if (database != null)
+ {
+ await database.DeleteAsync();
+ }
+ }
+
+ Assert.AreEqual(1, createItemCount, $"Request should use the custom handler, and it should only be used once. Count {createItemCount}");
+ }
+ }
+
  /// <summary>
  /// Tests that collection cache is refreshed when collection is recreated.
  /// The test just ensures that client retries and completes successfully for query - request which doesn't target single partition key.

Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.EmulatorTests/Utils/RequestHandlerHelper.cs

@@ -0,0 +1,22 @@
+//------------------------------------------------------------
+// Copyright (c) Microsoft Corporation. All rights reserved.
+//------------------------------------------------------------
+
+namespace Microsoft.Azure.Cosmos.SDK.EmulatorTests
+{
+ using System;
+ using System.Threading;
+ using System.Threading.Tasks;
+
+ public class RequestHandlerHelper : RequestHandler
+ {
+ public Action<RequestMessage> UpdateRequestMessage = null;
+
+ public override Task<ResponseMessage> SendAsync(RequestMessage request, CancellationToken cancellationToken)
+ {
+ this.UpdateRequestMessage?.Invoke(request);
+
+ return base.SendAsync(request, cancellationToken);
+ }
+ }
+}

changelog.md

@@ -30,7 +30,8 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 - [#1036](https://github.com/Azure/azure-cosmos-dotnet-v3/pull/1036) Fixed query responses to return null Content if it is a failure
 - [#1045](https://github.com/Azure/azure-cosmos-dotnet-v3/pull/1045) Added stack trace and innner exception to CosmosException
 - [#1050](https://github.com/Azure/azure-cosmos-dotnet-v3/pull/1050) Add mocking constructors to TransactionalBatchOperationResult
-
+- [#1070](https://github.com/Azure/azure-cosmos-dotnet-v3/pull/1070) CreateItem will only retry for auto-extracted partition key in-case of collection re-creation
+
 ## <a name="3.4.1"/> [3.4.1](https://www.nuget.org/packages/Microsoft.Azure.Cosmos/3.4.1) - 2019-11-06
 
 ### Fixed