PartitionKeyDelete: Adds API to delete all items in a partition key (#…

…1490)

Added DeleteItemsInPartitionKeyAsync API
Changes in GetStoreProxy in DocumentClient.cs, GetHttpMethod in RequestInvokerHandler.cs, PrepareRequestMessageAsync in GatewayStoreClient.cs for GatewayMode of PkDeletePurger to work.

Microsoft.Azure.Cosmos/src/DocumentClient.cs

@@ -6314,7 +6314,8 @@ internal IStoreModel GetStoreProxy(DocumentServiceRequest request)
                 (resourceType.IsScript() && operationType != OperationType.ExecuteJavaScript) ||
                 resourceType == ResourceType.PartitionKeyRange ||
                 resourceType == ResourceType.Snapshot ||
-                resourceType == ResourceType.ClientEncryptionKey)
+                resourceType == ResourceType.ClientEncryptionKey ||
+                (resourceType == ResourceType.PartitionKey && operationType == OperationType.Delete))
             {
                 return this.GatewayStoreModel;
             }

Microsoft.Azure.Cosmos/src/GatewayStoreClient.cs

@@ -249,7 +249,8 @@ private async ValueTask<HttpRequestMessage> PrepareRequestMessageAsync(
                 request.OperationType == OperationType.SqlQuery ||
                 request.OperationType == OperationType.Batch ||
                 request.OperationType == OperationType.ExecuteJavaScript ||
-                request.OperationType == OperationType.QueryPlan)
+                request.OperationType == OperationType.QueryPlan ||
+                (request.ResourceType == ResourceType.PartitionKey && request.OperationType == OperationType.Delete))
             {
                 httpMethod = HttpMethod.Post;
             }

Microsoft.Azure.Cosmos/src/Handler/RequestInvokerHandler.cs

@@ -144,7 +144,7 @@ public virtual async Task<ResponseMessage> SendAsync(
 
             try
             {
-                HttpMethod method = RequestInvokerHandler.GetHttpMethod(operationType);
+                HttpMethod method = RequestInvokerHandler.GetHttpMethod(resourceType, operationType);
                 RequestMessage request = new RequestMessage(
                     method,
                     resourceUriString,
@@ -284,6 +284,7 @@ public virtual async Task<ResponseMessage> SendAsync(
         }
 
         internal static HttpMethod GetHttpMethod(
+            ResourceType resourceType,
             OperationType operationType)
         {
             if (operationType == OperationType.Create ||
@@ -293,7 +294,8 @@ internal static HttpMethod GetHttpMethod(
                 operationType == OperationType.QueryPlan ||
                 operationType == OperationType.Batch ||
                 operationType == OperationType.ExecuteJavaScript ||
-                operationType == OperationType.CompleteUserTransaction)
+                operationType == OperationType.CompleteUserTransaction ||
+                (resourceType == ResourceType.PartitionKey && operationType == OperationType.Delete))
             {
                 return HttpMethod.Post;
             }

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

@@ -1173,6 +1173,24 @@ public abstract ChangeFeedEstimator GetChangeFeedEstimator(
         /// <returns>A new instance of <see cref="TransactionalBatch"/>.</returns>
         public abstract TransactionalBatch CreateTransactionalBatch(PartitionKey partitionKey);
 
+#if INTERNAL
+        /// <summary>
+        /// Deletes all items in the Container with the specified <see cref="PartitionKey"/> value.
+        /// Starts an asynchronous Cosmos DB background operation which deletes all items in the Container with the specified value. 
+        /// The asynchronous Cosmos DB background operation runs using a percentage of user RUs.
+        /// </summary>
+        /// <param name="partitionKey">The <see cref="PartitionKey"/> of the items to be deleted.</param>
+        /// <param name="requestOptions">(Optional) The options for the Partition Key Delete request.</param>
+        /// <param name="cancellationToken">(Optional) <see cref="CancellationToken"/> representing request cancellation.</param>
+        /// <returns>
+        /// A <see cref="Task"/> containing a <see cref="ResponseMessage"/>.
+        /// </returns>
+        public abstract Task<ResponseMessage> DeleteAllItemsByPartitionKeyStreamAsync(
+               Cosmos.PartitionKey partitionKey,
+               RequestOptions requestOptions = null,
+               CancellationToken cancellationToken = default(CancellationToken));
+#endif
+
 #if PREVIEW
         /// <summary>
         /// Obtains a list of <see cref="FeedRange"/> that can be used to parallelize Feed operations.

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

@@ -910,6 +910,40 @@ public override async Task<PartitionKey> GetPartitionKeyValueFromStreamAsync(
             }
         }
 
+        public Task<ResponseMessage> DeleteAllItemsByPartitionKeyStreamAsync(
+          Cosmos.PartitionKey partitionKey,
+          CosmosDiagnosticsContext diagnosticsContext,
+          ITrace trace,
+          RequestOptions requestOptions = null,
+          CancellationToken cancellationToken = default(CancellationToken))
+        {
+            Cosmos.PartitionKey? resultingPartitionKey;
+            if (requestOptions != null && requestOptions.IsEffectivePartitionKeyRouting)
+            {
+                resultingPartitionKey = null;
+            }
+            else
+            {
+                resultingPartitionKey = partitionKey;
+            }
+
+            ContainerCore.ValidatePartitionKey(resultingPartitionKey, requestOptions);
+
+            return this.ClientContext.ProcessResourceOperationStreamAsync(
+                resourceUri: this.LinkUri,
+                resourceType: ResourceType.PartitionKey,
+                operationType: OperationType.Delete,
+                requestOptions: requestOptions,
+                cosmosContainerCore: this,
+                partitionKey: resultingPartitionKey,
+                itemId: null,
+                streamPayload: null,
+                requestEnricher: null,
+                trace: trace,
+                diagnosticsContext: diagnosticsContext,
+                cancellationToken: cancellationToken);
+        }
+
         private static bool TryParseTokenListForElement(CosmosObject pathTraversal, IReadOnlyList<string> tokens, out CosmosElement result)
         {
             result = null;

Microsoft.Azure.Cosmos/src/Resource/Container/ContainerInlineCore.cs

@@ -454,5 +454,16 @@ public override IAsyncEnumerable<TryCatch<ReadFeedPage>> GetReadFeedAsyncEnumera
         {
             return base.GetReadFeedAsyncEnumerable(state, requestOptions);
         }
+
+        public override Task<ResponseMessage> DeleteAllItemsByPartitionKeyStreamAsync(
+          Cosmos.PartitionKey partitionKey,
+          RequestOptions requestOptions = null,
+          CancellationToken cancellationToken = default(CancellationToken))
+        {
+            return this.ClientContext.OperationHelperAsync(
+                nameof(DeleteAllItemsByPartitionKeyStreamAsync),
+                requestOptions,
+                (diagnostics, trace) => base.DeleteAllItemsByPartitionKeyStreamAsync(partitionKey, diagnostics, trace, requestOptions, cancellationToken));
+        }
     }
 }

Microsoft.Azure.Cosmos/src/Resource/Container/ContainerInternal.cs

@@ -129,6 +129,11 @@ public abstract Task<ItemResponse<T>> PatchItemAsync<T>(
             IReadOnlyList<PatchOperation> patchOperations,
             ItemRequestOptions requestOptions = null,
             CancellationToken cancellationToken = default);
+
+        public abstract Task<ResponseMessage> DeleteAllItemsByPartitionKeyStreamAsync(
+               Cosmos.PartitionKey partitionKey,
+               RequestOptions requestOptions = null,
+               CancellationToken cancellationToken = default(CancellationToken));
 #endif
 
 #if !PREVIEW

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

@@ -564,6 +564,59 @@ await feedIterator.ReadNextAsync(this.cancellationToken))
             Assert.AreEqual(itemIds.Count, 0);
         }
 
+        [TestMethod]
+        public async Task PartitionKeyDeleteTest()
+        {
+            string pKString = "PK1";
+            string pKString2 = "PK2";
+            dynamic testItem1 = new
+            {
+                id = "item1",
+                status = pKString
+            };
+
+            dynamic testItem2 = new
+            {
+                id = "item2",
+                status = pKString
+            };
+
+            dynamic testItem3 = new
+            {
+                id = "item3",
+                status = pKString2
+            };
+
+            ContainerInternal containerInternal = (ContainerInternal)this.Container;
+            ItemResponse<dynamic>  itemResponse = await this.Container.CreateItemAsync<dynamic>(testItem1);
+            ItemResponse<dynamic> itemResponse2 = await this.Container.CreateItemAsync<dynamic>(testItem2);
+            ItemResponse<dynamic> itemResponse3 = await this.Container.CreateItemAsync<dynamic>(testItem3);
+            Cosmos.PartitionKey partitionKey1 = new Cosmos.PartitionKey(pKString);
+            Cosmos.PartitionKey partitionKey2 = new Cosmos.PartitionKey(pKString2);
+            using (ResponseMessage pKDeleteResponse = await containerInternal.DeleteAllItemsByPartitionKeyStreamAsync(partitionKey1))
+            {
+                Assert.AreEqual(pKDeleteResponse.StatusCode, HttpStatusCode.OK);
+            }
+
+            using (ResponseMessage readResponse = await this.Container.ReadItemStreamAsync("item1", partitionKey1))
+            {
+                Assert.AreEqual(readResponse.StatusCode, HttpStatusCode.NotFound);
+                Assert.AreEqual(readResponse.Headers.SubStatusCode, SubStatusCodes.Unknown);
+            }
+
+            using (ResponseMessage readResponse = await this.Container.ReadItemStreamAsync("item2", partitionKey1))
+            {
+                Assert.AreEqual(readResponse.StatusCode, HttpStatusCode.NotFound);
+                Assert.AreEqual(readResponse.Headers.SubStatusCode, SubStatusCodes.Unknown);
+            }
+
+            //verify item with the other Partition Key is not deleted
+            using (ResponseMessage readResponse = await this.Container.ReadItemStreamAsync("item3", partitionKey2))
+            {
+                Assert.AreEqual(readResponse.StatusCode, HttpStatusCode.OK);
+            }
+        }
+
         [TestMethod]
         public async Task ItemCustomSerialzierTest()
         {

Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/CosmosItemUnitTests.cs

@@ -438,6 +438,17 @@ public async Task AllowBatchingRequestsSendsToExecutor_Patch()
             mockedExecutor.Verify(c => c.AddAsync(It.IsAny<ItemBatchOperation>(), It.IsAny<ItemRequestOptions>(), It.IsAny<CancellationToken>()), Times.Once);
         }
 
+        [TestMethod]
+        public async Task PartitionKeyDeleteUnitTest()
+        {
+            dynamic item = new
+            {
+                id = Guid.NewGuid().ToString(),
+                pk = "FF627B77-568E-4541-A47E-041EAC10E46F",
+            };
+            await this.VerifyPartitionKeyDeleteOperation(new Cosmos.PartitionKey(item.pk), "[\"FF627B77-568E-4541-A47E-041EAC10E46F\"]");
+        }
+
         [TestMethod]
         public async Task TestNestedPartitionKeyValueFromStreamAsync()
         {
@@ -870,6 +881,42 @@ private async Task VerifyItemOperations(
             Assert.AreEqual(10, testHandlerHitCount, "A stream operation did not make it to the handler");
         }
 
+        private async Task VerifyPartitionKeyDeleteOperation(
+            Cosmos.PartitionKey partitionKey,
+            string partitionKeySerialized,
+            RequestOptions requestOptions = null)
+        {
+            ResponseMessage response = null;
+            HttpStatusCode httpStatusCode = HttpStatusCode.OK;
+            int testHandlerHitCount = 0;
+            TestHandler testHandler = new TestHandler((request, cancellationToken) =>
+            {
+                Assert.IsTrue(request.RequestUri.OriginalString.StartsWith(@"dbs/testdb/colls/testcontainer"));
+                Assert.AreEqual(requestOptions, request.RequestOptions);
+                Assert.AreEqual(ResourceType.PartitionKey, request.ResourceType);
+                Assert.IsNotNull(request.Headers.PartitionKey);
+                Assert.AreEqual(partitionKeySerialized, request.Headers.PartitionKey);
+                testHandlerHitCount++;
+                response = new ResponseMessage(httpStatusCode, request, errorMessage: null);
+                response.Content = request.Content;
+                return Task.FromResult(response);
+            });
+
+            CosmosClient client = MockCosmosUtil.CreateMockCosmosClient(
+                (builder) => builder.AddCustomHandlers(testHandler));
+
+            Container container = client.GetDatabase("testdb")
+                                        .GetContainer("testcontainer");
+
+            ContainerInternal containerInternal = (ContainerInternal)container;
+            ResponseMessage responseMessage = await containerInternal.DeleteAllItemsByPartitionKeyStreamAsync(
+                partitionKey: partitionKey,
+                requestOptions: requestOptions);
+            Assert.IsNotNull(responseMessage);
+            Assert.AreEqual(httpStatusCode, responseMessage.StatusCode);
+            Assert.AreEqual(1, testHandlerHitCount, "The operation did not make it to the handler");
+        }
+
         private Mock<BatchAsyncContainerExecutor> GetMockedBatchExcecutor()
         {
             Mock<BatchAsyncContainerExecutor> mockedExecutor = new Mock<BatchAsyncContainerExecutor>();