CosmosQueryExecutionContextFactory Refactory

Microsoft.Azure.Cosmos/src/Query/Core/Aggregation/AverageAggregator.cs

@@ -58,7 +58,8 @@ public static TryCatch<IAggregator> TryCreate(string continuationToken)
             {
                 if (!AverageInfo.TryParse(continuationToken, out averageInfo))
                 {
-                    return TryCatch<IAggregator>.FromException(new ArgumentException($"Invalid continuation token: {continuationToken}"));
+                    return TryCatch<IAggregator>.FromException(
+                        new ArgumentException($"Invalid continuation token: {continuationToken}"));
                 }
             }
             else

Microsoft.Azure.Cosmos/src/Query/Core/AsyncLazy.cs

@@ -0,0 +1,53 @@
+// ------------------------------------------------------------
+// Copyright (c) Microsoft Corporation.  All rights reserved.
+// ------------------------------------------------------------
+
+namespace Microsoft.Azure.Cosmos.Query.Core
+{
+    using System;
+    using System.Threading;
+    using System.Threading.Tasks;
+
+    internal sealed class AsyncLazy<T>
+    {
+        private readonly Func<CancellationToken, Task<T>> valueFactory;
+        private T value;
+
+        public AsyncLazy(Func<CancellationToken, Task<T>> valueFactory)
+        {
+            if (valueFactory == null)
+            {
+                throw new ArgumentNullException(nameof(valueFactory));
+            }
+
+            this.valueFactory = valueFactory;
+        }
+
+        public bool ValueInitialized { get; private set; }
+
+        public async Task<T> GetValueAsync(CancellationToken cancellationToken)
+        {
+            cancellationToken.ThrowIfCancellationRequested();
+            if (!this.ValueInitialized)
+            {
+                this.value = await this.valueFactory(cancellationToken);
+                this.ValueInitialized = true;
+            }
+
+            return this.value;
+        }
+
+        public T Result
+        {
+            get
+            {
+                if (!this.ValueInitialized)
+                {
+                    throw new InvalidOperationException("Can not retrieve value before initialization.");
+                }
+
+                return this.value;
+            }
+        }
+    }
+}

Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/CosmosQueryExecutionContext.cs

@@ -26,10 +26,10 @@ public abstract bool IsDone
         /// <summary>
         /// Executes the context to feed the next page of results.
         /// </summary>
-        /// <param name="token">The cancellation token.</param>
+        /// <param name="cancellationToken">The cancellation token.</param>
         /// <returns>A task to await on, which in return provides a DoucmentFeedResponse of documents.</returns>
-        public abstract Task<QueryResponseCore> ExecuteNextAsync(CancellationToken token);
+        public abstract Task<QueryResponseCore> ExecuteNextAsync(CancellationToken cancellationToken);
 
-        public abstract bool TryGetContinuationToken(out string state);
+        public abstract bool TryGetContinuationToken(out string continuationToken);
     }
 }

Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/CosmosQueryExecutionContextWithNameCacheStaleRetry.cs

@@ -0,0 +1,67 @@
+// ------------------------------------------------------------
+// Copyright (c) Microsoft Corporation.  All rights reserved.
+// ------------------------------------------------------------
+
+namespace Microsoft.Azure.Cosmos.Query.Core.ExecutionContext
+{
+    using System;
+    using System.Threading;
+    using System.Threading.Tasks;
+
+    internal sealed class CosmosQueryExecutionContextWithNameCacheStaleRetry : CosmosQueryExecutionContext
+    {
+        private readonly CosmosQueryContext cosmosQueryContext;
+        private readonly Func<CosmosQueryExecutionContext> cosmosQueryExecutionContextFactory;
+        private CosmosQueryExecutionContext currentCosmosQueryExecutionContext;
+        private bool alreadyRetried;
+
+        public CosmosQueryExecutionContextWithNameCacheStaleRetry(
+            CosmosQueryContext cosmosQueryContext,
+            Func<CosmosQueryExecutionContext> cosmosQueryExecutionContextFactory)
+        {
+            this.cosmosQueryContext = cosmosQueryContext;
+            this.cosmosQueryExecutionContextFactory = cosmosQueryExecutionContextFactory;
+            this.currentCosmosQueryExecutionContext = cosmosQueryExecutionContextFactory();
+            this.alreadyRetried = false;
+        }
+
+        public override bool IsDone => this.currentCosmosQueryExecutionContext.IsDone;
+
+        public override void Dispose()
+        {
+            this.currentCosmosQueryExecutionContext.Dispose();
+        }
+
+        public override async Task<QueryResponseCore> ExecuteNextAsync(CancellationToken cancellationToken)
+        {
+            cancellationToken.ThrowIfCancellationRequested();
+
+            // If the cache is stale the entire execute context has incorrect values and should be recreated.
+            // This should only be done for the first execution.
+            // If results have already been pulled,
+            // then an error should be returned to the user,
+            // since it's not possible to combine query results from multiple containers.
+            QueryResponseCore queryResponse = await this.currentCosmosQueryExecutionContext.ExecuteNextAsync(cancellationToken);
+            if (
+                (queryResponse.StatusCode == System.Net.HttpStatusCode.Gone) &&
+                (queryResponse.SubStatusCode == Documents.SubStatusCodes.NameCacheIsStale) &&
+                !this.alreadyRetried)
+            {
+                this.currentCosmosQueryExecutionContext.Dispose();
+                await this.cosmosQueryContext.QueryClient.ForceRefreshCollectionCacheAsync(
+                    this.cosmosQueryContext.ResourceLink.OriginalString,
+                    cancellationToken);
+                this.currentCosmosQueryExecutionContext = this.cosmosQueryExecutionContextFactory();
+                this.alreadyRetried = true;
+                return await this.ExecuteNextAsync(cancellationToken);
+            }
+
+            return queryResponse;
+        }
+
+        public override bool TryGetContinuationToken(out string continuationToken)
+        {
+            return this.currentCosmosQueryExecutionContext.TryGetContinuationToken(out continuationToken);
+        }
+    }
+}

Microsoft.Azure.Cosmos/src/Query/Core/ExecutionContext/LazyCosmosQueryExecutionContext.cs

@@ -0,0 +1,100 @@
+// ------------------------------------------------------------
+// Copyright (c) Microsoft Corporation.  All rights reserved.
+// ------------------------------------------------------------
+
+namespace Microsoft.Azure.Cosmos.Query.Core.ExecutionContext
+{
+    using System;
+    using System.Threading;
+    using System.Threading.Tasks;
+    using Microsoft.Azure.Cosmos.Query.Core.Monads;
+
+    /// <summary>
+    /// Implementation of <see cref="CosmosQueryExecutionContext"/> that composes another context and defers it's initialization until the first read.
+    /// </summary>
+    internal sealed class LazyCosmosQueryExecutionContext : CosmosQueryExecutionContext
+    {
+        private readonly AsyncLazy<TryCatch<CosmosQueryExecutionContext>> lazyTryCreateCosmosQueryExecutionContext;
+
+        public LazyCosmosQueryExecutionContext(AsyncLazy<TryCatch<CosmosQueryExecutionContext>> lazyTryCreateCosmosQueryExecutionContext)
+        {
+            if (lazyTryCreateCosmosQueryExecutionContext == null)
+            {
+                throw new ArgumentNullException(nameof(lazyTryCreateCosmosQueryExecutionContext));
+            }
+
+            this.lazyTryCreateCosmosQueryExecutionContext = lazyTryCreateCosmosQueryExecutionContext;
+        }
+
+        public override bool IsDone
+        {
+            get
+            {
+                bool isDone;
+                if (this.lazyTryCreateCosmosQueryExecutionContext.ValueInitialized)
+                {
+                    TryCatch<CosmosQueryExecutionContext> tryCreateCosmosQueryExecutionContext = this.lazyTryCreateCosmosQueryExecutionContext.Result;
+                    if (tryCreateCosmosQueryExecutionContext.Succeeded)
+                    {
+                        isDone = tryCreateCosmosQueryExecutionContext.Result.IsDone;
+                    }
+                    else
+                    {
+                        isDone = true;
+                    }
+                }
+                else
+                {
+                    isDone = false;
+                }
+
+                return isDone;
+            }
+        }
+
+        public override void Dispose()
+        {
+            if (this.lazyTryCreateCosmosQueryExecutionContext.ValueInitialized)
+            {
+                TryCatch<CosmosQueryExecutionContext> tryCreateCosmosQueryExecutionContext = this.lazyTryCreateCosmosQueryExecutionContext.Result;
+                if (tryCreateCosmosQueryExecutionContext.Succeeded)
+                {
+                    tryCreateCosmosQueryExecutionContext.Result.Dispose();
+                }
+            }
+        }
+
+        public override async Task<QueryResponseCore> ExecuteNextAsync(CancellationToken cancellationToken)
+        {
+            cancellationToken.ThrowIfCancellationRequested();
+
+            TryCatch<CosmosQueryExecutionContext> tryCreateCosmosQueryExecutionContext = await this.lazyTryCreateCosmosQueryExecutionContext.GetValueAsync(cancellationToken);
+            if (!tryCreateCosmosQueryExecutionContext.Succeeded)
+            {
+                return QueryResponseFactory.CreateFromException(tryCreateCosmosQueryExecutionContext.Exception);
+            }
+
+            CosmosQueryExecutionContext cosmosQueryExecutionContext = tryCreateCosmosQueryExecutionContext.Result;
+            QueryResponseCore queryResponseCore = await cosmosQueryExecutionContext.ExecuteNextAsync(cancellationToken);
+            return queryResponseCore;
+        }
+
+        public override bool TryGetContinuationToken(out string state)
+        {
+            if (!this.lazyTryCreateCosmosQueryExecutionContext.ValueInitialized)
+            {
+                state = null;
+                return false;
+            }
+
+            TryCatch<CosmosQueryExecutionContext> tryCreateCosmosQueryExecutionContext = this.lazyTryCreateCosmosQueryExecutionContext.Result;
+            if (!tryCreateCosmosQueryExecutionContext.Succeeded)
+            {
+                state = null;
+                return false;
+            }
+
+            return tryCreateCosmosQueryExecutionContext.Result.TryGetContinuationToken(out state);
+        }
+    }
+}

Microsoft.Azure.Cosmos/src/Query/Core/Monads/TryCatch.cs

@@ -78,14 +78,7 @@ public TryCatch<T> Try<T>(
             TryCatch<T> matchResult;
             if (this.Succeeded)
             {
-                try
-                {
-                    matchResult = TryCatch<T>.FromResult(onSuccess(this.either.FromRight(default)));
-                }
-                catch (Exception ex)
-                {
-                    matchResult = TryCatch<T>.FromException(ex);
-                }
+                matchResult = TryCatch<T>.FromResult(onSuccess(this.either.FromRight(default)));
             }
             else
             {
@@ -101,14 +94,7 @@ public async Task<TryCatch<T>> TryAsync<T>(
             TryCatch<T> matchResult;
             if (this.Succeeded)
             {
-                try
-                {
-                    matchResult = TryCatch<T>.FromResult(await onSuccess(this.either.FromRight(default)));
-                }
-                catch (Exception ex)
-                {
-                    matchResult = TryCatch<T>.FromException(ex);
-                }
+                matchResult = TryCatch<T>.FromResult(await onSuccess(this.either.FromRight(default)));
             }
             else
             {

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

@@ -14,6 +14,7 @@ namespace Microsoft.Azure.Cosmos.Query
     internal struct QueryResponseCore
     {
         private static readonly IReadOnlyList<CosmosElement> EmptyList = new List<CosmosElement>().AsReadOnly();
+        internal static readonly string EmptyGuidString = Guid.Empty.ToString();
         internal static readonly IReadOnlyCollection<QueryPageDiagnostics> EmptyDiagnostics = new List<QueryPageDiagnostics>();
 
         private QueryResponseCore(

Microsoft.Azure.Cosmos/src/Query/Core/QueryResponseFactory.cs

@@ -0,0 +1,65 @@
+// ------------------------------------------------------------
+// Copyright (c) Microsoft Corporation.  All rights reserved.
+// ------------------------------------------------------------
+
+namespace Microsoft.Azure.Cosmos.Query.Core
+{
+    using System;
+
+    internal static class QueryResponseFactory
+    {
+        private static readonly string EmptyGuidString = Guid.Empty.ToString();
+
+        public static QueryResponseCore CreateFromException(Exception exception)
+        {
+            QueryResponseCore queryResponseCore;
+            if (exception is CosmosException cosmosException)
+            {
+                queryResponseCore = CreateFromCosmosException(cosmosException);
+            }
+            else if (exception is Microsoft.Azure.Documents.DocumentClientException documentClientException)
+            {
+                queryResponseCore = CreateFromDocumentClientException(documentClientException);
+            }
+            else
+            {
+                // Unknown exception type should become a 500
+                queryResponseCore = QueryResponseCore.CreateFailure(
+                    statusCode: System.Net.HttpStatusCode.InternalServerError,
+                    subStatusCodes: null,
+                    errorMessage: exception.ToString(),
+                    requestCharge: 0,
+                    activityId: QueryResponseCore.EmptyGuidString,
+                    diagnostics: QueryResponseCore.EmptyDiagnostics);
+            }
+
+            return queryResponseCore;
+        }
+
+        private static QueryResponseCore CreateFromCosmosException(CosmosException cosmosException)
+        {
+            QueryResponseCore queryResponseCore = QueryResponseCore.CreateFailure(
+                statusCode: cosmosException.StatusCode,
+                subStatusCodes: (Microsoft.Azure.Documents.SubStatusCodes)cosmosException.SubStatusCode,
+                errorMessage: cosmosException.Message,
+                requestCharge: 0,
+                activityId: cosmosException.ActivityId,
+                diagnostics: QueryResponseCore.EmptyDiagnostics);
+
+            return queryResponseCore;
+        }
+
+        private static QueryResponseCore CreateFromDocumentClientException(Microsoft.Azure.Documents.DocumentClientException documentClientException)
+        {
+            QueryResponseCore queryResponseCore = QueryResponseCore.CreateFailure(
+                statusCode: documentClientException.StatusCode.GetValueOrDefault(System.Net.HttpStatusCode.InternalServerError),
+                subStatusCodes: null,
+                errorMessage: documentClientException.Message,
+                requestCharge: 0,
+                activityId: documentClientException.ActivityId,
+                diagnostics: QueryResponseCore.EmptyDiagnostics);
+
+            return queryResponseCore;
+        }
+    }
+}

Microsoft.Azure.Cosmos/src/Query/v2Query/DefaultDocumentQueryExecutionContext.cs

@@ -22,6 +22,8 @@ namespace Microsoft.Azure.Cosmos.Query
     /// </summary>
     internal sealed class DefaultDocumentQueryExecutionContext : DocumentQueryExecutionContextBase
     {
+        private const string InternalPartitionKeyDefinitionProperty = "x-ms-query-partitionkey-definition";
+
         /// <summary>
         /// Whether or not a continuation is expected.
         /// </summary>
@@ -280,8 +282,9 @@ private static bool ServiceInteropAvailable()
                 {
                     FeedOptions feedOptions = this.GetFeedOptions(null);
                     PartitionKeyDefinition partitionKeyDefinition;
-                    object partitionKeyDefinitionObject;
-                    if (feedOptions.Properties != null && feedOptions.Properties.TryGetValue(CosmosQueryExecutionContextFactory.InternalPartitionKeyDefinitionProperty, out partitionKeyDefinitionObject))
+                    if ((feedOptions.Properties != null) && feedOptions.Properties.TryGetValue(
+                        DefaultDocumentQueryExecutionContext.InternalPartitionKeyDefinitionProperty,
+                        out object partitionKeyDefinitionObject))
                     {
                         if (partitionKeyDefinitionObject is PartitionKeyDefinition definition)
                         {
@@ -299,7 +302,6 @@ private static bool ServiceInteropAvailable()
                         partitionKeyDefinition = collection.PartitionKey;
                     }
 
-                    QueryInfo queryInfo;
                     providedRanges = PartitionRoutingHelper.GetProvidedPartitionKeyRanges(
                         (errorMessage) => new BadRequestException(errorMessage),
                         this.QuerySpec,
@@ -310,7 +312,7 @@ private static bool ServiceInteropAvailable()
                         partitionKeyDefinition,
                         queryPartitionProvider,
                         version,
-                        out queryInfo);
+                        out QueryInfo queryInfo);
                 }
                 else if (request.Properties != null && request.Properties.TryGetValue(
                     WFConstants.BackendHeaders.EffectivePartitionKeyString,
@@ -367,7 +369,7 @@ private async Task<DocumentServiceRequest> CreateRequestAsync()
             INameValueCollection requestHeaders = await this.CreateCommonHeadersAsync(
                     this.GetFeedOptions(this.ContinuationToken));
 
-            requestHeaders[HttpConstants.HttpHeaders.IsContinuationExpected] = isContinuationExpected.ToString();
+            requestHeaders[HttpConstants.HttpHeaders.IsContinuationExpected] = this.isContinuationExpected.ToString();
 
             return this.CreateDocumentServiceRequest(
                 requestHeaders,