Using ServiceInterop/Gateway to get QueryPlan for Specialized Pipeline

Microsoft.Azure.Cosmos/src/Query/Core/Monads/TryCatch{TResult}.cs

@@ -126,6 +126,22 @@ public async Task<TryCatch<T>> TryAsync<T>(
             return matchResult;
         }
 
+        public TryCatch<T> Try<T>(
+            Func<TResult, TryCatch<T>> onSuccess)
+        {
+            TryCatch<T> matchResult;
+            if (this.Succeeded)
+            {
+                matchResult = onSuccess(this.either.FromRight(default));
+            }
+            else
+            {
+                matchResult = TryCatch<T>.FromException(this.either.FromLeft(default));
+            }
+
+            return matchResult;
+        }
+
         public TryCatch<TResult> Catch(
             Action<Exception> onError)
         {

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

@@ -27,7 +27,6 @@ namespace Microsoft.Azure.Cosmos.Query.Core.ExecutionContext
     using Microsoft.Azure.Cosmos.SqlObjects;
     using Microsoft.Azure.Cosmos.SqlObjects.Visitors;
     using Microsoft.Azure.Cosmos.Tracing;
-    using static Microsoft.Azure.Cosmos.Query.Core.ExecutionContext.CosmosQueryExecutionContextFactory;
 
     internal static class CosmosQueryExecutionContextFactory
     {
@@ -138,12 +137,12 @@ private static async Task<TryCatch<IQueryPipelineStage>> TryCreateCoreContextAsy
                 cosmosQueryContext.ContainerResourceId = containerQueryProperties.ResourceId;
 
                 Documents.PartitionKeyRange targetRange = await GetTargetRangeOptimisticDirectExecutionAsync(
-                    inputParameters, 
-                    queryPlanFromContinuationToken, 
-                    cosmosQueryContext, 
-                    containerQueryProperties, 
+                    inputParameters,
+                    queryPlanFromContinuationToken,
+                    cosmosQueryContext,
+                    containerQueryProperties,
                     trace);
-     
+
                 if (targetRange != null)
                 {
                     // Test code added to confirm the correct pipeline is being utilized
@@ -229,33 +228,12 @@ private static async Task<TryCatch<IQueryPipelineStage>> TryCreateCoreContextAsy
                         }
                     }
 
-                    if (cosmosQueryContext.QueryClient.ByPassQueryParsing())
-                    {
-                        // For non-Windows platforms(like Linux and OSX) in .NET Core SDK, we cannot use ServiceInterop, so need to bypass in that case.
-                        // We are also now bypassing this for 32 bit host process running even on Windows as there are many 32 bit apps that will not work without this
-                        partitionedQueryExecutionInfo = await QueryPlanRetriever.GetQueryPlanThroughGatewayAsync(
-                            cosmosQueryContext,
-                            inputParameters.SqlQuerySpec,
-                            cosmosQueryContext.ResourceLink,
-                            inputParameters.PartitionKey,
-                            createQueryPipelineTrace,
-                            cancellationToken);
-                    }
-                    else
-                    {
-                        Documents.PartitionKeyDefinition partitionKeyDefinition = GetPartitionKeyDefinition(inputParameters, containerQueryProperties);
-
-                        partitionedQueryExecutionInfo = await QueryPlanRetriever.GetQueryPlanWithServiceInteropAsync(
-                            cosmosQueryContext.QueryClient,
-                            inputParameters.SqlQuerySpec,
-                            cosmosQueryContext.ResourceTypeEnum,
-                            partitionKeyDefinition,
-                            inputParameters.PartitionKey != null,
-                            containerQueryProperties.GeospatialType,
-                            cosmosQueryContext.UseSystemPrefix,
-                            createQueryPipelineTrace,
-                            cancellationToken);
-                    }
+                    partitionedQueryExecutionInfo = await GetPartitionedQueryExecutionInfoFromGatewayOrServiceInteropAsync(
+                        cosmosQueryContext,
+                        inputParameters,
+                        containerQueryProperties,
+                        createQueryPipelineTrace,
+                        cancellationToken);
                 }
 
                 return await TryCreateFromPartitionedQueryExecutionInfoAsync(
@@ -407,7 +385,7 @@ private static TryCatch<IQueryPipelineStage> TryCreateOptimisticDirectExecutionC
                 inputParameters: inputParameters,
                 targetRange: new FeedRangeEpk(targetRange.ToRange()),
                 queryPaginationOptions: new QueryPaginationOptions(pageSizeHint: inputParameters.MaxItemCount),
-                queryPipelineStage: (CosmosElement continuationToken) =>
+                fallbackQueryPipelineStageFactory: (continuationToken) =>
                 {
                     InputParameters updatedInputParameters = new InputParameters(
                       inputParameters.SqlQuerySpec,
@@ -447,44 +425,12 @@ public static async Task<TryCatch<IQueryPipelineStage>> TryCreateOdeFallbackPipe
             ITrace trace,
             CancellationToken cancellationToken) 
         {
-            PartitionedQueryExecutionInfo partitionedQueryExecutionInfo;
-            Documents.PartitionKeyDefinition partitionKeyDefinition = GetPartitionKeyDefinition(inputParameters, containerQueryProperties);
-
-            if (partitionKeyDefinition != null)
-            {
-                partitionedQueryExecutionInfo = await QueryPlanRetriever.GetQueryPlanWithServiceInteropAsync(
-                cosmosQueryContext.QueryClient,
-                inputParameters.SqlQuerySpec,
-                cosmosQueryContext.ResourceTypeEnum,
-                partitionKeyDefinition,
-                inputParameters.PartitionKey != null,
-                containerQueryProperties.GeospatialType,
-                cosmosQueryContext.UseSystemPrefix,
-                trace,
-                cancellationToken);
-            }
-            else
-            {
-                partitionedQueryExecutionInfo = new PartitionedQueryExecutionInfo()
-                {
-                    QueryInfo = new QueryInfo()
-                    {
-                        Aggregates = null,
-                        DistinctType = DistinctQueryType.None,
-                        GroupByAliases = null,
-                        GroupByAliasToAggregateType = null,
-                        GroupByExpressions = null,
-                        HasSelectValue = false,
-                        Limit = null,
-                        Offset = null,
-                        OrderBy = null,
-                        OrderByExpressions = null,
-                        RewrittenQuery = null,
-                        Top = null,
-                    },
-                    QueryRanges = new List<Documents.Routing.Range<string>>(),
-                };
-            }
+            PartitionedQueryExecutionInfo partitionedQueryExecutionInfo = await GetPartitionedQueryExecutionInfoFromGatewayOrServiceInteropAsync(
+               cosmosQueryContext,
+               inputParameters,
+               containerQueryProperties,
+               trace,
+               cancellationToken);
 
             List<Documents.PartitionKeyRange> targetRanges = await CosmosQueryExecutionContextFactory.GetTargetPartitionKeyRangesAsync(
                   cosmosQueryContext.QueryClient,
@@ -588,6 +534,45 @@ private static TryCatch<IQueryPipelineStage> TryCreateSpecializedDocumentQueryEx
                 requestCancellationToken: cancellationToken);
         }
 
+        private static async Task<PartitionedQueryExecutionInfo> GetPartitionedQueryExecutionInfoFromGatewayOrServiceInteropAsync(
+           CosmosQueryContext cosmosQueryContext,
+           InputParameters inputParameters,
+           ContainerQueryProperties containerQueryProperties,
+           ITrace trace,
+           CancellationToken cancellationToken)
+        {
+            PartitionedQueryExecutionInfo partitionedQueryExecutionInfo;
+            if (cosmosQueryContext.QueryClient.ByPassQueryParsing())
+            {
+                // For non-Windows platforms(like Linux and OSX) in .NET Core SDK, we cannot use ServiceInterop, so need to bypass in that case.
+                // We are also now bypassing this for 32 bit host process running even on Windows as there are many 32 bit apps that will not work without this
+                partitionedQueryExecutionInfo = await QueryPlanRetriever.GetQueryPlanThroughGatewayAsync(
+                    cosmosQueryContext,
+                    inputParameters.SqlQuerySpec,
+                    cosmosQueryContext.ResourceLink,
+                    inputParameters.PartitionKey,
+                    trace,
+                    cancellationToken);
+            }
+            else
+            {
+                Documents.PartitionKeyDefinition partitionKeyDefinition = GetPartitionKeyDefinition(inputParameters, containerQueryProperties);
+
+                partitionedQueryExecutionInfo = await QueryPlanRetriever.GetQueryPlanWithServiceInteropAsync(
+                    cosmosQueryContext.QueryClient,
+                    inputParameters.SqlQuerySpec,
+                    cosmosQueryContext.ResourceTypeEnum,
+                    partitionKeyDefinition,
+                    inputParameters.PartitionKey != null,
+                    containerQueryProperties.GeospatialType,
+                    cosmosQueryContext.UseSystemPrefix,
+                    trace,
+                    cancellationToken);
+            }
+
+            return partitionedQueryExecutionInfo;
+        }
+
         /// <summary>
         /// Gets the list of partition key ranges.
         /// 1. Check partition key range id

Microsoft.Azure.Cosmos/src/Query/Core/Pipeline/OptimisticDirectExecution/OptimisticDirectExecutionQueryPipelineStage.cs

@@ -6,7 +6,6 @@ namespace Microsoft.Azure.Cosmos.Query.Core.Pipeline.OptimisticDirectExecutionQu
 {
     using System;
     using System.Collections.Generic;
-    using System.Diagnostics;
     using System.Linq;
     using System.Threading;
     using System.Threading.Tasks;
@@ -18,41 +17,44 @@ namespace Microsoft.Azure.Cosmos.Query.Core.Pipeline.OptimisticDirectExecutionQu
     using Microsoft.Azure.Cosmos.Query.Core.Pipeline.CrossPartition;
     using Microsoft.Azure.Cosmos.Query.Core.Pipeline.CrossPartition.Parallel;
     using Microsoft.Azure.Cosmos.Query.Core.Pipeline.Pagination;
-    using Microsoft.Azure.Cosmos.Query.Core.QueryClient;
     using Microsoft.Azure.Cosmos.Tracing;
     using Microsoft.Azure.Documents;
-    using static Microsoft.Azure.Cosmos.Query.Core.ExecutionContext.CosmosQueryExecutionContextFactory;
-    using static Microsoft.Azure.Cosmos.Query.Core.Pipeline.CrossPartition.PartitionMapper;
 
     internal sealed class OptimisticDirectExecutionQueryPipelineStage : IQueryPipelineStage
     {
+        private enum ExecutionState
+        {
+            OptimisticDirectExecution,
+            SpecializedDocumentQueryExecution,
+        }
+
         private const string optimisticDirectExecutionToken = "OptimisticDirectExecutionToken";
-        private const string optimisticDirectExecution = "OptimisticDirectExecution";
-        private const string specializedDocumentQueryExecution = "SpecializedDocumentQueryExecution";
         public delegate Task<TryCatch<IQueryPipelineStage>> FallbackQueryPipelineStageFactory(CosmosElement continuationToken);
         private readonly FallbackQueryPipelineStageFactory queryPipelineStageFactory;
         private TryCatch<IQueryPipelineStage> inner;
         private CosmosElement continuationToken;
-        private string queryState;
-        
-        private OptimisticDirectExecutionQueryPipelineStage(TryCatch<IQueryPipelineStage> queryPipelineStage, FallbackQueryPipelineStageFactory queryPipelineStageFactory, CosmosElement continuationToken)
+        private ExecutionState executionState;
+
+        private OptimisticDirectExecutionQueryPipelineStage(TryCatch<IQueryPipelineStage> inner, FallbackQueryPipelineStageFactory queryPipelineStageFactory, CosmosElement continuationToken)
         {
-            this.inner = queryPipelineStage;
+            this.inner = inner;
             this.queryPipelineStageFactory = queryPipelineStageFactory;
             this.continuationToken = continuationToken;
+            this.executionState = ExecutionState.OptimisticDirectExecution;
         }
 
-        public TryCatch<QueryPage> Current => this.inner.Result.Current;
+        public TryCatch<QueryPage> Current => this.inner.Try<QueryPage>(pipelineStage => pipelineStage.Current);
 
         public ValueTask DisposeAsync()
         {
-            return this.inner.Result.DisposeAsync();
+            return this.inner.Try<ValueTask>(pipelineStage => pipelineStage.DisposeAsync()).Result;
         }
-       
+
         public async ValueTask<bool> MoveNextAsync(ITrace trace)
         {
-            if (this.inner.Result == null)
+            if (this.inner.Failed)
             {
+                this.inner = TryCatch<IQueryPipelineStage>.FromException(this.inner.Exception);
                 return false;
             }
 
@@ -64,44 +66,49 @@ public async ValueTask<bool> MoveNextAsync(ITrace trace)
 
             if (success)
             {
-                this.queryState = optimisticDirectExecution;
                 this.SaveContinuation(this.Current.Result.State?.Value);
             }
-            else if (isGoneException)
+            else if (isGoneException && this.executionState == ExecutionState.OptimisticDirectExecution)
             {
-                // Only Ode pipeline GoneException should be handled
-                Debug.Assert(this.queryState == optimisticDirectExecution);
+                this.inner = await this.queryPipelineStageFactory(TryUnwrapContinuationToken(this.continuationToken).Result);
+                this.executionState = ExecutionState.SpecializedDocumentQueryExecution;
 
-                this.inner = await this.queryPipelineStageFactory(this.UnwrapContinuationToken(this.continuationToken));
-                this.queryState = specializedDocumentQueryExecution;
-
-                // TODO: Failure check for this.inner
-                bool fallbackPipelineSuccess = await this.inner.Result.MoveNextAsync(trace);
+                if (this.inner.Failed)
+                {
+                    this.inner = TryCatch<IQueryPipelineStage>.FromException(this.inner.Exception);
+                    return false;
+                }
 
-                return fallbackPipelineSuccess;
+                success = await this.inner.Result.MoveNextAsync(trace);
             }
 
-            return success; 
+            return success;
         }
 
         public void SetCancellationToken(CancellationToken cancellationToken)
         {
-            this.inner.Result.SetCancellationToken(cancellationToken);
+            this.inner.Try<IQueryPipelineStage>((pipelineStage) => pipelineStage = this.inner.Result).Result.SetCancellationToken(cancellationToken);
         }
 
-        public CosmosElement UnwrapContinuationToken(CosmosElement continuationToken) 
+        public static TryCatch<CosmosElement> TryUnwrapContinuationToken(CosmosElement continuationToken)
         {
-            ((CosmosObject)continuationToken).TryGetValue(optimisticDirectExecutionToken, out CosmosElement specializedContinuationToken);
+            if (!((CosmosObject)continuationToken).TryGetValue(optimisticDirectExecutionToken, out CosmosElement specializedContinuationToken))
+            {
+                return TryCatch<CosmosElement>.FromException(
+                    new FormatException(
+                        $"Unable to convert CosmosObject '{optimisticDirectExecutionToken}' to CosmosElement."));
+            }
+
             CosmosArray cosmosElementFallbackToken = CosmosArray.Create(specializedContinuationToken);
-            return cosmosElementFallbackToken;
+            return TryCatch<CosmosElement>.FromResult(cosmosElementFallbackToken);
         }
 
         public static TryCatch<IQueryPipelineStage> MonadicCreate(
             DocumentContainer documentContainer,
-            InputParameters inputParameters,
+            CosmosQueryExecutionContextFactory.InputParameters inputParameters,
             FeedRangeEpk targetRange,
             QueryPaginationOptions queryPaginationOptions,
-            FallbackQueryPipelineStageFactory queryPipelineStage,
+            FallbackQueryPipelineStageFactory fallbackQueryPipelineStageFactory,
             CancellationToken cancellationToken)
         {
             TryCatch<IQueryPipelineStage> pipelineStage = OptimisticDirectExecutionQueryPipelineImpl.MonadicCreate(
@@ -118,7 +125,7 @@ public static TryCatch<IQueryPipelineStage> MonadicCreate(
                 return TryCatch<IQueryPipelineStage>.FromException(pipelineStage.Exception);
             }
 
-            OptimisticDirectExecutionQueryPipelineStage odePipelineStageMonadicCreate = new OptimisticDirectExecutionQueryPipelineStage(pipelineStage, queryPipelineStage, inputParameters.InitialUserContinuationToken);
+            OptimisticDirectExecutionQueryPipelineStage odePipelineStageMonadicCreate = new OptimisticDirectExecutionQueryPipelineStage(pipelineStage, fallbackQueryPipelineStageFactory, inputParameters.InitialUserContinuationToken);
 
             return TryCatch<IQueryPipelineStage>.FromResult(odePipelineStageMonadicCreate);
         }
@@ -262,7 +269,7 @@ private static TryCatch<FeedRangeState<QueryState>> MonadicExtractState(
                     return TryCatch<FeedRangeState<QueryState>>.FromException(tryCreateContinuationToken.Exception);
                 }
 
-                TryCatch<PartitionMapping<OptimisticDirectExecutionContinuationToken>> partitionMappingMonad = PartitionMapper.MonadicGetPartitionMapping(
+                TryCatch<PartitionMapper.PartitionMapping<OptimisticDirectExecutionContinuationToken>> partitionMappingMonad = PartitionMapper.MonadicGetPartitionMapping(
                     range,
                     tryCreateContinuationToken.Result);
 
@@ -272,7 +279,7 @@ private static TryCatch<FeedRangeState<QueryState>> MonadicExtractState(
                         partitionMappingMonad.Exception);
                 }
 
-                PartitionMapping<OptimisticDirectExecutionContinuationToken> partitionMapping = partitionMappingMonad.Result;
+                PartitionMapper.PartitionMapping<OptimisticDirectExecutionContinuationToken> partitionMapping = partitionMappingMonad.Result;
 
                 KeyValuePair<FeedRangeEpk, OptimisticDirectExecutionContinuationToken> kvpRange = new KeyValuePair<FeedRangeEpk, OptimisticDirectExecutionContinuationToken>(
                     partitionMapping.TargetMapping.Keys.First(),