CreateItem will only retry for auto-extracted partition key in-case o…

…f container re-creation (#1070)

* Fixed bug where if a user provided the wrong partition key value the request would be sent twice. The retry logic for the partition key extraction was always being used if it was a user provided partition key value or an extracted partition key value. The retry logic should only be used for SDK extracted partition key value in scenarios where the cache is stale.

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