Query: Add Optimistic Direct Execution configuration override support on the Client

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

@@ -9,7 +9,6 @@ namespace Microsoft.Azure.Cosmos.Query.Core.ExecutionContext
     using System.Linq;
     using System.Threading;
     using System.Threading.Tasks;
-    using global::Azure;
     using Microsoft.Azure.Cosmos;
     using Microsoft.Azure.Cosmos.CosmosElements;
     using Microsoft.Azure.Cosmos.Pagination;
@@ -33,7 +32,9 @@ namespace Microsoft.Azure.Cosmos.Query.Core.ExecutionContext
     internal static class CosmosQueryExecutionContextFactory
     {
         private const string InternalPartitionKeyDefinitionProperty = "x-ms-query-partitionkey-definition";
+        private const string AllowOptimisticDirectExecution = "allowOptimisticDirectExecution";
         private const string OptimisticDirectExecution = "OptimisticDirectExecution";
+        private const string ClientDisableOptimisticDirectExecution = "clientDisableOptimisticDirectExecution";
         private const string Passthrough = "Passthrough";
         private const string Specialized = "Specialized";
         private const int PageSizeFactorForTop = 5;
@@ -289,9 +290,9 @@ private static async Task<TryCatch<IQueryPipelineStage>> TryCreateFromPartitione
             TryCatch<IQueryPipelineStage> tryCreatePipelineStage;
 
             Documents.PartitionKeyRange targetRange = await TryGetTargetRangeOptimisticDirectExecutionAsync(
-                inputParameters, 
-                partitionedQueryExecutionInfo, 
-                cosmosQueryContext, 
+                inputParameters,
+                partitionedQueryExecutionInfo,
+                cosmosQueryContext,
                 containerQueryProperties,
                 trace);
 
@@ -761,13 +762,20 @@ private static Documents.PartitionKeyDefinition GetPartitionKeyDefinition(InputP
             ContainerQueryProperties containerQueryProperties,
             ITrace trace)
         {
-            if (!inputParameters.EnableOptimisticDirectExecution)
+            object queryConfigValue = await cosmosQueryContext.QueryClient.GetQueryEngineConfigurationValueAsync(ClientDisableOptimisticDirectExecution);
+            bool.TryParse(queryConfigValue.ToString(), out bool clientDisableOptimisticDirectExecution);
+
+            // Use the Ode code path only if ClientDisableOptimisticDirectExecution is false and EnableOptimisticDirectExecution is true
+            if (clientDisableOptimisticDirectExecution || !inputParameters.EnableOptimisticDirectExecution)
             {
-                if (inputParameters.InitialUserContinuationToken != null 
-                    && OptimisticDirectExecutionContinuationToken.IsOptimisticDirectExecutionContinuationToken(inputParameters.InitialUserContinuationToken))
+                if (inputParameters.InitialUserContinuationToken != null
+                          && OptimisticDirectExecutionContinuationToken.IsOptimisticDirectExecutionContinuationToken(inputParameters.InitialUserContinuationToken))
                 {
-                    throw new MalformedContinuationTokenException($"The continuation token supplied requires the Optimistic Direct Execution flag to be enabled in QueryRequestOptions for the query execution to resume. " +
-                            $"{inputParameters.InitialUserContinuationToken}");
+                    string errorMessage = "This query cannot be executed using the provided continuation token. " +
+                        "Please ensure that the EnableOptimisticDirectExecution flag is enabled in the QueryRequestOptions. " +
+                        "If after enabling this flag, you still see this error, contact the database administrator for assistance or retry the query without the continuation token.";
+
+                    throw new MalformedContinuationTokenException($"{errorMessage} Continuation Token: {inputParameters.InitialUserContinuationToken}");
                 }
 
                 return null;

Microsoft.Azure.Cosmos/src/Query/Core/QueryClient/CosmosQueryClient.cs

@@ -63,6 +63,8 @@ public abstract Task<TryCatch<QueryPage>> ExecuteItemQueryAsync(
             ITrace trace,
             CancellationToken cancellationToken);
 
+        public abstract Task<object> GetQueryEngineConfigurationValueAsync(string key);
+
         public abstract Task<PartitionedQueryExecutionInfo> ExecuteQueryPlanRequestAsync(
             string resourceUri,
             Documents.ResourceType resourceType,

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

@@ -44,6 +44,7 @@ internal sealed class QueryPartitionProvider : IDisposable
         private IntPtr serviceProvider;
         private bool disposed;
         private string queryengineConfiguration;
+        public IDictionary<string, object> QueryEngineConfigurationValues;
 
         public QueryPartitionProvider(IDictionary<string, object> queryengineConfiguration)
         {
@@ -59,6 +60,7 @@ public QueryPartitionProvider(IDictionary<string, object> queryengineConfigurati
 
             this.disposed = false;
             this.queryengineConfiguration = JsonConvert.SerializeObject(queryengineConfiguration);
+            this.QueryEngineConfigurationValues = queryengineConfiguration;
             this.serviceProvider = IntPtr.Zero;
 
             this.serviceProviderStateLock = new object();
@@ -91,6 +93,7 @@ public void Update(IDictionary<string, object> queryengineConfiguration)
                     if (!string.Equals(this.queryengineConfiguration, newConfiguration))
                     {
                         this.queryengineConfiguration = newConfiguration;
+                        this.QueryEngineConfigurationValues = queryengineConfiguration;
 
                         if (!this.disposed && this.serviceProvider != IntPtr.Zero)
                         {

Microsoft.Azure.Cosmos/src/Query/v3Query/CosmosQueryClientCore.cs

@@ -209,6 +209,22 @@ public override async Task<PartitionedQueryExecutionInfo> ExecuteQueryPlanReques
             return partitionedQueryExecutionInfo;
         }
 
+        public override async Task<object> GetQueryEngineConfigurationValueAsync(string key)
+        {
+            if (key == null)
+            {
+                throw new ArgumentNullException(nameof(key));
+            }
+
+            QueryPartitionProvider provider = await this.clientContext.DocumentClient.QueryPartitionProvider;
+            if (provider.QueryEngineConfigurationValues.TryGetValue(key, out object queryConfigProperty))
+            {
+                return queryConfigProperty;
+            }
+
+            throw new KeyNotFoundException($"The key '{key}' was not found in the QueryEngineConfiguration dictionary.");
+        }
+
         public override async Task<List<PartitionKeyRange>> GetTargetPartitionKeyRangeByFeedRangeAsync(
             string resourceLink,
             string collectionResourceId,

Microsoft.Azure.Cosmos/src/Query/v3Query/QueryIterator.cs

@@ -146,9 +146,9 @@ public static QueryIterator Create(
                 forcePassthrough: forcePassthrough,
                 enableOptimisticDirectExecution: queryRequestOptions.EnableOptimisticDirectExecution,
                 testInjections: queryRequestOptions.TestSettings);
-
+
             return new QueryIterator(
-                cosmosQueryContext,
+                cosmosQueryContext,
                 CosmosQueryExecutionContextFactory.Create(documentContainer, cosmosQueryContext, inputParameters, NoOpTrace.Singleton),
                 queryRequestOptions.CosmosSerializationFormatOptions,
                 queryRequestOptions,

Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.EmulatorTests/Query/OptimisticDirectExecutionQueryTests.cs

@@ -17,6 +17,7 @@ public sealed class OptimisticDirectExecutionQueryTests : QueryTestsBase
         private const string PartitionKeyField = "key";
         private const string NumberField = "numberField";
         private const string NullField = "nullField";
+        private const string ClientDisableOptimisticDirectExecution = "clientDisableOptimisticDirectExecution";
 
         private static class PageSizeOptions
         {
@@ -538,6 +539,19 @@ await this.CreateIngestQueryDeleteAsync(
                 documents,
                 (container, documents) => RunFailingTests(container, invalidQueries),
                 "/" + PartitionKeyField);
+        }
+
+        [TestMethod]
+        public async Task TestClientDisableOdeFlagInCosmosClient()
+        {
+            string authKey = Utils.ConfigurationManager.AppSettings["MasterKey"];
+            string endpoint = Utils.ConfigurationManager.AppSettings["GatewayEndpoint"];
+
+            CosmosClient client = new CosmosClient($"AccountEndpoint={endpoint};AccountKey={authKey}");
+            AccountProperties properties = await client.ReadAccountAsync();
+
+            Assert.IsTrue(properties.QueryEngineConfigurationString.Contains(ClientDisableOptimisticDirectExecution));
+            Assert.IsFalse(Convert.ToBoolean(properties.QueryEngineConfiguration[ClientDisableOptimisticDirectExecution]));
         }
 
         private static async Task RunTests(IEnumerable<DirectExecutionTestCase> testCases, Container container)