Implementing a custom QueryProvider with in-memory query

Question

I'm trying to create a wrapper around QueryableBase and INhQueryProvider that would receive a collection in the constructor and query it in-memory instead of going to a database. This is so I can mock the behavior of NHibernate's ToFuture() and properly unit test my classes.

The problem is that I'm facing a stack overflow due to infinite recursion and I'm struggling to find the reason.

Here's my implementation:

public class NHibernateQueryableProxy<T> : QueryableBase<T>, IOrderedQueryable<T>
{
    public NHibernateQueryableProxy(IQueryable<T> data) : base(new NhQueryProviderProxy<T>(data))
    {
    }

    public NHibernateQueryableProxy(IQueryParser queryParser, IQueryExecutor executor) : base(queryParser, executor)
    {
    }

    public NHibernateQueryableProxy(IQueryProvider provider) : base(provider)
    {
    }

    public NHibernateQueryableProxy(IQueryProvider provider, Expression expression) : base(provider, expression)
    {
    }

    public new IEnumerator<T> GetEnumerator()
    {
        return Provider.Execute<IEnumerable<T>>(Expression).GetEnumerator();
    }

    IEnumerator IEnumerable.GetEnumerator()
    {
        return GetEnumerator();
    }
}

internal class NhQueryProviderProxy<T> : INhQueryProvider
{
    private readonly IQueryProvider provider;

    public NhQueryProviderProxy(IQueryable<T> data)
    {
        provider = data.AsQueryable().Provider;
    }

    public IQueryable CreateQuery(Expression expression)
    {
        return new NHibernateQueryableProxy<T>(this, expression);
    }

    public IQueryable<TElement> CreateQuery<TElement>(Expression expression)
    {
        return new NHibernateQueryableProxy<TElement>(this, expression);
    }

    public object Execute(Expression expression)
    {
        return provider.Execute(expression);
    }

    public TResult Execute<TResult>(Expression expression)
    {
        return provider.Execute<TResult>(expression);
    }

    public object ExecuteFuture(Expression expression)
    {
        return provider.Execute(expression);
    }

    public void SetResultTransformerAndAdditionalCriteria(IQuery query, NhLinqExpression nhExpression, IDictionary<string, Tuple<object, IType>> parameters)
    {
        throw new NotImplementedException();
    }
}

Edit: I've kind of figured out the problem. One of the arguments to expression is my custom queryable. When this expression is executed by the provider, it causes an infinite call loop between CreateQuery and Execute. Is it possible to change all the references to my custom queryable to the queryable wrapped by this class?

Answer 1

After a while I decided to give it another try and I guess I've managed to mock it. I didn't test it with real case scenarios but I don't think many tweaks will be necessary. Most of this code is either taken from or based on this tutorial. There are some caveats related to IEnumerable when dealing with those queries.

We need to implement QueryableBase since NHibernate asserts the type when using ToFuture.

public class NHibernateQueryableProxy<T> : QueryableBase<T>
{
    public NHibernateQueryableProxy(IQueryable<T> data) : base(new NhQueryProviderProxy<T>(data))
    {
    }

    public NHibernateQueryableProxy(IQueryProvider provider, Expression expression) : base(provider, expression)
    {
    }
}

Now we need to mock a QueryProvider since that's what LINQ queries depend on and it needs to implement INhQueryProvider because ToFuture() also uses it.

public class NhQueryProviderProxy<T> : INhQueryProvider
{
    private readonly IQueryable<T> _data;

    public NhQueryProviderProxy(IQueryable<T> data)
    {
        _data = data;
    }

    // These two CreateQuery methods get called by LINQ extension methods to build up the query
    // and by ToFuture to return a queried collection and allow us to apply more filters
    public IQueryable CreateQuery(Expression expression)
    {
        Type elementType = TypeSystem.GetElementType(expression.Type);

        return (IQueryable)Activator.CreateInstance(typeof(NHibernateQueryableProxy<>)
                                    .MakeGenericType(elementType), new object[] { this, expression });
    }

    public IQueryable<TElement> CreateQuery<TElement>(Expression expression)
    {
        return new NHibernateQueryableProxy<TElement>(this, expression);
    }

    // Those two Execute methods are called by terminal methods like .ToList() and .ToArray()
    public object Execute(Expression expression)
    {
        return ExecuteInMemoryQuery(expression, false);
    }

    public TResult Execute<TResult>(Expression expression)
    {
        bool IsEnumerable = typeof(TResult).Name == "IEnumerable`1";
        return (TResult)ExecuteInMemoryQuery(expression, IsEnumerable);
    }

    public object ExecuteFuture(Expression expression)
    {
        // Here we need to return a NhQueryProviderProxy so we can add more queries
        // to the queryable and use another ToFuture if desired
        return CreateQuery(expression);
    }

    private object ExecuteInMemoryQuery(Expression expression, bool isEnumerable)
    {
        var newExpr = new ExpressionTreeModifier<T>(_data).Visit(expression);

        if (isEnumerable)
        {
            return _data.Provider.CreateQuery(newExpr);
        }

        return _data.Provider.Execute(newExpr);
    }

    public void SetResultTransformerAndAdditionalCriteria(IQuery query, NhLinqExpression nhExpression, IDictionary<string, Tuple<object, IType>> parameters)
    {
        throw new NotImplementedException();
    }
}

The expression tree visitor will change the type of the query for us:

internal class ExpressionTreeModifier<T> : ExpressionVisitor
{
    private IQueryable<T> _queryableData;

    internal ExpressionTreeModifier(IQueryable<T> queryableData)
    {
        _queryableData = queryableData;
    }

    protected override Expression VisitConstant(ConstantExpression c)
    {
        // Here the magic happens: the expression types are all NHibernateQueryableProxy,
        // so we replace them by the correct ones
        if (c.Type == typeof(NHibernateQueryableProxy<T>))
            return Expression.Constant(_queryableData);
        else
            return c;
    }
}

And we also need a helper (taken from the tutorial) to get the type being queried:

internal static class TypeSystem
{
    internal static Type GetElementType(Type seqType)
    {
        Type ienum = FindIEnumerable(seqType);
        if (ienum == null) return seqType;
        return ienum.GetGenericArguments()[0];
    }

    private static Type FindIEnumerable(Type seqType)
    {
        if (seqType == null || seqType == typeof(string))
            return null;

        if (seqType.IsArray)
            return typeof(IEnumerable<>).MakeGenericType(seqType.GetElementType());

        if (seqType.IsGenericType)
        {
            foreach (Type arg in seqType.GetGenericArguments())
            {
                Type ienum = typeof(IEnumerable<>).MakeGenericType(arg);
                if (ienum.IsAssignableFrom(seqType))
                {
                    return ienum;
                }
            }
        }

        Type[] ifaces = seqType.GetInterfaces();
        if (ifaces != null && ifaces.Length > 0)
        {
            foreach (Type iface in ifaces)
            {
                Type ienum = FindIEnumerable(iface);
                if (ienum != null) return ienum;
            }
        }

        if (seqType.BaseType != null && seqType.BaseType != typeof(object))
        {
            return FindIEnumerable(seqType.BaseType);
        }

        return null;
    }
}

To test the above code, I ran the following snippet:

var arr = new NHibernateQueryableProxy<int>(Enumerable.Range(1, 10000).AsQueryable());

var fluentQuery = arr.Where(x => x > 1 && x < 4321443)
            .Take(1000)
            .Skip(3)
            .Union(new[] { 4235, 24543, 52 })
            .GroupBy(x => x.ToString().Length)
            .ToFuture()
            .ToList();

var linqQuery = (from n in arr
                    where n > 40 && n < 50
                    select n.ToString())
                    .ToFuture()
                    .ToList();

As I said, no complex scenarios were tested but I guess only a few tweaks will be necessary for real-world usages.