CosmosQueryExecutionContextFactory Refactory

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

@@ -0,0 +1,55 @@
+// ------------------------------------------------------------
+// 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)
+        {
+            // Note that this class is not thread safe.
+            // if the valueFactory has side effects than this will have issues.
+            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/CatchAllCosmosQueryExecutionContext.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 CatchAllCosmosQueryExecutionContext : CosmosQueryExecutionContext
+    {
+        private readonly CosmosQueryExecutionContext cosmosQueryExecutionContext;
+        private bool hitException;
+
+        public CatchAllCosmosQueryExecutionContext(
+            CosmosQueryExecutionContext cosmosQueryExecutionContext)
+        {
+            if (cosmosQueryExecutionContext == null)
+            {
+                throw new ArgumentNullException(nameof(cosmosQueryExecutionContext));
+            }
+
+            this.cosmosQueryExecutionContext = cosmosQueryExecutionContext;
+        }
+
+        public override bool IsDone => this.hitException || this.cosmosQueryExecutionContext.IsDone;
+
+        public override void Dispose()
+        {
+            this.cosmosQueryExecutionContext.Dispose();
+        }
+
+        public override async Task<QueryResponseCore> ExecuteNextAsync(CancellationToken cancellationToken)
+        {
+            if (this.IsDone)
+            {
+                throw new InvalidOperationException(
+                    $"Can not {nameof(ExecuteNextAsync)} from a {nameof(CosmosQueryExecutionContext)} where {nameof(this.IsDone)}.");
+            }
+
+            cancellationToken.ThrowIfCancellationRequested();
+
+            QueryResponseCore queryResponseCore;
+            try
+            {
+                queryResponseCore = await this.cosmosQueryExecutionContext.ExecuteNextAsync(cancellationToken);
+            }
+            catch (Exception ex)
+            {
+                queryResponseCore = QueryResponseFactory.CreateFromException(ex);
+            }
+
+            if (!queryResponseCore.IsSuccess)
+            {
+                this.hitException = true;
+            }
+
+            return queryResponseCore;
+        }
+
+        public override bool TryGetContinuationToken(out string continuationToken)
+        {
+            return this.cosmosQueryExecutionContext.TryGetContinuationToken(out continuationToken);
+        }
+    }
+}

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