Fix batch paging translation errors with inheritance selectors (#8938)

Author: PHILLIPS71

src/GreenDonut/src/GreenDonut.Data.EntityFramework/Expressions/ExpressionHelpers.cs

@@ -122,6 +122,9 @@ public static (Expression<Func<T, bool>> WhereExpression, int Offset) BuildWhere
     /// <param name="forward">
     /// Defines how the dataset is sorted.
     /// </param>
+    /// <param name="selector">
+    /// Optional selector to apply to each item before materialization.
+    /// </param>
     /// <param name="requestedCount">
     /// The number of items that are requested.
     /// </param>
@@ -141,6 +144,7 @@ public static BatchExpression<TK, TV> BuildBatchExpression<TK, TV>(
         ReadOnlySpan<LambdaExpression> orderExpressions,
         ReadOnlySpan<string> orderMethods,
         bool forward,
+        Expression<Func<TV, TV>>? selector,
         ref int requestedCount)
     {
         if (keys.Length == 0)
@@ -177,6 +181,17 @@ public static BatchExpression<TK, TV> BuildBatchExpression<TK, TV>(
                 typedOrderExpression);
         }
 
+        // apply the selector to each item in the grouping after ordering
+        if (selector is not null)
+        {
+            var selectMethod = typeof(Enumerable)
+                .GetMethods(BindingFlags.Static | BindingFlags.Public)
+                .First(m => m.Name == nameof(Enumerable.Select) && m.GetParameters().Length == 2)
+                .MakeGenericMethod(typeof(TV), typeof(TV));
+
+            source = Expression.Call(selectMethod, source, selector);
+        }
+
         var offset = 0;
         var usesRelativeCursors = false;
         Cursor? cursor = null;

src/GreenDonut/src/GreenDonut.Data.EntityFramework/Expressions/QueryHelpers.cs

@@ -74,14 +74,21 @@ static List<MemberExpression> GetMemberExpressions(Expression<Func<T, TKey>> key
         }
     }
 
-    private static Expression<Func<T, T>>? ExtractCurrentSelector<T>(
+    public static Expression<Func<T, T>>? ExtractCurrentSelector<T>(
         IQueryable<T> query)
     {
         var visitor = new ExtractSelectExpressionVisitor();
         visitor.Visit(query.Expression);
         return visitor.Selector as Expression<Func<T, T>>;
     }
 
+    public static IQueryable<T> RemoveSelector<T>(IQueryable<T> query)
+    {
+        var visitor = new RemoveSelectorVisitor();
+        var newExpression = visitor.Visit(query.Expression);
+        return query.Provider.CreateQuery<T>(newExpression);
+    }
+
     private static Expression<Func<T, T>> AddPropertiesInSelector<T>(
         Expression<Func<T, T>> selector,
         List<MemberExpression> properties)
@@ -109,6 +116,51 @@ private static IQueryable<T> ReplaceSelector<T>(
         return query.Provider.CreateQuery<T>(newExpression);
     }
 
+    private sealed class RemoveSelectorVisitor : ExpressionVisitor
+    {
+        private const string SelectMethod = "Select";
+
+        protected override Expression VisitMethodCall(MethodCallExpression node)
+        {
+            if (node.Method.Name != SelectMethod || node.Arguments.Count != 2)
+            {
+                return base.VisitMethodCall(node);
+            }
+
+            var lambda = ConvertToLambda(node.Arguments[1]);
+            if (!lambda.Type.IsGenericType || lambda.Type.GetGenericTypeDefinition() != typeof(Func<,>))
+            {
+                return base.VisitMethodCall(node);
+            }
+
+            var genericArgs = lambda.Type.GetGenericArguments();
+            // remove selectors of type Func<T, T>
+            if (genericArgs[0] == genericArgs[1])
+            {
+                // return the source expression, effectively removing the Select
+                return Visit(node.Arguments[0]);
+            }
+
+            return base.VisitMethodCall(node);
+        }
+
+        private static LambdaExpression ConvertToLambda(Expression e)
+        {
+            while (e.NodeType == ExpressionType.Quote)
+            {
+                e = ((UnaryExpression)e).Operand;
+            }
+
+            if (e.NodeType != ExpressionType.MemberAccess)
+            {
+                return (LambdaExpression)e;
+            }
+
+            var typeArguments = e.Type.GetGenericArguments()[0].GetGenericArguments();
+            return Expression.Lambda(e, Expression.Parameter(typeArguments[0]));
+        }
+    }
+
     public class AddPropertiesVisitorRewriter : ExpressionVisitor
     {
         private readonly List<MemberExpression> _propertiesToAdd;
@@ -122,6 +174,38 @@ public AddPropertiesVisitorRewriter(
             _parameter = parameter;
         }
 
+        protected override Expression VisitConditional(ConditionalExpression node)
+        {
+            // recursively visit conditional branches to handle type switches in inheritance scenarios
+            var test = Visit(node.Test);
+            var ifTrue = Visit(node.IfTrue);
+            var ifFalse = Visit(node.IfFalse);
+
+            if (test != node.Test || ifTrue != node.IfTrue || ifFalse != node.IfFalse)
+            {
+                return Expression.Condition(test, ifTrue, ifFalse, node.Type);
+            }
+
+            return node;
+        }
+
+        protected override Expression VisitUnary(UnaryExpression node)
+        {
+            // handle convert expressions that wrap member init expressions in inheritance scenarios
+            if (node.NodeType != ExpressionType.Convert)
+            {
+                return base.VisitUnary(node);
+            }
+
+            var operand = Visit(node.Operand);
+            if (operand != node.Operand)
+            {
+                return Expression.Convert(operand, node.Type);
+            }
+
+            return base.VisitUnary(node);
+        }
+
         protected override Expression VisitMemberInit(MemberInitExpression node)
         {
             // Get existing bindings (properties in the current selector)

src/GreenDonut/src/GreenDonut.Data.EntityFramework/Extensions/PagingQueryableExtensions.cs

@@ -436,6 +436,18 @@ public static async ValueTask<Dictionary<TKey, Page<TValue>>> ToBatchPageAsync<T
         }
 
         source = QueryHelpers.EnsureOrderPropsAreSelected(source);
+
+        // extract the selector before ensuring group props are selected,
+        // as we need to remove it before grouping and re-apply it after
+        var selector = QueryHelpers.ExtractCurrentSelector(source);
+
+        // if we have a selector, remove it before grouping
+        // we'll re-apply it to the grouped items later
+        if (selector is not null)
+        {
+            source = QueryHelpers.RemoveSelector(source);
+        }
+
         source = QueryHelpers.EnsureGroupPropsAreSelected(source, keySelector);
 
         // we need to move the ordering into the select expression we are constructing
@@ -458,6 +470,7 @@ public static async ValueTask<Dictionary<TKey, Page<TValue>>> ToBatchPageAsync<T
                 ordering.OrderExpressions,
                 ordering.OrderMethods,
                 forward,
+                selector,
                 ref requestedCount);
         var map = new Dictionary<TKey, Page<TValue>>();
 

src/HotChocolate/Data/test/Data.PostgreSQL.Tests/__snapshots__/IntegrationTests.Query_Brands_First_2_And_Products_First_2_NET10_0.md

@@ -59,7 +59,7 @@ LIMIT @p
 
 ```sql
 -- @brandIds={ '11', '13' } (DbType = Object)
-SELECT p1."BrandId", p3."Id", p3."Name", p3."BrandId"
+SELECT p1."BrandId", p3."Id", p3."Name"
 FROM (
     SELECT p."BrandId"
     FROM "Products" AS p

src/HotChocolate/Data/test/Data.PostgreSQL.Tests/__snapshots__/IntegrationTests.Query_Brands_First_2_And_Products_First_2_NET8_0.md

@@ -59,7 +59,7 @@ LIMIT @__p_0
 
 ```sql
 -- @__brandIds_0={ '11', '13' } (DbType = Object)
-SELECT t."BrandId", t0."Id", t0."Name", t0."BrandId"
+SELECT t."BrandId", t0."Id", t0."Name"
 FROM (
     SELECT p."BrandId"
     FROM "Products" AS p

src/HotChocolate/Data/test/Data.PostgreSQL.Tests/__snapshots__/IntegrationTests.Query_Brands_First_2_And_Products_First_2_NET9_0.md

@@ -59,7 +59,7 @@ LIMIT @__p_0
 
 ```sql
 -- @__brandIds_0={ '11', '13' } (DbType = Object)
-SELECT p1."BrandId", p3."Id", p3."Name", p3."BrandId"
+SELECT p1."BrandId", p3."Id", p3."Name"
 FROM (
     SELECT p."BrandId"
     FROM "Products" AS p

src/HotChocolate/Data/test/Data.PostgreSQL.Tests/__snapshots__/IntegrationTests.Query_Brands_First_2_And_Products_First_2_Name_Desc_NET10_0.md

@@ -59,7 +59,7 @@ LIMIT @p
 
 ```sql
 -- @brandIds={ '11', '13' } (DbType = Object)
-SELECT p1."BrandId", p3."Id", p3."Name", p3."BrandId"
+SELECT p1."BrandId", p3."Id", p3."Name"
 FROM (
     SELECT p."BrandId"
     FROM "Products" AS p

src/HotChocolate/Data/test/Data.PostgreSQL.Tests/__snapshots__/IntegrationTests.Query_Brands_First_2_And_Products_First_2_Name_Desc_NET8_0.md

@@ -59,7 +59,7 @@ LIMIT @__p_0
 
 ```sql
 -- @__brandIds_0={ '11', '13' } (DbType = Object)
-SELECT t."BrandId", t0."Id", t0."Name", t0."BrandId"
+SELECT t."BrandId", t0."Id", t0."Name"
 FROM (
     SELECT p."BrandId"
     FROM "Products" AS p

src/HotChocolate/Data/test/Data.PostgreSQL.Tests/__snapshots__/IntegrationTests.Query_Brands_First_2_And_Products_First_2_Name_Desc_NET9_0.md

@@ -59,7 +59,7 @@ LIMIT @__p_0
 
 ```sql
 -- @__brandIds_0={ '11', '13' } (DbType = Object)
-SELECT p1."BrandId", p3."Id", p3."Name", p3."BrandId"
+SELECT p1."BrandId", p3."Id", p3."Name"
 FROM (
     SELECT p."BrandId"
     FROM "Products" AS p

src/HotChocolate/Data/test/Data.PostgreSQL.Tests/__snapshots__/IntegrationTests.Query_Brands_First_2_Products_First_2_ForwardCursors_NET10_0.md

@@ -62,7 +62,7 @@ GROUP BY p."BrandId"
 
 ```sql
 -- @brandIds={ '11' } (DbType = Object)
-SELECT p1."BrandId", p3."Name", p3."Id", p3."BrandId"
+SELECT p1."BrandId", p3."Name", p3."Id"
 FROM (
     SELECT p."BrandId"
     FROM "Products" AS p