Handle cacheable queries

Author: bahusoid

src/NHibernate/Async/Multi/QueryBatch.cs

@@ -139,33 +139,30 @@ protected async Task ExecuteBatchedAsync(CancellationToken cancellationToken)
 			}
 
 			var rowCount = 0;
+			CacheBatcher cacheBatcher = null;
 			try
 			{
 				if (resultSetsCommand.HasQueries)
 				{
 					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))
 								{
+									if(cacheBatcher == null)
+										cacheBatcher = new CacheBatcher(Session);
 									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)
 				{
 					//TODO 6.0: Replace with query.ProcessResults();
@@ -174,6 +171,17 @@ protected async Task ExecuteBatchedAsync(CancellationToken cancellationToken)
 					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)

src/NHibernate/Async/Multi/QueryBatchItemBase.cs

@@ -131,7 +131,9 @@ public async Task ProcessResultsAsync(CancellationToken cancellationToken)
 			cancellationToken.ThrowIfCancellationRequested();
 			ThrowIfNotInitialized();
 
-			await (InitializeEntitiesAndCollectionsAsync(_reader, _hydratedObjects, cancellationToken)).ConfigureAwait(false);
+			using (Session.SwitchCacheMode(_cacheMode))
+				await (InitializeEntitiesAndCollectionsAsync(_reader, _hydratedObjects, cancellationToken)).ConfigureAwait(false);
+
 			for (var i = 0; i < _queryInfos.Count; i++)
 			{
 				var queryInfo = _queryInfos[i];

src/NHibernate/Multi/QueryBatch.cs

@@ -151,33 +151,30 @@ protected void ExecuteBatched()
 			}
 
 			var rowCount = 0;
+			CacheBatcher cacheBatcher = null;
 			try
 			{
 				if (resultSetsCommand.HasQueries)
 				{
 					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))
 								{
+									if(cacheBatcher == null)
+										cacheBatcher = new CacheBatcher(Session);
 									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)
 				{
 					//TODO 6.0: Replace with query.ProcessResults();
@@ -186,6 +183,9 @@ protected void ExecuteBatched()
 					else
 						query.ProcessResults();
 				}
+
+				cacheBatcher?.ExecuteBatch();
+				PutCacheableResults();
 			}
 			catch (Exception sqle)
 			{

src/NHibernate/Multi/QueryBatchItemBase.cs

@@ -274,7 +274,9 @@ public void ProcessResults()
 		{
 			ThrowIfNotInitialized();
 
-			InitializeEntitiesAndCollections(_reader, _hydratedObjects);
+			using (Session.SwitchCacheMode(_cacheMode))
+				InitializeEntitiesAndCollections(_reader, _hydratedObjects);
+
 			for (var i = 0; i < _queryInfos.Count; i++)
 			{
 				var queryInfo = _queryInfos[i];