[Internal] Query : Fixes response parsing error for ReadFeed by limiting binary format to query operations

Microsoft.Azure.Cosmos/src/Query/v2Query/DocumentQueryExecutionContextBase.cs

@@ -32,8 +32,9 @@ namespace Microsoft.Azure.Cosmos.Query
 
     internal abstract class DocumentQueryExecutionContextBase : IDocumentQueryExecutionContext
     {
-        public static readonly string DefaultSupportedSerializationFormats = string.Join(",", SupportedSerializationFormats.JsonText, SupportedSerializationFormats.CosmosBinary);
-
+        // ISSUE-TODO-adityasa-2024/5/27 - Re-enable binary format. Currently encryption codepath is not able to handle binary response.
+        public static readonly string DefaultSupportedSerializationFormats = SupportedSerializationFormats.JsonText.ToString();
+
         public readonly struct InitParams
         {
             public IDocumentQueryClient Client { get; }

Microsoft.Azure.Cosmos/src/RequestOptions/QueryRequestOptions.cs

@@ -237,8 +237,14 @@ internal override void PopulateRequestOptions(RequestMessage request)
             {
                 request.Headers.Add(HttpConstants.HttpHeaders.ResponseContinuationTokenLimitInKB, this.ResponseContinuationTokenLimitInKb.ToString());
             }
-
-            request.Headers.CosmosMessageHeaders.SupportedSerializationFormats = this.SupportedSerializationFormats?.ToString() ?? DocumentQueryExecutionContextBase.DefaultSupportedSerializationFormats;
+
+            // All query APIs (GetItemQueryIterator, GetItemLinqQueryable and GetItemQueryStreamIterator) turn into ReadFeed operation if query text is null.
+            // In such a case, query pipelines are still involved (including QueryRequestOptions). In general backend only honors SupportedSerializationFormats
+            //  for OperationType Query but has a bug where it returns a binary response for ReadFeed API when partition key is also specified in the request.
+            if (request.OperationType == OperationType.Query)
+            {
+                request.Headers.CosmosMessageHeaders.SupportedSerializationFormats = this.SupportedSerializationFormats?.ToString() ?? DocumentQueryExecutionContextBase.DefaultSupportedSerializationFormats;
+            }
 
             if (this.StartId != null)
             {

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

@@ -389,7 +389,7 @@ private void SupportedSerializationFormatsPositiveCases(
             INameValueCollection headers = new RequestNameValueCollection();
 
             headers.Add(HttpConstants.HttpHeaders.SupportedSerializationFormats, supportedSerializationFormats);
-            DocumentServiceResponse response;                   
+            DocumentServiceResponse response;
             if(sqlQuerySpec!=null)
             {
                 response = this.QueryRequest(client, collection.ResourceId, sqlQuerySpec, headers);
@@ -440,34 +440,28 @@ private void ValidateSupportedSerializationFormatsQuery(DocumentClient client, D
                 expectedFormat: SupportedSerializationFormats.JsonText,
                 supportedSerializationFormats: "jsontext",
                 sqlQuerySpec: sqlQuerySpec);
-            /*
             this.SupportedSerializationFormatsPositiveCases(client, collection,
                 expectedFormat: SupportedSerializationFormats.CosmosBinary,
                 supportedSerializationFormats: "COSMOSBINARY",
                 sqlQuerySpec: sqlQuerySpec);
-            */
             this.SupportedSerializationFormatsPositiveCases(client, collection,
-                expectedFormat: SupportedSerializationFormats.JsonText,
+                expectedFormat: SupportedSerializationFormats.CosmosBinary,
                 supportedSerializationFormats: "JsonText, CosmosBinary",
                 sqlQuerySpec: sqlQuerySpec);
-            /*
             this.SupportedSerializationFormatsPositiveCases(client, collection,
                 expectedFormat: SupportedSerializationFormats.CosmosBinary,
                 supportedSerializationFormats: "CosmosBinary, HybridRow",
                 sqlQuerySpec: sqlQuerySpec);
-            */
             this.SupportedSerializationFormatsPositiveCases(client, collection,
-                expectedFormat: SupportedSerializationFormats.JsonText,
+                expectedFormat: SupportedSerializationFormats.CosmosBinary,
                 supportedSerializationFormats: "JsonText, CosmosBinary, HybridRow",
                 sqlQuerySpec: sqlQuerySpec);
-            /*
             this.SupportedSerializationFormatsPositiveCases(client, collection,
                 expectedFormat: SupportedSerializationFormats.CosmosBinary,
                 supportedSerializationFormats: "JsonText, CosmosBinary, HybridRow",
                 sqlQuerySpec: sqlQuerySpec);
-            */
             this.SupportedSerializationFormatsPositiveCases(client, collection,
-                expectedFormat: SupportedSerializationFormats.JsonText,
+                expectedFormat: SupportedSerializationFormats.CosmosBinary,
                 supportedSerializationFormats: "JsonText, CosmosBinary, HybridRow",
                 sqlQuerySpec: sqlQuerySpec);
         }

Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.EmulatorTests/Query/SupportedSerializationFormatsTest.cs

@@ -1,14 +1,13 @@
-namespace Microsoft.Azure.Cosmos.Query
+namespace Microsoft.Azure.Cosmos.EmulatorTests.Query
 {
-    using System;
     using System.Collections.Generic;
+    using System.IO;
     using System.Linq;
     using System.Threading.Tasks;
     using Microsoft.Azure.Cosmos.CosmosElements;
-    using Microsoft.Azure.Cosmos.EmulatorTests.Query;
-    using Microsoft.Azure.Cosmos.SDK.EmulatorTests.QueryOracle;
     using Microsoft.Azure.Documents;
     using Microsoft.VisualStudio.TestTools.UnitTesting;
+    using Newtonsoft.Json.Linq;
 
     [TestClass]
     [TestCategory("Query")]
@@ -36,66 +35,126 @@ await this.CreateIngestQueryDeleteAsync(
                 CollectionTypes.SinglePartition | CollectionTypes.MultiPartition,
                 inputDocuments,
                 this.TestSupportedSerializationFormatsHelper,
-                    "/id");
+                "/id");
         }
 
         private async Task TestSupportedSerializationFormatsHelper(Container container, IReadOnlyList<CosmosObject> documents)
         {
-            string[] expectedResults = new[] { "document_0", "document_1", "document_2", "document_3", "document_4", "document_5", "document_6", "document_7", "document_8", "document_9" };
-            string query = string.Format("SELECT c.name FROM c");
-            List<QueryRequestOptions> queryRequestOptionsList = new List<QueryRequestOptions>()
-            {
-                new QueryRequestOptions()
-                {
-                    SupportedSerializationFormats = SupportedSerializationFormats.CosmosBinary | SupportedSerializationFormats.HybridRow
-                },
-                new QueryRequestOptions()
-                {
-                    SupportedSerializationFormats = SupportedSerializationFormats.JsonText | SupportedSerializationFormats.CosmosBinary
-                },
-                new QueryRequestOptions()
-                {
-                    SupportedSerializationFormats = SupportedSerializationFormats.JsonText | SupportedSerializationFormats.HybridRow
-                },
-                new QueryRequestOptions()
+            List<(Cosmos.PartitionKey?, string[]) > partitionKeyAndExpectedResults = new List<(Cosmos.PartitionKey? partitionKey, string[] documents)>
                 {
-                    SupportedSerializationFormats = SupportedSerializationFormats.JsonText | SupportedSerializationFormats.CosmosBinary | SupportedSerializationFormats.HybridRow
-                },
-                new QueryRequestOptions()
-                {
-                    SupportedSerializationFormats = SupportedSerializationFormats.JsonText
-                },
-                new QueryRequestOptions()
+                    (
+                        partitionKey: null,
+                        documents: new[] { "document_0", "document_1", "document_2", "document_3", "document_4","document_5", "document_6", "document_7", "document_8", "document_9" }
+                    ),
+                    (
+                        partitionKey: new Cosmos.PartitionKey("0"),
+                        documents: new[] { "document_0" }
+                    )
+                };
+
+            foreach (string query in new string[]
                 {
-                    SupportedSerializationFormats = SupportedSerializationFormats.CosmosBinary
-                },
-                new QueryRequestOptions()
+                    null,   // With null query, this turns into a ReadFeed API, for which SupportedSerializationFormats should be ignored.
+                    "SELECT * FROM c"
+                })
+            {
+                List<QueryRequestOptions> queryRequestOptionsList = new List<QueryRequestOptions>()
                 {
-                    SupportedSerializationFormats = SupportedSerializationFormats.CosmosBinary
-                },
-                new QueryRequestOptions()
+                    new QueryRequestOptions()
+                    {
+                        SupportedSerializationFormats = SupportedSerializationFormats.CosmosBinary | SupportedSerializationFormats.HybridRow
+                    },
+                    new QueryRequestOptions()
+                    {
+                        SupportedSerializationFormats = SupportedSerializationFormats.JsonText | SupportedSerializationFormats.CosmosBinary
+                    },
+                    new QueryRequestOptions()
+                    {
+                        SupportedSerializationFormats = SupportedSerializationFormats.JsonText | SupportedSerializationFormats.HybridRow
+                    },
+                    new QueryRequestOptions()
+                    {
+                        SupportedSerializationFormats = SupportedSerializationFormats.JsonText | SupportedSerializationFormats.CosmosBinary | SupportedSerializationFormats.HybridRow
+                    },
+                    new QueryRequestOptions()
+                    {
+                        SupportedSerializationFormats = SupportedSerializationFormats.JsonText
+                    },
+                    new QueryRequestOptions()
+                    {
+                        SupportedSerializationFormats = SupportedSerializationFormats.CosmosBinary
+                    },
+                    new QueryRequestOptions()
+                    {
+                        SupportedSerializationFormats = SupportedSerializationFormats.CosmosBinary
+                    },
+                    new QueryRequestOptions()
+                    {
+                        SupportedSerializationFormats = SupportedSerializationFormats.JsonText
+                    }
+                };
+
+                // GetItemQueryIterator
+                foreach (QueryRequestOptions requestOptions in queryRequestOptionsList)
                 {
-                    SupportedSerializationFormats = SupportedSerializationFormats.JsonText
+                    QueryDefinition queryDefinition = query != null ? new QueryDefinition(query) : null;
+                    foreach ((Cosmos.PartitionKey? partitionKey, string[] expectedResults) in partitionKeyAndExpectedResults)
+                    {
+                        requestOptions.PartitionKey = partitionKey;
+
+                        List<JObject> queryResults = new List<JObject>();
+                        using (FeedIterator<JObject> feedIterator = container.GetItemQueryIterator<JObject>(queryDefinition, requestOptions: requestOptions))
+                        {
+                            while (feedIterator.HasMoreResults)
+                            {
+                                FeedResponse<JObject> response = await feedIterator.ReadNextAsync();
+                                queryResults.AddRange(response.ToList());
+                            }
+                        }
+
+                        string[] actualResults = queryResults
+                            .Select(doc => doc["name"].ToString())
+                            .ToArray();
+
+                        CollectionAssert.AreEquivalent(expectedResults, actualResults);
+                    }
                 }
-            };
 
-            foreach (QueryRequestOptions requestOptions in queryRequestOptionsList)
-            {
-                List<CosmosElement> results = new List<CosmosElement>();
-                using (FeedIterator<CosmosElement> feedIterator = container.GetItemQueryIterator<CosmosElement>(new QueryDefinition(query), requestOptions: requestOptions))
+                // GetItemQueryStreamIterator
+                foreach (QueryRequestOptions requestOptions in queryRequestOptionsList)
                 {
-                    while (feedIterator.HasMoreResults)
+                    QueryDefinition queryDefinition = query != null ? new QueryDefinition(query) : null;
+                    foreach ((Cosmos.PartitionKey? partitionKey, string[] expectedResults) in partitionKeyAndExpectedResults)
                     {
-                        FeedResponse<CosmosElement> response = await feedIterator.ReadNextAsync();
-                        results.AddRange(response.ToList());
+                        requestOptions.PartitionKey = partitionKey;
+
+                        List<CosmosElement> queryResults = new List<CosmosElement>();
+                        using (FeedIterator feedIterator = container.GetItemQueryStreamIterator(queryDefinition, requestOptions: requestOptions))
+                        {
+                            while (feedIterator.HasMoreResults)
+                            {
+                                ResponseMessage response = await feedIterator.ReadNextAsync();
+                                queryResults.AddRange(Deserialize(response.Content));
+                            }
+                        }
+
+                        string[] actualResults = queryResults
+                            .Select(doc => ((CosmosString)((CosmosObject)doc)["name"]).Value.ToString())
+                            .ToArray();
+
+                        CollectionAssert.AreEquivalent(expectedResults, actualResults);
                     }
                 }
+            }
+        }
 
-                string[] actualResults = results
-                    .Select(doc => ((CosmosString)(doc as CosmosObject)["name"]).Value.ToString())
-                    .ToArray();
-
-                CollectionAssert.AreEquivalent(expectedResults, actualResults);
+        private static IEnumerable<CosmosElement> Deserialize(Stream content)
+        {
+            string contentAsString = new StreamReader(content).ReadToEnd();
+            CosmosObject obj = CosmosObject.Parse(contentAsString);
+            foreach (CosmosElement element in (CosmosArray)obj["Documents"])
+            {
+                yield return element;
             }
         }
     }