Query cleanup

src/EFCore.Cosmos/Query/Internal/CosmosQueryContextFactory.cs

@@ -13,8 +13,9 @@ namespace Microsoft.EntityFrameworkCore.Cosmos.Query.Internal
     ///     any release. You should only use it directly in your code with extreme caution and knowing that
     ///     doing so can result in application failures when updating to a new Entity Framework Core release.
     /// </summary>
-    public class CosmosQueryContextFactory : QueryContextFactory
+    public class CosmosQueryContextFactory : IQueryContextFactory
     {
+        private readonly QueryContextDependencies _dependencies;
         private readonly CosmosClientWrapper _cosmosClient;
 
         /// <summary>
@@ -26,8 +27,8 @@ public class CosmosQueryContextFactory : QueryContextFactory
         public CosmosQueryContextFactory(
             [NotNull] QueryContextDependencies dependencies,
             [NotNull] CosmosClientWrapper cosmosClient)
-            : base(dependencies)
         {
+            _dependencies = dependencies;
             _cosmosClient = cosmosClient;
         }
 
@@ -37,7 +38,7 @@ public CosmosQueryContextFactory(
         ///     any release. You should only use it directly in your code with extreme caution and knowing that
         ///     doing so can result in application failures when updating to a new Entity Framework Core release.
         /// </summary>
-        public override QueryContext Create()
-            => new CosmosQueryContext(Dependencies, _cosmosClient);
+        public virtual QueryContext Create()
+            => new CosmosQueryContext(_dependencies, _cosmosClient);
     }
 }

src/EFCore.Cosmos/Query/Internal/CosmosShapedQueryCompilingExpressionVisitor.AsyncQueryingEnumerable.cs

@@ -0,0 +1,122 @@
+// Copyright (c) .NET Foundation. All rights reserved.
+// Licensed under the Apache License, Version 2.0. See License.txt in the project root for license information.
+
+using System;
+using System.Collections.Generic;
+using System.Threading;
+using System.Threading.Tasks;
+using Microsoft.EntityFrameworkCore.Diagnostics;
+using Microsoft.EntityFrameworkCore.Query;
+using Newtonsoft.Json.Linq;
+
+namespace Microsoft.EntityFrameworkCore.Cosmos.Query.Internal
+{
+    public partial class CosmosShapedQueryCompilingExpressionVisitor
+    {
+        private class AsyncQueryingEnumerable<T> : IAsyncEnumerable<T>
+        {
+            private readonly CosmosQueryContext _cosmosQueryContext;
+            private readonly SelectExpression _selectExpression;
+            private readonly Func<QueryContext, JObject, T> _shaper;
+            private readonly ISqlExpressionFactory _sqlExpressionFactory;
+            private readonly IQuerySqlGeneratorFactory _querySqlGeneratorFactory;
+            private readonly Type _contextType;
+            private readonly IDiagnosticsLogger<DbLoggerCategory.Query> _logger;
+
+            public AsyncQueryingEnumerable(
+                CosmosQueryContext cosmosQueryContext,
+                ISqlExpressionFactory sqlExpressionFactory,
+                IQuerySqlGeneratorFactory querySqlGeneratorFactory,
+                SelectExpression selectExpression,
+                Func<QueryContext, JObject, T> shaper,
+                Type contextType,
+                IDiagnosticsLogger<DbLoggerCategory.Query> logger)
+            {
+                _cosmosQueryContext = cosmosQueryContext;
+                _sqlExpressionFactory = sqlExpressionFactory;
+                _querySqlGeneratorFactory = querySqlGeneratorFactory;
+                _selectExpression = selectExpression;
+                _shaper = shaper;
+                _contextType = contextType;
+                _logger = logger;
+            }
+
+            public IAsyncEnumerator<T> GetAsyncEnumerator(CancellationToken cancellationToken = default)
+            {
+                return new AsyncEnumerator(this, cancellationToken);
+            }
+
+            private sealed class AsyncEnumerator : IAsyncEnumerator<T>
+            {
+                private IAsyncEnumerator<JObject> _enumerator;
+                private readonly CosmosQueryContext _cosmosQueryContext;
+                private readonly SelectExpression _selectExpression;
+                private readonly Func<QueryContext, JObject, T> _shaper;
+                private readonly ISqlExpressionFactory _sqlExpressionFactory;
+                private readonly IQuerySqlGeneratorFactory _querySqlGeneratorFactory;
+                private readonly Type _contextType;
+                private readonly IDiagnosticsLogger<DbLoggerCategory.Query> _logger;
+                private readonly CancellationToken _cancellationToken;
+
+                public AsyncEnumerator(AsyncQueryingEnumerable<T> queryingEnumerable, CancellationToken cancellationToken)
+                {
+                    _cosmosQueryContext = queryingEnumerable._cosmosQueryContext;
+                    _shaper = queryingEnumerable._shaper;
+                    _selectExpression = queryingEnumerable._selectExpression;
+                    _sqlExpressionFactory = queryingEnumerable._sqlExpressionFactory;
+                    _querySqlGeneratorFactory = queryingEnumerable._querySqlGeneratorFactory;
+                    _contextType = queryingEnumerable._contextType;
+                    _logger = queryingEnumerable._logger;
+                    _cancellationToken = cancellationToken;
+                }
+
+                public T Current { get; private set; }
+
+                public async ValueTask<bool> MoveNextAsync()
+                {
+                    try
+                    {
+                        using (_cosmosQueryContext.ConcurrencyDetector.EnterCriticalSection())
+                        {
+                            if (_enumerator == null)
+                            {
+                                var selectExpression = (SelectExpression)new InExpressionValuesExpandingExpressionVisitor(
+                                   _sqlExpressionFactory, _cosmosQueryContext.ParameterValues).Visit(_selectExpression);
+
+                                _enumerator = _cosmosQueryContext.CosmosClient
+                                    .ExecuteSqlQueryAsync(
+                                        _selectExpression.Container,
+                                        _querySqlGeneratorFactory.Create().GetSqlQuery(selectExpression, _cosmosQueryContext.ParameterValues))
+                                    .GetAsyncEnumerator(_cancellationToken);
+
+                            }
+
+                            var hasNext = await _enumerator.MoveNextAsync();
+
+                            Current
+                                = hasNext
+                                    ? _shaper(_cosmosQueryContext, _enumerator.Current)
+                                    : default;
+
+                            return hasNext;
+                        }
+                    }
+                    catch (Exception exception)
+                    {
+                        _logger.QueryIterationFailed(_contextType, exception);
+
+                        throw;
+                    }
+                }
+
+                public ValueTask DisposeAsync()
+                {
+                    _enumerator?.DisposeAsync();
+                    _enumerator = null;
+
+                    return default;
+                }
+            }
+        }
+    }
+}