Skip to content

Commit 4ecb25b

Browse files
neildshneildsh
committed
Do not use passthrough context for collections with HPK
1 parent 0b775de commit 4ecb25b

4 files changed

Lines changed: 253 additions & 96 deletions

File tree

‎Microsoft.Azure.Cosmos/src/Handler/RequestMessage.cs‎

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -385,9 +385,10 @@ private bool AssertPartitioningPropertiesAndHeaders()
385385

386386
bool partitionKeyRangeIdExists = !string.IsNullOrEmpty(this.Headers.PartitionKeyRangeId);
387387
if (partitionKeyRangeIdExists)
388-
{
388+
{
389+
OperationType op = this.OperationType;
389390
// Assert operation type is not write
390-
if (this.OperationType != OperationType.Query && this.OperationType != OperationType.ReadFeed && this.OperationType != OperationType.Batch)
391+
if (op != OperationType.Query && op != OperationType.QueryPlan && op != OperationType.ReadFeed && op != OperationType.Batch)
391392
{
392393
throw new ArgumentOutOfRangeException(RMResources.UnexpectedPartitionKeyRangeId);
393394
}

‎Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/CosmosQueryExecutionContextFactory.cs‎

Lines changed: 4 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -191,7 +191,7 @@ private static async Task<TryCatch<IQueryPipelineStage>> TryCreateCoreContextAsy
191191
// then try seeing if we can execute as a passthrough using client side only logic.
192192
// This is to short circuit the need to go to the gateway to get the query plan.
193193
if (cosmosQueryContext.QueryClient.BypassQueryParsing()
194-
&& inputParameters.PartitionKey.HasValue)
194+
&& inputParameters.PartitionKey.HasValue && containerQueryProperties.PartitionKeyDefinition.Paths.Count <= 1)
195195
{
196196
bool parsed;
197197
SqlQuery sqlQuery;
@@ -301,9 +301,10 @@ private static async Task<TryCatch<IQueryPipelineStage>> TryCreateFromPartitione
301301
}
302302
else
303303
{
304-
bool singleLogicalPartitionKeyQuery = (inputParameters.PartitionKey.HasValue && targetRanges.Count == 1)
304+
bool singleLogicalPartitionKeyQuery = ((inputParameters.PartitionKey.HasValue && targetRanges.Count == 1)
305305
|| ((partitionedQueryExecutionInfo.QueryRanges.Count == 1)
306-
&& partitionedQueryExecutionInfo.QueryRanges[0].IsSingleValue);
306+
&& partitionedQueryExecutionInfo.QueryRanges[0].IsSingleValue))
307+
&& containerQueryProperties.PartitionKeyDefinition.Paths.Count <= 1;
307308
bool serverStreamingQuery = !partitionedQueryExecutionInfo.QueryInfo.HasAggregates
308309
&& !partitionedQueryExecutionInfo.QueryInfo.HasDistinct
309310
&& !partitionedQueryExecutionInfo.QueryInfo.HasNonStreamingOrderBy

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

Lines changed: 115 additions & 25 deletions
Original file line numberDiff line numberDiff line change
@@ -2,44 +2,116 @@ namespace Microsoft.Azure.Cosmos.EmulatorTests.Query
22
{
33
using System;
44
using System.Collections.Generic;
5+
using System.Collections.ObjectModel;
56
using System.Linq;
67
using System.Threading.Tasks;
78
using Microsoft.Azure.Cosmos.CosmosElements;
9+
using Microsoft.Azure.Cosmos.Query.Core;
810
using Microsoft.VisualStudio.TestTools.UnitTesting;
911

1012
[TestClass]
1113
[TestCategory("Query")]
1214
public sealed class BypassQueryParsingTests : QueryTestsBase
1315
{
14-
[TestMethod]
15-
public async Task TestBypassQueryParsingWithNonePartitionKey()
16+
private const int DocumentCount = 400;
17+
18+
private static readonly Documents.PartitionKeyDefinition HierarchicalPartitionKeyDefinition = new Documents.PartitionKeyDefinition
1619
{
17-
int documentCount = 400;
20+
Paths = new Collection<string> { "/nullField", "/numberField" },
21+
Kind = Documents.PartitionKind.MultiHash,
22+
Version = Documents.PartitionKeyDefinitionVersion.V2
23+
};
1824

19-
string query = "SELECT VALUE r.numberField FROM r";
20-
IReadOnlyList<string> expectedOutput = Enumerable.Range(0, documentCount).Select(i => i.ToString()).ToList();
25+
private static readonly Documents.PartitionKeyDefinition SimplePartitionKeyDefinition = new Documents.PartitionKeyDefinition
26+
{
27+
Paths = new Collection<string> { "/UndefinedField" },
28+
Kind = Documents.PartitionKind.Hash
29+
};
2130

22-
await this.ValidateQueryBypassWithNonePartitionKey(documentCount, query, expectedOutput);
31+
[TestMethod]
32+
[TestCategory("Query")]
33+
public async Task TestBypassQueryParsing()
34+
{
35+
IReadOnlyList<QueryTestCase> testCases = new List<QueryTestCase>
36+
{
37+
new(
38+
PartitionKey.None,
39+
"SELECT VALUE r.numberField FROM r",
40+
Enumerable.Range(0, DocumentCount).Select(i => i.ToString()).ToList(),
41+
TestInjections.PipelineType.Passthrough),
42+
new(
43+
PartitionKey.None,
44+
@"SELECT VALUE { """" : r.numberField } FROM r",
45+
Enumerable.Range(0, DocumentCount).Select(i => String.Format("{{\"\":{0}}}", i)).ToList(),
46+
TestInjections.PipelineType.Passthrough),
47+
};
48+
49+
await this.ValidateQueryBypass(
50+
ConnectionModes.Direct | ConnectionModes.Gateway,
51+
CollectionTypes.NonPartitioned | CollectionTypes.SinglePartition | CollectionTypes.MultiPartition,
52+
SimplePartitionKeyDefinition,
53+
testCases);
2354
}
2455

2556
[TestMethod]
2657
[TestCategory("Query")]
27-
public async Task TestBypassQueryParsingWithNonePartitionKeyEmptyPropertyName()
58+
public async Task TestBypassQueryParsingWithHPK()
2859
{
29-
int documentCount = 400;
30-
31-
string query = @"SELECT VALUE { """" : r.numberField } FROM r";
32-
IReadOnlyList<string> expectedOutput = Enumerable.Range(0, documentCount).Select(i => String.Format("{{\"\":{0}}}", i)).ToList();
60+
PartitionKey partialPartitionKey = new PartitionKeyBuilder().AddNullValue().Build();
3361

34-
await this.ValidateQueryBypassWithNonePartitionKey(documentCount, query, expectedOutput);
62+
IReadOnlyList<QueryTestCase> testCases = new List<QueryTestCase>
63+
{
64+
new(
65+
partialPartitionKey,
66+
"SELECT VALUE r.numberField FROM r",
67+
Enumerable.Range(0, DocumentCount).Select(i => i.ToString()).ToList(),
68+
TestInjections.PipelineType.Passthrough), // Passthrough because it is a client streaming query
69+
new(
70+
partialPartitionKey,
71+
$"SELECT TOP {DocumentCount} VALUE r.numberField FROM r",
72+
Enumerable.Range(0, DocumentCount).Select(i => i.ToString()).ToList(),
73+
TestInjections.PipelineType.Specialized),
74+
new(
75+
partialPartitionKey,
76+
@"SELECT VALUE { """" : r.numberField } FROM r",
77+
Enumerable.Range(0, DocumentCount).Select(i => String.Format("{{\"\":{0}}}", i)).ToList(),
78+
TestInjections.PipelineType.Passthrough), // Passthrough because it is a client streaming query
79+
new(
80+
partialPartitionKey,
81+
"SELECT VALUE r.numberField FROM r ORDER BY r.numberField",
82+
Enumerable.Range(0, DocumentCount).Select(i => i.ToString()).ToList(),
83+
TestInjections.PipelineType.Specialized),
84+
};
85+
86+
await this.ValidateQueryBypass(
87+
ConnectionModes.Gateway,
88+
CollectionTypes.MultiPartition,
89+
HierarchicalPartitionKeyDefinition,
90+
testCases);
3591
}
3692

37-
private async Task ValidateQueryBypassWithNonePartitionKey(int documentCount, string query, IReadOnlyList<string> expectedOutput)
93+
private Task ValidateQueryBypass(ConnectionModes connectionModes, CollectionTypes collectionTypes, Documents.PartitionKeyDefinition partitionKeyDefinition, IReadOnlyList<QueryTestCase> testCases)
3894
{
39-
QueryRequestOptions feedOptions = new QueryRequestOptions { PartitionKey = PartitionKey.None };
95+
IReadOnlyList<string> documents = CreateDocuments(DocumentCount);
4096

41-
async Task ImplementationAsync(Container container, IReadOnlyList<CosmosObject> documents)
97+
return this.CreateIngestQueryDeleteAsync(
98+
connectionModes,
99+
collectionTypes,
100+
documents,
101+
query: (container, _) => RunTestsAsync(container, testCases),
102+
partitionKeyDefinition);
103+
}
104+
105+
private static async Task RunTestsAsync(Container container, IReadOnlyList<QueryTestCase> testCases)
106+
{
107+
foreach (QueryTestCase testCase in testCases)
42108
{
109+
QueryRequestOptions feedOptions = new QueryRequestOptions
110+
{
111+
PartitionKey = testCase.PartitionKey,
112+
TestSettings = new TestInjections(simulate429s: false, simulateEmptyPages: false, responseStats: new())
113+
};
114+
43115
ContainerInternal containerCore = container as ContainerInlineCore;
44116

45117
MockCosmosQueryClient cosmosQueryClientCore = new MockCosmosQueryClient(
@@ -53,20 +125,19 @@ async Task ImplementationAsync(Container container, IReadOnlyList<CosmosObject>
53125
containerCore.Id,
54126
cosmosQueryClientCore);
55127

56-
List<CosmosElement> items = await RunQueryAsync(containerWithBypassParsing, query, feedOptions);
128+
List<CosmosElement> items = await RunQueryAsync(containerWithBypassParsing, testCase.Query, feedOptions);
57129
string[] actualOutput = items.Select(x => x.ToString()).ToArray();
58130

59-
Assert.IsTrue(expectedOutput.SequenceEqual(actualOutput));
60-
}
131+
if(!testCase.ExpectedOutput.SequenceEqual(actualOutput))
132+
{
133+
System.Diagnostics.Trace.WriteLine($"Expected: [{string.Join(", ", testCase.ExpectedOutput)}]");
134+
System.Diagnostics.Trace.WriteLine($"Actual: [{string.Join(", ", actualOutput)}]");
61135

62-
IReadOnlyList<string> documents = CreateDocuments(documentCount);
136+
Assert.Fail("Query results do not match expected results.");
137+
}
63138

64-
await this.CreateIngestQueryDeleteAsync(
65-
ConnectionModes.Direct | ConnectionModes.Gateway,
66-
CollectionTypes.NonPartitioned | CollectionTypes.SinglePartition | CollectionTypes.MultiPartition,
67-
documents,
68-
ImplementationAsync,
69-
"/undefinedPartitionKey");
139+
Assert.AreEqual(testCase.ExpectedPipelineType, feedOptions.TestSettings.Stats.PipelineType);
140+
}
70141
}
71142

72143
private static IReadOnlyList<string> CreateDocuments(int documentCount)
@@ -80,5 +151,24 @@ private static IReadOnlyList<string> CreateDocuments(int documentCount)
80151

81152
return documents;
82153
}
154+
155+
private sealed class QueryTestCase
156+
{
157+
public PartitionKey PartitionKey { get; }
158+
159+
public string Query { get; }
160+
161+
public IReadOnlyList<string> ExpectedOutput { get; }
162+
163+
public TestInjections.PipelineType ExpectedPipelineType { get; }
164+
165+
public QueryTestCase(PartitionKey partitionKey, string query, IReadOnlyList<string> expectedOutput, TestInjections.PipelineType expectedPipelineType)
166+
{
167+
this.PartitionKey = partitionKey;
168+
this.Query = query ?? throw new ArgumentNullException(nameof(query));
169+
this.ExpectedOutput = expectedOutput ?? throw new ArgumentNullException(nameof(expectedOutput));
170+
this.ExpectedPipelineType = expectedPipelineType;
171+
}
172+
}
83173
}
84174
}

0 commit comments

Comments
 (0)