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