Query: Fixes ResponseMessage not parsing the IndexMetrics as text in …

…latest sdk (#4397) * Fix response message not parsing the index advice as text in latest sdk * add test coverage * address code review --------- Co-authored-by: Minh Le (from Dev Box) <leminh@microsoft.com>
Azure · Apr 11, 2024 · 0e06625 · 0e06625
1 parent cb7f923
commit 0e06625
Show file tree

Hide file tree

Showing 2 changed files with 48 additions and 1 deletion.
diff --git a/Microsoft.Azure.Cosmos/src/Handler/ResponseMessage.cs b/Microsoft.Azure.Cosmos/src/Handler/ResponseMessage.cs
@@ -83,7 +83,7 @@ internal ResponseMessage(
             this.CosmosException = cosmosException;
             this.Headers = headers ?? new Headers();
 
-            this.IndexUtilizationText = ResponseMessage.DecodeIndexMetrics(this.Headers, isBase64Encoded: true);
+            this.IndexUtilizationText = ResponseMessage.DecodeIndexMetrics(this.Headers, isBase64Encoded: false);
 
             if (requestMessage != null && requestMessage.Trace != null)
             {

diff --git a/...Azure.Cosmos/tests/Microsoft.Azure.Cosmos.EmulatorTests/IndexMetricsParserBaselineTest.cs b/...Azure.Cosmos/tests/Microsoft.Azure.Cosmos.EmulatorTests/IndexMetricsParserBaselineTest.cs
@@ -384,6 +384,18 @@ public override IndexMetricsParserTestOutput ExecuteTest(IndexMetricsParserTestI
 
             // Make sure ODE and non-ODE is consistent
             Assert.AreEqual(indexMetricsNonODE, indexMetricsODE);
+
+            // ----------------------------
+            // Test stream API
+            // ----------------------------
+            // Execute without ODE
+            string indexMetricsNonODEStreaming = RunStreamAPITest(input.Query, enableOptimisticDirectExecution: false);
+
+            // Execute with ODE
+            string indexMetricsODEStreaming = RunStreamAPITest(input.Query, enableOptimisticDirectExecution: true);
+
+            // Make sure ODE and non-ODE is consistent
+            Assert.AreEqual(indexMetricsNonODEStreaming, indexMetricsODEStreaming);
 
             return new IndexMetricsParserTestOutput(indexMetricsNonODE);
         }
@@ -403,7 +415,42 @@ private static string RunTest(string query, bool enableOptimisticDirectExecution
             {
                 FeedResponse<CosmosElement> page = itemQuery.ReadNextAsync().Result;
                 Assert.IsTrue(page.Headers.AllKeys().Length > 1);
+                if (roundTripCount > 0)
+                {
+                    if (page.IndexMetrics != null)
+                    {
+                        Assert.Fail("Expected only Index Metrics on first round trip. Current round trip %n", roundTripCount);
+                    }
+                }
+                else
+                {
+                    Assert.IsNotNull(page.Headers.Get(HttpConstants.HttpHeaders.IndexUtilization), "Expected index utilization headers for query");
+                    Assert.IsNotNull(page.IndexMetrics, "Expected index metrics response for query");
+
+                    indexMetrics = page.IndexMetrics;
+                }
+
+                roundTripCount++;
+            }
+
+            return indexMetrics;
+        }
+
+        private static string RunStreamAPITest(string query, bool enableOptimisticDirectExecution)
+        {
+            QueryRequestOptions requestOptions = new QueryRequestOptions() { PopulateIndexMetrics = true, EnableOptimisticDirectExecution = enableOptimisticDirectExecution };
 
+            using FeedIterator itemQuery = testContainer.GetItemQueryStreamIterator(
+                queryText: query,
+                requestOptions: requestOptions);
+
+            // Index Metrics is returned fully on the first page so no need to worry about result set
+            int roundTripCount = 0;
+            string indexMetrics = null;
+            while (itemQuery.HasMoreResults)
+            {
+                ResponseMessage page = itemQuery.ReadNextAsync().Result;
+                Assert.IsTrue(page.Headers.AllKeys().Length > 1);
                 if (roundTripCount > 0)
                 {
                     if (page.IndexMetrics != null)