fixup! Implement universal query batch

Expose afterLoad callback
Add ovverides with aggregator selector

Author: bahusoid

src/NHibernate.Test/Async/Futures/QueryBatchFixture.cs

@@ -153,33 +153,34 @@ public async Task CanFetchCollectionInBatchAsync()
 		}
 
 		[Test]
-		public async Task OnAfterLoadAsync()
+		public async Task AfterLoadCallbackAsync()
 		{
 			using (var session = OpenSession())
 			{
 				var batch = session.CreateQueryBatch();
 				IList<EntityComplex> results = null;
-				var batchItem = LinqBatchItem.Create(session.Query<EntityComplex>().WithOptions(o => o.SetCacheable(true)));
-				batchItem.OnAfterLoad += list => results = list;
-				batch.Add(batchItem);
-
+				int count = 0;
+				batch.Add(session.Query<EntityComplex>().WithOptions(o => o.SetCacheable(true)), r => results = r);
+				batch.Add(session.Query<EntityComplex>().WithOptions(o => o.SetCacheable(true)), ec => ec.Count(), r => count = r);
 				await (batch.ExecuteAsync(CancellationToken.None));
 
 				Assert.That(results, Is.Not.Null);
+				Assert.That(count, Is.GreaterThan(0));
 			}
 
 			using (var sqlLog = new SqlLogSpy())
 			using (var session = OpenSession())
 			{
 				var batch = session.CreateQueryBatch();
 				IList<EntityComplex> results = null;
-				var batchItem = LinqBatchItem.Create(session.Query<EntityComplex>().WithOptions(o => o.SetCacheable(true)));
-				batchItem.OnAfterLoad += list => results = list;
-				batch.Add(batchItem);
+				int count = 0;
+				batch.Add(session.Query<EntityComplex>().WithOptions(o => o.SetCacheable(true)), r => results = r);
+				batch.Add(session.Query<EntityComplex>().WithOptions(o => o.SetCacheable(true)), ec => ec.Count(), r => count = r);
 
 				await (batch.ExecuteAsync(CancellationToken.None));
 
 				Assert.That(results, Is.Not.Null);
+				Assert.That(count, Is.GreaterThan(0));
 				Assert.That(sqlLog.Appender.GetEvents().Length, Is.EqualTo(0), "Query is expected to be retrieved from cache");
 			}
 		}

src/NHibernate.Test/Futures/QueryBatchFixture.cs

@@ -141,33 +141,34 @@ public void CanFetchCollectionInBatch()
 		}
 
 		[Test]
-		public void OnAfterLoad()
+		public void AfterLoadCallback()
 		{
 			using (var session = OpenSession())
 			{
 				var batch = session.CreateQueryBatch();
 				IList<EntityComplex> results = null;
-				var batchItem = LinqBatchItem.Create(session.Query<EntityComplex>().WithOptions(o => o.SetCacheable(true)));
-				batchItem.OnAfterLoad += list => results = list;
-				batch.Add(batchItem);
-
+				int count = 0;
+				batch.Add(session.Query<EntityComplex>().WithOptions(o => o.SetCacheable(true)), r => results = r);
+				batch.Add(session.Query<EntityComplex>().WithOptions(o => o.SetCacheable(true)), ec => ec.Count(), r => count = r);
 				batch.Execute();
 
 				Assert.That(results, Is.Not.Null);
+				Assert.That(count, Is.GreaterThan(0));
 			}
 
 			using (var sqlLog = new SqlLogSpy())
 			using (var session = OpenSession())
 			{
 				var batch = session.CreateQueryBatch();
 				IList<EntityComplex> results = null;
-				var batchItem = LinqBatchItem.Create(session.Query<EntityComplex>().WithOptions(o => o.SetCacheable(true)));
-				batchItem.OnAfterLoad += list => results = list;
-				batch.Add(batchItem);
+				int count = 0;
+				batch.Add(session.Query<EntityComplex>().WithOptions(o => o.SetCacheable(true)), r => results = r);
+				batch.Add(session.Query<EntityComplex>().WithOptions(o => o.SetCacheable(true)), ec => ec.Count(), r => count = r);
 
 				batch.Execute();
 
 				Assert.That(results, Is.Not.Null);
+				Assert.That(count, Is.GreaterThan(0));
 				Assert.That(sqlLog.Appender.GetEvents().Length, Is.EqualTo(0), "Query is expected to be retrieved from cache");
 			}
 		}

src/NHibernate/Multi/IQueryBatchItem.cs

@@ -16,9 +16,9 @@ public interface IQueryBatchItem<TResult> : IQueryBatchItem
 		/// Must be called only after <see cref="IQueryBatch.Execute"/>.
 		/// </summary>
 		IList<TResult> GetResults();
-		
+
 		/// <summary>
-		/// Event is executed after results are loaded by batch.
+		/// Callback is executed after results are loaded by batch.
 		/// Loaded results are provided in action parameter.
 		/// </summary>
 		event Action<IList<TResult>> OnAfterLoad;

src/NHibernate/Multi/QueryBatchExtensions.cs

@@ -17,13 +17,13 @@ public static partial class QueryBatchExtensions
 		/// </summary>
 		/// <param name="batch">The batch.</param>
 		/// <param name="query">The query.</param>
+		/// <param name="afterLoad">Callback to execute when query is loaded. Loaded results are provided as action parameter.</param>
 		/// <exception cref="InvalidOperationException">Thrown if the batch has already been executed.</exception>
 		/// <exception cref="ArgumentNullException">Thrown if <paramref name="query"/> is <see langword="null"/>.</exception>
 		/// <returns>The batch instance for method chain.</returns>
-		public static IQueryBatch Add<TResult>(this IQueryBatch batch, IQueryOver query)
+		public static IQueryBatch Add<TResult>(this IQueryBatch batch, IQueryOver query, Action<IList<TResult>> afterLoad = null)
 		{
-			batch.Add(For<TResult>(query));
-			return batch;
+			return batch.Add(For<TResult>(query), afterLoad);
 		}
 
 		/// <summary>
@@ -48,13 +48,13 @@ public static IQueryBatch Add<TResult>(this IQueryBatch batch, string key, IQuer
 		/// </summary>
 		/// <param name="batch">The batch.</param>
 		/// <param name="query">The query.</param>
+		/// <param name="afterLoad">Callback to execute when query is loaded. Loaded results are provided as action parameter.</param>
 		/// <exception cref="InvalidOperationException">Thrown if the batch has already been executed.</exception>
 		/// <exception cref="ArgumentNullException">Thrown if <paramref name="query"/> is <see langword="null"/>.</exception>
 		/// <returns>The batch instance for method chain.</returns>
-		public static IQueryBatch Add<TResult>(this IQueryBatch batch, IQueryOver<TResult> query)
+		public static IQueryBatch Add<TResult>(this IQueryBatch batch, IQueryOver<TResult> query, Action<IList<TResult>> afterLoad = null)
 		{
-			batch.Add(For<TResult>(query));
-			return batch;
+			return batch.Add(For<TResult>(query), afterLoad);
 		}
 
 		/// <summary>
@@ -77,13 +77,13 @@ public static IQueryBatch Add<TResult>(this IQueryBatch batch, string key, IQuer
 		/// </summary>
 		/// <param name="batch">The batch.</param>
 		/// <param name="query">The query.</param>
+		/// <param name="afterLoad">Callback to execute when query is loaded. Loaded results are provided as action parameter.</param>
 		/// <exception cref="InvalidOperationException">Thrown if the batch has already been executed.</exception>
 		/// <exception cref="ArgumentNullException">Thrown if <paramref name="query"/> is <see langword="null"/>.</exception>
 		/// <returns>The batch instance for method chain.</returns>
-		public static IQueryBatch Add<TResult>(this IQueryBatch batch, ICriteria query)
+		public static IQueryBatch Add<TResult>(this IQueryBatch batch, ICriteria query, Action<IList<TResult>> afterLoad = null)
 		{
-			batch.Add(For<TResult>(query));
-			return batch;
+			return batch.Add(For<TResult>(query), afterLoad);
 		}
 
 		/// <summary>
@@ -106,13 +106,13 @@ public static IQueryBatch Add<TResult>(this IQueryBatch batch, string key, ICrit
 		/// </summary>
 		/// <param name="batch">The batch.</param>
 		/// <param name="query">The query.</param>
+		/// <param name="afterLoad">Callback to execute when query is loaded. Loaded results are provided as action parameter.</param>
 		/// <exception cref="InvalidOperationException">Thrown if the batch has already been executed.</exception>
 		/// <exception cref="ArgumentNullException">Thrown if <paramref name="query"/> is <see langword="null"/>.</exception>
 		/// <returns>The batch instance for method chain.</returns>
-		public static IQueryBatch Add<TResult>(this IQueryBatch batch, DetachedCriteria query)
+		public static IQueryBatch Add<TResult>(this IQueryBatch batch, DetachedCriteria query, Action<IList<TResult>> afterLoad = null)
 		{
-			batch.Add(For<TResult>(query));
-			return batch;
+			return batch.Add(For<TResult>(query), afterLoad);
 		}
 
 		/// <summary>
@@ -135,13 +135,13 @@ public static IQueryBatch Add<TResult>(this IQueryBatch batch, string key, Detac
 		/// </summary>
 		/// <param name="batch">The batch.</param>
 		/// <param name="query">The query.</param>
+		/// <param name="afterLoad">Callback to execute when query is loaded. Loaded results are provided as action parameter.</param>
 		/// <exception cref="InvalidOperationException">Thrown if the batch has already been executed.</exception>
 		/// <exception cref="ArgumentNullException">Thrown if <paramref name="query"/> is <see langword="null"/>.</exception>
 		/// <returns>The batch instance for method chain.</returns>
-		public static IQueryBatch Add<TResult>(this IQueryBatch batch, IQuery query)
+		public static IQueryBatch Add<TResult>(this IQueryBatch batch, IQuery query, Action<IList<TResult>> afterLoad = null)
 		{
-			batch.Add(For<TResult>(query));
-			return batch;
+			return batch.Add(For<TResult>(query), afterLoad);
 		}
 
 		/// <summary>
@@ -164,13 +164,13 @@ public static IQueryBatch Add<TResult>(this IQueryBatch batch, string key, IQuer
 		/// </summary>
 		/// <param name="batch">The batch.</param>
 		/// <param name="query">The query.</param>
+		/// <param name="afterLoad">Callback to execute when query is loaded. Loaded results are provided as action parameter.</param>
 		/// <exception cref="InvalidOperationException">Thrown if the batch has already been executed.</exception>
 		/// <exception cref="ArgumentNullException">Thrown if <paramref name="query"/> is <see langword="null"/>.</exception>
 		/// <returns>The batch instance for method chain.</returns>
-		public static IQueryBatch Add<TResult>(this IQueryBatch batch, IQueryable<TResult> query)
+		public static IQueryBatch Add<TResult>(this IQueryBatch batch, IQueryable<TResult> query, Action<IList<TResult>> afterLoad = null)
 		{
-			batch.Add(For(query));
-			return batch;
+			return batch.Add(For<TResult>(query), afterLoad);
 		}
 
 		/// <summary>
@@ -188,6 +188,37 @@ public static IQueryBatch Add<TResult>(this IQueryBatch batch, string key, IQuer
 			return batch;
 		}
 
+		/// <summary>
+		/// Adds a query to the batch.
+		/// </summary>
+		/// <param name="batch">The batch.</param>
+		/// <param name="query">The query.</param>
+		/// <param name="selector">An aggregation function to apply to <paramref name="query"/>.</param>
+		/// <param name="afterLoad">Callback to execute when query is loaded. Loaded results are provided as action parameter.</param>
+		/// <exception cref="InvalidOperationException">Thrown if the batch has already been executed.</exception>
+		/// <exception cref="ArgumentNullException">Thrown if <paramref name="query"/> is <see langword="null"/>.</exception>
+		/// <returns>The batch instance for method chain.</returns>
+		public static IQueryBatch Add<TSource, TResult>(this IQueryBatch batch, IQueryable<TSource> query, Expression<Func<IQueryable<TSource>, TResult>> selector, Action<TResult> afterLoad = null)
+		{
+			return batch.Add(For(query, selector), afterLoad == null ? (Action<IList<TResult>>) null : list => afterLoad(list.FirstOrDefault()));
+		}
+
+		/// <summary>
+		/// Adds a query to the batch.
+		/// </summary>
+		/// <param name="batch">The batch.</param>
+		/// <param name="key">A key for retrieval of the query result.</param>
+		/// <param name="query">The query.</param>
+		/// <param name="selector">An aggregation function to apply to <paramref name="query"/>.</param>
+		/// <exception cref="InvalidOperationException">Thrown if the batch has already been executed.</exception>
+		/// <exception cref="ArgumentNullException">Thrown if <paramref name="query"/> is <see langword="null"/>.</exception>
+		/// <returns>The batch instance for method chain.</returns>
+		public static IQueryBatch Add<TSource, TResult>(this IQueryBatch batch, string key, IQueryable<TSource> query, Expression<Func<IQueryable<TSource>, TResult>> selector)
+		{
+			batch.Add(key, For(query, selector));
+			return batch;
+		}
+
 		/// <summary>
 		/// Sets the timeout in seconds for the underlying ADO.NET query.
 		/// </summary>
@@ -324,6 +355,16 @@ private static CriteriaBatchItem<TResult> For<TResult>(IQueryOver query)
 			return For<TResult>(query.RootCriteria);
 		}
 
+		private static IQueryBatch Add<TResult>(this IQueryBatch batch, IQueryBatchItem<TResult> query, Action<IList<TResult>> afterLoad)
+		{
+			if (afterLoad != null)
+			{
+				query.OnAfterLoad += afterLoad;
+			}
+			batch.Add(query);
+			return batch;
+		}
+
 		#region Helper classes
 
 		partial class FutureValue<TResult> : IFutureValue<TResult>