Modifying LINQ to SQL command text without Reflection

Following my post last year about modifying LINQ to SQL command text (evil, as it calls private methods through reflection) here’s an equally evil, but faster version that pre-compiles most of its work through an expression tree:

public delegate string ModifyCommandDelegate( string commandText, IDictionary<string, object> parameters );

public class DataContextInterceptor
{
    private DataContext dc;
    private object oldProvider;
    private Type providerType;
    private ModifyCommandDelegate modifyCommand;

    private static Func<object, ModifyCommandDelegate, DataContext, Func<Expression, object>> CompileFactory = CreateCompileFactory().Compile();
    private static Func<object, ModifyCommandDelegate, Func<Expression, IExecuteResult>> ExecuteFactory = CreateExecuteFactory().Compile();

    public static TDataContext Intercept<TDataContext>( TDataContext dc, ModifyCommandDelegate modifyCommand )
        where TDataContext : DataContext
    {
        new DataContextInterceptor( dc, modifyCommand );

        //            typeof( Expression ).GetProperty( "DebugView", BindingFlags.Instance | BindingFlags.NonPublic ).GetValue( executeFactoryExp, null ).Dump();

        return dc;
    }

    public DataContextInterceptor( DataContext dc, ModifyCommandDelegate modifyCommand )
    {
        this.dc = dc;
        this.modifyCommand = modifyCommand;

        FieldInfo providerField = typeof( DataContext ).GetField( "provider", BindingFlags.Instance | BindingFlags.NonPublic );
        var existingProvider = providerField.GetValue( dc );

        if ( existingProvider is IProviderProxy )
        {
            // System.Diagnostics.Trace.WriteLine( string.Format( "DataContext {0} already intercepted", dc.GetHashCode() ) );
        }
        else
        {
            oldProvider = existingProvider;

            var proxy = new ProviderProxy( this, oldProvider ).GetTransparentProxy();

            providerField.SetValue( dc, proxy );

            // System.Diagnostics.Trace.WriteLine( string.Format( "DataContext {0} intercepted", dc.GetHashCode() ) );
        }
    }

    public static MethodCallExpression MakeMethodCall( Type type, string methodName, params Expression[] arguments )
    {
        // ( "Making MethodCallExpression for " + methodName ).Dump();

        var methodInfo = type.GetMethod(
            methodName,
            BindingFlags.Static | BindingFlags.Public | BindingFlags.NonPublic,
            null,
            arguments.Select( a => a.Type ).ToArray(),
            null );

        if ( methodInfo == null ) throw new ArgumentException( string.Format( "Unable to find method {0}.{1}", type.Name, methodName ), "methodName" );

        return Expression.Call( methodInfo, arguments );
    }

    public static MethodCallExpression MakeMethodCall( Expression instance, string methodName, params Expression[] arguments )
    {
        // ( "Making MethodCallExpression for " + instance.ToString() + "." + methodName ).Dump();

        var methodInfo = instance.Type.GetMethod(
            methodName,
            BindingFlags.Instance | BindingFlags.Public | BindingFlags.NonPublic,
            null,
            arguments.Select( a => a.Type ).ToArray(),
            null );

        if ( methodInfo == null ) throw new ArgumentException( string.Format( "Unable to find method {0}.{1}", instance.Type.Name, methodName ), "methodName" );

        return Expression.Call( instance, methodInfo, arguments );
    }

    protected internal object Compile( Expression query )
    {
        return DataContextInterceptor.CompileFactory( oldProvider, modifyCommand, dc )( query );
    }

    protected internal virtual IExecuteResult Execute( Expression query )
    {
        return DataContextInterceptor.ExecuteFactory( oldProvider, modifyCommand )( query );
    }

    public static Expression<Func<object, ModifyCommandDelegate, DataContext, Func<Expression, object>>> CreateCompileFactory()
    {
        var assembly = typeof( SqlProvider ).Assembly;

        var providerType = assembly.GetType( "System.Data.Linq.SqlClient.SqlProvider" );
        var providerParam = Expression.Parameter( typeof( object ), "provider" );
        var modifyCommandParam = Expression.Parameter( typeof( ModifyCommandDelegate ), "modifyCommand" );
        var dataContextParam = Expression.Parameter( typeof( DataContext ), "dataContext" );

        return Expression.Lambda<Func<object, ModifyCommandDelegate, DataContext, Func<Expression, object>>>(
            CreateCompileMethod( Expression.Convert( providerParam, providerType ), modifyCommandParam, dataContextParam ),
            providerParam,
            modifyCommandParam,
            dataContextParam );
    }

    public static Expression<Func<Expression, object>> CreateCompileMethod( Expression oldProvider, Expression modifyCommand, Expression dataContextParam )
    {
        var assembly = typeof( SqlProvider ).Assembly;

        var providerType = assembly.GetType( "System.Data.Linq.Provider.IProvider" );
        var funcletizerType = assembly.GetType( "System.Data.Linq.SqlClient.Funcletizer" );
        var annotationsType = assembly.GetType( "System.Data.Linq.SqlClient.SqlNodeAnnotations" );
        var queryInfoType = assembly.GetType( "System.Data.Linq.SqlClient.SqlProvider" ).GetNestedType( "QueryInfo", BindingFlags.NonPublic );
        var readerFactoryType = assembly.GetType( "System.Data.Linq.SqlClient.IObjectReaderFactory" );
        var resultShapeType = assembly.GetType( "System.Data.Linq.SqlClient.SqlProvider" ).GetNestedType( "ResultShape", BindingFlags.NonPublic );
        var compiledSubQueryType = assembly.GetType( "System.Data.Linq.SqlClient.ICompiledSubQuery" );
        var typeSystemType = assembly.GetType( "System.Data.Linq.SqlClient.TypeSystem" );
        var compiledQueryType = assembly.GetType( "System.Data.Linq.SqlClient.SqlProvider+CompiledQuery" );

        var queryParam = Expression.Variable( typeof( Expression ), "query" );

        var providerField = typeof( DataContext ).GetField( "provider", BindingFlags.Instance | BindingFlags.NonPublic );

        var annotationsVar = Expression.Variable( annotationsType, "annotations" );
        var queriesVar = Expression.Variable( queryInfoType.MakeArrayType(), "queries" );
        var infoVar = Expression.Variable( queryInfoType, "info" );
        var lambdaVar = Expression.Variable( typeof( LambdaExpression ), "lambda" );
        var readerFactoryVar = Expression.Variable( readerFactoryType, "readerFactory" );
        var subQueriesVar = Expression.Variable( compiledSubQueryType.MakeArrayType(), "subQueries" );
        var resultShapeVar = Expression.Variable( resultShapeType, "resultShape" );
        var returnTarget = Expression.Label( compiledQueryType );

        var getQuery = MakeMethodCall( infoVar, "get_Query" );
        var getResultType = MakeMethodCall( infoVar, "get_ResultType" );
        var getResultElementType = MakeMethodCall( typeSystemType, "GetElementType", getResultType );
        var modifySubQueries = CreateModifySubQueriesMethod();
        var compiledQueryConstructor = compiledQueryType.GetConstructors( BindingFlags.NonPublic | BindingFlags.Instance ).First();

        return Expression.Lambda<Func<Expression, object>>(
            Expression.Block(
                new[] { annotationsVar, queriesVar, infoVar, lambdaVar, readerFactoryVar, subQueriesVar, resultShapeVar },

                // this.InitializeProviderMode();
                MakeMethodCall( oldProvider, "InitializeProviderMode" ),

                // SqlNodeAnnotations annotations = new SqlNodeAnnotations();
                Expression.Assign( annotationsVar, Expression.New( annotationsType ) ),

                // QueryInfo[] queries = this.BuildQuery( query, annotations );
                Expression.Assign( queriesVar, MakeMethodCall( oldProvider, "BuildQuery", queryParam, annotationsVar ) ),

                // var info = ModifyQueries( (IEnumerable)queries );
                Expression.Assign( infoVar, Expression.Invoke( CreateModifyQueriesMethod(), queriesVar, modifyCommand ) ),

                // this.CheckSqlCompatibility(queries, annotations);
                MakeMethodCall( oldProvider, "CheckSqlCompatibility", queriesVar, annotationsVar ),

                // var lambda = query as LambdaExpression;
                Expression.Assign( lambdaVar, Expression.TypeAs( queryParam, lambdaVar.Type ) ),

                // if ( lambda != null )
                Expression.IfThen(
                    Expression.NotEqual( lambdaVar, Expression.Constant( null ) ),
                    // query = lambda.Body;
                    Expression.Assign( queryParam, Expression.Property( lambdaVar, "Body" ) ) ),

                // IObjectReaderFactory readerFactory = null;
                // ICompiledSubQuery[] subQueries = null;

                // QueryInfo info = queries[queries.Length - 1];
                // info defined above
                // var resultShape = (int)Invoke( info, "get_ResultShape" );
                Expression.Assign( resultShapeVar, MakeMethodCall( infoVar, "get_ResultShape" ) ),

                Expression.Switch(
                    Expression.Convert( resultShapeVar, typeof( int ) ),
                    // if ( resultShape == 1 /* ResultShape.Singleton */ )
                    Expression.SwitchCase(
                        Expression.Block(
                            // subQueries = this.CompileSubQueries( info.Query );
                            Expression.Assign( subQueriesVar, MakeMethodCall( oldProvider, "CompileSubQueries", getQuery ) ),
                            // ModifySubQueries( (IEnumerable)subQueries );
                            Expression.Invoke( modifySubQueries, subQueriesVar, modifyCommand ),
                            // readerFactory = this.GetReaderFactory( info.Query, info.ResultType );
                            Expression.Assign( readerFactoryVar, MakeMethodCall( oldProvider, "GetReaderFactory", getQuery, getResultType ) ),
                            Expression.Empty() ),
                        Expression.Constant( 1 ) ),
                    // else if ( resultShape == 2 /* ResultShape.Sequence */ )
                    Expression.SwitchCase(
                        Expression.Block(
                            // subQueries = this.CompileSubQueries( info.Query );
                            Expression.Assign( subQueriesVar, MakeMethodCall( oldProvider, "CompileSubQueries", getQuery ) ),
                            // ModifySubQueries( (IEnumerable)subQueries );
                            Expression.Invoke( modifySubQueries, subQueriesVar, modifyCommand ),
                            // readerFactory = this.GetReaderFactory( info.Query, TypeSystem.GetElementType( info.ResultType ) );
                            Expression.Assign( readerFactoryVar, MakeMethodCall( oldProvider, "GetReaderFactory", getQuery, getResultElementType ) ),
                            Expression.Empty() ),
                        Expression.Constant( 2 ) ) ),

                // dc.provider = oldProvider;    // (Unfortunately needed to ensure compilation runs)
                Expression.Assign( Expression.MakeMemberAccess( dataContextParam, providerField ), oldProvider ),

                // return new CompiledQuery( oldProvider, query, queries, readerFactory, subQueries );
                Expression.Return( returnTarget, Expression.New( compiledQueryConstructor, oldProvider, queryParam, queriesVar, readerFactoryVar, subQueriesVar ) ),
                // --- Expression.Throw( Expression.New( typeof( Exception ).GetConstructor( new[] { typeof( string ) } ), Expression.Constant( "boo4" ) ) ),
                Expression.Label( returnTarget, Expression.Constant( null, compiledQueryType ) ) ),
            queryParam );
    }

    public static Expression<Func<object, ModifyCommandDelegate, Func<Expression, IExecuteResult>>> CreateExecuteFactory()
    {
        var assembly = typeof( SqlProvider ).Assembly;

        var providerType = assembly.GetType( "System.Data.Linq.SqlClient.SqlProvider" );
        var providerParam = Expression.Parameter( typeof( object ), "provider" );
        var modifyCommandParam = Expression.Parameter( typeof( ModifyCommandDelegate ), "modifyCommand" );

        return Expression.Lambda<Func<object, ModifyCommandDelegate, Func<Expression, IExecuteResult>>>(
            CreateExecuteMethod( Expression.Convert( providerParam, providerType ), modifyCommandParam ),
            providerParam,
            modifyCommandParam );
    }

    public static Expression<Func<Expression, IExecuteResult>> CreateExecuteMethod( Expression oldProvider, Expression modifyCommand )
    {
        var assembly = typeof( SqlProvider ).Assembly;

        var providerType = assembly.GetType( "System.Data.Linq.Provider.IProvider" );
        var funcletizerType = assembly.GetType( "System.Data.Linq.SqlClient.Funcletizer" );
        var annotationsType = assembly.GetType( "System.Data.Linq.SqlClient.SqlNodeAnnotations" );
        var queryInfoType = assembly.GetType( "System.Data.Linq.SqlClient.SqlProvider" ).GetNestedType( "QueryInfo", BindingFlags.NonPublic );
        var readerFactoryType = assembly.GetType( "System.Data.Linq.SqlClient.IObjectReaderFactory" );
        var resultShapeType = assembly.GetType( "System.Data.Linq.SqlClient.SqlProvider" ).GetNestedType( "ResultShape", BindingFlags.NonPublic );
        var compiledSubQueryType = assembly.GetType( "System.Data.Linq.SqlClient.ICompiledSubQuery" );
        var typeSystemType = assembly.GetType( "System.Data.Linq.SqlClient.TypeSystem" );

        var queryParam = Expression.Variable( typeof( Expression ), "query" );

        var cachedResultVar = Expression.Variable( typeof( IExecuteResult ), "cachedResult" );
        var annotationsVar = Expression.Variable( annotationsType, "annotations" );
        var queriesVar = Expression.Variable( queryInfoType.MakeArrayType(), "queries" );
        var infoVar = Expression.Variable( queryInfoType, "info" );
        var lambdaVar = Expression.Variable( typeof( LambdaExpression ), "lambda" );
        var readerFactoryVar = Expression.Variable( readerFactoryType, "readerFactory" );
        var subQueriesVar = Expression.Variable( compiledSubQueryType.MakeArrayType(), "subQueries" );
        var resultShapeVar = Expression.Variable( resultShapeType, "resultShape" );
        var returnTarget = Expression.Label( typeof( IExecuteResult ) );

        var getQuery = MakeMethodCall( infoVar, "get_Query" );
        var getResultType = MakeMethodCall( infoVar, "get_ResultType" );
        var getResultElementType = MakeMethodCall( typeSystemType, "GetElementType", getResultType );
        var modifySubQueriesMethod = CreateModifySubQueriesMethod();

        return Expression.Lambda<Func<Expression, IExecuteResult>>(
            Expression.Block(
                new[] { annotationsVar, queriesVar, infoVar, lambdaVar, readerFactoryVar, subQueriesVar, resultShapeVar },

                // this.InitializeProviderMode();
                MakeMethodCall( oldProvider, "InitializeProviderMode" ),

                // query = Funcletizer.Funcletize(query);
                Expression.Assign( queryParam, MakeMethodCall( funcletizerType, "Funcletize", queryParam ) ),

                // if ( this.EnableCacheLookup )
                Expression.IfThen(
                    MakeMethodCall( oldProvider, "get_EnableCacheLookup" ),
                    Expression.Block(
                        new[] { cachedResultVar },
                        // IExecuteResult cachedResult = this.GetCachedResult(query);
                        Expression.Assign( cachedResultVar, MakeMethodCall( oldProvider, "GetCachedResult", queryParam ) ),
                        // if ( cachedResult != null )
                        Expression.IfThen(
                            Expression.NotEqual( cachedResultVar, Expression.Constant( null ) ),
                            // return cachedResult;
                            Expression.Return( returnTarget, cachedResultVar ) ) ) ),

                // SqlNodeAnnotations annotations = new SqlNodeAnnotations();
                Expression.Assign( annotationsVar, Expression.New( annotationsType ) ),

                // QueryInfo[] queries = this.BuildQuery(query, annotations);
                Expression.Assign( queriesVar, MakeMethodCall( oldProvider, "BuildQuery", queryParam, annotationsVar ) ),

                // var info = ModifyQueries( (IEnumerable)queries );
                Expression.Assign( infoVar, Expression.Invoke( CreateModifyQueriesMethod(), queriesVar, modifyCommand ) ),

                // this.CheckSqlCompatibility(queries, annotations);
                MakeMethodCall( oldProvider, "CheckSqlCompatibility", queriesVar, annotationsVar ),

                // var lambda = query as LambdaExpression;
                Expression.Assign( lambdaVar, Expression.TypeAs( queryParam, lambdaVar.Type ) ),

                // if ( lambda != null )
                Expression.IfThen(
                    Expression.NotEqual( lambdaVar, Expression.Constant( null ) ),
                    // query = lambda.Body;
                    Expression.Assign( queryParam, Expression.Property( lambdaVar, "Body" ) ) ),

                // IObjectReaderFactory readerFactory = null;
                // ICompiledSubQuery[] subQueries = null;

                // QueryInfo info = queries[queries.Length - 1];
                // info defined above
                // var resultShape = (int)Invoke( info, "get_ResultShape" );
                Expression.Assign( resultShapeVar, MakeMethodCall( infoVar, "get_ResultShape" ) ),

                Expression.Switch(
                    Expression.Convert( resultShapeVar, typeof( int ) ),
                    // if ( resultShape == 1 /* ResultShape.Singleton */ )
                    Expression.SwitchCase(
                        Expression.Block(
                            // subQueries = this.CompileSubQueries( info.Query );
                            Expression.Assign( subQueriesVar, MakeMethodCall( oldProvider, "CompileSubQueries", getQuery ) ),
                            // ModifySubQueries( (IEnumerable)subQueries );
                            Expression.Invoke( modifySubQueriesMethod, subQueriesVar, modifyCommand ),
                            // readerFactory = this.GetReaderFactory( info.Query, info.ResultType );
                            Expression.Assign( readerFactoryVar, MakeMethodCall( oldProvider, "GetReaderFactory", getQuery, getResultType ) ),
                            Expression.Empty() ),
                        Expression.Constant( 1 ) ),
                    // else if ( resultShape == 2 /* ResultShape.Sequence */ )
                    Expression.SwitchCase(
                        Expression.Block(
                            // subQueries = this.CompileSubQueries( info.Query );
                            Expression.Assign( subQueriesVar, MakeMethodCall( oldProvider, "CompileSubQueries", getQuery ) ),
                            // ModifySubQueries( (IEnumerable)subQueries );
                            Expression.Invoke( modifySubQueriesMethod, subQueriesVar, modifyCommand ),
                            // readerFactory = this.GetReaderFactory( info.Query, TypeSystem.GetElementType( info.ResultType ) );
                            Expression.Assign( readerFactoryVar, MakeMethodCall( oldProvider, "GetReaderFactory", getQuery, getResultElementType ) ),
                            Expression.Empty() ),
                        Expression.Constant( 2 ) ) ),

                // return this.ExecuteAll(query, queries, readerFactory, null, subQueries);
                Expression.Return( returnTarget, MakeMethodCall( oldProvider, "ExecuteAll", queryParam, queriesVar, readerFactoryVar, Expression.Constant( null, typeof( object[] ) ), subQueriesVar ) ),
                // --- Expression.Throw( Expression.New( typeof( Exception ).GetConstructor( new[] { typeof( string ) } ), Expression.Constant( "boo4" ) ) ),
                Expression.Label( returnTarget, Expression.Constant( null, typeof( IExecuteResult ) ) ) ),
            queryParam );
    }

    private static LambdaExpression CreateModifyQueriesMethod()
    {
        var assembly = typeof( SqlProvider ).Assembly;

        var queryInfoType = assembly.GetType( "System.Data.Linq.SqlClient.SqlProvider" ).GetNestedType( "QueryInfo", BindingFlags.NonPublic );
        var queriesParam = Expression.Parameter( queryInfoType.MakeArrayType(), "queries" );
        var modifyCommandParam = Expression.Parameter( typeof( ModifyCommandDelegate ), "modifyCommand" );

        var returnLabel = Expression.Label();
        var indexVar = Expression.Variable( typeof( int ), "index" );
        var queryVar = Expression.Variable( queryInfoType, "query" );

        return Expression.Lambda(
            Expression.Block(
                new[] { indexVar, queryVar },
                // var index = 0;
                Expression.Assign( indexVar, Expression.Constant( 0 ) ),
                Expression.Loop(
                    Expression.Block(
                        // if ( index >= queries.Length ) break;
                        Expression.IfThen(
                            Expression.GreaterThanOrEqual( indexVar, Expression.ArrayLength( queriesParam ) ),
                            Expression.Break( returnLabel ) ),
                        // query = queries[ index ];
                        Expression.Assign( queryVar, Expression.ArrayIndex( queriesParam, indexVar ) ),
                        // ModifyQuery( query );
                        Expression.Invoke( CreateModifyQueryMethod(), queryVar, modifyCommandParam ),
                        // ++ index;
                        Expression.PreIncrementAssign( indexVar ) ),
                    returnLabel ),
                queryVar ),
            queriesParam,
            modifyCommandParam );
    }

    private static LambdaExpression CreateModifySubQueriesMethod()
    {
        var assembly = typeof( SqlProvider ).Assembly;

        var iCompiledSubQueryType = assembly.GetType( "System.Data.Linq.SqlClient.ICompiledSubQuery" );
        var compiledSubQueryType = assembly.GetType( "System.Data.Linq.SqlClient.SqlProvider" ).GetNestedType( "CompiledSubQuery", BindingFlags.NonPublic );
        var queryInfoType = assembly.GetType( "System.Data.Linq.SqlClient.SqlProvider" ).GetNestedType( "QueryInfo", BindingFlags.NonPublic );
        var subQueriesType = iCompiledSubQueryType.MakeArrayType();

        var queryInfoField = compiledSubQueryType.GetField( "queryInfo", BindingFlags.Instance | BindingFlags.Public | BindingFlags.NonPublic );
        var subQueriesField = compiledSubQueryType.GetField( "subQueries", BindingFlags.Instance | BindingFlags.Public | BindingFlags.NonPublic );

        var subQueriesParam = Expression.Parameter( subQueriesType, "subQueries" );
        var modifyCommandParam = Expression.Parameter( typeof( ModifyCommandDelegate ), "modifyCommand" );

        var indexVar = Expression.Variable( typeof( int ), "index" );
        var subQueryVar = Expression.Variable( compiledSubQueryType, "subQuery" );
        var queryInfoVar = Expression.Variable( queryInfoType, "queryInfo" );
        var nestedSubQueriesVar = Expression.Variable( iCompiledSubQueryType.MakeArrayType(), "nestedSubQueries" );
        var modifySubQueriesDelegateType = typeof( Action<> ).MakeGenericType( iCompiledSubQueryType.MakeArrayType() );
        var modifySubQueriesDelegateVar = Expression.Variable( modifySubQueriesDelegateType, "modifySubQueries" );
        var loopTarget = Expression.Label();

        var innerLambda = Expression.Lambda(
            modifySubQueriesDelegateType,
            Expression.Block(
                new[] { indexVar },
            // var index = 0;
                Expression.Assign( indexVar, Expression.Constant( 0 ) ),
                Expression.Loop(
                    Expression.Block(
                        new[] { subQueryVar },
                        // if ( index >= subQueries.Length ) break;
                        Expression.IfThen(
                            Expression.GreaterThanOrEqual( indexVar, Expression.ArrayLength( subQueriesParam ) ),
                            Expression.Break( loopTarget ) ),
                        // var subQuery = subQueries[ index ];
                        Expression.Assign( subQueryVar, Expression.TypeAs( Expression.ArrayIndex( subQueriesParam, indexVar ), compiledSubQueryType ) ),
                        // if ( subQuery != null )
                        Expression.IfThen(
                            Expression.NotEqual( subQueryVar, Expression.Constant( null ) ),
                            Expression.Block(
                                new[] { queryInfoVar, nestedSubQueriesVar },
                                // var queryInfo = subQuery.queryInfo;
                                Expression.Assign( queryInfoVar, Expression.MakeMemberAccess( subQueryVar, queryInfoField ) ),
                                // ModifyQuery( queryInfo, modifyCommand );
                                Expression.Invoke( CreateModifyQueryMethod(), queryInfoVar, modifyCommandParam ),
                                // var nestedSubQueries = subQuery.subQueries;
                                Expression.Assign( nestedSubQueriesVar, Expression.MakeMemberAccess( subQueryVar, subQueriesField ) ),
                                // if ( nestedSubQueries != null )
                                Expression.IfThen(
                                    Expression.NotEqual( nestedSubQueriesVar, Expression.Constant( null ) ),
                                    // ModifySubQueries( nestedSubQueries, modifyCommand );
                                    Expression.Invoke( modifySubQueriesDelegateVar, nestedSubQueriesVar ) ) ) ),
                        // ++ index;
                        Expression.PreIncrementAssign( indexVar ) ),
                    loopTarget ) ),
                subQueriesParam );

        return Expression.Lambda(
            Expression.Block(
                new[] { modifySubQueriesDelegateVar },
                Expression.Assign( modifySubQueriesDelegateVar, innerLambda ),
                Expression.Invoke( modifySubQueriesDelegateVar, subQueriesParam ) ),
            subQueriesParam,
            modifyCommandParam );
    }

    private static LambdaExpression CreateModifyQueryMethod()
    {
        var assembly = typeof( SqlProvider ).Assembly;

        var queryInfoType = assembly.GetType( "System.Data.Linq.SqlClient.SqlProvider" ).GetNestedType( "QueryInfo", BindingFlags.NonPublic );
        var sqlParameterInfoType = assembly.GetType( "System.Data.Linq.SqlClient.SqlParameterInfo" );
        var sqlParameterType = assembly.GetType( "System.Data.Linq.SqlClient.SqlParameter" );

        var queryInfoParam = Expression.Parameter( queryInfoType, "query" );
        var modifyCommandParam = Expression.Parameter( typeof( ModifyCommandDelegate ), "modifyCommand" );

        var commandTextField = queryInfoType.GetField( "commandText", BindingFlags.Instance | BindingFlags.Public | BindingFlags.NonPublic );
        var parametersField = queryInfoType.GetField( "parameters", BindingFlags.Instance | BindingFlags.Public | BindingFlags.NonPublic );

        var commandTextVar = Expression.Variable( typeof( string ), "commandText" );
        var parameterInfosVar = Expression.Variable( typeof( ReadOnlyCollection<> ).MakeGenericType( sqlParameterInfoType ), "parameterInfos" );
        var parametersVar = Expression.Variable( typeof( Dictionary<string, object> ), "parameters" );
        var indexVar = Expression.Variable( typeof( int ), "index" );
        var paramInfoVar = Expression.Variable( sqlParameterInfoType, "paramInfo" );
        var paramVar = Expression.Variable( sqlParameterType, "param" );
        var loopTarget = Expression.Label();

        return Expression.Lambda(
            Expression.IfThen(
                Expression.NotEqual( modifyCommandParam, Expression.Constant( null ) ),
                Expression.Block(
                    new[] { commandTextVar, parameterInfosVar, parametersVar, indexVar },
                    // var commandText = queryInfo.CommandText;
                    Expression.Assign( commandTextVar, Expression.MakeMemberAccess( queryInfoParam, commandTextField ) ),
                    // var parameterInfos = queryInfo.Parameters;
                    Expression.Assign( parameterInfosVar, Expression.MakeMemberAccess( queryInfoParam, parametersField ) ),
                    // var parameters = new Dictionary<string, object>();
                    Expression.Assign( parametersVar, Expression.New( typeof( Dictionary<string, object> ) ) ),
                    // var index = 0;
                    Expression.Assign( indexVar, Expression.Constant( 0 ) ),
                    Expression.Loop(
                        Expression.Block(
                            new[] { paramInfoVar, paramVar },
                            // if ( index >= parameterInfos.Count ) break;
                            Expression.IfThen(
                                Expression.GreaterThanOrEqual( indexVar, Expression.Property( parameterInfosVar, "Count" ) ),
                                Expression.Break( loopTarget ) ),
                            // var paramInfo = parameterInfos[ index ];
                            Expression.Assign( paramInfoVar, MakeMethodCall( parameterInfosVar, "get_Item", indexVar ) ),
                            // var param = paramInfo.Parameter;
                            Expression.Assign( paramVar, MakeMethodCall( paramInfoVar, "get_Parameter" ) ),
                            // parameters[ param.Name ] = paramInfo.Value;
                            MakeMethodCall( parametersVar, "set_Item", MakeMethodCall( paramVar, "get_Name" ), MakeMethodCall( paramInfoVar, "get_Value" ) ),
                            // ++ index;
                            Expression.PreIncrementAssign( indexVar ) ),
                        loopTarget ),
                    // queryInfo.CommandText = modifyCommand( commandText, parameters );
                    Expression.Assign( Expression.MakeMemberAccess( queryInfoParam, commandTextField ), Expression.Invoke( modifyCommandParam, commandTextVar, parametersVar ) ) ) ),
            queryInfoParam,
            modifyCommandParam );
    }

    internal interface IProviderProxy
    {
        DataContextInterceptor Interceptor { get; }
        object OldProvider { get; }
    }

    public class ProviderProxy : RealProxy, IRemotingTypeInfo, IProviderProxy
    {
        public DataContextInterceptor Interceptor { get; private set; }
        public object OldProvider { get; private set; }

        internal ProviderProxy( DataContextInterceptor extender, object oldProvider )
            : base( typeof( ContextBoundObject ) )
        {
            this.Interceptor = extender;
            this.OldProvider = oldProvider;
        }

        public override IMessage Invoke( IMessage msg )
        {
            var call = msg as IMethodCallMessage;

            if ( call != null && OldProvider != null )
            {
                try
                {
                    if ( call.MethodBase.DeclaringType.Name == "IProvider" && call.MethodBase.DeclaringType.IsInterface )
                    {
                        Interceptor.providerType = call.MethodBase.DeclaringType;

                        switch ( call.MethodName )
                        {
                            case "Compile": return new ReturnMessage( Interceptor.Compile( call.Args.Cast<Expression>().First() ), null, 0, null, call );
                            case "Execute": return new ReturnMessage( Interceptor.Execute( call.Args.Cast<Expression>().First() ), null, 0, null, call );
                        }
                    }

                    return new ReturnMessage( call.MethodBase.Invoke( OldProvider, call.Args ), null, 0, null, call );
                }
                catch ( TargetInvocationException e )
                {
                    return new ReturnMessage( e.InnerException, call );
                }
            }

            throw new NotImplementedException();
        }

        public bool CanCastTo( Type fromType, object o )
        {
            return true;
        }

        public string TypeName
        {
            get { return this.GetType().Name; }
            set { }
        }
    }
}

To use it, simply call DataContextInterceptor.Intercept on your DataContext before use (passing in a delegate to modify command text as needed):

DataContextInterceptor.Intercept( this, ( c, p ) => c.Replace( "[t0].[cName]", "NULL" ).Dump( "test") );

 

As it’s still calling private methods, you’ll need to be running in a full-trust environment.

About these ads