Fix auto flush for Enumerable and AsyncEnumerable methods

src/NHibernate.Test/Async/GenericTest/Methods/Fixture.cs

@@ -24,7 +24,7 @@ protected override string[] Mappings
 		{
 			get
 			{
-				return new string[] { "One.hbm.xml", "Many.hbm.xml" };
+				return new string[] { "One.hbm.xml", "Many.hbm.xml", "Simple.hbm.xml" };
 			}
 		}
 
@@ -49,12 +49,15 @@ protected override void OnSetUp()
 			many2.One = one;
 			one.Manies.Add( many2 );
 
-			using( ISession s = OpenSession() )
+			var simple = new Simple(1) {Count = 1};
+
+			using ( ISession s = OpenSession() )
 			using( ITransaction t = s.BeginTransaction() )
 			{
 				s.Save( one );
 				s.Save( many1 );
 				s.Save( many2 );
+				s.Save(simple, 1);
 				t.Commit();
 			}
 		}
@@ -66,6 +69,7 @@ protected override void OnTearDown()
 			{
 				session.Delete( "from Many" );
 				session.Delete( "from One" );
+				session.Delete("from Simple");
 				tx.Commit();
 			}
 			base.OnTearDown();

src/NHibernate.Test/GenericTest/Methods/Fixture.cs

@@ -14,7 +14,7 @@ protected override string[] Mappings
 		{
 			get
 			{
-				return new string[] { "One.hbm.xml", "Many.hbm.xml" };
+				return new string[] { "One.hbm.xml", "Many.hbm.xml", "Simple.hbm.xml" };
 			}
 		}
 
@@ -39,12 +39,15 @@ protected override void OnSetUp()
 			many2.One = one;
 			one.Manies.Add( many2 );
 
-			using( ISession s = OpenSession() )
+			var simple = new Simple(1) {Count = 1};
+
+			using ( ISession s = OpenSession() )
 			using( ITransaction t = s.BeginTransaction() )
 			{
 				s.Save( one );
 				s.Save( many1 );
 				s.Save( many2 );
+				s.Save(simple, 1);
 				t.Commit();
 			}
 		}
@@ -56,6 +59,7 @@ protected override void OnTearDown()
 			{
 				session.Delete( "from Many" );
 				session.Delete( "from One" );
+				session.Delete("from Simple");
 				tx.Commit();
 			}
 			base.OnTearDown();
@@ -106,6 +110,40 @@ public void QueryEnumerable()
 			}
 		}
 
+		[Test]
+		public void AutoFlushQueryEnumerable()
+		{
+			using (var s = OpenSession())
+			using (var t = s.BeginTransaction())
+			{
+				Assert.That(s.FlushMode, Is.EqualTo(FlushMode.Auto));
+				var results = s.CreateQuery("from Simple").Enumerable<Simple>();
+
+				var id = 2;
+				var simple = new Simple(id) {Count = id};
+				s.Save(simple, id);
+				var enumerator = results.GetEnumerator();
+
+				Assert.That(enumerator.MoveNext(), Is.True);
+				Assert.That(enumerator.MoveNext(), Is.True);
+				Assert.That(enumerator.MoveNext(), Is.False);
+				enumerator.Dispose();
+
+				id++;
+				simple = new Simple(id) {Count = id};
+				s.Save(simple, id);
+				enumerator = results.GetEnumerator();
+
+				Assert.That(enumerator.MoveNext(), Is.True);
+				Assert.That(enumerator.MoveNext(), Is.True);
+				Assert.That(enumerator.MoveNext(), Is.True);
+				Assert.That(enumerator.MoveNext(), Is.False);
+				enumerator.Dispose();
+
+				t.Rollback();
+			}
+		}
+
 		[Test]
 		public async Task QueryEnumerableAsync()
 		{
@@ -120,6 +158,40 @@ public async Task QueryEnumerableAsync()
 			}
 		}
 
+		[Test]
+		public async Task AutoFlushQueryEnumerableAsync()
+		{
+			using (var s = OpenSession())
+			using (var t = s.BeginTransaction())
+			{
+				Assert.That(s.FlushMode, Is.EqualTo(FlushMode.Auto));
+				var results = s.CreateQuery("from Simple").AsyncEnumerable<Simple>();
+
+				var id = 2;
+				var simple = new Simple(id) {Count = id};
+				s.Save(simple, id);
+				var enumerator = results.GetAsyncEnumerator();
+
+				Assert.That(await enumerator.MoveNextAsync(), Is.True);
+				Assert.That(await enumerator.MoveNextAsync(), Is.True);
+				Assert.That(await enumerator.MoveNextAsync(), Is.False);
+				await enumerator.DisposeAsync();
+
+				id++;
+				simple = new Simple(id) {Count = id};
+				s.Save(simple, id);
+				enumerator = results.GetAsyncEnumerator();
+
+				Assert.That(await enumerator.MoveNextAsync(), Is.True);
+				Assert.That(await enumerator.MoveNextAsync(), Is.True);
+				Assert.That(await enumerator.MoveNextAsync(), Is.True);
+				Assert.That(await enumerator.MoveNextAsync(), Is.False);
+				await enumerator.DisposeAsync();
+
+				await t.RollbackAsync();
+			}
+		}
+
 		[Test]
 		public void Filter()
 		{

src/NHibernate/Async/Impl/SessionImpl.cs

@@ -302,38 +302,38 @@ public override async Task<IQueryTranslator[]> GetQueriesAsync(IQueryExpression
 		/// <inheritdoc />
 		// Since v5.3
 		[Obsolete("Use AsyncEnumerable extension method instead.")]
-		public override async Task<IEnumerable<T>> EnumerableAsync<T>(IQueryExpression queryExpression, QueryParameters queryParameters, CancellationToken cancellationToken)
+		public override Task<IEnumerable<T>> EnumerableAsync<T>(IQueryExpression queryExpression, QueryParameters queryParameters, CancellationToken cancellationToken)
 		{
-			cancellationToken.ThrowIfCancellationRequested();
-			using (BeginProcess())
+			if (cancellationToken.IsCancellationRequested)
 			{
-				queryParameters.ValidateParameters();
-				var plan = GetHQLQueryPlan(queryExpression, true);
-				await (AutoFlushIfRequiredAsync(plan.QuerySpaces, cancellationToken)).ConfigureAwait(false);
-
-				using (SuspendAutoFlush()) //stops flush being called multiple times if this method is recursively called
-				{
-					return plan.PerformIterate<T>(queryParameters, this);
-				}
+				return Task.FromCanceled<IEnumerable<T>>(cancellationToken);
+			}
+			try
+			{
+				return Task.FromResult<IEnumerable<T>>(Enumerable<T>(queryExpression, queryParameters));
+			}
+			catch (Exception ex)
+			{
+				return Task.FromException<IEnumerable<T>>(ex);
 			}
 		}
 
 		/// <inheritdoc />
 		// Since v5.3
 		[Obsolete("Use AsyncEnumerable extension method instead.")]
-		public override async Task<IEnumerable> EnumerableAsync(IQueryExpression queryExpression, QueryParameters queryParameters, CancellationToken cancellationToken)
+		public override Task<IEnumerable> EnumerableAsync(IQueryExpression queryExpression, QueryParameters queryParameters, CancellationToken cancellationToken)
 		{
-			cancellationToken.ThrowIfCancellationRequested();
-			using (BeginProcess())
+			if (cancellationToken.IsCancellationRequested)
 			{
-				queryParameters.ValidateParameters();
-				var plan = GetHQLQueryPlan(queryExpression, true);
-				await (AutoFlushIfRequiredAsync(plan.QuerySpaces, cancellationToken)).ConfigureAwait(false);
-
-				using (SuspendAutoFlush()) //stops flush being called multiple times if this method is recursively called
-				{
-					return plan.PerformIterate(queryParameters, this);
-				}
+				return Task.FromCanceled<IEnumerable>(cancellationToken);
+			}
+			try
+			{
+				return Task.FromResult<IEnumerable>(Enumerable(queryExpression, queryParameters));
+			}
+			catch (Exception ex)
+			{
+				return Task.FromException<IEnumerable>(ex);
 			}
 		}
 

src/NHibernate/Async/Loader/Hql/QueryLoader.cs

@@ -90,22 +90,26 @@ protected override async Task<object[]> GetResultRowAsync(object[] row, DbDataRe
 		internal async Task<InitializeEnumerableResult> InitializeEnumerableAsync(QueryParameters queryParameters, ISessionImplementor session, CancellationToken cancellationToken)
 		{
 			cancellationToken.ThrowIfCancellationRequested();
-			Stopwatch stopWatch = null;
-			if (session.Factory.Statistics.IsStatisticsEnabled)
+			await (session.AutoFlushIfRequiredAsync(_queryTranslator.QuerySpaces, cancellationToken)).ConfigureAwait(false);
+			using (session.SuspendAutoFlush())
 			{
-				stopWatch = Stopwatch.StartNew();
-			}
+				Stopwatch stopWatch = null;
+				if (session.Factory.Statistics.IsStatisticsEnabled)
+				{
+					stopWatch = Stopwatch.StartNew();
+				}
 
-			var command = await (PrepareQueryCommandAsync(queryParameters, false, session, cancellationToken)).ConfigureAwait(false);
-			var dataReader = await (GetResultSetAsync(command, queryParameters, session, null, cancellationToken)).ConfigureAwait(false);
-			if (stopWatch != null)
-			{
-				stopWatch.Stop();
-				session.Factory.StatisticsImplementor.QueryExecuted("HQL: " + _queryTranslator.QueryString, 0, stopWatch.Elapsed);
-				session.Factory.StatisticsImplementor.QueryExecuted(QueryIdentifier, 0, stopWatch.Elapsed);
-			}
+				var command = await (PrepareQueryCommandAsync(queryParameters, false, session, cancellationToken)).ConfigureAwait(false);
+				var dataReader = await (GetResultSetAsync(command, queryParameters, session, null, cancellationToken)).ConfigureAwait(false);
+				if (stopWatch != null)
+				{
+					stopWatch.Stop();
+					session.Factory.StatisticsImplementor.QueryExecuted("HQL: " + _queryTranslator.QueryString, 0, stopWatch.Elapsed);
+					session.Factory.StatisticsImplementor.QueryExecuted(QueryIdentifier, 0, stopWatch.Elapsed);
+				}
 
-			return new InitializeEnumerableResult(command, dataReader);
+				return new InitializeEnumerableResult(command, dataReader);
+			}
 		}
 	}
 }

src/NHibernate/Engine/ISessionImplementor.cs

@@ -64,6 +64,11 @@ internal static void AutoFlushIfRequired(this ISessionImplementor implementor, I
 			(implementor as AbstractSessionImpl)?.AutoFlushIfRequired(querySpaces);
 		}
 
+		internal static IDisposable SuspendAutoFlush(this ISessionImplementor implementor)
+		{
+			return (implementor as IEventSource)?.SuspendAutoFlush();
+		}
+
 		/// <summary>
 		/// Returns an <see cref="IAsyncEnumerable{T}" /> which can be enumerated asynchronously.
 		/// </summary>

src/NHibernate/Impl/SessionImpl.cs

@@ -585,12 +585,9 @@ public override IEnumerable<T> Enumerable<T>(IQueryExpression queryExpression, Q
 			{
 				queryParameters.ValidateParameters();
 				var plan = GetHQLQueryPlan(queryExpression, true);
-				AutoFlushIfRequired(plan.QuerySpaces);
 
-				using (SuspendAutoFlush()) //stops flush being called multiple times if this method is recursively called
-				{
-					return plan.PerformIterate<T>(queryParameters, this);
-				}
+				// AutoFlushIfRequired will be called when iterating through the enumerable
+				return plan.PerformIterate<T>(queryParameters, this);
 			}
 		}
 
@@ -600,12 +597,9 @@ public override IAsyncEnumerable<T> AsyncEnumerable<T>(IQueryExpression queryExp
 			{
 				queryParameters.ValidateParameters();
 				var plan = GetHQLQueryPlan(queryExpression, true);
-				AutoFlushIfRequired(plan.QuerySpaces);
 
-				using (SuspendAutoFlush()) //stops flush being called multiple times if this method is recursively called
-				{
-					return plan.PerformAsyncIterate<T>(queryParameters, this);
-				}
+				// AutoFlushIfRequired will be called when iterating through the enumerable
+				return plan.PerformAsyncIterate<T>(queryParameters, this);
 			}
 		}
 
@@ -615,12 +609,9 @@ public override IEnumerable Enumerable(IQueryExpression queryExpression, QueryPa
 			{
 				queryParameters.ValidateParameters();
 				var plan = GetHQLQueryPlan(queryExpression, true);
-				AutoFlushIfRequired(plan.QuerySpaces);
 
-				using (SuspendAutoFlush()) //stops flush being called multiple times if this method is recursively called
-				{
-					return plan.PerformIterate(queryParameters, this);
-				}
+				// AutoFlushIfRequired will be called when iterating through the enumerable
+				return plan.PerformIterate(queryParameters, this);
 			}
 		}
 

src/NHibernate/Loader/Hql/QueryLoader.cs

@@ -486,22 +486,26 @@ internal AsyncEnumerableImpl<T> GetAsyncEnumerable<T>(QueryParameters queryParam
 
 		internal InitializeEnumerableResult InitializeEnumerable(QueryParameters queryParameters, ISessionImplementor session)
 		{
-			Stopwatch stopWatch = null;
-			if (session.Factory.Statistics.IsStatisticsEnabled)
+			session.AutoFlushIfRequired(_queryTranslator.QuerySpaces);
+			using (session.SuspendAutoFlush())
 			{
-				stopWatch = Stopwatch.StartNew();
-			}
+				Stopwatch stopWatch = null;
+				if (session.Factory.Statistics.IsStatisticsEnabled)
+				{
+					stopWatch = Stopwatch.StartNew();
+				}
 
-			var command = PrepareQueryCommand(queryParameters, false, session);
-			var dataReader = GetResultSet(command, queryParameters, session, null);
-			if (stopWatch != null)
-			{
-				stopWatch.Stop();
-				session.Factory.StatisticsImplementor.QueryExecuted("HQL: " + _queryTranslator.QueryString, 0, stopWatch.Elapsed);
-				session.Factory.StatisticsImplementor.QueryExecuted(QueryIdentifier, 0, stopWatch.Elapsed);
-			}
+				var command = PrepareQueryCommand(queryParameters, false, session);
+				var dataReader = GetResultSet(command, queryParameters, session, null);
+				if (stopWatch != null)
+				{
+					stopWatch.Stop();
+					session.Factory.StatisticsImplementor.QueryExecuted("HQL: " + _queryTranslator.QueryString, 0, stopWatch.Elapsed);
+					session.Factory.StatisticsImplementor.QueryExecuted(QueryIdentifier, 0, stopWatch.Elapsed);
+				}
 
-			return new InitializeEnumerableResult(command, dataReader);
+				return new InitializeEnumerableResult(command, dataReader);
+			}
 		}
 
 		protected override void ResetEffectiveExpectedType(IEnumerable<IParameterSpecification> parameterSpecs, QueryParameters queryParameters)