Combining Mapped Database Functions on Multiple DataContexts

The LINQ to SQL designer provides some great flexibility for calling scalar and table-valued functions and stored procedures in your database.  You usually just need to drag the function from the Server Explorer into the DBML designer’s method pane and indicate how the output of the function should be materialized.  I’m a fan of keeping as much complex logic out of the database as possible, but sometimes you just can’t efficiently get the data you need from your database through LINQ into your middle tier.  Table-valued functions are especially useful for queries that can’t be expressed easily (or at all) using standard LINQ expressions.  Queries that can be made more efficient using ranking and windowing functions or recursive CTEs are good candidates since they often have no direct translation to LINQ.  Best of all, since the generated methods for table-valued functions return an instance of IQueryable, you can use function results within more complex LINQ expressions and still have the entire query executed at the database.

As you have come to expect form LINQ to SQL, pretty much anything you get out of a DataContext is tied to that DataContext instance.  This includes any instances of IQueryable you get from the generated methods for table-valued functions.  In my last post, I discussed using the LINQ to SQL O/RM behind the repositories in domain-driven design.  In this context, and in many other instances, it can be helpful to maintain just one DataContext instance for a given thread of execution or “unit of work.”  A singleton instance per thread.  This can be difficult to design, however, if some of your code relies on a particular derived DataContext type generated by the DBML designer to access its mapped functions.  Worse yet, you might have taken a divide and conquer approach to the database design and used the designer to produce more than one type of derived DataContext.  These DataContext classes could exist in separate assemblies and model different entities within the database (for example, maybe you have developed a reusable library for authentication and user management that has its own DataContext to interact with tables in the database).

When we look at the code generated for a derived DataContext, in a simple case we see properties defined for tables:

public System.Data.Linq.Table<Account> Accounts
{
	get
	{
		return this.GetTable<Account>();
	}
}

The method GetTable<T> is defined by the DataContext class, so regardless of the particular DataContext type you have, it’s clear that we can easily request a table of any entity type.  The table properties generated by the designer are purely a convenience.  If your code is only ever concerned with basic CRUD operations on tables (i.e. you never use mapped database functions), then you should be able to use an appropriately constructed instance of any DataContext, even the DataContext class itself, as your singleton DataContext instance.

Where you’ll get into trouble is when you need to use a mapped database function.  For now I’m only going to tackle mapped table-valued functions like the one below.  Stored procedures look a little different, but the approach I’m about to discuss can still be applied.  Let’s look at the code that the designer generates for us for a simple mapped function (this mapped function uses a recursive CTE defined in the database to get the ancestors of a particular node in a hierarchy):

[Function(Name = "dbo.GetNodePath", IsComposable = true)]
public IQueryable<NodeParent> GetNodePath([Parameter(Name = "NodeID", DbType = "Int")] System.Nullable<int> nodeID)
{
    return this.CreateMethodCallQuery<NodeParent>(this, ((MethodInfo)(MethodInfo.GetCurrentMethod())), nodeID);
}

There isn’t very much going on here.  We have a FunctionAttribute that tells LINQ to SQL that this is a mapped database function.  It indicates the function name and that it is composable (meaning we can include this in larger LINQ expressions).  The name of the method isn’t really important and the return type is simply whatever you select in the designer (LINQ to SQL tries its best to materialize the result as whatever type you request, or even an auto-generated type).  The parameters to the method are each decorated with a ParameterAttribute to help translate the method parameters into SQL parameters.  The body of the method calls a method inherited from DataContext with the current instance, the current method, and each of the parameters passed to the method. 

The designer generates this code for us within a derived DataContext type.  Ideally, we’d like to define this mapped function anywhere so that we can use it against any DataContext instance (especially our singleton instance we’re trying to keep).

Here is the metadata documentation from CreateMethodCallQuery:

//
// Summary:
//     Executes the table-valued database function associated with the specified
//     CLR method.
//
// Parameters:
//   instance:
//  The instance of the method invocation (the current object).
//
//   methodInfo:
//     Identifies the CLR method that corresponds to a database method.
//
//   parameters:
//     The array of parameters to be passed to the command.
//
// Type parameters:
//   TResult:
//     The type of the elements in the returned collection.
//
// Returns:
//     A collection of resultant values returned by the database query.
protected internal IQueryable<TResult> CreateMethodCallQuery<TResult>(object instance, MethodInfo methodInfo, params object[] parameters);

So far so good.  We see that the method is protected internal, so if we want to call this from outside a derived DataContext, we’re going to have to make a proxy method.  The proxy method would accept the same arguments as CreateMethodCallQuery and simply pass them along to the inherited method.  Here is the first attempt:

public class DataContextProxy : DataContext
{
    public DataContextProxy(string connectionString) : base(connectionString) { }

    public new IQueryable<TResult> CreateMethodCallQuery<TResult>(object instance, MethodInfo methodInfo, params object[] parameters)
    {
        return base.CreateMethodCallQuery<TResult>(instance, methodInfo, parameters);
    }
}

public class DatabaseFunctions
{
    private DataContextProxy _Proxy;
    public DatabaseFunctions(DataContextProxy proxy) { _Proxy = proxy; }

    [Function(Name = "dbo.GetNodePath", IsComposable = true)]
    public IQueryable<NodeParent> GetNodePath([Parameter(Name = "NodeID", DbType = "Int")] System.Nullable<int> nodeID)
    {
        return _Proxy.CreateMethodCallQuery<NodeParent>(this, ((MethodInfo)(MethodInfo.GetCurrentMethod())), nodeID);
    }
}


class Program
{
    static void Main(string[] args)
    {
        using (DataContextProxy dc = new DataContextProxy(Settings.Default.DefaultConnection))
        {
            DatabaseFunctions dbf = new DatabaseFunctions(dc);

            dbf.GetNodePath(0).ToList();    //call the method and force execution by copying to a list
        }
    }
}

Again, this is the direction we’re going in: define your mapped functions in a way that allows you to execute them against any DataContext.  Here our function is defined in a class called DatabaseFunctions, which receives an instance of our DataContextProxy class.

Executing this code produces the following exception:

System.InvalidOperationException was unhandled
  Message="The method 'GetNodePath' is not mapped as a stored procedure or user-defined function."
  Source="System.Data.Linq"

Well, it sure looks like that method is mapped!  Time to break out Reflector!  (Aside: if you’ve ever been curious about how a .NET assembly was implemented, and you’ve never heard of Reflector, go download it right now!)

We’re using an AttributeMappingSource, by default, to tell LINQ to SQL how to map our entities and functions to SQL statements.  Derived instances of MappingSource are responsible for creating an instance of MetaModel that gives LINQ to SQL this mapping.  There is also an internal class, System.Data.Linq.SqlClient.QueryConverter, that visits LINQ expressions and helps with the actual translation to SQL.  A quick browse through the disassembled DataContext and QueryConverter classes reveals some important facts:

• The CreateModel method of any mapping source is intended to examine types derived from DataContext when building a meta model.
• The DataContext class creates a MetaModel by examining its own type: MetaModel model = mapping.GetModel(base.GetType());.
• The QueryConverter VisitMethodCall method will only proceed to VisitMappedFunctionCall if the method is an instance method defined on a class derived from DataContext.

Our mapped function isn’t defined on a class derived from DataContext, so the QueryConverter won’t treat it as a mapped function.  However, we don’t even get to that point yet since our mapped function is defined outside of the DataContextProxy class and the AttributedMetaModel isn’t aware of it.  To take care of both, we’ll have DatabaseFunctions subclass DataContext and we’ll define our own MappingSource and MetaModel classes that dynamically build models as necessary based on the reflected type of the MethodInfo passed to CreateMethodCallQuery.  First, the DynamicAttributeMappingSource:

internal class DynamicAttributeMappingSource : MappingSource
{
    private static AttributeMappingSource AttributeMapping = new AttributeMappingSource();

    private Dictionary<Type, MetaModel> _ModelDictionary = new Dictionary<Type, MetaModel>();
    private ReaderWriterLockSlim _ModelLock = new ReaderWriterLockSlim();

    public MetaModel Primary { get { return GetModel(typeof(DataContextProxy)); } }

    protected override MetaModel CreateModel(Type dataContextType)
    {
        MetaModel Result;

        _ModelLock.EnterReadLock();
        try
        {
            if (_ModelDictionary.TryGetValue(dataContextType, out Result))
                return Result;
        }
        finally
        {
            _ModelLock.ExitReadLock();
        }

        _ModelLock.EnterWriteLock();
        try
        {
            if (_ModelDictionary.TryGetValue(dataContextType, out Result))
                return Result;

            Result = new DynamicAttributedMetaModel(this, AttributeMapping.GetModel(dataContextType));
            _ModelDictionary.Add(dataContextType, Result);
        }
        finally
        {
            _ModelLock.ExitWriteLock();
        }

        return Result;
    }
}

There’s nothing fancy really going on here.  A dictionary keeps track of models we’ve created for particular types and, in CreateModel, we have some thread-safe code for creating a model: we either get it out of the dictionary, or create a new one and put it into the dictionary.  When we create our MetaModel instances, we give them a MetaModel generated by an AttributeMappingSource as a starting point.  Here is the DynamicAttributedMetaModel:

internal class DynamicAttributedMetaModel : MetaModel
{
    public DynamicAttributeMappingSource DynamicSource { get; private set; }
    public MetaModel Surrogate { get; private set; }

    public DynamicAttributedMetaModel(DynamicAttributeMappingSource dynamicSource, MetaModel surrogate)
    {
        DynamicSource = dynamicSource;
        Surrogate = surrogate;
    }

    public override Type ContextType
    {
        get { return Surrogate.ContextType; }
    }

    public override string DatabaseName
    {
        get { return Surrogate.DatabaseName; }
    }

    public override MetaFunction GetFunction(System.Reflection.MethodInfo method)
    {
        Type MethodType = method.ReflectedType;

        if (Surrogate.ContextType == MethodType)
            return Surrogate.GetFunction(method);
        else
            return DynamicSource.GetModel(MethodType).GetFunction(method);
    }

    public override IEnumerable<MetaFunction> GetFunctions()
    {
        return Surrogate.GetFunctions();
    }

    public override MetaType GetMetaType(Type type)
    {
        return Surrogate.GetMetaType(type);
    }

    public override MetaTable GetTable(Type rowType)
    {
        return Surrogate.GetTable(rowType);
    }

    public override IEnumerable<MetaTable> GetTables()
    {
        return Surrogate.GetTables();
    }

    public override MappingSource MappingSource
    {
        get { return DynamicSource; }
    }

    public override Type ProviderType
    {
        get { return Surrogate.ProviderType; }
    }
}

This class is a proxy for a MetaModel, passing calls along to the surrogate model.  The MappingSource property refers back to the DynamicAttributeMappingSource (since that is in fact our mapping source), and there is some extra logic in the GetFunction method.  When the mapping for a function is requested, we check to see if the method is defined on a type that matches our model’s type.  If it is, then it has already been mapped and we can easily get the MetaFunction instance describing that mapped function from our surrogate model.  Otherwise, we have to get the MetaModel corresponding to the method’s type.  Our (dynamic) MappingSource can provide that for us easily.

Let’s change our DataContextProxy to use our new mapping source by default.  MappingSource instances provided here and by the framework are thread-safe.  This allows us to use one static instance and avoid the overhead of using reflection to re-build our MetaModel instances every time we create a new DataContext (the MappingSource class even implements caching of models in much the same way the above code does).

private static MappingSource _Mapping = new DynamicAttributeMappingSource();

public DataContextProxy(string connectionString) : base(connectionString, _Mapping) { }

And for our DatabaseFunctions class:

public class DatabaseFunctions : DataContext
{
    public DatabaseFunctions(DataContextProxy proxy) : base(string.Empty) { _Proxy = proxy; }

…

}

Okay let’s give it a go with the same driver code as before!

System.InvalidOperationException was unhandled
  Message="The query contains references to items defined on a different data context."
  Source="System.Data.Linq"

Well I suppose we should have expected this.  We have two DataContext types now: DataContextProxy and DatabaseFunctions.  The stack trace for the exception reveals that the check is made by the QueryConverter class:

private SqlNode VisitMappedFunctionCall(MethodCallExpression mc)
{
    MetaFunction function = this.services.Model.GetFunction(mc.Method);
    this.CheckContext(this.VisitExpression(mc.Object));
    if (!function.IsComposable)
    {
        return this.TranslateStoredProcedureCall(mc, function);
    }
    if (function.ResultRowTypes.Count > 0)
    {
        return this.TranslateTableValuedFunction(mc, function);
    }
    ProviderType sqlType = ((function.ReturnParameter != null) && !string.IsNullOrEmpty(function.ReturnParameter.DbType)) ? this.typeProvider.Parse(function.ReturnParameter.DbType) : this.typeProvider.From(mc.Method.ReturnType);
    List<SqlExpression> functionParameters = this.GetFunctionParameters(mc, function);
    return this.sql.FunctionCall(mc.Method.ReturnType, sqlType, function.MappedName, functionParameters, mc);
}

private void CheckContext(SqlExpression expr)
{
    SqlValue value2 = expr as SqlValue;
    if (value2 != null)
    {
        DataContext context = value2.Value as DataContext;
        if ((context != null) && (context != this.services.Context))
        {
            throw Error.WrongDataContext();
        }
    }
}

OK, let’s think here.  What is the QueryConverter actually working with?  It’s trying to translate an expression tree associated with an IQueryable object into an SQL statement.  The expression tree isn’t being executed, it’s just being parsed.  We need an instance of an object to execute an instance method on that object, but there’s no reason we can’t have an expression tree that represents calling such a method on a null reference.  If I can convince the QueryConverter that the context variable is null, then the CheckContext method will return successfully!  The context variable is the instance of the derived DataContext that was used to call our mapped function.  In our case, this is an instance of the class DatabaseFunctions.

So where is this part of the expression tree generated?  By the DataContext, of course:

protected internal IQueryable<TResult> CreateMethodCallQuery<TResult>(object instance, MethodInfo methodInfo, params object[] parameters)
{
    this.CheckDispose();
    if (instance == null)
    {
        throw Error.ArgumentNull("instance");
    }
    if (methodInfo == null)
    {
        throw Error.ArgumentNull("methodInfo");
    }
    if (parameters == null)
    {
        throw Error.ArgumentNull("parameters");
    }
    if (!typeof(IQueryable<TResult>).IsAssignableFrom(methodInfo.ReturnType))
    {
        throw Error.ExpectedQueryableArgument("methodInfo", typeof(IQueryable<TResult>));
    }
    return new DataQuery<TResult>(this, this.GetMethodCall(instance, methodInfo, parameters));
}

private Expression GetMethodCall(object instance, MethodInfo methodInfo, params object[] parameters)
{
    this.CheckDispose();
    if (parameters.Length <= 0)
    {
        return Expression.Call(Expression.Constant(instance), methodInfo);
    }
    ParameterInfo[] infoArray = methodInfo.GetParameters();
    List<Expression> arguments = new List<Expression>(parameters.Length);
    int index = 0;
    int length = parameters.Length;
    while (index < length)
    {
        Type parameterType = infoArray[index].ParameterType;
        if (parameterType.IsByRef)
        {
            parameterType = parameterType.GetElementType();
        }
        arguments.Add(Expression.Constant(parameters[index], parameterType));
        index++;
    }
    return Expression.Call(Expression.Constant(instance), methodInfo, arguments);
}

The GetMethodCall method is responsible for generating the expression tree, which always uses the instance passed to CreateMethodCallQuery in Expression.Call(…).  We can’t pass an instance of any type other than a type that declares our mapped function to CreateMethodCallQuery, since Expression.Call verifies that the instance type supports the specified method.

To circumvent this, we need to use our own GetMethodCall method and construct an instance of the DataQuery class (which is internal, unfortunately).

Here is the completed DataContextProxy class:

public class DataContextProxy : DataContext
{
    private static MappingSource _Mapping = new DynamicAttributeMappingSource();

    private static bool _ReflectionInit = false;
    private static ReaderWriterLockSlim _ReflectionLock = new ReaderWriterLockSlim();
    private static Type _DataQueryType;
    private static FieldInfo _DataContextProviderMember;
    private static MethodInfo _IProviderExecuteMethod;

    /// <summary>
    /// Construct a new instance of the data context proxy using a provided connection string and dynamic attribute mapping.
    /// </summary>
    /// <param name="connectionString">The connection string used to connect to the underlying database.</param>
    public DataContextProxy(string connectionString)
        : base(connectionString, _Mapping)
    {
    }

    #region Proxy Methods for Cross-Context Calls

    private static void InitReflection()
    {
        _ReflectionLock.EnterReadLock();
        try
        {
            if (_ReflectionInit)
                return;
        }
        finally
        {
            _ReflectionLock.ExitReadLock();
        }

        _ReflectionLock.EnterWriteLock();
        try
        {
            if (_ReflectionInit)
                return;

            _DataQueryType = typeof(DataContext).Assembly.GetType("System.Data.Linq.DataQuery`1", true);
            _DataContextProviderMember = typeof(DataContext).GetField("provider", BindingFlags.Instance | BindingFlags.NonPublic);
            _IProviderExecuteMethod = typeof(DataContext).Assembly.GetType("System.Data.Linq.Provider.IProvider").GetMethod("Execute");

            _ReflectionInit = true;
        }
        finally
        {
            _ReflectionLock.ExitWriteLock();
        }
    }

    /// <summary>
    /// Uses the mapping attributes on a method definition to execute a stored procedure.
    /// </summary>
    /// <param name="instance">The instance of the class that exposes the method.</param>
    /// <param name="method">The method invoked.</param>
    /// <param name="parameters">The parameters to the method invoked.</param>
    /// <returns>The execution result.</returns>
    public IExecuteResult ExecuteMethodCall(DataContext instance, MethodInfo method, params object[] parameters)
    {
        InitReflection();

        object Provider = _DataContextProviderMember.GetValue(this);
        return (IExecuteResult)_IProviderExecuteMethod.Invoke(Provider, new object[] { GetMethodCall(instance, method, parameters) });
    }

    /// <summary>
    /// Uses the mapping attributes on a method definition to query a table-valued function.
    /// </summary>
    /// <typeparam name="TElement">The type of element queried.</typeparam>
    /// <param name="instance">The instance of the class that exposes the method.</param>
    /// <param name="method">The method invoked.</param>
    /// <param name="parameters">The parameters to the method invoked.</param>
    /// <returns>The queryable result.</returns>
    public IQueryable<TElement> CreateMethodCallQuery<TElement>(DataContext instance, MethodInfo method, params object[] parameters)
    {
        InitReflection();

        Type DataQueryType = _DataQueryType.MakeGenericType(typeof(TElement));

        return (IQueryable<TElement>)Activator.CreateInstance(DataQueryType, this, GetMethodCall(instance, method, parameters));
    }

    /// <remarks>This is copied almost verbatim from System.DataLinq.DataContext:GetMethodCall.
    /// However, note that the instance is null but of the correct type.
    /// This exploits two loop-holes in System.Data.Linq.SqlClient.QueryConverter:
    /// (1) An instance inheriting DataContext is required to map functions, and
    /// (2) A null instance passes the CheckContext method by default.</remarks>
    private static Expression GetMethodCall(DataContext instance, MethodInfo methodInfo, params object[] parameters)
    {
        if (parameters.Length <= 0)
        {
            return Expression.Call(Expression.Constant(null, methodInfo.DeclaringType), methodInfo);
        }
        ParameterInfo[] infoArray = methodInfo.GetParameters();
        List<Expression> arguments = new List<Expression>(parameters.Length);
        int index = 0;
        int length = parameters.Length;
        while (index < length)
        {
            Type parameterType = infoArray[index].ParameterType;
            if (parameterType.IsByRef)
            {
                parameterType = parameterType.GetElementType();
            }
            arguments.Add(Expression.Constant(parameters[index], parameterType));
            index++;
        }
        return Expression.Call(Expression.Constant(null, methodInfo.DeclaringType), methodInfo, arguments);
    }

    #endregion
}

Notice how GetMethodCall passes null and a type to Expression.Constant, instead of the instance parameter.  This circumvents the context type check performed by QueryConverter.  We also have to use reflection to dig into some of the internal classes used by DataContext.  Obviously if LINQ to SQL changes significantly, this may not work anymore (but hopefully they’ll have added this or similar functionality in).

The above code includes both the CreateMethodCallQuery method to handle table-valued functions and the ExecuteMethodCall method to handle stored procedures.  Using either method appropriately, you can now define your mapped database functions outside of the designer generated code and use them against any data context!