Fix future queries with non lazy associations (#2174)

bahusoid · hazzik · commit a8f4ea352407 · 2019-07-16T09:50:31.000Z
diff --git a/src/AsyncGenerator.yml b/src/AsyncGenerator.yml
@@ -5,6 +5,10 @@
   applyChanges: true
   analyzation:
     methodConversion:
+#TODO 6.0: Remove ignore rule for IQueryBatchItem.ProcessResults  
+    - conversion: Ignore
+      name: ProcessResults
+      containingTypeName: IQueryBatchItem
     - conversion: Ignore
       name: PostProcessInsert
       containingTypeName: HqlSqlWalker
diff --git a/src/NHibernate.Test/Async/Futures/QueryBatchFixture.cs b/src/NHibernate.Test/Async/Futures/QueryBatchFixture.cs
@@ -498,6 +498,35 @@ public async Task CacheModeWorksWithFutureAsync()
 			}
 		}
 
+		//GH-2173
+		[Test]
+		public async Task CanFetchNonLazyEntitiesInSubsequentQueryAsync()
+		{
+			Sfi.Statistics.IsStatisticsEnabled = true;
+			using (var s = OpenSession())
+			using (var t = s.BeginTransaction())
+			{
+				await (s.SaveAsync(
+					new EntityEager
+					{
+						Name = "EagerManyToOneAssociation",
+						EagerEntity = new EntityEagerChild {Name = "association"}
+					}));
+				await (t.CommitAsync());
+			}
+
+			using (var s = OpenSession())
+			{
+				Sfi.Statistics.Clear();
+				//EntityEager.EagerEntity is lazy initialized instead of being loaded by the second query 
+				s.QueryOver<EntityEager>().Fetch(SelectMode.Skip, x => x.EagerEntity).Future();
+				await (s.QueryOver<EntityEager>().Fetch(SelectMode.Fetch, x => x.EagerEntity).Future().GetEnumerableAsync());
+
+				if(SupportsMultipleQueries)
+					Assert.That(Sfi.Statistics.PrepareStatementCount, Is.EqualTo(1));
+			}
+		}
+
 		#region Test Setup
 
 		protected override HbmMapping GetMappings()
@@ -543,6 +572,10 @@ protected override HbmMapping GetMappings()
 					rc.Id(x => x.Id, m => m.Generator(Generators.GuidComb));
 					rc.Property(x => x.Name);
 
+					rc.ManyToOne(x => x.EagerEntity, m =>
+					{
+						m.Cascade(Mapping.ByCode.Cascade.Persist);
+					});
 					rc.Bag(ep => ep.ChildrenListSubselect,
 							m =>
 							{
@@ -560,6 +593,14 @@ protected override HbmMapping GetMappings()
 							},
 							a => a.OneToMany());
 				});
+			mapper.Class<EntityEagerChild>(
+				rc =>
+				{
+					rc.Lazy(false);
+
+					rc.Id(x => x.Id, m => m.Generator(Generators.GuidComb));
+					rc.Property(x => x.Name);
+				});
 			mapper.Class<EntitySubselectChild>(
 				rc =>
 				{
diff --git a/src/NHibernate.Test/Futures/Entities.cs b/src/NHibernate.Test/Futures/Entities.cs
@@ -36,11 +36,18 @@ public class EntitySubselectChild
 		public virtual EntityEager Parent { get; set; }
 	}
 
+	public class EntityEagerChild
+	{
+		public Guid Id { get; set; }
+		public string Name { get; set; }
+	}
+	
 	public class EntityEager
 	{
 		public Guid Id { get; set; }
 		public string Name { get; set; }
 
+		public EntityEagerChild EagerEntity { get; set; }
 		public IList<EntitySubselectChild> ChildrenListSubselect { get; set; }
 		public IList<EntitySimpleChild> ChildrenListEager { get; set; } //= new HashSet<EntitySimpleChild>();
 	}
diff --git a/src/NHibernate.Test/Futures/QueryBatchFixture.cs b/src/NHibernate.Test/Futures/QueryBatchFixture.cs
@@ -486,6 +486,35 @@ public void CacheModeWorksWithFuture()
 			}
 		}
 
+		//GH-2173
+		[Test]
+		public void CanFetchNonLazyEntitiesInSubsequentQuery()
+		{
+			Sfi.Statistics.IsStatisticsEnabled = true;
+			using (var s = OpenSession())
+			using (var t = s.BeginTransaction())
+			{
+				s.Save(
+					new EntityEager
+					{
+						Name = "EagerManyToOneAssociation",
+						EagerEntity = new EntityEagerChild {Name = "association"}
+					});
+				t.Commit();
+			}
+
+			using (var s = OpenSession())
+			{
+				Sfi.Statistics.Clear();
+				//EntityEager.EagerEntity is lazy initialized instead of being loaded by the second query 
+				s.QueryOver<EntityEager>().Fetch(SelectMode.Skip, x => x.EagerEntity).Future();
+				s.QueryOver<EntityEager>().Fetch(SelectMode.Fetch, x => x.EagerEntity).Future().GetEnumerable();
+
+				if(SupportsMultipleQueries)
+					Assert.That(Sfi.Statistics.PrepareStatementCount, Is.EqualTo(1));
+			}
+		}
+
 		#region Test Setup
 
 		protected override HbmMapping GetMappings()
@@ -531,6 +560,10 @@ protected override HbmMapping GetMappings()
 					rc.Id(x => x.Id, m => m.Generator(Generators.GuidComb));
 					rc.Property(x => x.Name);
 
+					rc.ManyToOne(x => x.EagerEntity, m =>
+					{
+						m.Cascade(Mapping.ByCode.Cascade.Persist);
+					});
 					rc.Bag(ep => ep.ChildrenListSubselect,
 							m =>
 							{
@@ -548,6 +581,14 @@ protected override HbmMapping GetMappings()
 							},
 							a => a.OneToMany());
 				});
+			mapper.Class<EntityEagerChild>(
+				rc =>
+				{
+					rc.Lazy(false);
+
+					rc.Id(x => x.Id, m => m.Generator(Generators.GuidComb));
+					rc.Property(x => x.Name);
+				});
 			mapper.Class<EntitySubselectChild>(
 				rc =>
 				{
diff --git a/src/NHibernate/Async/Multi/IQueryBatchItem.cs b/src/NHibernate/Async/Multi/IQueryBatchItem.cs
@@ -36,4 +36,9 @@ public partial interface IQueryBatchItem
 		/// <param name="cancellationToken">A cancellation token that can be used to cancel the work</param>
 		Task ExecuteNonBatchedAsync(CancellationToken cancellationToken);
 	}
+
+	internal partial interface IQueryBatchItemWithAsyncProcessResults
+	{
+		Task ProcessResultsAsync(CancellationToken cancellationToken);
+	}
 }
diff --git a/src/NHibernate/Async/Multi/QueryBatch.cs b/src/NHibernate/Async/Multi/QueryBatch.cs
@@ -127,37 +127,48 @@ protected async Task ExecuteBatchedAsync(CancellationToken cancellationToken)
 			}
 
 			var rowCount = 0;
+			CacheBatcher cacheBatcher = null;
 			try
 			{
 				if (resultSetsCommand.HasQueries)
 				{
+					cacheBatcher = new CacheBatcher(Session);
 					using (var reader = await (resultSetsCommand.GetReaderAsync(Timeout, cancellationToken)).ConfigureAwait(false))
 					{
-						var cacheBatcher = new CacheBatcher(Session);
 						foreach (var query in _queries)
 						{
 							if (query.CachingInformation != null)
 							{
-								foreach (var cachingInfo in query.CachingInformation.Where(ci => ci.IsCacheable))
+								foreach (var cachingInfo in query.CachingInformation)
 								{
 									cachingInfo.SetCacheBatcher(cacheBatcher);
 								}
 							}
 
 							rowCount += await (query.ProcessResultsSetAsync(reader, cancellationToken)).ConfigureAwait(false);
 						}
-						await (cacheBatcher.ExecuteBatchAsync(cancellationToken)).ConfigureAwait(false);
 					}
 				}
 
-				// Query cacheable results must be cached untransformed: the put does not need to wait for
-				// the ProcessResults.
-				await (PutCacheableResultsAsync(cancellationToken)).ConfigureAwait(false);
-
 				foreach (var query in _queries)
 				{
-					query.ProcessResults();
+					//TODO 6.0: Replace with query.ProcessResults();
+					if (query is IQueryBatchItemWithAsyncProcessResults q)
+						await (q.ProcessResultsAsync(cancellationToken)).ConfigureAwait(false);
+					else
+						query.ProcessResults();
 				}
+
+				var executeBatchTask = cacheBatcher?.ExecuteBatchAsync(cancellationToken);
+
+				if (executeBatchTask != null)
+
+				{
+
+					await (executeBatchTask).ConfigureAwait(false);
+
+				}
+				await (PutCacheableResultsAsync(cancellationToken)).ConfigureAwait(false);
 			}
 			catch (OperationCanceledException) { throw; }
 			catch (Exception sqle)
diff --git a/src/NHibernate/Async/Multi/QueryBatchItemBase.cs b/src/NHibernate/Async/Multi/QueryBatchItemBase.cs
@@ -23,7 +23,7 @@ namespace NHibernate.Multi
 {
 	using System.Threading.Tasks;
 	using System.Threading;
-	public abstract partial class QueryBatchItemBase<TResult> : IQueryBatchItem<TResult>
+	public abstract partial class QueryBatchItemBase<TResult> : IQueryBatchItem<TResult>, IQueryBatchItemWithAsyncProcessResults
 	{
 
 		/// <inheritdoc />
@@ -103,17 +103,46 @@ public async Task<int> ProcessResultsSetAsync(DbDataReader reader, CancellationT
 
 					queryInfo.Result = tmpResults;
 					if (queryInfo.CanPutToCache)
-						queryInfo.ResultToCache = tmpResults;
+						queryInfo.ResultToCache = new List<object>(tmpResults);
 
 					await (reader.NextResultAsync(cancellationToken)).ConfigureAwait(false);
 				}
 
-				await (InitializeEntitiesAndCollectionsAsync(reader, hydratedObjects, cancellationToken)).ConfigureAwait(false);
-
+				StopLoadingCollections(reader);
+				_reader = reader;
+				_hydratedObjects = hydratedObjects;
 				return rowCount;
 			}
 		}
 
+		/// <inheritdoc cref="IQueryBatchItem.ProcessResults" />
+		public async Task ProcessResultsAsync(CancellationToken cancellationToken)
+		{
+			cancellationToken.ThrowIfCancellationRequested();
+			ThrowIfNotInitialized();
+
+			using (Session.SwitchCacheMode(_cacheMode))
+				await (InitializeEntitiesAndCollectionsAsync(_reader, _hydratedObjects, cancellationToken)).ConfigureAwait(false);
+
+			for (var i = 0; i < _queryInfos.Count; i++)
+			{
+				var queryInfo = _queryInfos[i];
+				if (_subselectResultKeys[i] != null)
+				{
+					queryInfo.Loader.CreateSubselects(_subselectResultKeys[i], queryInfo.Parameters, Session);
+				}
+
+				if (queryInfo.IsCacheable)
+				{
+					// This transformation must not be applied to ResultToCache.
+					queryInfo.Result =
+						queryInfo.Loader.TransformCacheableResults(
+							queryInfo.Parameters, queryInfo.CacheKey.ResultTransformer, queryInfo.Result);
+				}
+			}
+			AfterLoadCallback?.Invoke(GetResults());
+		}
+
 		/// <inheritdoc />
 		public async Task ExecuteNonBatchedAsync(CancellationToken cancellationToken)
 		{
diff --git a/src/NHibernate/Engine/Loading/CollectionLoadContext.cs b/src/NHibernate/Engine/Loading/CollectionLoadContext.cs
@@ -124,6 +124,8 @@ public IPersistentCollection GetLoadingCollection(ICollectionPersister persister
 			}
 			else
 			{
+				if (loadingCollectionEntry.StopLoading)
+					return null;
 				if (loadingCollectionEntry.ResultSet == resultSet)
 				{
 					log.Debug("found loading collection bound to current result set processing; reading row");
@@ -398,5 +400,17 @@ public override string ToString()
 		{
 			return base.ToString() + "<rs=" + ResultSet + ">";
 		}
+
+		internal void StopLoadingCollections(ICollectionPersister[] collectionPersisters)
+		{
+			foreach (var collectionKey in localLoadingCollectionKeys)
+			{
+				var loadingCollectionEntry = LoadContext.LocateLoadingCollectionEntry(collectionKey);
+				if (loadingCollectionEntry != null && Array.IndexOf(collectionPersisters, loadingCollectionEntry.Persister) >= 0)
+				{
+					loadingCollectionEntry.StopLoading = true;
+				}
+			}
+		}
 	}
 }
diff --git a/src/NHibernate/Engine/Loading/LoadingCollectionEntry.cs b/src/NHibernate/Engine/Loading/LoadingCollectionEntry.cs
@@ -44,6 +44,8 @@ public IPersistentCollection Collection
 			get { return collection; }
 		}
 
+		public bool StopLoading { get; set; }
+
 		public override string ToString()
 		{
 			return GetType().FullName + "<rs=" + ResultSet + ", coll=" + MessageHelper.InfoString(Persister.Role, Key) + ">@" + Convert.ToString(GetHashCode(), 16);
diff --git a/src/NHibernate/Loader/Loader.cs b/src/NHibernate/Loader/Loader.cs
@@ -672,6 +672,18 @@ internal void InitializeEntitiesAndCollections(
 			}
 		}
 
+		/// <summary>
+		/// Stops further collection population without actual collection initialization.
+		/// </summary>
+		internal void StopLoadingCollections(ISessionImplementor session, DbDataReader reader)
+		{
+			var collectionPersisters = CollectionPersisters;
+			if (collectionPersisters == null || collectionPersisters.Length == 0)
+				return;
+
+			session.PersistenceContext.LoadContexts.GetCollectionLoadContext(reader).StopLoadingCollections(collectionPersisters);
+		}
+
 		private void EndCollectionLoad(DbDataReader reader, ISessionImplementor session, ICollectionPersister collectionPersister)
 		{
 			//this is a query and we are loading multiple instances of the same collection role
diff --git a/src/NHibernate/Multi/IQueryBatchItem.cs b/src/NHibernate/Multi/IQueryBatchItem.cs
@@ -80,4 +80,10 @@ public partial interface IQueryBatchItem
 		/// </summary>
 		void ExecuteNonBatched();
 	}
+
+	//TODO 6.0: Remove along with ignore rule for IQueryBatchItem.ProcessResults in AsyncGenerator.yml
+	internal partial interface IQueryBatchItemWithAsyncProcessResults
+	{
+		void ProcessResults();
+	}
 }
diff --git a/src/NHibernate/Multi/QueryBatch.cs b/src/NHibernate/Multi/QueryBatch.cs
@@ -153,37 +153,40 @@ protected void ExecuteBatched()
 			}
 
 			var rowCount = 0;
+			CacheBatcher cacheBatcher = null;
 			try
 			{
 				if (resultSetsCommand.HasQueries)
 				{
+					cacheBatcher = new CacheBatcher(Session);
 					using (var reader = resultSetsCommand.GetReader(Timeout))
 					{
-						var cacheBatcher = new CacheBatcher(Session);
 						foreach (var query in _queries)
 						{
 							if (query.CachingInformation != null)
 							{
-								foreach (var cachingInfo in query.CachingInformation.Where(ci => ci.IsCacheable))
+								foreach (var cachingInfo in query.CachingInformation)
 								{
 									cachingInfo.SetCacheBatcher(cacheBatcher);
 								}
 							}
 
 							rowCount += query.ProcessResultsSet(reader);
 						}
-						cacheBatcher.ExecuteBatch();
 					}
 				}
 
-				// Query cacheable results must be cached untransformed: the put does not need to wait for
-				// the ProcessResults.
-				PutCacheableResults();
-
 				foreach (var query in _queries)
 				{
-					query.ProcessResults();
+					//TODO 6.0: Replace with query.ProcessResults();
+					if (query is IQueryBatchItemWithAsyncProcessResults q)
+						q.ProcessResults();
+					else
+						query.ProcessResults();
 				}
+
+				cacheBatcher?.ExecuteBatch();
+				PutCacheableResults();
 			}
 			catch (Exception sqle)
 			{
diff --git a/src/NHibernate/Multi/QueryBatchItemBase.cs b/src/NHibernate/Multi/QueryBatchItemBase.cs

Original file line number	Diff line number	Diff line change
`@@ -36,11 +36,18 @@ public class EntitySubselectChild`
`36`	`36`	`public virtual EntityEager Parent { get; set; }`
`37`	`37`	`}`
`38`	`38`
	`39`	`+ public class EntityEagerChild`
	`40`	`+ {`
	`41`	`+ public Guid Id { get; set; }`
	`42`	`+ public string Name { get; set; }`
	`43`	`+ }`
	`44`	`+`
`39`	`45`	`public class EntityEager`
`40`	`46`	`{`
`41`	`47`	`public Guid Id { get; set; }`
`42`	`48`	`public string Name { get; set; }`
`43`	`49`
	`50`	`+ public EntityEagerChild EagerEntity { get; set; }`
`44`	`51`	`public IList<EntitySubselectChild> ChildrenListSubselect { get; set; }`
`45`	`52`	`public IList<EntitySimpleChild> ChildrenListEager { get; set; } //= new HashSet<EntitySimpleChild>();`
`46`	`53`	`}`
Original file line number	Diff line number	Diff line change
`@@ -36,4 +36,9 @@ public partial interface IQueryBatchItem`
`36`	`36`	`/// <param name="cancellationToken">A cancellation token that can be used to cancel the work</param>`
`37`	`37`	`Task ExecuteNonBatchedAsync(CancellationToken cancellationToken);`
`38`	`38`	`}`
	`39`	`+`
	`40`	`+ internal partial interface IQueryBatchItemWithAsyncProcessResults`
	`41`	`+ {`
	`42`	`+ Task ProcessResultsAsync(CancellationToken cancellationToken);`
	`43`	`+ }`
`39`	`44`	`}`
Original file line number	Diff line number	Diff line change
`@@ -127,37 +127,48 @@ protected async Task ExecuteBatchedAsync(CancellationToken cancellationToken)`
`127`	`127`	`}`
`128`	`128`
`129`	`129`	`var rowCount = 0;`
	`130`	`+ CacheBatcher cacheBatcher = null;`
`130`	`131`	`try`
`131`	`132`	`{`
`132`	`133`	`if (resultSetsCommand.HasQueries)`
`133`	`134`	`{`
	`135`	`+ cacheBatcher = new CacheBatcher(Session);`
`134`	`136`	`using (var reader = await (resultSetsCommand.GetReaderAsync(Timeout, cancellationToken)).ConfigureAwait(false))`
`135`	`137`	`{`
`136`		`- var cacheBatcher = new CacheBatcher(Session);`
`137`	`138`	`foreach (var query in _queries)`
`138`	`139`	`{`
`139`	`140`	`if (query.CachingInformation != null)`
`140`	`141`	`{`
`141`		`- foreach (var cachingInfo in query.CachingInformation.Where(ci => ci.IsCacheable))`
	`142`	`+ foreach (var cachingInfo in query.CachingInformation)`
`142`	`143`	`{`
`143`	`144`	`cachingInfo.SetCacheBatcher(cacheBatcher);`
`144`	`145`	`}`
`145`	`146`	`}`
`146`	`147`
`147`	`148`	`rowCount += await (query.ProcessResultsSetAsync(reader, cancellationToken)).ConfigureAwait(false);`
`148`	`149`	`}`
`149`		`- await (cacheBatcher.ExecuteBatchAsync(cancellationToken)).ConfigureAwait(false);`
`150`	`150`	`}`
`151`	`151`	`}`
`152`	`152`
`153`		`- // Query cacheable results must be cached untransformed: the put does not need to wait for`
`154`		`- // the ProcessResults.`
`155`		`- await (PutCacheableResultsAsync(cancellationToken)).ConfigureAwait(false);`
`156`		`-`
`157`	`153`	`foreach (var query in _queries)`
`158`	`154`	`{`
`159`		`- query.ProcessResults();`
	`155`	`+ //TODO 6.0: Replace with query.ProcessResults();`
	`156`	`+ if (query is IQueryBatchItemWithAsyncProcessResults q)`
	`157`	`+ await (q.ProcessResultsAsync(cancellationToken)).ConfigureAwait(false);`
	`158`	`+ else`
	`159`	`+ query.ProcessResults();`
`160`	`160`	`}`
	`161`	`+`
	`162`	`+ var executeBatchTask = cacheBatcher?.ExecuteBatchAsync(cancellationToken);`
	`163`	`+`
	`164`	`+ if (executeBatchTask != null)`
	`165`	`+`
	`166`	`+ {`
	`167`	`+`
	`168`	`+ await (executeBatchTask).ConfigureAwait(false);`
	`169`	`+`
	`170`	`+ }`
	`171`	`+ await (PutCacheableResultsAsync(cancellationToken)).ConfigureAwait(false);`
`161`	`172`	`}`
`162`	`173`	`catch (OperationCanceledException) { throw; }`
`163`	`174`	`catch (Exception sqle)`
Original file line number	Diff line number	Diff line change
`@@ -124,6 +124,8 @@ public IPersistentCollection GetLoadingCollection(ICollectionPersister persister`
`124`	`124`	`}`
`125`	`125`	`else`
`126`	`126`	`{`
	`127`	`+ if (loadingCollectionEntry.StopLoading)`
	`128`	`+ return null;`
`127`	`129`	`if (loadingCollectionEntry.ResultSet == resultSet)`
`128`	`130`	`{`
`129`	`131`	`log.Debug("found loading collection bound to current result set processing; reading row");`
`@@ -398,5 +400,17 @@ public override string ToString()`
`398`	`400`	`{`
`399`	`401`	`return base.ToString() + "<rs=" + ResultSet + ">";`
`400`	`402`	`}`
	`403`	`+`
	`404`	`+ internal void StopLoadingCollections(ICollectionPersister[] collectionPersisters)`
	`405`	`+ {`
	`406`	`+ foreach (var collectionKey in localLoadingCollectionKeys)`
	`407`	`+ {`
	`408`	`+ var loadingCollectionEntry = LoadContext.LocateLoadingCollectionEntry(collectionKey);`
	`409`	`+ if (loadingCollectionEntry != null && Array.IndexOf(collectionPersisters, loadingCollectionEntry.Persister) >= 0)`
	`410`	`+ {`
	`411`	`+ loadingCollectionEntry.StopLoading = true;`
	`412`	`+ }`
	`413`	`+ }`
	`414`	`+ }`
`401`	`415`	`}`
`402`	`416`	`}`
Original file line number	Diff line number	Diff line change
`@@ -44,6 +44,8 @@ public IPersistentCollection Collection`
`44`	`44`	`get { return collection; }`
`45`	`45`	`}`
`46`	`46`
	`47`	`+ public bool StopLoading { get; set; }`
	`48`	`+`
`47`	`49`	`public override string ToString()`
`48`	`50`	`{`
`49`	`51`	`return GetType().FullName + "<rs=" + ResultSet + ", coll=" + MessageHelper.InfoString(Persister.Role, Key) + ">@" + Convert.ToString(GetHashCode(), 16);`
Original file line number	Diff line number	Diff line change
`@@ -672,6 +672,18 @@ internal void InitializeEntitiesAndCollections(`
`672`	`672`	`}`
`673`	`673`	`}`
`674`	`674`
	`675`	`+ /// <summary>`
	`676`	`+ /// Stops further collection population without actual collection initialization.`
	`677`	`+ /// </summary>`
	`678`	`+ internal void StopLoadingCollections(ISessionImplementor session, DbDataReader reader)`
	`679`	`+ {`
	`680`	`+ var collectionPersisters = CollectionPersisters;`
	`681`	`+ if (collectionPersisters == null \|\| collectionPersisters.Length == 0)`
	`682`	`+ return;`
	`683`	`+`
	`684`	`+ session.PersistenceContext.LoadContexts.GetCollectionLoadContext(reader).StopLoadingCollections(collectionPersisters);`
	`685`	`+ }`
	`686`	`+`
`675`	`687`	`private void EndCollectionLoad(DbDataReader reader, ISessionImplementor session, ICollectionPersister collectionPersister)`
`676`	`688`	`{`
`677`	`689`	`//this is a query and we are loading multiple instances of the same collection role`