Query: Auto-buffer data readers in SqlServer when split query & MARS off

Resolves #21420

src/EFCore.SqlServer/Extensions/SqlServerServiceCollectionExtensions.cs

@@ -73,8 +73,8 @@ public static IServiceCollection AddEntityFrameworkSqlServer([NotNull] this ISer
                 .TryAdd<IHistoryRepository, SqlServerHistoryRepository>()
                 .TryAdd<IExecutionStrategyFactory, SqlServerExecutionStrategyFactory>()
                 .TryAdd<IRelationalQueryStringFactory, SqlServerQueryStringFactory>()
-
-                // New Query Pipeline
+                .TryAdd<ICompiledQueryCacheKeyGenerator, SqlServerCompiledQueryCacheKeyGenerator>()
+                .TryAdd<IQueryCompilationContextFactory, SqlServerQueryCompilationContextFactory>()
                 .TryAdd<IMethodCallTranslatorProvider, SqlServerMethodCallTranslatorProvider>()
                 .TryAdd<IMemberTranslatorProvider, SqlServerMemberTranslatorProvider>()
                 .TryAdd<IQuerySqlGeneratorFactory, SqlServerQuerySqlGeneratorFactory>()

src/EFCore.SqlServer/Query/Internal/SqlServerCompiledQueryCacheKeyGenerator.cs

@@ -0,0 +1,75 @@
+// 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.Linq.Expressions;
+using JetBrains.Annotations;
+using Microsoft.EntityFrameworkCore.Query;
+using Microsoft.EntityFrameworkCore.SqlServer.Storage.Internal;
+using Microsoft.EntityFrameworkCore.Utilities;
+
+namespace Microsoft.EntityFrameworkCore.SqlServer.Query.Internal
+{
+    /// <summary>
+    ///     This is an internal API that supports the Entity Framework Core infrastructure and not subject to
+    ///     the same compatibility standards as public APIs. It may be changed or removed without notice in
+    ///     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 SqlServerCompiledQueryCacheKeyGenerator : RelationalCompiledQueryCacheKeyGenerator
+    {
+        private readonly ISqlServerConnection _sqlServerConnection;
+
+        /// <summary>
+        ///     This is an internal API that supports the Entity Framework Core infrastructure and not subject to
+        ///     the same compatibility standards as public APIs. It may be changed or removed without notice in
+        ///     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 SqlServerCompiledQueryCacheKeyGenerator(
+            [NotNull] CompiledQueryCacheKeyGeneratorDependencies dependencies,
+            [NotNull] RelationalCompiledQueryCacheKeyGeneratorDependencies relationalDependencies,
+            [NotNull] ISqlServerConnection sqlServerConnection)
+            : base(dependencies, relationalDependencies)
+        {
+            Check.NotNull(sqlServerConnection, nameof(sqlServerConnection));
+
+            _sqlServerConnection = sqlServerConnection;
+        }
+
+        /// <summary>
+        ///     This is an internal API that supports the Entity Framework Core infrastructure and not subject to
+        ///     the same compatibility standards as public APIs. It may be changed or removed without notice in
+        ///     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 object GenerateCacheKey(Expression query, bool async)
+            => new SqlServerCompiledQueryCacheKey(
+                GenerateCacheKeyCore(query, async),
+                _sqlServerConnection.IsMultipleActiveResultSetsEnabled);
+
+        private readonly struct SqlServerCompiledQueryCacheKey
+        {
+            private readonly RelationalCompiledQueryCacheKey _relationalCompiledQueryCacheKey;
+            private readonly bool _multipleActiveResultSetsEnabled;
+
+            public SqlServerCompiledQueryCacheKey(
+                RelationalCompiledQueryCacheKey relationalCompiledQueryCacheKey, bool multipleActiveResultSetsEnabled)
+            {
+                _relationalCompiledQueryCacheKey = relationalCompiledQueryCacheKey;
+                _multipleActiveResultSetsEnabled = multipleActiveResultSetsEnabled;
+            }
+
+            public override bool Equals(object obj)
+                => !(obj is null)
+                   && obj is SqlServerCompiledQueryCacheKey sqlServerCompiledQueryCacheKey
+                   && Equals(sqlServerCompiledQueryCacheKey);
+
+            private bool Equals(SqlServerCompiledQueryCacheKey other)
+                => _relationalCompiledQueryCacheKey.Equals(other._relationalCompiledQueryCacheKey)
+                   && _multipleActiveResultSetsEnabled == other._multipleActiveResultSetsEnabled;
+
+            public override int GetHashCode() => HashCode.Combine(_relationalCompiledQueryCacheKey, _multipleActiveResultSetsEnabled);
+        }
+    }
+}

src/EFCore.SqlServer/Query/Internal/SqlServerQueryCompilationContext.cs

@@ -0,0 +1,46 @@
+// 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 JetBrains.Annotations;
+using Microsoft.EntityFrameworkCore.Query;
+using Microsoft.EntityFrameworkCore.Storage;
+
+namespace Microsoft.EntityFrameworkCore.SqlServer.Query.Internal
+{
+    /// <summary>
+    ///     This is an internal API that supports the Entity Framework Core infrastructure and not subject to
+    ///     the same compatibility standards as public APIs. It may be changed or removed without notice in
+    ///     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 SqlServerQueryCompilationContext : RelationalQueryCompilationContext
+    {
+        private readonly bool _multipleActiveResultSetsEnabled;
+
+        /// <summary>
+        ///     This is an internal API that supports the Entity Framework Core infrastructure and not subject to
+        ///     the same compatibility standards as public APIs. It may be changed or removed without notice in
+        ///     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 SqlServerQueryCompilationContext(
+            [NotNull] QueryCompilationContextDependencies dependencies,
+            [NotNull] RelationalQueryCompilationContextDependencies relationalDependencies,
+            bool async,
+            bool multipleActiveResultSetsEnabled)
+            : base(dependencies, relationalDependencies, async)
+        {
+            _multipleActiveResultSetsEnabled = multipleActiveResultSetsEnabled;
+        }
+
+        /// <summary>
+        ///     This is an internal API that supports the Entity Framework Core infrastructure and not subject to
+        ///     the same compatibility standards as public APIs. It may be changed or removed without notice in
+        ///     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 bool IsBuffering => base.IsBuffering
+            || (QuerySplittingBehavior == EntityFrameworkCore.QuerySplittingBehavior.SplitQuery
+                && !_multipleActiveResultSetsEnabled);
+    }
+}

src/EFCore.SqlServer/Query/Internal/SqlServerQueryCompilationContextFactory.cs

@@ -0,0 +1,61 @@
+// 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 JetBrains.Annotations;
+using Microsoft.EntityFrameworkCore.Query;
+using Microsoft.EntityFrameworkCore.SqlServer.Storage.Internal;
+using Microsoft.EntityFrameworkCore.Utilities;
+using Microsoft.Extensions.DependencyInjection;
+
+namespace Microsoft.EntityFrameworkCore.SqlServer.Query.Internal
+{
+    /// <summary>
+    ///     <para>
+    ///         This is an internal API that supports the Entity Framework Core infrastructure and not subject to
+    ///         the same compatibility standards as public APIs. It may be changed or removed without notice in
+    ///         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.
+    ///     </para>
+    ///     <para>
+    ///         The service lifetime is <see cref="ServiceLifetime.Scoped" />. This means that each
+    ///         <see cref="DbContext" /> instance will use its own instance of this service.
+    ///         The implementation may depend on other services registered with any lifetime.
+    ///         The implementation does not need to be thread-safe.
+    ///     </para>
+    /// </summary>
+    public class SqlServerQueryCompilationContextFactory : IQueryCompilationContextFactory
+    {
+        private readonly QueryCompilationContextDependencies _dependencies;
+        private readonly RelationalQueryCompilationContextDependencies _relationalDependencies;
+        private readonly ISqlServerConnection _sqlServerConnection;
+
+        /// <summary>
+        ///     This is an internal API that supports the Entity Framework Core infrastructure and not subject to
+        ///     the same compatibility standards as public APIs. It may be changed or removed without notice in
+        ///     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 SqlServerQueryCompilationContextFactory(
+            [NotNull] QueryCompilationContextDependencies dependencies,
+            [NotNull] RelationalQueryCompilationContextDependencies relationalDependencies,
+            [NotNull] ISqlServerConnection sqlServerConnection)
+        {
+            Check.NotNull(dependencies, nameof(dependencies));
+            Check.NotNull(relationalDependencies, nameof(relationalDependencies));
+
+            _dependencies = dependencies;
+            _relationalDependencies = relationalDependencies;
+            _sqlServerConnection = sqlServerConnection;
+        }
+
+        /// <summary>
+        ///     This is an internal API that supports the Entity Framework Core infrastructure and not subject to
+        ///     the same compatibility standards as public APIs. It may be changed or removed without notice in
+        ///     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 virtual QueryCompilationContext Create(bool async)
+            => new SqlServerQueryCompilationContext(
+                _dependencies, _relationalDependencies, async, _sqlServerConnection.IsMultipleActiveResultSetsEnabled);
+    }
+}

src/EFCore.SqlServer/Storage/Internal/ISqlServerConnection.cs

@@ -29,5 +29,13 @@ public interface ISqlServerConnection : IRelationalConnection
         ///     doing so can result in application failures when updating to a new Entity Framework Core release.
         /// </summary>
         ISqlServerConnection CreateMasterConnection();
+
+        /// <summary>
+        ///     This is an internal API that supports the Entity Framework Core infrastructure and not subject to
+        ///     the same compatibility standards as public APIs. It may be changed or removed without notice in
+        ///     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>
+        bool IsMultipleActiveResultSetsEnabled { get; }
     }
 }

src/EFCore.SqlServer/Storage/Internal/SqlServerConnection.cs

@@ -1,6 +1,7 @@
 // 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.Collections.Concurrent;
 using System.Data.Common;
 using JetBrains.Annotations;
 using Microsoft.Data.SqlClient;
@@ -27,6 +28,8 @@ public class SqlServerConnection : RelationalConnection, ISqlServerConnection
     {
         // Compensate for slow SQL Server database creation
         private const int DefaultMasterConnectionCommandTimeout = 60;
+        private static readonly ConcurrentDictionary<string, bool> _multipleActiveResultSetsEnabledMap =
+            new ConcurrentDictionary<string, bool>();
 
         /// <summary>
         ///     This is an internal API that supports the Entity Framework Core infrastructure and not subject to
@@ -87,6 +90,16 @@ public virtual ISqlServerConnection CreateMasterConnection()
             return new SqlServerConnection(Dependencies.With(contextOptions));
         }
 
+        /// <summary>
+        ///     This is an internal API that supports the Entity Framework Core infrastructure and not subject to
+        ///     the same compatibility standards as public APIs. It may be changed or removed without notice in
+        ///     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 virtual bool IsMultipleActiveResultSetsEnabled
+            => _multipleActiveResultSetsEnabledMap.GetOrAdd(
+                ConnectionString, cs => new SqlConnectionStringBuilder(cs).MultipleActiveResultSets);
+
         /// <summary>
         ///     Indicates whether the store connection supports ambient transactions
         /// </summary>

test/EFCore.SqlServer.FunctionalTests/Query/QueryBugsTest.cs

@@ -7423,7 +7423,7 @@ ORDER BY [p].[Id]"
         }
 
         [ConditionalFact]
-        public virtual void Using_AsSingleQuery_withouth_context_configuration_does_not_throw_warning()
+        public virtual void Using_AsSingleQuery_without_context_configuration_does_not_throw_warning()
         {
             var (options, testSqlLoggerFactory) = CreateOptions21355(null);
             using var context = new BugContext21355(options);
@@ -7442,16 +7442,16 @@ FROM [Parents] AS [p]
         }
 
         [ConditionalFact]
-        public virtual void Using_AsSplitQuery_withouth_context_configuration_does_not_throw_warning()
+        public virtual void Using_AsSplitQuery_without_context_configuration_does_not_throw_warning()
         {
             var (options, testSqlLoggerFactory) = CreateOptions21355(null);
             using var context = new BugContext21355(options);
 
             context.Parents.Include(p => p.Children1).Include(p => p.Children2).AsSplitQuery().ToList();
 
             testSqlLoggerFactory.AssertBaseline(
-                            new[]
-                            {
+                new[]
+                {
                     @"SELECT [p].[Id]
 FROM [Parents] AS [p]
 ORDER BY [p].[Id]",
@@ -7465,7 +7465,7 @@ FROM [Parents] AS [p]
 FROM [Parents] AS [p]
 INNER JOIN [AnotherChild21355] AS [a] ON [p].[Id] = [a].[ParentId]
 ORDER BY [p].[Id]"
-                            });
+                });
         }
 
         [ConditionalFact]
@@ -7524,6 +7524,21 @@ public virtual async Task SplitQuery_disposes_inner_data_readers_single_async()
             Assert.Equal(ConnectionState.Closed, dbConnection.State);
         }
 
+        [ConditionalFact]
+        public virtual void Using_AsSplitQuery_without_multiple_active_result_sets_work()
+        {
+            var (options, testSqlLoggerFactory) = CreateOptions21355(null, mars: true);
+            using var context = new BugContext21355(options);
+
+            context.Parents.Include(p => p.Children1).Include(p => p.Children2).AsSplitQuery().ToList();
+
+            var connectionStringWithoutMars = SqlServerTestStore.CreateConnectionString("QueryBugsTest", multipleActiveResultSets: false);
+            var connection = context.GetService<IRelationalConnection>();
+            connection.ConnectionString = connectionStringWithoutMars;
+
+            context.Parents.Include(p => p.Children1).Include(p => p.Children2).AsSplitQuery().ToList();
+        }
+
         private class Parent21355
         {
             public string Id { get; set; }
@@ -7545,9 +7560,9 @@ private class AnotherChild21355
             public Parent21355 Parent { get; set; }
         }
 
-        private (DbContextOptions, TestSqlLoggerFactory) CreateOptions21355(QuerySplittingBehavior? querySplittingBehavior)
+        private (DbContextOptions, TestSqlLoggerFactory) CreateOptions21355(QuerySplittingBehavior? querySplittingBehavior, bool mars = true)
         {
-            var testStore = SqlServerTestStore.CreateInitialized("QueryBugsTest", multipleActiveResultSets: true);
+            var testStore = SqlServerTestStore.CreateInitialized("QueryBugsTest", multipleActiveResultSets: mars);
             var testSqlLoggerFactory = new TestSqlLoggerFactory();
             var serviceProvider = new ServiceCollection().AddSingleton<ILoggerFactory>(testSqlLoggerFactory).BuildServiceProvider();