ReadMany: Fixes missing headers and message on failure scenarios and …

…AddRequestHeaders request option (Azure#2567) Fixes custom headers via AddRequestHeaders in ReadManyRequestOptions Fixes the ResponseMessage in failure scenarios to have the exception and headers.
juraj-blazek · Jun 21, 2021 · 275e020 · 275e020
1 parent 1797aaa
commit 275e020
Show file tree

Hide file tree

Showing 3 changed files with 85 additions and 5 deletions.
diff --git a/Microsoft.Azure.Cosmos/src/ReadManyQueryHelper.cs b/Microsoft.Azure.Cosmos/src/ReadManyQueryHelper.cs
@@ -160,10 +160,12 @@ private ResponseMessage CombineStreamsFromQueryResponses(List<ResponseMessage>[]
                     {
                         if (!responseMessage.IsSuccessStatusCode)
                         {
-                            return new ResponseMessage(responseMessage.StatusCode)
-                            {
-                                Trace = trace
-                            };
+                            return new ResponseMessage(
+                                responseMessage.StatusCode,
+                                requestMessage: null,
+                                responseMessage.Headers,
+                                responseMessage.CosmosException,
+                                trace);
                         }
 
                         if (responseMessage is QueryResponse queryResponse)

diff --git a/Microsoft.Azure.Cosmos/src/ReadManyRequestOptions.cs b/Microsoft.Azure.Cosmos/src/ReadManyRequestOptions.cs
@@ -67,7 +67,8 @@ internal QueryRequestOptions ConvertToQueryRequestOptions()
                 SessionToken = this.SessionToken,
                 IfMatchEtag = this.IfMatchEtag,
                 IfNoneMatchEtag = this.IfNoneMatchEtag,
-                Properties = this.Properties
+                Properties = this.Properties,
+                AddRequestHeaders = this.AddRequestHeaders
             };
         }
     }

diff --git a/...osoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.EmulatorTests/CosmosReadManyItemsTests.cs b/...osoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.EmulatorTests/CosmosReadManyItemsTests.cs
@@ -252,6 +252,83 @@ public async Task MultipleQueriesToSamePartitionTest()
             Assert.AreEqual(feedResponse.Count, 2500);
         }
 
+        [TestMethod]
+        public async Task ReadManyTestWithIncorrectIntendedContainerRid()
+        {
+            for (int i = 0; i < 2; i++)
+            {
+                await this.Container.CreateItemAsync<ToDoActivity>(ToDoActivity.CreateRandomToDoActivity("pk", i.ToString()));
+            }
+
+            List<(string, PartitionKey)> itemList = new List<(string, PartitionKey)>();
+            for (int i = 0; i < 2; i++)
+            {
+                itemList.Add((i.ToString(), new PartitionKey("pk")));
+            }           
+
+            // pass incorrect Rid.
+            ReadManyRequestOptions readManyRequestOptions = new ReadManyRequestOptions
+            {
+                AddRequestHeaders = (headers) =>
+                {
+                    headers["x-ms-cosmos-is-client-encrypted"] = bool.TrueString;
+                    headers["x-ms-cosmos-intended-collection-rid"] = "iCoRrecTrID=";
+                }
+            };
+
+            FeedResponse<ToDoActivity> feedResponse;
+            try
+            {
+                feedResponse = await this.Container.ReadManyItemsAsync<ToDoActivity>(itemList, readManyRequestOptions);
+                Assert.Fail("ReadManyItemsAsync execution should have failed. ");
+            }
+            catch(CosmosException ex)
+            {
+                if (ex.StatusCode != HttpStatusCode.BadRequest || ex.SubStatusCode != 1024)
+                {
+                    Assert.Fail("ReadManyItemsAsync execution should have failed with the StatusCode: BadRequest and SubStatusCode: 1024. ");
+                }
+            }
+
+            using (ResponseMessage responseMessage = await this.Container.ReadManyItemsStreamAsync(itemList , readManyRequestOptions))
+            {
+                if(responseMessage.StatusCode != HttpStatusCode.BadRequest ||
+                    !string.Equals(responseMessage.Headers.Get("x-ms-substatus"), "1024"))
+                {
+                    Assert.Fail("ReadManyItemsStreamAsync execution should have failed with the StatusCode: BadRequest and SubStatusCode: 1024. ");
+                }
+            }
+
+            // validate by passing correct Rid.
+            ContainerInlineCore containerInternal = (ContainerInlineCore)this.Container;
+            string rid = await containerInternal.GetCachedRIDAsync(forceRefresh: false, NoOpTrace.Singleton, cancellationToken: default);
+
+            readManyRequestOptions = new ReadManyRequestOptions
+            {
+                AddRequestHeaders = (headers) =>
+                {
+                    headers["x-ms-cosmos-is-client-encrypted"] = bool.TrueString;
+                    headers["x-ms-cosmos-intended-collection-rid"] = rid;
+                }
+            };
+
+            feedResponse = await this.Container.ReadManyItemsAsync<ToDoActivity>(itemList, readManyRequestOptions);
+            Assert.AreEqual(feedResponse.Count, 2);
+
+            using (ResponseMessage responseMessage = await this.Container.ReadManyItemsStreamAsync(itemList, readManyRequestOptions))
+            {
+                Assert.AreEqual(responseMessage.StatusCode, HttpStatusCode.OK);
+
+                Assert.IsNotNull(responseMessage);
+                Assert.IsTrue(responseMessage.Headers.RequestCharge > 0);
+                Assert.IsNotNull(responseMessage.Diagnostics);
+
+                ToDoActivity[] items = this.cosmosClient.ClientContext.SerializerCore.FromFeedStream<ToDoActivity>(
+                                        CosmosFeedResponseSerializer.GetStreamWithoutServiceEnvelope(responseMessage.Content));
+                Assert.AreEqual(items.Length, 2);
+            }
+        }
+
         [TestMethod]
         public async Task ReadMany404ExceptionTest()
         {