From 7982369eb4551f5c4ef5c55a8702920a969d1a6b Mon Sep 17 00:00:00 2001 From: Michael Staib Date: Tue, 1 Oct 2024 21:54:17 +0200 Subject: [PATCH] Improved selector when using batching --- .../ExtractOrderPropertiesVisitor.cs | 53 +++++++++++++ .../ExtractSelectExpressionVisitor.cs | 41 ++++++++++ .../Expressions/QueryHelpers.cs | 75 +++++++++++++++++++ .../Expressions/ReplaceSelectorVisitor.cs | 25 +++++++ .../Expressions/ReplacerParameterVisitor.cs | 14 ++++ .../Extensions/PagingQueryableExtensions.cs | 4 + 6 files changed, 212 insertions(+) create mode 100644 src/HotChocolate/Pagination/src/Pagination.EntityFramework/Expressions/ExtractOrderPropertiesVisitor.cs create mode 100644 src/HotChocolate/Pagination/src/Pagination.EntityFramework/Expressions/ExtractSelectExpressionVisitor.cs create mode 100644 src/HotChocolate/Pagination/src/Pagination.EntityFramework/Expressions/QueryHelpers.cs create mode 100644 src/HotChocolate/Pagination/src/Pagination.EntityFramework/Expressions/ReplaceSelectorVisitor.cs create mode 100644 src/HotChocolate/Pagination/src/Pagination.EntityFramework/Expressions/ReplacerParameterVisitor.cs diff --git a/src/HotChocolate/Pagination/src/Pagination.EntityFramework/Expressions/ExtractOrderPropertiesVisitor.cs b/src/HotChocolate/Pagination/src/Pagination.EntityFramework/Expressions/ExtractOrderPropertiesVisitor.cs new file mode 100644 index 00000000000..72393194fd1 --- /dev/null +++ b/src/HotChocolate/Pagination/src/Pagination.EntityFramework/Expressions/ExtractOrderPropertiesVisitor.cs @@ -0,0 +1,53 @@ +using System.Linq.Expressions; + +namespace HotChocolate.Pagination.Expressions; + +internal sealed class ExtractOrderPropertiesVisitor : ExpressionVisitor +{ + private const string _orderByMethod = "OrderBy"; + private const string _thenByMethod = "ThenBy"; + private const string _orderByDescendingMethod = "OrderByDescending"; + private const string _thenByDescendingMethod = "ThenByDescending"; + private bool _isOrderScope; + + public List OrderProperties { get; } = []; + + protected override Expression VisitMethodCall(MethodCallExpression node) + { + if (node.Method.Name == _orderByMethod || + node.Method.Name == _thenByMethod || + node.Method.Name == _orderByDescendingMethod || + node.Method.Name == _thenByDescendingMethod) + { + _isOrderScope = true; + + var lambda = StripQuotes(node.Arguments[1]); + Visit(lambda.Body); + + _isOrderScope = false; + } + + return base.VisitMethodCall(node); + } + + protected override Expression VisitMember(MemberExpression node) + { + if (_isOrderScope) + { + // we only collect members that are within an order method. + OrderProperties.Add(node); + } + + return base.VisitMember(node); + } + + private static LambdaExpression StripQuotes(Expression expression) + { + while (expression.NodeType == ExpressionType.Quote) + { + expression = ((UnaryExpression)expression).Operand; + } + + return (LambdaExpression)expression; + } +} diff --git a/src/HotChocolate/Pagination/src/Pagination.EntityFramework/Expressions/ExtractSelectExpressionVisitor.cs b/src/HotChocolate/Pagination/src/Pagination.EntityFramework/Expressions/ExtractSelectExpressionVisitor.cs new file mode 100644 index 00000000000..e1f6950918d --- /dev/null +++ b/src/HotChocolate/Pagination/src/Pagination.EntityFramework/Expressions/ExtractSelectExpressionVisitor.cs @@ -0,0 +1,41 @@ +using System.Linq.Expressions; + +namespace HotChocolate.Pagination.Expressions; + +internal sealed class ExtractSelectExpressionVisitor : ExpressionVisitor +{ + private const string _selectMethod = "Select"; + + public LambdaExpression? Selector { get; private set; } + + protected override Expression VisitMethodCall(MethodCallExpression node) + { + if (node.Method.Name == _selectMethod && node.Arguments.Count == 2) + { + var lambda = StripQuotes(node.Arguments[1]); + if (lambda.Type.IsGenericType + && lambda.Type.GetGenericTypeDefinition() == typeof(Func<,>)) + { + // we make sure that the selector is of type Expression> + // otherwise we are not interested in it. + var genericArgs = lambda.Type.GetGenericArguments(); + if (genericArgs[0] == genericArgs[1]) + { + Selector = lambda; + } + } + } + + return base.VisitMethodCall(node); + } + + private static LambdaExpression StripQuotes(Expression e) + { + while (e.NodeType == ExpressionType.Quote) + { + e = ((UnaryExpression)e).Operand; + } + + return (LambdaExpression)e; + } +} diff --git a/src/HotChocolate/Pagination/src/Pagination.EntityFramework/Expressions/QueryHelpers.cs b/src/HotChocolate/Pagination/src/Pagination.EntityFramework/Expressions/QueryHelpers.cs new file mode 100644 index 00000000000..ae6680628f4 --- /dev/null +++ b/src/HotChocolate/Pagination/src/Pagination.EntityFramework/Expressions/QueryHelpers.cs @@ -0,0 +1,75 @@ +using System.Linq.Expressions; + +namespace HotChocolate.Pagination.Expressions; + +internal static class QueryHelpers +{ + public static IQueryable EnsureOrderPropsAreSelected( + IQueryable query) + { + var selector = ExtractCurrentSelector(query); + if (selector is null) + { + return query; + } + + var orderByProperties = ExtractOrderProperties(query); + if(orderByProperties.Count == 0) + { + return query; + } + + var updatedSelector = AddPropertiesInSelector(selector, orderByProperties); + return ReplaceSelector(query, updatedSelector); + } + + private static Expression>? ExtractCurrentSelector( + IQueryable query) + { + var visitor = new ExtractSelectExpressionVisitor(); + visitor.Visit(query.Expression); + return visitor.Selector as Expression>; + } + + private static Expression> AddPropertiesInSelector( + Expression> selector, + List properties) + { + var parameter = selector.Parameters[0]; + var bindings = ((MemberInitExpression)selector.Body).Bindings.Cast().ToList(); + + foreach (var property in properties) + { + var propertyName = property.Member.Name; + if(property.Expression is not ParameterExpression parameterExpression + || bindings.Any(b => b.Member.Name == propertyName)) + { + continue; + } + + var replacer = new ReplacerParameterVisitor(parameterExpression, parameter); + var rewrittenProperty = (MemberExpression)replacer.Visit(property); + bindings.Add(Expression.Bind(rewrittenProperty.Member, rewrittenProperty)); + } + + var newBody = Expression.MemberInit(Expression.New(typeof(T)), bindings); + return Expression.Lambda>(newBody, parameter); + } + + private static List ExtractOrderProperties( + IQueryable query) + { + var visitor = new ExtractOrderPropertiesVisitor(); + visitor.Visit(query.Expression); + return visitor.OrderProperties; + } + + private static IQueryable ReplaceSelector( + IQueryable query, + Expression> newSelector) + { + var visitor = new ReplaceSelectorVisitor(newSelector); + var newExpression = visitor.Visit(query.Expression); + return query.Provider.CreateQuery(newExpression); + } +} diff --git a/src/HotChocolate/Pagination/src/Pagination.EntityFramework/Expressions/ReplaceSelectorVisitor.cs b/src/HotChocolate/Pagination/src/Pagination.EntityFramework/Expressions/ReplaceSelectorVisitor.cs new file mode 100644 index 00000000000..6c3c325a2c4 --- /dev/null +++ b/src/HotChocolate/Pagination/src/Pagination.EntityFramework/Expressions/ReplaceSelectorVisitor.cs @@ -0,0 +1,25 @@ +using System.Linq.Expressions; + +namespace HotChocolate.Pagination.Expressions; + +internal sealed class ReplaceSelectorVisitor( + Expression> newSelector) + : ExpressionVisitor +{ + private const string _selectMethod = "Select"; + + protected override Expression VisitMethodCall(MethodCallExpression node) + { + if (node.Method.Name == _selectMethod && node.Arguments.Count == 2) + { + return Expression.Call( + node.Method.DeclaringType!, + node.Method.Name, + [typeof(T), typeof(T)], + node.Arguments[0], + newSelector); + } + + return base.VisitMethodCall(node); + } +} diff --git a/src/HotChocolate/Pagination/src/Pagination.EntityFramework/Expressions/ReplacerParameterVisitor.cs b/src/HotChocolate/Pagination/src/Pagination.EntityFramework/Expressions/ReplacerParameterVisitor.cs new file mode 100644 index 00000000000..5007c3a3891 --- /dev/null +++ b/src/HotChocolate/Pagination/src/Pagination.EntityFramework/Expressions/ReplacerParameterVisitor.cs @@ -0,0 +1,14 @@ +using System.Linq.Expressions; + +namespace HotChocolate.Pagination.Expressions; + +internal sealed class ReplacerParameterVisitor( + ParameterExpression oldParameter, + ParameterExpression newParameter) + : ExpressionVisitor +{ + protected override Expression VisitParameter(ParameterExpression node) + => node == oldParameter + ? newParameter + : base.VisitParameter(node); +} diff --git a/src/HotChocolate/Pagination/src/Pagination.EntityFramework/Extensions/PagingQueryableExtensions.cs b/src/HotChocolate/Pagination/src/Pagination.EntityFramework/Extensions/PagingQueryableExtensions.cs index 572d74580df..9fa20a89a70 100644 --- a/src/HotChocolate/Pagination/src/Pagination.EntityFramework/Extensions/PagingQueryableExtensions.cs +++ b/src/HotChocolate/Pagination/src/Pagination.EntityFramework/Extensions/PagingQueryableExtensions.cs @@ -72,6 +72,8 @@ public static async ValueTask> ToPageAsync( throw new ArgumentNullException(nameof(source)); } + source = QueryHelpers.EnsureOrderPropsAreSelected(source); + var keys = ParseDataSetKeys(source); if (keys.Length == 0) @@ -215,6 +217,8 @@ public static async ValueTask>> ToBatchPageAsync