Query: Adds an environment config to suppress sending NonStreamingOrd…

…erBy in the list of query features sent to the gateway (#4492)

* Introduce an environment config to suppress sending NonStreamingOrderBy in the list of query features sent to the gateway

* 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

* Incorporate code review feedback

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

@@ -408,6 +408,7 @@ private static TryCatch<IQueryPipelineStage> TryCreateSpecializedDocumentQueryEx
                     inputParameters.ExecutionEnvironment,
                     inputParameters.ReturnResultsInDeterministicOrder,
                     inputParameters.EnableOptimisticDirectExecution,
+                    inputParameters.IsNonStreamingOrderByQueryFeatureDisabled,
                     inputParameters.TestInjections);
             }
 
@@ -593,6 +594,7 @@ private static async Task<PartitionedQueryExecutionInfo> 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);
             }
         }

Microsoft.Azure.Cosmos/src/Query/Core/QueryPlan/QueryPlanRetriever.cs

@@ -28,9 +28,23 @@ internal static class QueryPlanRetriever
             | QueryFeatures.OrderBy
             | QueryFeatures.Top
             | QueryFeatures.NonValueAggregate
-            | QueryFeatures.DCount;
-
-        private static readonly string SupportedQueryFeaturesString = SupportedQueryFeatures.ToString();
+            | QueryFeatures.DCount
+            | QueryFeatures.NonStreamingOrderBy;
+
+        private static readonly QueryFeatures SupportedQueryFeaturesWithoutNonStreamingOrderBy =
+            SupportedQueryFeatures & (~QueryFeatures.NonStreamingOrderBy);
+
+        private static readonly string SupportedQueryFeaturesString = SupportedQueryFeatures.ToString();
+
+        private static readonly string SupportedQueryFeaturesWithoutNonStreamingOrderByString =
+            SupportedQueryFeaturesWithoutNonStreamingOrderBy.ToString();
+
+        private static string GetSupportedQueryFeaturesString(bool isNonStreamingOrderByQueryFeatureDisabled)
+        {
+            return isNonStreamingOrderByQueryFeatureDisabled ?
+                SupportedQueryFeaturesWithoutNonStreamingOrderByString :
+                SupportedQueryFeaturesString;
+        }
 
         public static async Task<PartitionedQueryExecutionInfo> GetQueryPlanWithServiceInteropAsync(
             CosmosQueryClient queryClient,
@@ -94,7 +108,8 @@ public static Task<PartitionedQueryExecutionInfo> GetQueryPlanThroughGatewayAsyn
             CosmosQueryContext queryContext,
             SqlQuerySpec sqlQuerySpec,
             string resourceLink,
-            PartitionKey? partitionKey,
+            PartitionKey? partitionKey,
+            bool isNonStreamingOrderByQueryFeatureDisabled,
             ITrace trace,
             CancellationToken cancellationToken = default)
         {
@@ -130,7 +145,7 @@ public static Task<PartitionedQueryExecutionInfo> GetQueryPlanThroughGatewayAsyn
                     OperationType.QueryPlan,
                     sqlQuerySpec,
                     partitionKey,
-                    QueryPlanRetriever.SupportedQueryFeaturesString,
+                    GetSupportedQueryFeaturesString(isNonStreamingOrderByQueryFeatureDisabled),
                     trace,
                     cancellationToken);
             }

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(

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

@@ -50,7 +50,7 @@ public class QueryRequestOptions : RequestOptions
         /// <value>
         /// Direct (optimistic) execution offers improved performance for several kinds of queries such as a single partition streaming query.
         /// </value>
-        public bool EnableOptimisticDirectExecution { get; set; } = ConfigurationManager.IsOptimisticDirectExecutionEnabled(defaultValue: false);
+        public bool EnableOptimisticDirectExecution { get; set; } = ConfigurationManager.IsOptimisticDirectExecutionEnabled(defaultValue: false);
 
         /// <summary>
         /// Gets or sets the maximum number of items that can be buffered client side during 
@@ -189,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);
 
         /// <summary>
         /// Fill the CosmosRequestMessage headers with the set properties

Microsoft.Azure.Cosmos/src/Util/ConfigurationManager.cs

@@ -25,7 +25,12 @@ internal static class ConfigurationManager
         /// <summary>
         /// Environment variable name for overriding optimistic direct execution of queries.
         /// </summary>
-        internal static readonly string OptimisticDirectExecutionEnabled = "AZURE_COSMOS_OPTIMISTIC_DIRECT_EXECUTION_ENABLED";
+        internal static readonly string OptimisticDirectExecutionEnabled = "AZURE_COSMOS_OPTIMISTIC_DIRECT_EXECUTION_ENABLED";
+
+        /// <summary>
+        /// Environment variable name to disable sending non streaming order by query feature falg to the gateway.
+        /// </summary>
+        internal static readonly string NonStreamingOrderByQueryFeatureDisabled = "AZURE_COSMOS_NON_STREAMING_ORDER_BY_FLAG_DISABLED";
 
         public static T GetEnvironmentVariable<T>(string variable, T defaultValue)
         {
@@ -88,6 +93,19 @@ public static bool IsOptimisticDirectExecutionEnabled(
                     .GetEnvironmentVariable(
                         variable: OptimisticDirectExecutionEnabled,
                         defaultValue: defaultValue);
+        }
+
+        /// <summary>
+        /// 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.
+        /// </summary>
+        public static bool IsNonStreamingOrderByQueryFeatureDisabled(
+            bool defaultValue)
+        {
+            return ConfigurationManager
+                    .GetEnvironmentVariable(
+                        variable: NonStreamingOrderByQueryFeatureDisabled,
+                        defaultValue: defaultValue);
         }
     }
 }