From c94150ca29696235e0cbe0c07ca5fbc7203de253 Mon Sep 17 00:00:00 2001 From: Neil Deshpande Date: Tue, 14 May 2024 12:09:23 -0700 Subject: [PATCH 1/3] Introduce an environment config to suppress sending NonStreamingOrderBy in the list of query features sent to the gateway --- .../CosmosQueryExecutionContextFactory.cs | 6 + .../Core/QueryPlan/QueryPlanRetriever.cs | 34 ++- .../src/Query/v3Query/QueryIterator.cs | 3 +- .../src/RequestOptions/QueryRequestOptions.cs | 4 +- .../src/Util/ConfigurationManager.cs | 20 +- .../Query/NonStreamingOrderByQueryTests.cs | 245 ++++++++++++++++++ ...misticDirectExecutionQueryBaselineTests.cs | 3 +- .../Query/Pipeline/FullPipelineTests.cs | 3 +- .../Query/SubpartitionTests.cs | 3 +- 9 files changed, 312 insertions(+), 9 deletions(-) diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/CosmosQueryExecutionContextFactory.cs b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/CosmosQueryExecutionContextFactory.cs index 7676270416..042c8a7606 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/CosmosQueryExecutionContextFactory.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/CosmosQueryExecutionContextFactory.cs @@ -408,6 +408,7 @@ private static TryCatch TryCreateSpecializedDocumentQueryEx inputParameters.ExecutionEnvironment, inputParameters.ReturnResultsInDeterministicOrder, inputParameters.EnableOptimisticDirectExecution, + inputParameters.IsNonStreamingOrderByQueryFeatureDisabled, inputParameters.TestInjections); } @@ -593,6 +594,7 @@ private static async Task GetPartitionedQueryExec inputParameters.SqlQuerySpec, cosmosQueryContext.ResourceLink, inputParameters.PartitionKey, + inputParameters.IsNonStreamingOrderByQueryFeatureDisabled, trace, cancellationToken); } @@ -841,6 +843,7 @@ public InputParameters( ExecutionEnvironment? executionEnvironment, bool? returnResultsInDeterministicOrder, bool enableOptimisticDirectExecution, + bool isNonStreamingOrderByQueryFeatureDisabled, TestInjections testInjections) { this.SqlQuerySpec = sqlQuerySpec ?? throw new ArgumentNullException(nameof(sqlQuerySpec)); @@ -874,6 +877,7 @@ public InputParameters( this.ExecutionEnvironment = executionEnvironment.GetValueOrDefault(InputParameters.DefaultExecutionEnvironment); this.ReturnResultsInDeterministicOrder = returnResultsInDeterministicOrder.GetValueOrDefault(InputParameters.DefaultReturnResultsInDeterministicOrder); this.EnableOptimisticDirectExecution = enableOptimisticDirectExecution; + this.IsNonStreamingOrderByQueryFeatureDisabled = isNonStreamingOrderByQueryFeatureDisabled; this.TestInjections = testInjections; } @@ -890,6 +894,7 @@ public InputParameters( public bool ReturnResultsInDeterministicOrder { get; } public TestInjections TestInjections { get; } public bool EnableOptimisticDirectExecution { get; } + public bool IsNonStreamingOrderByQueryFeatureDisabled { get; } public InputParameters WithContinuationToken(CosmosElement token) { @@ -906,6 +911,7 @@ public InputParameters WithContinuationToken(CosmosElement token) this.ExecutionEnvironment, this.ReturnResultsInDeterministicOrder, this.EnableOptimisticDirectExecution, + this.IsNonStreamingOrderByQueryFeatureDisabled, this.TestInjections); } } diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/QueryPlan/QueryPlanRetriever.cs b/Microsoft.Azure.Cosmos/src/Query/Core/QueryPlan/QueryPlanRetriever.cs index 09ad98f395..9cbf6918fc 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/QueryPlan/QueryPlanRetriever.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/QueryPlan/QueryPlanRetriever.cs @@ -19,6 +19,19 @@ namespace Microsoft.Azure.Cosmos.Query.Core.QueryPlan internal static class QueryPlanRetriever { private static readonly QueryFeatures SupportedQueryFeatures = + QueryFeatures.Aggregate + | QueryFeatures.Distinct + | QueryFeatures.GroupBy + | QueryFeatures.MultipleOrderBy + | QueryFeatures.MultipleAggregates + | QueryFeatures.OffsetAndLimit + | QueryFeatures.OrderBy + | QueryFeatures.Top + | QueryFeatures.NonValueAggregate + | QueryFeatures.DCount + | QueryFeatures.NonStreamingOrderBy; + + private static readonly QueryFeatures LegacySupportedQueryFeatures = QueryFeatures.Aggregate | QueryFeatures.Distinct | QueryFeatures.GroupBy @@ -30,7 +43,21 @@ internal static class QueryPlanRetriever | QueryFeatures.NonValueAggregate | QueryFeatures.DCount; - private static readonly string SupportedQueryFeaturesString = SupportedQueryFeatures.ToString(); + private static readonly string SupportedQueryFeaturesString = SupportedQueryFeatures.ToString(); + + private static readonly string LegacySupportedQueryFeaturesString = LegacySupportedQueryFeatures.ToString(); + + private static string GetSupportedQueryFeaturesString(bool isNonStreamingOrderByQueryFeatureDisabled) + { + if (isNonStreamingOrderByQueryFeatureDisabled) + { + return LegacySupportedQueryFeaturesString; + } + else + { + return SupportedQueryFeaturesString; + } + } public static async Task GetQueryPlanWithServiceInteropAsync( CosmosQueryClient queryClient, @@ -94,7 +121,8 @@ public static Task GetQueryPlanThroughGatewayAsyn CosmosQueryContext queryContext, SqlQuerySpec sqlQuerySpec, string resourceLink, - PartitionKey? partitionKey, + PartitionKey? partitionKey, + bool isNonStreamingOrderByQueryFeatureDisabled, ITrace trace, CancellationToken cancellationToken = default) { @@ -130,7 +158,7 @@ public static Task GetQueryPlanThroughGatewayAsyn OperationType.QueryPlan, sqlQuerySpec, partitionKey, - QueryPlanRetriever.SupportedQueryFeaturesString, + GetSupportedQueryFeaturesString(isNonStreamingOrderByQueryFeatureDisabled), trace, cancellationToken); } diff --git a/Microsoft.Azure.Cosmos/src/Query/v3Query/QueryIterator.cs b/Microsoft.Azure.Cosmos/src/Query/v3Query/QueryIterator.cs index fe540f222e..abfe5fbb44 100644 --- a/Microsoft.Azure.Cosmos/src/Query/v3Query/QueryIterator.cs +++ b/Microsoft.Azure.Cosmos/src/Query/v3Query/QueryIterator.cs @@ -144,7 +144,8 @@ public static QueryIterator Create( partitionedQueryExecutionInfo: partitionedQueryExecutionInfo, executionEnvironment: queryRequestOptions.ExecutionEnvironment, returnResultsInDeterministicOrder: queryRequestOptions.ReturnResultsInDeterministicOrder, - enableOptimisticDirectExecution: queryRequestOptions.EnableOptimisticDirectExecution, + enableOptimisticDirectExecution: queryRequestOptions.EnableOptimisticDirectExecution, + isNonStreamingOrderByQueryFeatureDisabled: queryRequestOptions.IsNonStreamingOrderByQueryFeatureDisabled, testInjections: queryRequestOptions.TestSettings); return new QueryIterator( diff --git a/Microsoft.Azure.Cosmos/src/RequestOptions/QueryRequestOptions.cs b/Microsoft.Azure.Cosmos/src/RequestOptions/QueryRequestOptions.cs index ad15e0cd13..fd697c09ad 100644 --- a/Microsoft.Azure.Cosmos/src/RequestOptions/QueryRequestOptions.cs +++ b/Microsoft.Azure.Cosmos/src/RequestOptions/QueryRequestOptions.cs @@ -50,7 +50,9 @@ public class QueryRequestOptions : RequestOptions /// /// Direct (optimistic) execution offers improved performance for several kinds of queries such as a single partition streaming query. /// - public bool EnableOptimisticDirectExecution { get; set; } = ConfigurationManager.IsOptimisticDirectExecutionEnabled(defaultValue: false); + public bool EnableOptimisticDirectExecution { get; set; } = ConfigurationManager.IsOptimisticDirectExecutionEnabled(defaultValue: false); + + internal bool IsNonStreamingOrderByQueryFeatureDisabled { get; set; } = ConfigurationManager.IsNonStreamingOrderByQueryFeatureDisabled(defaultValue: false); /// /// Gets or sets the maximum number of items that can be buffered client side during diff --git a/Microsoft.Azure.Cosmos/src/Util/ConfigurationManager.cs b/Microsoft.Azure.Cosmos/src/Util/ConfigurationManager.cs index 1be83b2430..a49b3f093f 100644 --- a/Microsoft.Azure.Cosmos/src/Util/ConfigurationManager.cs +++ b/Microsoft.Azure.Cosmos/src/Util/ConfigurationManager.cs @@ -25,7 +25,12 @@ internal static class ConfigurationManager /// /// Environment variable name for overriding optimistic direct execution of queries. /// - internal static readonly string OptimisticDirectExecutionEnabled = "AZURE_COSMOS_OPTIMISTIC_DIRECT_EXECUTION_ENABLED"; + internal static readonly string OptimisticDirectExecutionEnabled = "AZURE_COSMOS_OPTIMISTIC_DIRECT_EXECUTION_ENABLED"; + + /// + /// Environment variable name to disable sending non streaming order by query feature falg to the gateway. + /// + internal static readonly string NonStreamingOrderByQueryFeatureDisabled = "AZURE_COSMOS_NON_STREAMING_ORDER_BY_FLAG_DISABLED"; public static T GetEnvironmentVariable(string variable, T defaultValue) { @@ -88,6 +93,19 @@ public static bool IsOptimisticDirectExecutionEnabled( .GetEnvironmentVariable( variable: OptimisticDirectExecutionEnabled, defaultValue: defaultValue); + } + + /// + /// Gets the boolean value indicating whether the non streaming order by query feature flag should be sent to the gateway + /// based on the environment variable override. + /// + public static bool IsNonStreamingOrderByQueryFeatureDisabled( + bool defaultValue) + { + return ConfigurationManager + .GetEnvironmentVariable( + variable: NonStreamingOrderByQueryFeatureDisabled, + defaultValue: defaultValue); } } } diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.EmulatorTests/Query/NonStreamingOrderByQueryTests.cs b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.EmulatorTests/Query/NonStreamingOrderByQueryTests.cs index 355e6b0fac..19659fc6b1 100644 --- a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.EmulatorTests/Query/NonStreamingOrderByQueryTests.cs +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.EmulatorTests/Query/NonStreamingOrderByQueryTests.cs @@ -4,9 +4,16 @@ namespace Microsoft.Azure.Cosmos.EmulatorTests.Query using System.Collections.Generic; using System.Diagnostics; using System.Linq; + using System.Threading; using System.Threading.Tasks; using Microsoft.Azure.Cosmos.CosmosElements; + using Microsoft.Azure.Cosmos.Query.Core.Monads; + using Microsoft.Azure.Cosmos.Query.Core.Pipeline.Pagination; + using Microsoft.Azure.Cosmos.Query.Core.QueryPlan; + using Microsoft.Azure.Cosmos.Query.Core; using Microsoft.VisualStudio.TestTools.UnitTesting; + using Microsoft.Azure.Cosmos.Query.Core.QueryClient; + using Microsoft.Azure.Cosmos.Query; [TestClass] [TestCategory("Query")] @@ -73,6 +80,133 @@ await this.CreateIngestQueryDeleteAsync( RunVectorDistanceTestsAsync); } + [TestMethod] + public async Task QueryFeatureFlagTests() + { + static async Task ImplementationAsync(Container container, IReadOnlyList _) + { + ContainerInlineCore containerCore = container as ContainerInlineCore; + + SupportedQueryFeaturesValidator validator = new SupportedQueryFeaturesValidator + { + ExpectNonStreamingOrderBy = true + }; + + CosmosQueryClient cosmosQueryClient = new MockCosmosQueryClient( + containerCore.ClientContext, + containerCore, + bypassQueryParsing: true, + validator); + + ContainerProperties containerProperties = await containerCore.GetCachedContainerPropertiesAsync( + forceRefresh: false, + Cosmos.Tracing.NoOpTrace.Singleton, + cancellationToken: default); + + IReadOnlyList testCases = new List + { + MakeQueryFeaturesTest( + environmentVariable: null, + requestOption: null, + expectNonStreamingOrderBy: true), + MakeQueryFeaturesTest( + environmentVariable: true, + requestOption: null, + expectNonStreamingOrderBy: false), + MakeQueryFeaturesTest( + environmentVariable: null, + requestOption: true, + expectNonStreamingOrderBy: false), + MakeQueryFeaturesTest( + environmentVariable: false, + requestOption: null, + expectNonStreamingOrderBy: true), + MakeQueryFeaturesTest( + environmentVariable: null, + requestOption: false, + expectNonStreamingOrderBy: true), + }; + + int validationCount = validator.InvocationCount; + foreach (QueryFeaturesTest testCase in testCases) + { + if (testCase.EnvironmentVariable.HasValue) + { + Environment.SetEnvironmentVariable( + ConfigurationManager.NonStreamingOrderByQueryFeatureDisabled, + testCase.EnvironmentVariable.Value.ToString()); + } + else + { + Environment.SetEnvironmentVariable( + ConfigurationManager.NonStreamingOrderByQueryFeatureDisabled, + null); + } + + QueryRequestOptions requestOptions = new QueryRequestOptions + { + EnableOptimisticDirectExecution = false, + }; + + if (testCase.RequestOption.HasValue) + { + requestOptions.IsNonStreamingOrderByQueryFeatureDisabled = testCase.RequestOption.Value; + } + + validator.ExpectNonStreamingOrderBy = testCase.ExpectNonStreamingOrderBy; + ++validationCount; + + DebugTraceHelpers.TraceSupportedFeaturesTestCase(testCase); + + QueryIterator iterator = QueryIterator.Create( + containerCore, + cosmosQueryClient, + containerCore.ClientContext, + new SqlQuerySpec("SELECT * FROM c"), + continuationToken: null, + feedRangeInternal: FeedRangeEpk.FullRange, + requestOptions, + resourceLink: containerProperties.SelfLink, + isContinuationExpected: true, + allowNonValueAggregateQuery: true, + partitionedQueryExecutionInfo: null, + resourceType: Azure.Documents.ResourceType.Document); + + Assert.IsTrue(iterator.HasMoreResults); + ResponseMessage responseMessage = await iterator.ReadNextAsync(); + Assert.AreEqual(System.Net.HttpStatusCode.OK, responseMessage.StatusCode); + Assert.AreEqual(validationCount, validator.InvocationCount); + } + } + + await this.CreateIngestQueryDeleteAsync( + ConnectionModes.Direct, + CollectionTypes.MultiPartition, + new List(), + ImplementationAsync); + } + + private static QueryFeaturesTest MakeQueryFeaturesTest(bool? environmentVariable, bool? requestOption, bool expectNonStreamingOrderBy) + { + return new QueryFeaturesTest(environmentVariable, requestOption, expectNonStreamingOrderBy); + } + + private sealed class QueryFeaturesTest + { + public bool? EnvironmentVariable { get; } + + public bool? RequestOption { get; } + + public bool ExpectNonStreamingOrderBy { get; } + + public QueryFeaturesTest(bool? environmentVariable, bool? requestOption, bool expectNonStreamingOrderBy) + { + this.EnvironmentVariable = environmentVariable; + this.RequestOption = requestOption; + this.ExpectNonStreamingOrderBy = expectNonStreamingOrderBy; + } + } + private static async Task RunVectorDistanceTestsAsync(Container container, IReadOnlyList _) { foreach (TestCase testCase in VectorDistanceTestCases) @@ -174,5 +308,116 @@ private sealed class Document public string Id { get; set; } public string Word { get; set; } } + + private static class DebugTraceHelpers + { + private const bool Enabled = false; + +#pragma warning disable CS0162 // Unreachable code detected + public static void TraceSupportedFeaturesString(string supportedQueryFeatures) + { + if (Enabled) + { + Trace.WriteLine($"Supported query features header: {supportedQueryFeatures}"); + } + } + + public static void TraceSupportedFeaturesTestCase(QueryFeaturesTest testCase) + { + if (Enabled) + { + Trace.WriteLine($"EnvironmentVariable: {testCase.EnvironmentVariable}, RequestOption: {testCase.RequestOption}, ExpectNonStreamingOrderBy: {testCase.ExpectNonStreamingOrderBy}"); + } + } +#pragma warning restore CS0162 // Unreachable code detected + } + + private sealed class SupportedQueryFeaturesValidator + { + public int InvocationCount { get; private set; } + + public bool ExpectNonStreamingOrderBy { get; set;} + + public void Validate(string supportedQueryFeatures) + { + DebugTraceHelpers.TraceSupportedFeaturesString(supportedQueryFeatures); + ++this.InvocationCount; + bool hasNonStreamingOrderBy = supportedQueryFeatures.Contains(QueryFeatures.NonStreamingOrderBy.ToString()); + Assert.AreEqual(this.ExpectNonStreamingOrderBy, hasNonStreamingOrderBy); + } + } + + private sealed class MockCosmosQueryClient : CosmosQueryClientCore + { + private readonly bool bypassQueryParsing; + + private readonly SupportedQueryFeaturesValidator validator; + + public MockCosmosQueryClient( + CosmosClientContext clientContext, + ContainerInternal cosmosContainerCore, + bool bypassQueryParsing, + SupportedQueryFeaturesValidator validateSupportedQueryFeatures) + : base(clientContext, cosmosContainerCore) + { + this.bypassQueryParsing = bypassQueryParsing; + this.validator = validateSupportedQueryFeatures; + } + + public override bool BypassQueryParsing() + { + return this.bypassQueryParsing; + } + + public override Task ExecuteQueryPlanRequestAsync( + string resourceUri, + Documents.ResourceType resourceType, + Documents.OperationType operationType, + SqlQuerySpec sqlQuerySpec, + Cosmos.PartitionKey? partitionKey, + string supportedQueryFeatures, + Guid clientQueryCorrelationId, + Cosmos.Tracing.ITrace trace, + CancellationToken cancellationToken) + { + this.validator.Validate(supportedQueryFeatures); + return base.ExecuteQueryPlanRequestAsync( + resourceUri, + resourceType, + operationType, + sqlQuerySpec, + partitionKey, + supportedQueryFeatures, + clientQueryCorrelationId, + trace, + cancellationToken); + } + + public override Task> ExecuteItemQueryAsync( + string resourceUri, + Documents.ResourceType resourceType, + Documents.OperationType operationType, + FeedRange feedRange, + QueryRequestOptions requestOptions, + AdditionalRequestHeaders additionalRequestHeaders, + SqlQuerySpec sqlQuerySpec, + string continuationToken, + int pageSize, + Cosmos.Tracing.ITrace trace, + CancellationToken cancellationToken) + { + QueryPage queryPage = new QueryPage( + documents: new List(), + requestCharge: 0, + activityId: Guid.NewGuid().ToString(), + cosmosQueryExecutionInfo: null, + distributionPlanSpec: null, + disallowContinuationTokenMessage: null, + additionalHeaders: null, + state: null, + streaming: false); + return Task.FromResult(TryCatch.FromResult(queryPage)); + } + } } } diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/OptimisticDirectExecutionQueryBaselineTests.cs b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/OptimisticDirectExecutionQueryBaselineTests.cs index d723a57df7..03f6807e1d 100644 --- a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/OptimisticDirectExecutionQueryBaselineTests.cs +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/OptimisticDirectExecutionQueryBaselineTests.cs @@ -1021,7 +1021,8 @@ public override OptimisticDirectExecutionTestOutput ExecuteTest(OptimisticDirect partitionedQueryExecutionInfo: null, executionEnvironment: null, returnResultsInDeterministicOrder: null, - enableOptimisticDirectExecution: queryRequestOptions.EnableOptimisticDirectExecution, + enableOptimisticDirectExecution: queryRequestOptions.EnableOptimisticDirectExecution, + isNonStreamingOrderByQueryFeatureDisabled: false, testInjections: queryRequestOptions.TestSettings); List targetPkRanges = new (); 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 c8ca339530..298387bbff 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 @@ -342,7 +342,8 @@ private async Task TestPageSizeAsync(string query, int expectedPageSize, int exp partitionedQueryExecutionInfo: null, executionEnvironment: null, returnResultsInDeterministicOrder: null, - enableOptimisticDirectExecution: queryRequestOptions.EnableOptimisticDirectExecution, + enableOptimisticDirectExecution: queryRequestOptions.EnableOptimisticDirectExecution, + isNonStreamingOrderByQueryFeatureDisabled: queryRequestOptions.IsNonStreamingOrderByQueryFeatureDisabled, testInjections: queryRequestOptions.TestSettings); string databaseId = "db1234"; diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/SubpartitionTests.cs b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/SubpartitionTests.cs index 35869ebbc8..8253e973e0 100644 --- a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/SubpartitionTests.cs +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Query/SubpartitionTests.cs @@ -130,7 +130,8 @@ public override SubpartitionTestOutput ExecuteTest(SubpartitionTestInput input) partitionedQueryExecutionInfo: null, executionEnvironment: null, returnResultsInDeterministicOrder: null, - enableOptimisticDirectExecution: queryRequestOptions.EnableOptimisticDirectExecution, + enableOptimisticDirectExecution: queryRequestOptions.EnableOptimisticDirectExecution, + isNonStreamingOrderByQueryFeatureDisabled: queryRequestOptions.IsNonStreamingOrderByQueryFeatureDisabled, testInjections: queryRequestOptions.TestSettings); string databaseId = "db1234"; From c19d05ab5233095fb1f7f92e964ad7111d536918 Mon Sep 17 00:00:00 2001 From: Neil Deshpande Date: Tue, 14 May 2024 14:16:38 -0700 Subject: [PATCH 2/3] Make sure that we dont leak state from the integration test. Add the suppression environment variable for the test pipeline since we run against an old emulator --- .../Query/NonStreamingOrderByQueryTests.cs | 20 +++++++++++++++++++ templates/build-test.yml | 9 +++++++++ 2 files changed, 29 insertions(+) diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.EmulatorTests/Query/NonStreamingOrderByQueryTests.cs b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.EmulatorTests/Query/NonStreamingOrderByQueryTests.cs index 19659fc6b1..9d88d2e003 100644 --- a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.EmulatorTests/Query/NonStreamingOrderByQueryTests.cs +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.EmulatorTests/Query/NonStreamingOrderByQueryTests.cs @@ -83,6 +83,9 @@ await this.CreateIngestQueryDeleteAsync( [TestMethod] public async Task QueryFeatureFlagTests() { + using EnvironmentVariable nonStreamingOrderByDisabled = new EnvironmentVariable( + ConfigurationManager.NonStreamingOrderByQueryFeatureDisabled); + static async Task ImplementationAsync(Container container, IReadOnlyList _) { ContainerInlineCore containerCore = container as ContainerInlineCore; @@ -332,6 +335,23 @@ public static void TraceSupportedFeaturesTestCase(QueryFeaturesTest testCase) #pragma warning restore CS0162 // Unreachable code detected } + private sealed class EnvironmentVariable : IDisposable + { + private readonly string name; + private readonly string value; + + public EnvironmentVariable(string name) + { + this.name = name; + this.value = Environment.GetEnvironmentVariable(name); + } + + public void Dispose() + { + Environment.SetEnvironmentVariable(this.name, this.value); + } + } + private sealed class SupportedQueryFeaturesValidator { public int InvocationCount { get; private set; } diff --git a/templates/build-test.yml b/templates/build-test.yml index f4074ecce1..0749412e24 100644 --- a/templates/build-test.yml +++ b/templates/build-test.yml @@ -118,6 +118,8 @@ jobs: nugetConfigPath: NuGet.config publishTestResults: true testRunTitle: Microsoft.Azure.Cosmos.EmulatorTests + env: + AZURE_COSMOS_NON_STREAMING_ORDER_BY_FLAG_DISABLED: true - job: displayName: EmulatorTests ${{ parameters.BuildConfiguration }} - ${{ parameters.EmulatorPipeline2CategoryListName }} @@ -149,6 +151,8 @@ jobs: nugetConfigPath: NuGet.config publishTestResults: true testRunTitle: Microsoft.Azure.Cosmos.EmulatorTests + env: + AZURE_COSMOS_NON_STREAMING_ORDER_BY_FLAG_DISABLED: true - job: displayName: Encryption EmulatorTests ${{ parameters.BuildConfiguration }} @@ -180,6 +184,8 @@ jobs: nugetConfigPath: NuGet.config publishTestResults: true testRunTitle: Microsoft.Azure.Cosmos.Encryption.EmulatorTests + env: + AZURE_COSMOS_NON_STREAMING_ORDER_BY_FLAG_DISABLED: true - job: displayName: Encryption.Custom EmulatorTests ${{ parameters.BuildConfiguration }} @@ -211,6 +217,8 @@ jobs: nugetConfigPath: NuGet.config publishTestResults: true testRunTitle: Microsoft.Azure.Cosmos.Encryption.Custom.EmulatorTests + env: + AZURE_COSMOS_NON_STREAMING_ORDER_BY_FLAG_DISABLED: true - job: displayName: EmulatorTests ${{ parameters.BuildConfiguration }} - ${{ parameters.EmulatorPipeline3CategoryListName }} @@ -242,3 +250,4 @@ jobs: testRunTitle: Microsoft.Azure.Cosmos.EmulatorTests env: COSMOSDB_MULTI_REGION: ${{ parameters.MultiRegionConnectionString }} + AZURE_COSMOS_NON_STREAMING_ORDER_BY_FLAG_DISABLED: true From 8c1e1a6b91f95bc18d982419f8d41586495a91a7 Mon Sep 17 00:00:00 2001 From: Neil Deshpande Date: Tue, 14 May 2024 16:56:20 -0700 Subject: [PATCH 3/3] Incorporate code review feedback --- .../Core/QueryPlan/QueryPlanRetriever.cs | 27 +++------- .../src/RequestOptions/QueryRequestOptions.cs | 6 +-- .../Query/NonStreamingOrderByQueryTests.cs | 53 +++++++++++++++---- 3 files changed, 52 insertions(+), 34 deletions(-) diff --git a/Microsoft.Azure.Cosmos/src/Query/Core/QueryPlan/QueryPlanRetriever.cs b/Microsoft.Azure.Cosmos/src/Query/Core/QueryPlan/QueryPlanRetriever.cs index 9cbf6918fc..5ac7f41043 100644 --- a/Microsoft.Azure.Cosmos/src/Query/Core/QueryPlan/QueryPlanRetriever.cs +++ b/Microsoft.Azure.Cosmos/src/Query/Core/QueryPlan/QueryPlanRetriever.cs @@ -31,32 +31,19 @@ internal static class QueryPlanRetriever | QueryFeatures.DCount | QueryFeatures.NonStreamingOrderBy; - private static readonly QueryFeatures LegacySupportedQueryFeatures = - QueryFeatures.Aggregate - | QueryFeatures.Distinct - | QueryFeatures.GroupBy - | QueryFeatures.MultipleOrderBy - | QueryFeatures.MultipleAggregates - | QueryFeatures.OffsetAndLimit - | QueryFeatures.OrderBy - | QueryFeatures.Top - | QueryFeatures.NonValueAggregate - | QueryFeatures.DCount; + private static readonly QueryFeatures SupportedQueryFeaturesWithoutNonStreamingOrderBy = + SupportedQueryFeatures & (~QueryFeatures.NonStreamingOrderBy); private static readonly string SupportedQueryFeaturesString = SupportedQueryFeatures.ToString(); - private static readonly string LegacySupportedQueryFeaturesString = LegacySupportedQueryFeatures.ToString(); + private static readonly string SupportedQueryFeaturesWithoutNonStreamingOrderByString = + SupportedQueryFeaturesWithoutNonStreamingOrderBy.ToString(); private static string GetSupportedQueryFeaturesString(bool isNonStreamingOrderByQueryFeatureDisabled) { - if (isNonStreamingOrderByQueryFeatureDisabled) - { - return LegacySupportedQueryFeaturesString; - } - else - { - return SupportedQueryFeaturesString; - } + return isNonStreamingOrderByQueryFeatureDisabled ? + SupportedQueryFeaturesWithoutNonStreamingOrderByString : + SupportedQueryFeaturesString; } public static async Task GetQueryPlanWithServiceInteropAsync( diff --git a/Microsoft.Azure.Cosmos/src/RequestOptions/QueryRequestOptions.cs b/Microsoft.Azure.Cosmos/src/RequestOptions/QueryRequestOptions.cs index fd697c09ad..88db939c96 100644 --- a/Microsoft.Azure.Cosmos/src/RequestOptions/QueryRequestOptions.cs +++ b/Microsoft.Azure.Cosmos/src/RequestOptions/QueryRequestOptions.cs @@ -51,8 +51,6 @@ public class QueryRequestOptions : RequestOptions /// Direct (optimistic) execution offers improved performance for several kinds of queries such as a single partition streaming query. /// public bool EnableOptimisticDirectExecution { get; set; } = ConfigurationManager.IsOptimisticDirectExecutionEnabled(defaultValue: false); - - internal bool IsNonStreamingOrderByQueryFeatureDisabled { get; set; } = ConfigurationManager.IsNonStreamingOrderByQueryFeatureDisabled(defaultValue: false); /// /// Gets or sets the maximum number of items that can be buffered client side during @@ -191,7 +189,9 @@ public ConsistencyLevel? ConsistencyLevel internal TestInjections TestSettings { get; set; } - internal FeedRange FeedRange { get; set; } + internal FeedRange FeedRange { get; set; } + + internal bool IsNonStreamingOrderByQueryFeatureDisabled { get; set; } = ConfigurationManager.IsNonStreamingOrderByQueryFeatureDisabled(defaultValue: false); /// /// Fill the CosmosRequestMessage headers with the set properties diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.EmulatorTests/Query/NonStreamingOrderByQueryTests.cs b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.EmulatorTests/Query/NonStreamingOrderByQueryTests.cs index 9d88d2e003..c21f76d610 100644 --- a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.EmulatorTests/Query/NonStreamingOrderByQueryTests.cs +++ b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.EmulatorTests/Query/NonStreamingOrderByQueryTests.cs @@ -1,4 +1,4 @@ -namespace Microsoft.Azure.Cosmos.EmulatorTests.Query +namespace Microsoft.Azure.Cosmos.EmulatorTests.Query { using System; using System.Collections.Generic; @@ -15,9 +15,9 @@ namespace Microsoft.Azure.Cosmos.EmulatorTests.Query using Microsoft.Azure.Cosmos.Query.Core.QueryClient; using Microsoft.Azure.Cosmos.Query; - [TestClass] - [TestCategory("Query")] - public sealed class NonStreamingOrderByQueryTests : QueryTestsBase + [TestClass] + [TestCategory("Query")] + public sealed class NonStreamingOrderByQueryTests : QueryTestsBase { private const string Embedding = "[-0.008861,0.097097,0.100236,0.070044,-0.079279,0.000923,-0.012829,0.064301,-0.029405,-0.009858,-0.017753,0.063115,0.033623,0.019805,0.052704,-0.100458,0.089387,-0.040792,-0.088936,0.110212,-0.044749,0.077675,-0.017062,-0.063745,-0.009502,-0.079371,0.066952,-0.070209,0.063761,-0.038194,-0.046252,0.049983,-0.094985,-0.086341,0.024665,-0.112857,-0.038358,-0.007008,-0.010063,-0.000183,0.068841,0.024942,-0.042561,-0.044576,0.010776,0.006323,0.088285,-0.062522,0.028216,0.088291,0.033231,-0.033732,-0.002995,0.118994,0.000453,0.158588,-0.044475,-0.137629,0.066080,0.062824,-0.128369,-0.087959,0.028080,0.070063,0.046700,-0.083278,-0.118428,0.071118,0.100757,0.017944,0.026296,0.017282,-0.082127,-0.006148,0.002967,-0.032857,-0.076493,-0.072842,-0.055179,-0.081703,0.011437,-0.038698,-0.062540,-0.027899,0.087635,0.031870,0.029164,0.000524,-0.039895,-0.055559,0.024582,-0.030595,0.003942,-0.034500,0.003012,-0.023863,0.033831,0.061476,-0.090183,-0.039206,-0.026586,-0.042763,0.049835,-0.052496,-0.020044,0.073703,0.096775,0.033063,0.000313,-0.022581,-0.141154,0.032095,0.077733,-0.063739,-0.055647,-0.017604,0.044639,-0.062925,-0.001960,0.024665,-0.009416,-0.021381,0.082724,-0.031026,0.027255,0.066198,0.000845,0.008393,0.039434,0.054104,-0.060255,0.034266,0.079435,0.043624,-0.015871,-0.038030,-0.030374,-0.020542,0.007132,0.008708,0.087840,0.017351,-0.089493,0.030182,0.026961,-0.071212,-0.004854,0.007389,0.067203,-0.026351,-0.011460,-0.058723,0.013153,-0.020313,-0.051170,0.002242,0.088222,-0.004267,-0.073523,-0.021874,-0.033585,-0.048553,-0.019119,-0.025310,0.053096,0.111063,0.035042,-0.082811,-0.073749,-0.010048,0.012265,-0.023893,-0.125340,0.026611,0.043258,-0.010473,-0.044428,-0.039251,-0.046891,-0.013008,0.062219,0.078732,-0.086303,0.016901,0.010331,-0.043754,-0.057733,-0.037964,0.024907,0.068143,-0.019992,-0.035030,0.038854,0.034345,-0.048839,-0.105419,0.043013,-0.023374,-0.077629,-0.076465,0.078564,-0.024519,0.041293,-0.032088,-0.007053,0.022618,-0.004657,-0.093970,-0.000199,0.004813,-0.044789,-0.127900,-0.033516,-0.043816,0.033056,-0.057619,0.004901,0.018863,0.039752,0.000739,-0.136350,-0.067819,-0.014856,0.058351,-0.014275,-0.000873,-0.039388,-0.017191,-0.051184,-0.046863,0.006143,-0.075998,-0.064695,0.046676,-0.020558,0.082474,0.160449,-0.027475,0.009541,-0.021876,0.027416,0.078049,0.089309,0.032928,-0.033272,0.048905,0.061164,0.054811,0.024527,-0.034978,-0.018083,-0.077601,0.034112,-0.021121,0.098856,0.019585,-0.058928,-0.016126,-0.011748,0.031588,0.003205,-0.077483,-0.002372,-0.113548,0.047445,-0.027094,-0.032843,0.042378,-0.074703,0.057001,0.012020,0.131156,0.002080,-0.065770,0.112443,0.047786,0.024492,-0.108401,0.016836,0.001478,0.041542,-0.067801,0.102876,-0.052808,-0.136035,0.073852,0.079966,-0.000586,0.034055,-0.053040,0.050461,-0.021550,0.014827,0.077605,-0.024783,-0.082388,0.074410,-0.033689,-0.010982,0.043733]"; @@ -151,9 +151,9 @@ static async Task ImplementationAsync(Container container, IReadOnlyList _) { - foreach (TestCase testCase in VectorDistanceTestCases) + foreach (TestCase testCase in VectorDistanceTestCases) { Trace.WriteLine("Executing test case: "); - Trace.WriteLine(testCase.Query); + Trace.WriteLine(testCase.Query); FeedIterator iterator = container.GetItemQueryIterator(testCase.Query); List documents = new List(); @@ -236,7 +246,7 @@ private static async Task RunVectorDistanceTestsAsync(Container container, IRead } Assert.Fail(); - } + } } } @@ -369,6 +379,20 @@ public void Validate(string supportedQueryFeatures) private sealed class MockCosmosQueryClient : CosmosQueryClientCore { + private static readonly QueryFeatures SupportedQueryFeaturesWithoutNonStreamingOrderBy = + QueryFeatures.Aggregate + | QueryFeatures.Distinct + | QueryFeatures.GroupBy + | QueryFeatures.MultipleOrderBy + | QueryFeatures.MultipleAggregates + | QueryFeatures.OffsetAndLimit + | QueryFeatures.OrderBy + | QueryFeatures.Top + | QueryFeatures.NonValueAggregate + | QueryFeatures.DCount; + + private static readonly string SupportedQueryFeaturesString = SupportedQueryFeaturesWithoutNonStreamingOrderBy.ToString(); + private readonly bool bypassQueryParsing; private readonly SupportedQueryFeaturesValidator validator; @@ -401,6 +425,13 @@ public override Task ExecuteQueryPlanRequestAsync CancellationToken cancellationToken) { this.validator.Validate(supportedQueryFeatures); + + if (this.validator.ExpectNonStreamingOrderBy) + { + // older emulator in the github repo does not support non-streaming order by, and will send back 400 + supportedQueryFeatures = SupportedQueryFeaturesString; + } + return base.ExecuteQueryPlanRequestAsync( resourceUri, resourceType,