Search code examples
c#.netlinqexpression-trees

ExpressionTree ExpressionVisitor to change/replace query OrderBy Field


I need to be able to change the orderby field on the expression tree before it gets converted to sql. For example if the query contains an orderby "className.PropertyA", I need to modify this to be orderby "className.PropertyB" So far my idea has been to write an expression visitor to change the orderby method node on the tree. My code looks like this

public class ClassName
{
    public string PropertyA { get; set; }

    public string PropertyB { get; set; }

    public string PropertyC { get; set; }

    public string PropertyD { get; set; }
}

public class ChangeOrderByVisitor : ExpressionVisitor
{

    protected override Expression VisitMethodCall(MethodCallExpression node)
    {
        if (node.Method.DeclaringType == typeof(Queryable) &&
            (node.Method.Name == "OrderBy" || node.Method.Name == "OrderByDescending"))
        {
            //Only if ordering by className.PropertyA
                    //Somehow change the order by arguments.operands.body from className.PropertyA to className.PropertyB
        }

        return base.VisitMethodCall(node);
    }

}

At some point this expression tree will get converted to sql and should be ordering by className.PropertyB when executed against the database.

Thanks.


Solution

  • Took this as a nice programming exercise :-)

    Here's a class that deals with reflection and expressions:

    using System;
    using System.Linq;
    using System.Linq.Expressions;
    using System.Reflection;
    
    public static class MyExtensions
    {
        /// <summary>
        /// Helper method to reflect the return type.
        /// </summary>
        public static Expression<Func<T1, TResult>> FuncX<T1, TResult>(Expression<Func<T1, TResult>> lambda)
            => lambda;
    
        /// <summary>
        /// Helper method to get a <see cref="MemberInfo"/>.
        /// </summary>
        public static MemberInfo GetMember<TSource, TResult>(this Expression<Func<TSource, TResult>> lambda)
            => (lambda?.Body as MemberExpression)?.Member ?? throw new ArgumentException($"Not a {nameof(MemberExpression)}.");
    
        /// <summary>
        /// Helper method to get a <see cref="MethodInfo"/>.
        /// </summary>
        public static MethodInfo GetMethod<TSource, TResult>(this Expression<Func<TSource, TResult>> lambda)
            => (lambda?.Body as MethodCallExpression)?.Method ?? throw new ArgumentException($"Not a {nameof(MethodCallExpression)}.");
    
        /// <summary>
        /// <see cref="Queryable.OrderBy{TSource,TKey}(System.Linq.IQueryable{TSource},System.Linq.Expressions.Expression{System.Func{TSource,TKey}})"/>
        /// </summary>
        private static readonly MethodInfo _miOrderBy = GetMethod((IQueryable<int> q) => q.OrderBy(x => x)).GetGenericMethodDefinition();
    
        /// <summary>
        /// Replace occurrencies of OrderBy(<paramref name="origKeySelector"/>) with OrderBy(<paramref name="newKeySelector"/>).
        /// </summary>
        /// <exception cref="ArgumentException"><paramref name="origKeySelector"/>'s body is not a <see cref="MemberExpression"/>.</exception>
        public static IQueryable<TQueryable> ChangeOrder<TQueryable, TOrdered, TOrigOrder, TNewOrder>(this IQueryable<TQueryable> queryable, Expression<Func<TOrdered, TOrigOrder>> origKeySelector, Expression<Func<TOrdered, TNewOrder>> newKeySelector)
        {
            var changed = new ChangeOrderVisitor<TOrdered, TOrigOrder, TNewOrder>(origKeySelector, newKeySelector).Visit(queryable.Expression);
            return queryable.Provider.CreateQuery<TQueryable>(changed);
        }
    
        private sealed class ChangeOrderVisitor<TOrdered, TOrigOrder, TNewOrder> : ExpressionVisitor
        {
            private static readonly MethodInfo _miOrigOrderBy = _miOrderBy.MakeGenericMethod(typeof(TOrdered), typeof(TOrigOrder));
            private static readonly MethodInfo _miNewOrderBy = _miOrderBy.MakeGenericMethod(typeof(TOrdered), typeof(TNewOrder));
    
            private readonly MemberInfo _origMember;
            private readonly Expression<Func<TOrdered, TNewOrder>> _newKeySelector;
    
            public ChangeOrderVisitor(Expression<Func<TOrdered, TOrigOrder>> origKeySelector, Expression<Func<TOrdered, TNewOrder>> newKeySelector)
            {
                _origMember = origKeySelector.GetMember();
                _newKeySelector = newKeySelector;
            }
    
            protected override Expression VisitMethodCall(MethodCallExpression node)
            {
                if (node.Method == _miOrigOrderBy)
                {
                    if (node.Arguments[1] is UnaryExpression u &&
                        u.Operand is LambdaExpression lambda &&
                        lambda.Body is MemberExpression mx &&
                        mx.Member == _origMember)
                        return Expression.Call(_miNewOrderBy, Visit(node.Arguments[0]), _newKeySelector);
                }
                return base.VisitMethodCall(node);
            }
        }
    }
    

    And here's the test code:

    var origOrder = MyExtensions.FuncX((Person p) => p.FirstName);
    var qOrig = new[]
    {
        new Person{ FirstName = "Elon", LastName = "Musk", ShoeSize = 44 },
        new Person{ FirstName = "Jeff", LastName = "Who?", ShoeSize = 40 }
    }
        .AsQueryable()
        .OrderBy(origOrder)
        .Select(p => p.LastName);
    var qChanged = qOrig.ChangeOrder(origOrder, x => x.ShoeSize); // <string, Person, string, int>
    var result = qChanged.ToList(); // "Who?", "Musk"