Query: Adds an environment config to suppress sending NonStreamingOrderBy in the list of query features sent to the gateway

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);
  }
  }
 }