Skip to content

Commit b1f3f57

Browse files
neildshneildsh
committed
Query: Fixes incorrect FeedResponse.Count when result contains undefined elements (#3574)
* Do not maintain an independent count on QueryResponse that can go out of sync * Add more test coverage for QueryResponse<T>.Count * Output the correct count from CosmosElementSerializer when the input contains CosmosUndefined * Add untyped tests for CosmosUndefined * Remove commented code
1 parent 5175a24 commit b1f3f57

File tree

4 files changed

+98
-6
lines changed

4 files changed

+98
-6
lines changed

‎Microsoft.Azure.Cosmos/src/Query/v3Query/QueryResponse.cs‎

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -163,6 +163,7 @@ class QueryResponse<T> : FeedResponse<T>
163163
{
164164
private readonly CosmosSerializerCore serializerCore;
165165
private readonly CosmosSerializationFormatOptions serializationOptions;
166+
private readonly IReadOnlyList<T> resource;
166167

167168
private QueryResponse(
168169
HttpStatusCode httpStatusCode,
@@ -178,8 +179,7 @@ private QueryResponse(
178179
this.serializerCore = serializerCore;
179180
this.serializationOptions = serializationOptions;
180181
this.StatusCode = httpStatusCode;
181-
this.Count = cosmosElements.Count;
182-
this.Resource = CosmosElementSerializer.GetResources<T>(
182+
this.resource = CosmosElementSerializer.GetResources<T>(
183183
cosmosArray: cosmosElements,
184184
serializerCore: serializerCore);
185185

@@ -197,7 +197,7 @@ private QueryResponse(
197197

198198
public override CosmosDiagnostics Diagnostics { get; }
199199

200-
public override int Count { get; }
200+
public override int Count => this.resource.Count;
201201

202202
internal CosmosQueryResponseMessageHeaders QueryHeaders { get; }
203203

@@ -210,7 +210,7 @@ public override IEnumerator<T> GetEnumerator()
210210
return this.Resource.GetEnumerator();
211211
}
212212

213-
public override IEnumerable<T> Resource { get; }
213+
public override IEnumerable<T> Resource => this.resource;
214214

215215
internal override RequestMessage RequestMessage { get; }
216216

‎Microsoft.Azure.Cosmos/src/Serializer/CosmosElementSerializer.cs‎

Lines changed: 5 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -74,8 +74,11 @@ internal static MemoryStream ToStream(
7474
jsonWriter.WriteArrayStart();
7575
foreach (CosmosElement element in cosmosElements)
7676
{
77-
count++;
78-
element.WriteTo(jsonWriter);
77+
if (element is not CosmosUndefined)
78+
{
79+
count++;
80+
element.WriteTo(jsonWriter);
81+
}
7982
}
8083

8184
jsonWriter.WriteArrayEnd();

‎Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.EmulatorTests/Query/CosmosUndefinedQueryTests.cs‎

Lines changed: 65 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -5,6 +5,7 @@ namespace Microsoft.Azure.Cosmos.EmulatorTests.Query
55
using System.Collections.ObjectModel;
66
using System.Linq;
77
using System.Threading.Tasks;
8+
using Azure;
89
using Microsoft.Azure.Cosmos.CosmosElements;
910
using Microsoft.Azure.Cosmos.CosmosElements.Numbers;
1011
using Microsoft.Azure.Cosmos.Json;
@@ -50,6 +51,70 @@ private static async Task RunTests(Container container, IReadOnlyList<CosmosObje
5051
{
5152
await OrderByTests(container);
5253
await GroupByTests(container);
54+
await UntypedTests(container);
55+
}
56+
57+
private static async Task UntypedTests(Container container)
58+
{
59+
UndefinedProjectionTestCase[] undefinedProjectionTestCases = new[]
60+
{
61+
MakeUndefinedProjectionTest(
62+
query: "SELECT VALUE c.AlwaysUndefinedField FROM c",
63+
expectedCount: 0),
64+
MakeUndefinedProjectionTest(
65+
query: "SELECT VALUE c.AlwaysUndefinedField FROM c ORDER BY c.AlwaysUndefinedField",
66+
expectedCount: 0),
67+
MakeUndefinedProjectionTest(
68+
query: "SELECT c.AlwaysUndefinedField FROM c ORDER BY c.AlwaysUndefinedField",
69+
expectedCount: DocumentCount),
70+
MakeUndefinedProjectionTest(
71+
query: "SELECT VALUE c.AlwaysUndefinedField FROM c GROUP BY c.AlwaysUndefinedField",
72+
expectedCount: 0),
73+
MakeUndefinedProjectionTest(
74+
query: $"SELECT VALUE SUM(c.{nameof(MixedTypeDocument.MixedTypeField)}) FROM c",
75+
expectedCount: 0),
76+
MakeUndefinedProjectionTest(
77+
query: $"SELECT VALUE AVG(c.{nameof(MixedTypeDocument.MixedTypeField)}) FROM c",
78+
expectedCount: 0),
79+
};
80+
81+
foreach (UndefinedProjectionTestCase testCase in undefinedProjectionTestCases)
82+
{
83+
foreach (int pageSize in PageSizes)
84+
{
85+
IAsyncEnumerable<ResponseMessage> results = RunSimpleQueryAsync(
86+
container,
87+
testCase.Query,
88+
new QueryRequestOptions { MaxItemCount = pageSize });
89+
90+
long actualCount = 0;
91+
await foreach (ResponseMessage responseMessage in results)
92+
{
93+
Assert.IsTrue(responseMessage.IsSuccessStatusCode);
94+
95+
string content = responseMessage.Content.ReadAsString();
96+
IJsonNavigator navigator = JsonNavigator.Create(System.Text.Encoding.UTF8.GetBytes(content));
97+
IJsonNavigatorNode rootNode = navigator.GetRootNode();
98+
Assert.IsTrue(navigator.TryGetObjectProperty(rootNode, "_count", out ObjectProperty countProperty));
99+
100+
long count = Number64.ToLong(navigator.GetNumber64Value(countProperty.ValueNode));
101+
actualCount += count;
102+
103+
Assert.IsTrue(navigator.TryGetObjectProperty(rootNode, "Documents", out ObjectProperty documentsProperty));
104+
int documentCount = navigator.GetArrayItemCount(documentsProperty.ValueNode);
105+
Assert.AreEqual(count, documentCount);
106+
107+
for (int index= 0; index < documentCount; ++index)
108+
{
109+
IJsonNavigatorNode documentNode = navigator.GetArrayItemAt(documentsProperty.ValueNode, index);
110+
int propertyCount = navigator.GetObjectPropertyCount(documentNode);
111+
Assert.AreEqual(0, propertyCount);
112+
}
113+
}
114+
115+
Assert.AreEqual(testCase.ExpectedResultCount, actualCount);
116+
}
117+
}
53118
}
54119

55120
private static async Task OrderByTests(Container container)

‎Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.EmulatorTests/Query/QueryTestsBase.cs‎

Lines changed: 24 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -622,6 +622,7 @@ internal static async Task<List<T>> QueryWithContinuationTokensAsync<T>(
622622
}
623623

624624
List<T> resultsFromContinuationToken = new List<T>();
625+
int resultCount = 0;
625626
string continuationToken = null;
626627
do
627628
{
@@ -645,6 +646,7 @@ internal static async Task<List<T>> QueryWithContinuationTokensAsync<T>(
645646
}
646647

647648
resultsFromContinuationToken.AddRange(cosmosQueryResponse);
649+
resultCount += cosmosQueryResponse.Count;
648650
continuationToken = cosmosQueryResponse.ContinuationToken;
649651
break;
650652
}
@@ -664,6 +666,7 @@ internal static async Task<List<T>> QueryWithContinuationTokensAsync<T>(
664666
}
665667
} while (continuationToken != null);
666668

669+
Assert.AreEqual(resultsFromContinuationToken.Count, resultCount);
667670
return resultsFromContinuationToken;
668671
}
669672

@@ -677,6 +680,7 @@ internal static async Task<List<T>> QueryWithoutContinuationTokensAsync<T>(
677680
queryRequestOptions = new QueryRequestOptions();
678681
}
679682

683+
int resultCount = 0;
680684
List<T> results = new List<T>();
681685
FeedIterator<T> itemQuery = container.GetItemQueryIterator<T>(
682686
queryText: query,
@@ -690,6 +694,7 @@ internal static async Task<List<T>> QueryWithoutContinuationTokensAsync<T>(
690694
{
691695
FeedResponse<T> page = await itemQuery.ReadNextAsync();
692696
results.AddRange(page);
697+
resultCount += page.Count;
693698

694699
if (queryRequestOptions.MaxItemCount.HasValue)
695700
{
@@ -723,6 +728,7 @@ internal static async Task<List<T>> QueryWithoutContinuationTokensAsync<T>(
723728
{
724729
// The query failed and we don't have a save point, so just restart the whole thing.
725730
results = new List<T>();
731+
resultCount = 0;
726732
}
727733
}
728734
}
@@ -732,6 +738,7 @@ internal static async Task<List<T>> QueryWithoutContinuationTokensAsync<T>(
732738
itemQuery.Dispose();
733739
}
734740

741+
Assert.AreEqual(results.Count, resultCount);
735742
return results;
736743
}
737744

@@ -884,6 +891,23 @@ internal static async IAsyncEnumerable<FeedResponse<T>> RunSimpleQueryWithNewIte
884891
}
885892
}
886893

894+
internal static async IAsyncEnumerable<ResponseMessage> RunSimpleQueryAsync(
895+
Container container,
896+
string query,
897+
QueryRequestOptions requestOptions = null)
898+
{
899+
using FeedIterator resultSetIterator = container.GetItemQueryStreamIterator(
900+
query,
901+
null,
902+
requestOptions: requestOptions);
903+
904+
while (resultSetIterator.HasMoreResults)
905+
{
906+
ResponseMessage response = await resultSetIterator.ReadNextAsync();
907+
yield return response;
908+
}
909+
}
910+
887911
internal async Task<List<T>> RunSinglePartitionQuery<T>(
888912
Container container,
889913
string query,

0 commit comments

Comments
 (0)