From e100f1ce581b60c8bdc1238c3e4188a4ed19bc28 Mon Sep 17 00:00:00 2001 From: j82w Date: Tue, 3 May 2022 07:54:39 -0700 Subject: [PATCH] Query: Fixes query plan via service interop to use custom serializer (#3154) Customer can pass a customer serializer in which can apply and data transformation necessary. For example it could convert a int or double to a string. Query has 3 different ways to get the query plan. 1. ServiceInterop.dll which requires Windows x64 2. Antlr parser 3. Gateway ### The custom serializer is used in query for the following scenarios. 1. ServiceInterop.dll-> Using a random [JsonConvert ](https://github.com/Azure/azure-cosmos-dotnet-v3/blob/c8935ac2f864fb829f5d941dda07c74aec86a677/Microsoft.Azure.Cosmos/src/Query/Core/QueryPlan/QueryPartitionProvider.cs#L169)instead of the [standard serialization contract](https://github.com/Azure/azure-cosmos-dotnet-v3/blob/c8935ac2f864fb829f5d941dda07c74aec86a677/Microsoft.Azure.Cosmos/src/Query/v3Query/CosmosQueryClientCore.cs#L123) 2. Antlr parser -> Does not take in parameters so serialization is not necessary. It just looks at query text. 3. Gateway -> Applies the[ correct serialization ](https://github.com/Azure/azure-cosmos-dotnet-v3/blob/c8935ac2f864fb829f5d941dda07c74aec86a677/Microsoft.Azure.Cosmos/src/Query/v3Query/CosmosQueryClientCore.cs#L163) 4. Executing the query(sending it to gateway or a partition): Applies the[ correct serialization ](https://github.com/Azure/azure-cosmos-dotnet-v3/blob/c8935ac2f864fb829f5d941dda07c74aec86a677/Microsoft.Azure.Cosmos/src/Query/v3Query/CosmosQueryClientCore.cs#L123) ### Why didn't testing catch this? There are existing[ tests which validate this scenario](https://github.com/Azure/azure-cosmos-dotnet-v3/blob/1d1d4c753cae896e6d96a98ef07a276cf1e4f130/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.EmulatorTests/CosmosItemTests.cs#L839), but they are not calculating the correct expected count. The test assumed the query pipeline would only serialize the SqlQuerySpec once. The implementation serializes it for every page request to the backend and to generate the query plan. ### Who is impacted? Only customers that are running on Windows x64 and have custom serialization that filters on the partition key value only in the query. If the partition key is provided in the request options then it is handled correctly because request option overrides the query text. ### Impact: This could result in the query being sent to the wrong partitions' which would cause it query returning possibly less results than expected because it would be routed to the wrong partition. ### Solution: Based on the current models and contract the serialization logic should not be in the query code. The query pipeline should not know or care about how the parameters are serialized. To keep this abstraction in place the serialization logic will be moved the CosmosQueryClientCore.cs like all of the other places that currently handle the custom serialization. This keeps all the serialization logic in the same file, and keeps it the same for all the different methods to get the query plan. This will cause the serialized string to be passed down instead of the SqlQuerySpec. This is a better contract because the service interop only requires the serialized string. All the contracts now match what is actually needed for it to execute getting the query plan. This follows the same model as getting it from gateway or sending it to be executed. --- .../CosmosQueryExecutionContextFactory.cs | 1 + .../Core/QueryClient/CosmosQueryClient.cs | 1 + .../Core/QueryPlan/QueryPartitionProvider.cs | 14 +++--- .../Query/Core/QueryPlan/QueryPlanHandler.cs | 6 +++ .../Core/QueryPlan/QueryPlanRetriever.cs | 2 + .../DefaultDocumentQueryExecutionContext.cs | 2 +- .../DocumentQueryExecutionContextBase.cs | 2 +- .../Query/v3Query/CosmosQueryClientCore.cs | 15 +++++- .../Resource/Container/ContainerCore.Items.cs | 1 + .../src/Routing/PartitionRoutingHelper.cs | 8 ++-- .../CosmosItemTests.cs | 46 +++++++++++++++---- .../CosmosJsonSerializerTests.cs | 9 ++-- .../Utils/ToDoActivity.cs | 26 +++++++++-- .../Query/ParsingBenchmark.cs | 3 +- .../Query/PassThroughQueryBaselineTests.cs | 16 +++++-- .../Query/Pipeline/FullPipelineTests.cs | 3 +- .../Query/QueryPlanBaselineTests.cs | 2 +- .../Query/QueryPlanRetrieverTests.cs | 5 ++ .../Routing/PartitionRoutingHelperTest.cs | 2 +- .../Tracing/TraceWriterBaselineTests.cs | 2 +- 20 files changed, 125 insertions(+), 41 deletions(-) diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/CosmosQueryExecutionContextFactory.cs b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/CosmosQueryExecutionContextFactory.cs index 5c97571e01..3d411f01e6 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/CosmosQueryExecutionContextFactory.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/CosmosQueryExecutionContextFactory.cs @@ -232,6 +232,7 @@ private static async Task> TryCreateCoreContextAsy partitionedQueryExecutionInfo = await QueryPlanRetriever.GetQueryPlanWithServiceInteropAsync( cosmosQueryContext.QueryClient, inputParameters.SqlQuerySpec, + cosmosQueryContext.ResourceTypeEnum, partitionKeyDefinition, inputParameters.PartitionKey != null, createQueryPipelineTrace, diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/QueryClient/CosmosQueryClient.cs b/Microsoft.Azure.Cosmos/src/Query/Core/QueryClient/CosmosQueryClient.cs index 342a887c60..0eedd2e162 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/QueryClient/CosmosQueryClient.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/QueryClient/CosmosQueryClient.cs @@ -39,6 +39,7 @@ public abstract Task GetCachedContainerQueryProperties public abstract Task> TryGetPartitionedQueryExecutionInfoAsync( SqlQuerySpec sqlQuerySpec, + Documents.ResourceType resourceType, Documents.PartitionKeyDefinition partitionKeyDefinition, bool requireFormattableOrderByQuery, bool isContinuationExpected, diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/QueryPlan/QueryPartitionProvider.cs b/Microsoft.Azure.Cosmos/src/Query/Core/QueryPlan/QueryPartitionProvider.cs index 5042b6f396..8d4fa07f71 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/QueryPlan/QueryPartitionProvider.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/QueryPlan/QueryPartitionProvider.cs @@ -104,7 +104,7 @@ public void Update(IDictionary queryengineConfiguration) } public TryCatch TryGetPartitionedQueryExecutionInfo( - SqlQuerySpec querySpec, + string querySpecJsonString, PartitionKeyDefinition partitionKeyDefinition, bool requireFormattableOrderByQuery, bool isContinuationExpected, @@ -113,7 +113,7 @@ public TryCatch TryGetPartitionedQueryExecutionIn bool allowDCount) { TryCatch tryGetInternalQueryInfo = this.TryGetPartitionedQueryExecutionInfoInternal( - querySpec: querySpec, + querySpecJsonString: querySpecJsonString, partitionKeyDefinition: partitionKeyDefinition, requireFormattableOrderByQuery: requireFormattableOrderByQuery, isContinuationExpected: isContinuationExpected, @@ -153,7 +153,7 @@ internal PartitionedQueryExecutionInfo ConvertPartitionedQueryExecutionInfo( } internal TryCatch TryGetPartitionedQueryExecutionInfoInternal( - SqlQuerySpec querySpec, + string querySpecJsonString, PartitionKeyDefinition partitionKeyDefinition, bool requireFormattableOrderByQuery, bool isContinuationExpected, @@ -161,13 +161,11 @@ internal TryCatch TryGetPartitionedQueryE bool hasLogicalPartitionKey, bool allowDCount) { - if (querySpec == null || partitionKeyDefinition == null) + if (querySpecJsonString == null || partitionKeyDefinition == null) { return TryCatch.FromResult(DefaultInfoInternal); } - string queryText = JsonConvert.SerializeObject(querySpec); - List paths = new List(partitionKeyDefinition.Paths); List> pathPartsList = new List>(paths.Count); uint[] partsLengths = new uint[paths.Count]; @@ -205,7 +203,7 @@ internal TryCatch TryGetPartitionedQueryE { errorCode = ServiceInteropWrapper.GetPartitionKeyRangesFromQuery2( this.serviceProvider, - queryText, + querySpecJsonString, requireFormattableOrderByQuery, isContinuationExpected, allowNonValueAggregateQuery, @@ -230,7 +228,7 @@ internal TryCatch TryGetPartitionedQueryE { errorCode = ServiceInteropWrapper.GetPartitionKeyRangesFromQuery2( this.serviceProvider, - queryText, + querySpecJsonString, requireFormattableOrderByQuery, isContinuationExpected, allowNonValueAggregateQuery, diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/QueryPlan/QueryPlanHandler.cs b/Microsoft.Azure.Cosmos/src/Query/Core/QueryPlan/QueryPlanHandler.cs index 8325050d92..ae18a2a101 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/QueryPlan/QueryPlanHandler.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/QueryPlan/QueryPlanHandler.cs @@ -24,6 +24,7 @@ public QueryPlanHandler(CosmosQueryClient queryClient) public async Task> TryGetQueryPlanAsync( SqlQuerySpec sqlQuerySpec, + Documents.ResourceType resourceType, PartitionKeyDefinition partitionKeyDefinition, QueryFeatures supportedQueryFeatures, bool hasLogicalPartitionKey, @@ -43,6 +44,7 @@ public async Task> TryGetQueryPlanAsync( TryCatch tryGetQueryInfo = await this.TryGetQueryInfoAsync( sqlQuerySpec, + resourceType, partitionKeyDefinition, hasLogicalPartitionKey, cancellationToken); @@ -68,6 +70,7 @@ public async Task> TryGetQueryPlanAsync( public async Task> TryGetQueryInfoAndIfSupportedAsync( QueryFeatures supportedQueryFeatures, SqlQuerySpec sqlQuerySpec, + Documents.ResourceType resourceType, PartitionKeyDefinition partitionKeyDefinition, bool hasLogicalPartitionKey, CancellationToken cancellationToken = default) @@ -86,6 +89,7 @@ public async Task> TryGetQueryPlanAsync( TryCatch tryGetQueryInfo = await this.TryGetQueryInfoAsync( sqlQuerySpec, + resourceType, partitionKeyDefinition, hasLogicalPartitionKey, cancellationToken); @@ -102,6 +106,7 @@ public async Task> TryGetQueryPlanAsync( private Task> TryGetQueryInfoAsync( SqlQuerySpec sqlQuerySpec, + Documents.ResourceType resourceType, PartitionKeyDefinition partitionKeyDefinition, bool hasLogicalPartitionKey, CancellationToken cancellationToken = default) @@ -110,6 +115,7 @@ private Task> TryGetQueryInfoAsync( return this.queryClient.TryGetPartitionedQueryExecutionInfoAsync( sqlQuerySpec: sqlQuerySpec, + resourceType: resourceType, partitionKeyDefinition: partitionKeyDefinition, requireFormattableOrderByQuery: true, isContinuationExpected: false, diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/QueryPlan/QueryPlanRetriever.cs b/Microsoft.Azure.Cosmos/src/Query/Core/QueryPlan/QueryPlanRetriever.cs index 24a56d82cd..947046b894 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/QueryPlan/QueryPlanRetriever.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/QueryPlan/QueryPlanRetriever.cs @@ -34,6 +34,7 @@ internal static class QueryPlanRetriever public static async Task GetQueryPlanWithServiceInteropAsync( CosmosQueryClient queryClient, SqlQuerySpec sqlQuerySpec, + Documents.ResourceType resourceType, PartitionKeyDefinition partitionKeyDefinition, bool hasLogicalPartitionKey, ITrace trace, @@ -62,6 +63,7 @@ public static async Task GetQueryPlanWithServiceI TryCatch tryGetQueryPlan = await queryPlanHandler.TryGetQueryPlanAsync( sqlQuerySpec, + resourceType, partitionKeyDefinition, QueryPlanRetriever.SupportedQueryFeatures, hasLogicalPartitionKey, diff --git a/Microsoft.Azure.Cosmos/src/Query/v2Query/DefaultDocumentQueryExecutionContext.cs b/Microsoft.Azure.Cosmos/src/Query/v2Query/DefaultDocumentQueryExecutionContext.cs index 1f6d86f87c..86ccaf0334 100644 --- a/Microsoft.Azure.Cosmos/src/Query/v2Query/DefaultDocumentQueryExecutionContext.cs +++ b/Microsoft.Azure.Cosmos/src/Query/v2Query/DefaultDocumentQueryExecutionContext.cs @@ -308,7 +308,7 @@ private static bool ServiceInteropAvailable() } providedRanges = PartitionRoutingHelper.GetProvidedPartitionKeyRanges( - querySpec: this.QuerySpec, + querySpecJsonString: JsonConvert.SerializeObject(this.QuerySpec), enableCrossPartitionQuery: enableCrossPartitionQuery, parallelizeCrossPartitionQuery: false, isContinuationExpected: this.isContinuationExpected, diff --git a/Microsoft.Azure.Cosmos/src/Query/v2Query/DocumentQueryExecutionContextBase.cs b/Microsoft.Azure.Cosmos/src/Query/v2Query/DocumentQueryExecutionContextBase.cs index 055869358e..c2c0333f74 100644 --- a/Microsoft.Azure.Cosmos/src/Query/v2Query/DocumentQueryExecutionContextBase.cs +++ b/Microsoft.Azure.Cosmos/src/Query/v2Query/DocumentQueryExecutionContextBase.cs @@ -167,7 +167,7 @@ public async Task GetPartitionedQueryExecutionInf QueryPartitionProvider queryPartitionProvider = await this.Client.GetQueryPartitionProviderAsync(); TryCatch tryGetPartitionedQueryExecutionInfo = queryPartitionProvider.TryGetPartitionedQueryExecutionInfo( - querySpec: this.QuerySpec, + querySpecJsonString: JsonConvert.SerializeObject(this.QuerySpec), partitionKeyDefinition: partitionKeyDefinition, requireFormattableOrderByQuery: requireFormattableOrderByQuery, isContinuationExpected: isContinuationExpected, diff --git a/Microsoft.Azure.Cosmos/src/Query/v3Query/CosmosQueryClientCore.cs b/Microsoft.Azure.Cosmos/src/Query/v3Query/CosmosQueryClientCore.cs index 04e71d02ba..b194dfd59e 100644 --- a/Microsoft.Azure.Cosmos/src/Query/v3Query/CosmosQueryClientCore.cs +++ b/Microsoft.Azure.Cosmos/src/Query/v3Query/CosmosQueryClientCore.cs @@ -79,6 +79,7 @@ public override async Task GetCachedContainerQueryProp public override async Task> TryGetPartitionedQueryExecutionInfoAsync( SqlQuerySpec sqlQuerySpec, + ResourceType resourceType, PartitionKeyDefinition partitionKeyDefinition, bool requireFormattableOrderByQuery, bool isContinuationExpected, @@ -87,8 +88,20 @@ public override async Task> TryGetPartit bool allowDCount, CancellationToken cancellationToken) { + string queryString = null; + if (sqlQuerySpec != null) + { + using (Stream stream = this.clientContext.SerializerCore.ToStreamSqlQuerySpec(sqlQuerySpec, resourceType)) + { + using (StreamReader reader = new StreamReader(stream)) + { + queryString = reader.ReadToEnd(); + } + } + } + return (await this.documentClient.QueryPartitionProvider).TryGetPartitionedQueryExecutionInfo( - querySpec: sqlQuerySpec, + querySpecJsonString: queryString, partitionKeyDefinition: partitionKeyDefinition, requireFormattableOrderByQuery: requireFormattableOrderByQuery, isContinuationExpected: isContinuationExpected, diff --git a/Microsoft.Azure.Cosmos/src/Resource/Container/ContainerCore.Items.cs b/Microsoft.Azure.Cosmos/src/Resource/Container/ContainerCore.Items.cs index d88714a724..5cbd91fb9f 100644 --- a/Microsoft.Azure.Cosmos/src/Resource/Container/ContainerCore.Items.cs +++ b/Microsoft.Azure.Cosmos/src/Resource/Container/ContainerCore.Items.cs @@ -414,6 +414,7 @@ public override async Task TryExecuteQueryAsync( TryCatch<(PartitionedQueryExecutionInfo queryPlan, bool supported)> tryGetQueryInfoAndIfSupported = await queryPlanHandler.TryGetQueryInfoAndIfSupportedAsync( supportedQueryFeatures, queryDefinition.ToSqlQuerySpec(), + ResourceType.Document, partitionKeyDefinition, requestOptions.PartitionKey.HasValue, cancellationToken); diff --git a/Microsoft.Azure.Cosmos/src/Routing/PartitionRoutingHelper.cs b/Microsoft.Azure.Cosmos/src/Routing/PartitionRoutingHelper.cs index 04a33b3871..60a6b56a20 100644 --- a/Microsoft.Azure.Cosmos/src/Routing/PartitionRoutingHelper.cs +++ b/Microsoft.Azure.Cosmos/src/Routing/PartitionRoutingHelper.cs @@ -26,7 +26,7 @@ namespace Microsoft.Azure.Cosmos.Routing internal class PartitionRoutingHelper { public static IReadOnlyList> GetProvidedPartitionKeyRanges( - SqlQuerySpec querySpec, + string querySpecJsonString, bool enableCrossPartitionQuery, bool parallelizeCrossPartitionQuery, bool isContinuationExpected, @@ -38,9 +38,9 @@ public static IReadOnlyList> GetProvidedPartitionKeyRanges( string clientApiVersion, out QueryInfo queryInfo) { - if (querySpec == null) + if (querySpecJsonString == null) { - throw new ArgumentNullException(nameof(querySpec)); + throw new ArgumentNullException(nameof(querySpecJsonString)); } if (partitionKeyDefinition == null) @@ -54,7 +54,7 @@ public static IReadOnlyList> GetProvidedPartitionKeyRanges( } TryCatch tryGetPartitionQueryExecutionInfo = queryPartitionProvider.TryGetPartitionedQueryExecutionInfo( - querySpec: querySpec, + querySpecJsonString: querySpecJsonString, partitionKeyDefinition: partitionKeyDefinition, requireFormattableOrderByQuery: VersionUtility.IsLaterThan(clientApiVersion, HttpConstants.VersionDates.v2016_11_14), isContinuationExpected: isContinuationExpected, diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.EmulatorTests/CosmosItemTests.cs b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.EmulatorTests/CosmosItemTests.cs index 7da4502a5e..f5b3d2c700 100644 --- a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.EmulatorTests/CosmosItemTests.cs +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.EmulatorTests/CosmosItemTests.cs @@ -54,6 +54,7 @@ public async Task TestInitialize() this.containerSettings = new ContainerProperties(id: Guid.NewGuid().ToString(), partitionKeyPath: PartitionKey); ContainerResponse response = await this.database.CreateContainerAsync( this.containerSettings, + throughput: 15000, cancellationToken: this.cancellationToken); Assert.IsNotNull(response); Assert.IsNotNull(response.Container); @@ -882,7 +883,7 @@ public async Task ItemCustomSerialzierTest() // Each parameter in query spec should be a call to the custom serializer int parameterCount = queryDefinition.ToSqlQuerySpec().Parameters.Count; - Assert.AreEqual(parameterCount, toStreamCount, $"missing to stream call. Expected: {parameterCount}, Actual: {toStreamCount} for query:{queryDefinition.ToSqlQuerySpec().QueryText}"); + Assert.AreEqual((parameterCount*pageCount)+parameterCount, toStreamCount, $"missing to stream call. Expected: {(parameterCount * pageCount) + parameterCount}, Actual: {toStreamCount} for query:{queryDefinition.ToSqlQuerySpec().QueryText}"); Assert.AreEqual(pageCount, fromStreamCount); } } @@ -1094,7 +1095,15 @@ public async Task QueryStreamValueTest() } Assert.AreEqual(2, allItems.Count, $"missing query results. Only found: {allItems.Count} items for query:{queryDefinition.ToSqlQuerySpec().QueryText}"); - Assert.AreEqual(pageCount, 1); + if (queryDefinition.QueryText.Contains("pk")) + { + Assert.AreEqual(1, pageCount); + } + else + { + Assert.AreEqual(3, pageCount); + } + IReadOnlyList<(string Name, object Value)> parameters1 = queryDefinition.GetQueryParameters(); @@ -1240,13 +1249,20 @@ public async Task ItemMultiplePartitionQuery() sql, requestOptions: requestOptions); + bool found = false; while (feedIterator.HasMoreResults) { FeedResponse iter = await feedIterator.ReadNextAsync(); - Assert.AreEqual(1, iter.Count()); - ToDoActivity response = iter.First(); - Assert.AreEqual(find.id, response.id); + Assert.IsTrue(iter.Count() <= 1); + if(iter.Count() == 1) + { + found = true; + ToDoActivity response = iter.First(); + Assert.AreEqual(find.id, response.id); + } } + + Assert.IsTrue(found); } /// @@ -1263,7 +1279,11 @@ public async Task ItemMultiplePartitionOrderByQueryStream() System.Globalization.CultureInfo.GetCultureInfo("fr-FR") }; - IList deleteList = await ToDoActivity.CreateRandomItems(this.Container, 300, randomPartitionKey: true); + IList deleteList = await ToDoActivity.CreateRandomItems( + this.Container, + 300, + randomPartitionKey: true, + randomTaskNumber: true); try { @@ -1349,8 +1369,8 @@ public async Task ItemMultiplePartitionQueryStream() Assert.AreEqual(deleteList.Count, resultList.Count); Assert.IsTrue(totalRequstCharge > 0); - List verifiedOrderBy = deleteList.OrderBy(x => x.taskNum).ToList(); - resultList = resultList.OrderBy(x => x.taskNum).ToList(); + List verifiedOrderBy = deleteList.OrderBy(x => x.id).ToList(); + resultList = resultList.OrderBy(x => x.id).ToList(); for (int i = 0; i < verifiedOrderBy.Count(); i++) { Assert.AreEqual(verifiedOrderBy[i].taskNum, resultList[i].taskNum); @@ -1432,6 +1452,11 @@ public async Task EpkPointReadTest() .InternalKey .GetEffectivePartitionKeyString(this.containerSettings.PartitionKey); + properties = new Dictionary() + { + { WFConstants.BackendHeaders.EffectivePartitionKeyString, epk }, + }; + QueryRequestOptions queryRequestOptions = new QueryRequestOptions { IsEffectivePartitionKeyRouting = true, @@ -1511,7 +1536,10 @@ public async Task ItemEpkQuerySingleKeyRangeValidation() [TestMethod] public async Task ItemQueryStreamSerializationSetting() { - IList deleteList = await ToDoActivity.CreateRandomItems(this.Container, 101, randomPartitionKey: true); + IList deleteList = await ToDoActivity.CreateRandomItems( + container: this.Container, + pkCount: 101, + randomTaskNumber: true); QueryDefinition sql = new QueryDefinition("SELECT * FROM toDoActivity t ORDER BY t.taskNum"); diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.EmulatorTests/CosmosJsonSerializerTests.cs b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.EmulatorTests/CosmosJsonSerializerTests.cs index 235116006e..0be44fbc1c 100644 --- a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.EmulatorTests/CosmosJsonSerializerTests.cs +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.EmulatorTests/CosmosJsonSerializerTests.cs @@ -27,7 +27,10 @@ public async Task TestInitialize() await base.TestInit(); string PartitionKey = "/pk"; ContainerResponse response = await this.database.CreateContainerAsync( - new ContainerProperties(id: Guid.NewGuid().ToString(), partitionKeyPath: PartitionKey), + new ContainerProperties( + id: Guid.NewGuid().ToString(), + partitionKeyPath: PartitionKey), + throughput: 15000, cancellationToken: this.cancellationToken); Assert.IsNotNull(response); this.container = response; @@ -102,7 +105,7 @@ public async Task TestQueryWithCustomJsonSerializer() items.AddRange(await itemIterator.ReadNextAsync()); } - Assert.AreEqual(1, toStreamCount); + Assert.AreEqual(2, toStreamCount); Assert.AreEqual(1, fromStreamCount); toStreamCount = 0; @@ -122,7 +125,7 @@ public async Task TestQueryWithCustomJsonSerializer() } } - Assert.AreEqual(1, toStreamCount); + Assert.AreEqual(2, toStreamCount); Assert.AreEqual(0, fromStreamCount); } diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.EmulatorTests/Utils/ToDoActivity.cs b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.EmulatorTests/Utils/ToDoActivity.cs index 3b075f85a6..c2800c02d2 100644 --- a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.EmulatorTests/Utils/ToDoActivity.cs +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.EmulatorTests/Utils/ToDoActivity.cs @@ -43,7 +43,12 @@ public override int GetHashCode() return base.GetHashCode(); } - public static async Task> CreateRandomItems(Container container, int pkCount, int perPKItemCount = 1, bool randomPartitionKey = true) + public static async Task> CreateRandomItems( + Container container, + int pkCount, + int perPKItemCount = 1, + bool randomPartitionKey = true, + bool randomTaskNumber = false) { Assert.IsFalse(!randomPartitionKey && pkCount > 1); @@ -58,7 +63,10 @@ public static async Task> CreateRandomItems(Container contai for (int j = 0; j < perPKItemCount; j++) { - ToDoActivity temp = ToDoActivity.CreateRandomToDoActivity(pk); + ToDoActivity temp = ToDoActivity.CreateRandomToDoActivity( + pk: pk, + id: null, + randomTaskNumber: randomTaskNumber); createdList.Add(temp); @@ -69,7 +77,10 @@ public static async Task> CreateRandomItems(Container contai return createdList; } - public static ToDoActivity CreateRandomToDoActivity(string pk = null, string id = null) + public static ToDoActivity CreateRandomToDoActivity( + string pk = null, + string id = null, + bool randomTaskNumber = false) { if (string.IsNullOrEmpty(pk)) { @@ -79,12 +90,19 @@ public static ToDoActivity CreateRandomToDoActivity(string pk = null, string id { id = Guid.NewGuid().ToString(); } + + int taskNum = 42; + if (randomTaskNumber) + { + taskNum = Random.Shared.Next(); + } + return new ToDoActivity() { id = id, description = "CreateRandomToDoActivity", pk = pk, - taskNum = 42, + taskNum = taskNum, cost = double.MaxValue, CamelCase = "camelCase", children = new ToDoActivity[] diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Performance.Tests/Query/ParsingBenchmark.cs b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Performance.Tests/Query/ParsingBenchmark.cs index cc24ddc86e..5ef8965a36 100644 --- a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Performance.Tests/Query/ParsingBenchmark.cs +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Performance.Tests/Query/ParsingBenchmark.cs @@ -14,6 +14,7 @@ namespace Microsoft.Azure.Cosmos.Query using Microsoft.Azure.Cosmos.Query.Core.QueryPlan; using Microsoft.Azure.Cosmos.SqlObjects; using Microsoft.Azure.Documents; + using Newtonsoft.Json; [MemoryDiagnoser] public class ParsingBenchmark @@ -118,7 +119,7 @@ private static void ParseUsingMangedParser(SqlQuerySpec sqlQuerySpec) private static void ParseUsingNativeParser(SqlQuerySpec sqlQuerySpec) { TryCatch tryGetQueryPlan = QueryPartitionProvider.TryGetPartitionedQueryExecutionInfo( - querySpec: sqlQuerySpec, + querySpecJsonString: JsonConvert.SerializeObject(sqlQuerySpec), partitionKeyDefinition: PartitionKeyDefinition, requireFormattableOrderByQuery: true, isContinuationExpected: false, diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/PassThroughQueryBaselineTests.cs b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/PassThroughQueryBaselineTests.cs index 3f1ffd6e87..29ee795ffc 100644 --- a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/PassThroughQueryBaselineTests.cs +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/PassThroughQueryBaselineTests.cs @@ -29,6 +29,7 @@ using System.Threading; using System.Linq; using Microsoft.Azure.Cosmos.Query.Core.Pipeline.Pagination; + using System.IO; [TestClass] public class PassThroughQueryBaselineTests : BaselineTests @@ -116,10 +117,10 @@ private static PassThroughQueryTestInput CreateInput( return new PassThroughQueryTestInput(description, query, new SqlQuerySpec(query), expectedPassThrough, partitionKeyPath, partitionKeyValue); } - private static PartitionedQueryExecutionInfo GetPartitionedQueryExecutionInfo(SqlQuerySpec sqlQuerySpec, PartitionKeyDefinition pkDefinition) + private static PartitionedQueryExecutionInfo GetPartitionedQueryExecutionInfo(string querySpecJsonString, PartitionKeyDefinition pkDefinition) { TryCatch tryGetQueryPlan = QueryPartitionProviderTestInstance.Object.TryGetPartitionedQueryExecutionInfo( - querySpec: sqlQuerySpec, + querySpecJsonString: querySpecJsonString, partitionKeyDefinition: pkDefinition, requireFormattableOrderByQuery: true, isContinuationExpected: false, @@ -153,8 +154,13 @@ public override PassThroughQueryTestOutput ExecuteTest(PassThroughQueryTestInput // gets input parameters QueryRequestOptions queryRequestOptions = new QueryRequestOptions(); - PartitionedQueryExecutionInfo partitionedQueryExecutionInfo = GetPartitionedQueryExecutionInfo(sqlQuerySpec, input.PartitionKeyDefinition); - if (input.PartitionKeyValue == null || input.PartitionKeyValue == Cosmos.PartitionKey.None) + + CosmosSerializerCore serializerCore = new(); + using StreamReader streamReader = new(serializerCore.ToStreamSqlQuerySpec(sqlQuerySpec, Documents.ResourceType.Document)); + string sqlQuerySpecJsonString = streamReader.ReadToEnd(); + + PartitionedQueryExecutionInfo partitionedQueryExecutionInfo = GetPartitionedQueryExecutionInfo(sqlQuerySpecJsonString, input.PartitionKeyDefinition); + if (input.PartitionKeyValue == default || input.PartitionKeyValue == Cosmos.PartitionKey.None) { input.PartitionKeyValue = Cosmos.PartitionKey.Null; } @@ -331,7 +337,7 @@ public override Task> TryGetOverlappingRangesAs throw new NotImplementedException(); } - public override Task> TryGetPartitionedQueryExecutionInfoAsync(SqlQuerySpec sqlQuerySpec, PartitionKeyDefinition partitionKeyDefinition, bool requireFormattableOrderByQuery, bool isContinuationExpected, bool allowNonValueAggregateQuery, bool hasLogicalPartitionKey, bool allowDCount, CancellationToken cancellationToken) + public override Task> TryGetPartitionedQueryExecutionInfoAsync(SqlQuerySpec sqlQuerySpec, ResourceType resourceType, PartitionKeyDefinition partitionKeyDefinition, bool requireFormattableOrderByQuery, bool isContinuationExpected, bool allowNonValueAggregateQuery, bool hasLogicalPartitionKey, bool allowDCount, CancellationToken cancellationToken) { throw new NotImplementedException(); } diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/Pipeline/FullPipelineTests.cs b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/Pipeline/FullPipelineTests.cs index 7b7f86db75..79fef002a8 100644 --- a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/Pipeline/FullPipelineTests.cs +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/Pipeline/FullPipelineTests.cs @@ -20,6 +20,7 @@ namespace Microsoft.Azure.Cosmos.Tests.Query.Pipeline using Microsoft.Azure.Cosmos.Tracing; using Microsoft.Azure.Documents; using Microsoft.VisualStudio.TestTools.UnitTesting; + using Newtonsoft.Json; [TestClass] public class FullPipelineTests @@ -388,7 +389,7 @@ private static async Task CreatePipelineAsync( private static QueryInfo GetQueryPlan(string query) { TryCatch info = QueryPartitionProviderTestInstance.Object.TryGetPartitionedQueryExecutionInfoInternal( - new SqlQuerySpec(query), + JsonConvert.SerializeObject(new SqlQuerySpec(query)), partitionKeyDefinition, requireFormattableOrderByQuery: true, isContinuationExpected: false, diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/QueryPlanBaselineTests.cs b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/QueryPlanBaselineTests.cs index fc98d4c148..1108a7cfb7 100644 --- a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/QueryPlanBaselineTests.cs +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/QueryPlanBaselineTests.cs @@ -1379,7 +1379,7 @@ private static QueryPlanBaselineTestInput Range( public override QueryPlanBaselineTestOutput ExecuteTest(QueryPlanBaselineTestInput input) { TryCatch info = QueryPartitionProviderTestInstance.Object.TryGetPartitionedQueryExecutionInfoInternal( - input.SqlQuerySpec, + JsonConvert.SerializeObject(input.SqlQuerySpec), input.PartitionKeyDefinition, requireFormattableOrderByQuery: true, isContinuationExpected: false, diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/QueryPlanRetrieverTests.cs b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/QueryPlanRetrieverTests.cs index 853f705264..55c3e871de 100644 --- a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/QueryPlanRetrieverTests.cs +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/QueryPlanRetrieverTests.cs @@ -15,6 +15,7 @@ namespace Microsoft.Azure.Cosmos.Tests.Query using Microsoft.Azure.Cosmos.Query.Core.QueryClient; using Microsoft.Azure.Cosmos.Query.Core.QueryPlan; using Microsoft.Azure.Cosmos.Tracing; + using Microsoft.Azure.Documents; using Microsoft.VisualStudio.TestTools.UnitTesting; using Moq; @@ -29,6 +30,7 @@ public async Task ServiceInterop_BadRequestContainsInnerException() queryClient.Setup(c => c.TryGetPartitionedQueryExecutionInfoAsync( It.IsAny(), + It.IsAny(), It.IsAny(), It.IsAny(), It.IsAny(), @@ -41,6 +43,7 @@ public async Task ServiceInterop_BadRequestContainsInnerException() CosmosException cosmosException = await Assert.ThrowsExceptionAsync(() => QueryPlanRetriever.GetQueryPlanWithServiceInteropAsync( queryClient.Object, new SqlQuerySpec("selectttttt * from c"), + ResourceType.Document, new Documents.PartitionKeyDefinition() { Paths = new Collection() { "/id" } }, hasLogicalPartitionKey: false, trace.Object, @@ -60,6 +63,7 @@ public async Task ServiceInterop_BadRequestContainsOriginalCosmosException() queryClient.Setup(c => c.TryGetPartitionedQueryExecutionInfoAsync( It.IsAny(), + It.IsAny(), It.IsAny(), It.IsAny(), It.IsAny(), @@ -72,6 +76,7 @@ public async Task ServiceInterop_BadRequestContainsOriginalCosmosException() CosmosException cosmosException = await Assert.ThrowsExceptionAsync(() => QueryPlanRetriever.GetQueryPlanWithServiceInteropAsync( queryClient.Object, new SqlQuerySpec("selectttttt * from c"), + ResourceType.Document, new Documents.PartitionKeyDefinition() { Paths = new Collection() { "/id" } }, hasLogicalPartitionKey: false, trace.Object, diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Routing/PartitionRoutingHelperTest.cs b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Routing/PartitionRoutingHelperTest.cs index 81cc135671..775b179685 100644 --- a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Routing/PartitionRoutingHelperTest.cs +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Routing/PartitionRoutingHelperTest.cs @@ -337,7 +337,7 @@ FROM c try { IReadOnlyList> _ = PartitionRoutingHelper.GetProvidedPartitionKeyRanges( - querySpec: new Cosmos.Query.Core.SqlQuerySpec(testcase.Query), + querySpecJsonString: JsonConvert.SerializeObject(new Cosmos.Query.Core.SqlQuerySpec(testcase.Query)), enableCrossPartitionQuery: testcase.EnableCrossPartitionQuery, parallelizeCrossPartitionQuery: false, isContinuationExpected: true, diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Tracing/TraceWriterBaselineTests.cs b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Tracing/TraceWriterBaselineTests.cs index 174dcf7992..ede11316e9 100644 --- a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Tracing/TraceWriterBaselineTests.cs +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Tracing/TraceWriterBaselineTests.cs @@ -806,7 +806,7 @@ private static IQueryPipelineStage CreatePipeline(IDocumentContainer documentCon private static QueryInfo GetQueryPlan(string query) { TryCatch info = QueryPartitionProviderTestInstance.Object.TryGetPartitionedQueryExecutionInfoInternal( - new SqlQuerySpec(query), + Newtonsoft.Json.JsonConvert.SerializeObject(new SqlQuerySpec(query)), partitionKeyDefinition, requireFormattableOrderByQuery: true, isContinuationExpected: false,