FixCountExpression

OData/src/System.Web.OData/Extensions/HttpRequestMessageProperties.cs

@@ -33,10 +33,11 @@ public class HttpRequestMessageProperties
         private const string RoutingConventionsKey = "System.Web.OData.RoutingConventions";
         private const string SelectExpandClauseKey = "System.Web.OData.SelectExpandClause";
         private const string TotalCountKey = "System.Web.OData.TotalCount";
+        private const string TotalCountFuncKey = "System.Web.OData.TotalCountFunc";
 
         internal const string ODataServiceVersionHeader = "OData-Version";
         internal const string ODataMaxServiceVersionHeader = "OData-MaxVersion";
-        
+
         internal const ODataVersion DefaultODataVersion = ODataVersion.V4;
 
         private HttpRequestMessage _request;
@@ -47,6 +48,24 @@ internal HttpRequestMessageProperties(HttpRequestMessage request)
             _request = request;
         }
 
+        internal Func<long> TotalCountFunc
+        {
+            get
+            {
+                object totalCountFunc;
+                if (_request.Properties.TryGetValue(TotalCountFuncKey, out totalCountFunc))
+                {
+                    return (Func<long>)totalCountFunc;
+                }
+
+                return null;
+            }
+            set
+            {
+                _request.Properties[TotalCountFuncKey] = value;
+            }
+        }
+
         /// <summary>
         /// Gets or sets the EDM model associated with the request.
         /// </summary>
@@ -146,6 +165,11 @@ public Routing.ODataPath Path
                     return (long)totalCount;
                 }
 
+                if (this.TotalCountFunc != null)
+                {
+                    return this.TotalCountFunc();
+                }
+
                 return null;
             }
             set

OData/src/System.Web.OData/OData/ExpressionHelpers.cs

@@ -13,10 +13,11 @@ namespace System.Web.OData
 {
     internal static class ExpressionHelpers
     {
-        public static long Count(IQueryable query, Type type)
+        public static Func<long> Count(IQueryable query, Type type)
         {
             MethodInfo countMethod = ExpressionHelperMethods.QueryableCountGeneric.MakeGenericMethod(type);
-            return (long)countMethod.Invoke(null, new object[] { query });
+            Expression<Func<long>> func = Expression.Lambda<Func<long>>(Expression.Call(null, countMethod, query.Expression));
+            return func.Compile();
         }
 
         public static IQueryable Skip(IQueryable query, int count, Type type, bool parameterize)

OData/src/System.Web.OData/OData/Query/CountQueryOption.cs

@@ -4,6 +4,7 @@
 using System.Collections.Generic;
 using System.Diagnostics.Contracts;
 using System.Linq;
+using System.Linq.Expressions;
 using System.Web.Http;
 using System.Web.OData.Properties;
 using System.Web.OData.Query.Validators;
@@ -123,7 +124,7 @@ public void Validate(ODataValidationSettings validationSettings)
         }
 
         /// <summary>
-        /// Gets the number of entities that satify the given query if the response should include a count query option, or <c>null</c> otherwise.
+        /// Gets the number of entities that satisfy the given query if the response should include a count query option, or <c>null</c> otherwise.
         /// </summary>
         /// <param name="query">The query to compute the count for.</param>
         /// <returns>The number of entities that satisfy the specified query if the response should include a count query option, or <c>null</c> otherwise.</returns>
@@ -134,6 +135,28 @@ public void Validate(ODataValidationSettings validationSettings)
                 throw Error.NotSupported(SRResources.ApplyToOnUntypedQueryOption, "GetEntityCount");
             }
 
+            if (Value)
+            {
+                return ExpressionHelpers.Count(query, Context.ElementClrType)();
+            }
+            else
+            {
+                return null;
+            }
+        }
+
+        /// <summary>
+        /// Gets the Func of entities number that satisfy the given query if the response should include a count query option, or <c>null</c> otherwise.
+        /// </summary>
+        /// <param name="query">The query to compute the count for.</param>
+        /// <returns>The the Func of entities number that satisfy the specified query if the response should include a count query option, or <c>null</c> otherwise.</returns>
+        internal Func<long> GetEntityCountFunc(IQueryable query)
+        {
+            if (Context.ElementClrType == null)
+            {
+                throw Error.NotSupported(SRResources.ApplyToOnUntypedQueryOption, "GetEntityCount");
+            }
+
             if (Value)
             {
                 return ExpressionHelpers.Count(query, Context.ElementClrType);

OData/src/System.Web.OData/OData/Query/ODataQueryOptions.cs

@@ -261,12 +261,12 @@ public virtual IQueryable ApplyTo(IQueryable query, ODataQuerySettings querySett
 
             if (Count != null)
             {
-                if (Request.ODataProperties().TotalCount == null)
+                if (Request.ODataProperties().TotalCountFunc == null)
                 {
-                    long? count = Count.GetEntityCount(result);
-                    if (count.HasValue)
+                    Func<long> countFunc = Count.GetEntityCountFunc(result);
+                    if (countFunc != null)
                     {
-                        Request.ODataProperties().TotalCount = count.Value;
+                        Request.ODataProperties().TotalCountFunc = countFunc;
                     }
                 }
 

OData/test/System.Web.OData.Test/OData/Query/CountQueryOptionTest.cs

@@ -133,6 +133,26 @@ public void GetEntityCount_ReturnsNull_IfValueIsFalse()
             Assert.Null(countOption.GetEntityCount(_customers));
         }
 
+        [Fact]
+        public void GetEntityCountFunc_ReturnsFunc_IfValueIsTrue()
+        {
+            // Arrange
+            var countOption = new CountQueryOption("true", _context);
+
+            // Act & Assert
+            Assert.Equal(3, countOption.GetEntityCountFunc(_customers)());
+        }
+
+        [Fact]
+        public void GetEntityCountFunc_ReturnsNull_IfValueIsFalse()
+        {
+            // Arrange
+            var countOption = new CountQueryOption("false", _context);
+
+            // Act & Assert
+            Assert.Null(countOption.GetEntityCountFunc(_customers));
+        }
+
         [Fact]
         public void Property_Value_WorksWithUnTypedContext()
         {