Reduce cast usage for COUNT aggregate and add support for Mssql count_big (#2061)

Co-authored-by: Alexander Zaytsev <hazzik@gmail.com>

Author: maca88

src/NHibernate.Test/Async/Linq/ByMethod/CountTests.cs

@@ -11,6 +11,7 @@
 using System;
 using System.Linq;
 using NHibernate.Cfg;
+using NHibernate.Dialect;
 using NUnit.Framework;
 using NHibernate.Linq;
 
@@ -110,5 +111,50 @@ into temp
 
 			Assert.That(result.Count, Is.EqualTo(77));
 		}
+
+		[Test]
+		public async Task CheckSqlFunctionNameLongCountAsync()
+		{
+			var name = Dialect is MsSql2000Dialect ? "count_big" : "count";
+			using (var sqlLog = new SqlLogSpy())
+			{
+				var result = await (db.Orders.LongCountAsync());
+				Assert.That(result, Is.EqualTo(830));
+
+				var log = sqlLog.GetWholeLog();
+				Assert.That(log, Does.Contain($"{name}("));
+			}
+		}
+
+		[Test]
+		public async Task CheckSqlFunctionNameForCountAsync()
+		{
+			using (var sqlLog = new SqlLogSpy())
+			{
+				var result = await (db.Orders.CountAsync());
+				Assert.That(result, Is.EqualTo(830));
+
+				var log = sqlLog.GetWholeLog();
+				Assert.That(log, Does.Contain("count("));
+			}
+		}
+
+		[Test]
+		public async Task CheckMssqlCountCastAsync()
+		{
+			if (!(Dialect is MsSql2000Dialect))
+			{
+				Assert.Ignore();
+			}
+
+			using (var sqlLog = new SqlLogSpy())
+			{
+				var result = await (db.Orders.CountAsync());
+				Assert.That(result, Is.EqualTo(830));
+
+				var log = sqlLog.GetWholeLog();
+				Assert.That(log, Does.Not.Contain("cast("));
+			}
+		}
 	}
 }

src/NHibernate.Test/Linq/ByMethod/CountTests.cs

@@ -1,6 +1,7 @@
 using System;
 using System.Linq;
 using NHibernate.Cfg;
+using NHibernate.Dialect;
 using NUnit.Framework;
 
 namespace NHibernate.Test.Linq.ByMethod
@@ -98,5 +99,50 @@ into temp
 
 			Assert.That(result.Count, Is.EqualTo(77));
 		}
+
+		[Test]
+		public void CheckSqlFunctionNameLongCount()
+		{
+			var name = Dialect is MsSql2000Dialect ? "count_big" : "count";
+			using (var sqlLog = new SqlLogSpy())
+			{
+				var result = db.Orders.LongCount();
+				Assert.That(result, Is.EqualTo(830));
+
+				var log = sqlLog.GetWholeLog();
+				Assert.That(log, Does.Contain($"{name}("));
+			}
+		}
+
+		[Test]
+		public void CheckSqlFunctionNameForCount()
+		{
+			using (var sqlLog = new SqlLogSpy())
+			{
+				var result = db.Orders.Count();
+				Assert.That(result, Is.EqualTo(830));
+
+				var log = sqlLog.GetWholeLog();
+				Assert.That(log, Does.Contain("count("));
+			}
+		}
+
+		[Test]
+		public void CheckMssqlCountCast()
+		{
+			if (!(Dialect is MsSql2000Dialect))
+			{
+				Assert.Ignore();
+			}
+
+			using (var sqlLog = new SqlLogSpy())
+			{
+				var result = db.Orders.Count();
+				Assert.That(result, Is.EqualTo(830));
+
+				var log = sqlLog.GetWholeLog();
+				Assert.That(log, Does.Not.Contain("cast("));
+			}
+		}
 	}
 }

src/NHibernate/Dialect/Dialect.cs

@@ -55,6 +55,7 @@ public abstract partial class Dialect
 		static Dialect()
 		{
 			StandardAggregateFunctions["count"] = new CountQueryFunctionInfo();
+			StandardAggregateFunctions["count_big"] = new CountQueryFunctionInfo();
 			StandardAggregateFunctions["avg"] = new AvgQueryFunctionInfo();
 			StandardAggregateFunctions["max"] = new ClassicAggregateFunction("max", false);
 			StandardAggregateFunctions["min"] = new ClassicAggregateFunction("min", false);

src/NHibernate/Dialect/Function/ClassicAggregateFunction.cs

@@ -1,15 +1,15 @@
 using System;
 using System.Collections;
-using System.Text;
+using System.Collections.Generic;
+using System.Linq;
 using NHibernate.Engine;
 using NHibernate.SqlCommand;
 using NHibernate.Type;
-using NHibernate.Util;
 
 namespace NHibernate.Dialect.Function
 {
 	[Serializable]
-	public class ClassicAggregateFunction : ISQLFunction, IFunctionGrammar
+	public class ClassicAggregateFunction : ISQLFunction, IFunctionGrammar, ISQLFunctionExtended
 	{
 		private IType returnType = null;
 		private readonly string name;
@@ -45,6 +45,15 @@ public virtual IType ReturnType(IType columnType, IMapping mapping)
 			return returnType ?? columnType;
 		}
 
+		/// <inheritdoc />
+		public virtual IType GetEffectiveReturnType(IEnumerable<IType> argumentTypes, IMapping mapping, bool throwOnError)
+		{
+			return ReturnType(argumentTypes.FirstOrDefault(), mapping);
+		}
+
+		/// <inheritdoc />
+		public string FunctionName => name;
+
 		public bool HasArguments
 		{
 			get { return true; }

src/NHibernate/Dialect/Function/ISQLFunction.cs

@@ -1,4 +1,6 @@
 using System.Collections;
+using System.Collections.Generic;
+using System.Linq;
 using NHibernate.Engine;
 using NHibernate.SqlCommand;
 using NHibernate.Type;
@@ -41,4 +43,45 @@ public interface ISQLFunction
 		/// <returns>SQL fragment for the function.</returns>
 		SqlString Render(IList args, ISessionFactoryImplementor factory);
 	}
+
+	// 6.0 TODO: Remove
+	internal static class SQLFunctionExtensions
+	{
+		/// <summary>
+		/// Get the type that will be effectively returned by the underlying database.
+		/// </summary>
+		/// <param name="sqlFunction">The sql function.</param>
+		/// <param name="argumentTypes">The types of arguments.</param>
+		/// <param name="mapping">The mapping for retrieving the argument sql types.</param>
+		/// <param name="throwOnError">Whether to throw when the number of arguments is invalid or they are not supported.</param>
+		/// <returns>The type returned by the underlying database or <see langword="null"/> when the number of arguments
+		/// is invalid or they are not supported.</returns>
+		/// <exception cref="QueryException">When <paramref name="throwOnError"/> is set to <see langword="true"/> and the
+		/// number of arguments is invalid or they are not supported.</exception>
+		public static IType GetEffectiveReturnType(
+			this ISQLFunction sqlFunction,
+			IEnumerable<IType> argumentTypes,
+			IMapping mapping,
+			bool throwOnError)
+		{
+			if (!(sqlFunction is ISQLFunctionExtended extendedSqlFunction))
+			{
+				try
+				{
+					return sqlFunction.ReturnType(argumentTypes.FirstOrDefault(), mapping);
+				}
+				catch (QueryException)
+				{
+					if (throwOnError)
+					{
+						throw;
+					}
+
+					return null;
+				}
+			}
+
+			return extendedSqlFunction.GetEffectiveReturnType(argumentTypes, mapping, throwOnError);
+		}
+	}
 }

src/NHibernate/Dialect/Function/ISQLFunctionExtended.cs

@@ -0,0 +1,27 @@
+using System.Collections.Generic;
+using NHibernate.Engine;
+using NHibernate.Type;
+
+namespace NHibernate.Dialect.Function
+{
+	// 6.0 TODO: Merge into ISQLFunction
+	internal interface ISQLFunctionExtended : ISQLFunction
+	{
+		/// <summary>
+		/// The function name or <see langword="null"/> when multiple functions/operators/statements are used.
+		/// </summary>
+		string FunctionName { get; }
+
+		/// <summary>
+		/// Get the type that will be effectively returned by the underlying database.
+		/// </summary>
+		/// <param name="argumentTypes">The types of arguments.</param>
+		/// <param name="mapping">The mapping for retrieving the argument sql types.</param>
+		/// <param name="throwOnError">Whether to throw when the number of arguments is invalid or they are not supported.</param>
+		/// <returns>The type returned by the underlying database or <see langword="null"/> when the number of arguments
+		/// is invalid or they are not supported.</returns>
+		/// <exception cref="QueryException">When <paramref name="throwOnError"/> is set to <see langword="true"/> and the
+		/// number of arguments is invalid or they are not supported.</exception>
+		IType GetEffectiveReturnType(IEnumerable<IType> argumentTypes, IMapping mapping, bool throwOnError);
+	}
+}

src/NHibernate/Dialect/MsSql2000Dialect.cs

@@ -286,7 +286,8 @@ protected virtual void RegisterKeywords()
 
 		protected virtual void RegisterFunctions()
 		{
-			RegisterFunction("count", new CountBigQueryFunction());
+			RegisterFunction("count", new CountQueryFunction());
+			RegisterFunction("count_big", new CountBigQueryFunction());
 
 			RegisterFunction("abs", new StandardSQLFunction("abs"));
 			RegisterFunction("absval", new StandardSQLFunction("absval"));
@@ -704,11 +705,15 @@ protected virtual string GetSelectExistingObject(string catalog, string schema,
 		[Serializable]
 		protected class CountBigQueryFunction : ClassicAggregateFunction
 		{
-			public CountBigQueryFunction() : base("count_big", true) { }
+			public CountBigQueryFunction() : base("count_big", true, NHibernateUtil.Int64) { }
+		}
 
-			public override IType ReturnType(IType columnType, IMapping mapping)
+		[Serializable]
+		private class CountQueryFunction : CountQueryFunctionInfo
+		{
+			public override IType GetEffectiveReturnType(IEnumerable<IType> argumentTypes, IMapping mapping, bool throwOnError)
 			{
-				return NHibernateUtil.Int64;
+				return NHibernateUtil.Int32;
 			}
 		}
 

src/NHibernate/Hql/Ast/ANTLR/SqlGenerator.cs

@@ -310,6 +310,12 @@ private void EndFunctionTemplate(IASTNode m)
 			}
 		}
 
+		private void OutAggregateFunctionName(IASTNode m)
+		{
+			var aggregateNode = (AggregateNode) m;
+			Out(aggregateNode.FunctionName);
+		}
+
 		private void CommaBetweenParameters(String comma)
 		{
 			writer.CommaBetweenParameters(comma);

src/NHibernate/Hql/Ast/ANTLR/SqlGenerator.g

@@ -150,7 +150,7 @@ selectExpr
 	;
 
 count
-	: ^(COUNT { Out("count("); }  ( distinctOrAll ) ? countExpr { Out(")"); } )
+	: ^(c=COUNT { OutAggregateFunctionName(c); Out("("); }  ( distinctOrAll ) ? countExpr { Out(")"); } )
 	;
 
 distinctOrAll
@@ -344,7 +344,7 @@ caseExpr
 	;
 
 aggregate
-	: ^(a=AGGREGATE { Out(a); Out("("); }  expr { Out(")"); } )
+	: ^(a=AGGREGATE { OutAggregateFunctionName(a); Out("("); }  expr { Out(")"); } )
 	;
 
 

src/NHibernate/Hql/Ast/ANTLR/Tree/AggregateNode.cs

@@ -1,5 +1,6 @@
 using System;
 using Antlr.Runtime;
+using NHibernate.Dialect.Function;
 using NHibernate.Type;
 using NHibernate.Hql.Ast.ANTLR.Util;
 
@@ -19,6 +20,19 @@ public AggregateNode(IToken token)
 		{
 		}
 
+		public string FunctionName
+		{
+			get
+			{
+				if (SessionFactoryHelper.FindSQLFunction(Text) is ISQLFunctionExtended sqlFunction)
+				{
+					return sqlFunction.FunctionName;
+				}
+
+				return Text;
+			}
+		}
+
 		public override IType DataType
 		{
 			get
@@ -31,6 +45,7 @@ public override IType DataType
 				base.DataType = value;
 			}
 		}
+		
 		public override void SetScalarColumnText(int i)
 		{
 			ColumnHelper.GenerateSingleScalarColumn(ASTFactory, this, i);

src/NHibernate/Hql/Ast/ANTLR/Tree/CountNode.cs

@@ -1,4 +1,5 @@
 using Antlr.Runtime;
+using NHibernate.Dialect.Function;
 using NHibernate.Hql.Ast.ANTLR.Util;
 using NHibernate.Type;
 
@@ -9,7 +10,7 @@ namespace NHibernate.Hql.Ast.ANTLR.Tree
 	/// Author: josh
 	/// Ported by: Steve Strong
 	/// </summary>
-	class CountNode : AbstractSelectExpression, ISelectExpression
+	class CountNode : AggregateNode, ISelectExpression
 	{
 		public CountNode(IToken token) : base(token)
 		{
@@ -26,9 +27,5 @@ public override IType DataType
 				base.DataType = value;
 			}
 		}
-		public override void SetScalarColumnText(int i)
-		{
-			ColumnHelper.GenerateSingleScalarColumn(ASTFactory, this, i);
-		}
 	}
 }

src/NHibernate/Hql/Ast/HqlTreeBuilder.cs

@@ -307,6 +307,11 @@ public HqlCount Count(HqlExpression child)
 			return new HqlCount(_factory, child);
 		}
 
+		public HqlCountBig CountBig(HqlExpression child)
+		{
+			return new HqlCountBig(_factory, child);
+		}
+
 		public HqlRowStar RowStar()
 		{
 			return new HqlRowStar(_factory);

src/NHibernate/Hql/Ast/HqlTreeNode.cs

@@ -697,6 +697,19 @@ public HqlCount(IASTFactory factory, HqlExpression child)
 		}
 	}
 
+	public class HqlCountBig : HqlExpression
+	{
+		public HqlCountBig(IASTFactory factory)
+			: base(HqlSqlWalker.COUNT, "count_big", factory)
+		{
+		}
+
+		public HqlCountBig(IASTFactory factory, HqlExpression child)
+			: base(HqlSqlWalker.COUNT, "count_big", factory, child)
+		{
+		}
+	}
+
 	public class HqlAs : HqlExpression
 	{
 		public HqlAs(IASTFactory factory, HqlExpression expression, System.Type type)