Customizing SQL Generation in Entity Framework Core

CustomMethodInfoBasedNodeTypeRegistryFactory.cs

using Microsoft.EntityFrameworkCore.Query.Internal;
using Remotion.Linq.Parsing.Structure;

namespace MyApp.EFCoreExtensions
{
  internal class CustomMethodInfoBasedNodeTypeRegistryFactory : DefaultMethodInfoBasedNodeTypeRegistryFactory
  {
    public override INodeTypeProvider Create()
    {
      RegisterMethods(WithSqlTweaksExpressionNode.SupportedMethods, typeof(WithSqlTweaksExpressionNode));
      base.Create();
    }
  }
}

CustomSelectExpression.cs

using Microsoft.EntityFrameworkCore.Query;
using Microsoft.EntityFrameworkCore.Query.Expressions;
using System.Linq;

namespace MyApp.EFCoreExtensions
{
  internal class CustomSelectExpression : SelectExpression
  {
    public bool UseSqlTweaks { get; set; }
    
    public CustomSelectExpression(
      SelectExpressionDependencies dependencies,
      RelationalQueryCompilationContext queryCompilationContext): base(dependencies, queryCompilationContext)
    {
      SetCustomSelectExpressionProperties(queryCompilationContext);
    }
    
    public CustomSelectExpression(
      SelectExpressionDependencies dependencies,
      RelationalQueryCompilationContext queryCompilationContext,
      string alias): base(dependencies, queryCompilationContext, alias)
    {
      SetCustomSelectExpressionProperties(queryCompilationContext);
    }
    
    private void SetCustomSelectExpressionProperties(RelationalQueryCompilationContext queryCompilationContext)
    {
      // If the WithSqlTweaksResultOperator query annotation exists then set the property to true.
      if(queryCompilationContext.QueryAnnotations.Any(a => a.GetType() == typeof(WithSqlTweaksResultOperator)))
      {
        UseSqlTweaks = true;
      }
    }
  }
  
  internal class CustomSelectExpressionFactory : SelectExpressionFactory
  {
    public CustomSelectExpressionFactory(SelectExpressionDependencies dependencies)
      : base(dependencies)
    {
    }
    
    public override SelectExpression Create(RelationalQueryCompilationContext queryCompilationContext)
      => new CustomSelectExpression(Dependencies, queryCompilationContext);
      
    public override SelectExpression Create(RelationalQueryCompilationContext queryCompilationContext, string alias)
      => new CustomSelectExpression(Dependencies, queryCompilationContext, alias);
  }
}

CustomSqlServerQuerySqlGenerator.cs

using Microsoft.EntityFrameworkCore.Infrastructure.Internal;
using Microsoft.EntityFrameworkCore.Query.Expressions;
using Microsoft.EntityFrameworkCore.Query.Sql;
using Microsoft.EntityFrameworkCore.Query.Sql.Internal;
using Microsoft.EntityFrameworkCore.Storage;
using System;
using System.Collections.Generic;
using System.Linq;
using System.Linq.Expressions;

namespace MyApp.EfCoreExtensions
{
  internal class CustomSqlServerQuerySqlGenerator : SqlServerQuerySqlGenerator
  {
    public CustomSqlServerQuerySqlGenerator(
      QuerySqlGeneratorDependencies dependencies,
      SelectExpression selectExpression,
      bool rowNumberPagingEnabled)
      : base(dependencies, selectExpression, rowNumberPagingEnabled)
    {
    }
    
    public override Expresssion VisitSelect(SelectExpression selectExpression)
    {
      // other code left out for simplicity
      if(selectExpression is CustomSelectExpression)
      {
        if(((CustomSelectExpression)selectExpression).UseSqlTweaks)
        {
          // Do SQL tweaks here!
        }
      }
      // other code left out for simplicity
    }
  }
  
  internal class CustomSqlServerQuerySqlGeneratorFactory : QuerySqlGeneratorFactoryBase
  {
    private readonly ISqlServerOptions _sqlServerOptions;
    
    public CustomSqlServerQuerySqlGeneratorFactory(
      QuerySqlGeneratorDependencies dependencies, 
      ISqlServerOptions sqlServerOptions) : base(dependencies)
    {
      _sqlServerOptions = sqlServerOptions;
    }
    
    public override IQuerySqlGenerator CreateDefault(SelectExpression selectExpression)
      => new CustomSqlServerQuerySqlGenerator(
        Dependencies,
        selectExpression,
        _sqlServerOptions.RowNumberPagingEnabled);
  }
}

IQueryableExtensions.cs

using Microsoft.EntityFrameworkCore.Query.Internal;
using System.Linq;
using System.Linq.Expressions;
using System.Reflection;

namespace MyApp.EFCoreExtensions
{
  public static class IQueryableExtensions
  {
    internal static readonly MethodInfo WithSqlTweaksMethodInfo 
      = typeof(IQueryableExtensions).GetTypeInfo().GetDeclaredMethod(nameof(WithSqlTeaks));
      
    public static IQueryable<TEntity> WithSqlTweaks<TEntity>(this IQueryable<TEntity> source) where TEntity : class
    {
      return
        source.Provider is EntityQueryProvider
          ? source.Provider.CreateQuery<TEntity>(
            Expression.Call(
              instance: null,
              method: WithSqlTweaksMethodInfo.MakeGenericMethod(typeof(TEntity)),
              arguments: source.Expression))
          : source;
    }
  }
}

MyDbContext-2.cs

using Microsoft.EntityFrameworkCore;
using Microsoft.EntityFrameworkCore.Query.Expressions;
using Microsoft.EntityFrameworkCore.Query.Internal;
using Microsoft.EntityFrameworkCore.Query.Sql;
using MyApp.EfCoreExtensions;
using System.Collections.Generic;

namespace MyApp
{
  public class MyDbContext : DbContext
  {
    public DbSet<Person> People { get; set; }
    
    protected override void OnConfiguring(DbContextOptionsBuilder optionsBuilder)
    {
      optionsBuilder
        .UseSqlServer("connectionString")
        .ReplaceService<INodeTypeProviderFactory, CustomMethodInfoBasedNodeTypeRegistryFactory>()
        .ReplaceService<ISelectExpressionFactory, CustomSelectExpressionFactory>();
      base.OnConfiguring(optionsBuilder);
    }
  }
}

MyDbContext-3.cs

using Microsoft.EntityFrameworkCore;
using Microsoft.EntityFrameworkCore.Query.Expressions;
using Microsoft.EntityFrameworkCore.Query.Internal;
using Microsoft.EntityFrameworkCore.Query.Sql;
using MyApp.EfCoreExtensions;
using System.Collections.Generic;

namespace MyApp
{
  public class MyDbContext : DbContext
  {
    public DbSet<Person> People { get; set; }
    
    protected override void OnConfiguring(DbContextOptionsBuilder optionsBuilder)
    {
      optionsBuilder
        .UseSqlServer("connectionString")
        .ReplaceService<INodeTypeProviderFactory, CustomMethodInfoBasedNodeTypeRegistryFactory>()
        .ReplaceService<ISelectExpressionFactory, CustomSelectExpressionFactory>();
        .ReplaceService<IQuerySqlGeneratorFactory, CustomSqlServerQuerySqlGeneratorFactory>();
      base.OnConfiguring(optionsBuilder);
    }
  }
}

MyDbContext.cs

using Microsoft.EntityFrameworkCore;
using Microsoft.EntityFrameworkCore.Query.Expressions;
using Microsoft.EntityFrameworkCore.Query.Internal;
using Microsoft.EntityFrameworkCore.Query.Sql;
using MyApp.EfCoreExtensions;
using System.Collections.Generic;

namespace MyApp
{
  public class MyDbContext : DbContext
  {
    public DbSet<Person> People { get; set; }
    
    protected override void OnConfiguring(DbContextOptionsBuilder optionsBuilder)
    {
      optionsBuilder
        .UseSqlServer("connectionString")
        .ReplaceService<INodeTypeProviderFactory, CustomMethodInfoBasedNodeTypeRegistryFactory>();
      base.OnConfiguring(optionsBuilder);
    }
  }
}

WithSqlTweaksExpressionNode.cs

using Remotion.Linq.Clauses;
using Remotion.Linq.Parsing.Structure.IntermediateModel;
using System.Collections.Generic;
using System.Linq.Expressions;
using System.Reflection;

namespace MyApp.EFCoreExtensions
{
  internal class WithSqlTweaksExpressionNode : ResultOperatorExpressionNodeBase
  {
    public static readonly IReadOnlyCollection<MethodInfo> SupportedMethods = new[]
    {
      IQueryableExtensions.WithSqlTweaksMethodInfo
    };
    
    public WithSqlTweaksExpressionNode(MethodCallExpressionParseInfo parseInfo)
      : base(parseInfo, null, null)
    {
    }
    
    protected override ResultOperatorBase CreateResultOperator(ClauseGenerationContext clauseGenerationContext)
      => new WithSqlTweaksResultOperator();
      
    public override Expression Resolve(
      ParameterExpression inputParameter,
      Expression expressionToBeResolved,
      ClauseGenerationContext clauseGenerationContext)
      => Source.Resolve(inputParameter, expressionToBeResolved, clauseGenerationContext);
  }
}

WithSqlTweaksResultOperator.cs

using Microsoft.EntityFrameworkCore.Query.ResultOperators;
using Remotion.Linq;
using Remotion.Linq.Clauses;
using Remotion.Linq.Clauses.ResultOperators;
using Remotion.Linq.Clauses.StreamedData;
using System;
using System.Linq.Expressions;

namespace MyApp.EFCoreExtensions
{
  internal class WithSqlTweaksResultOperator : SequenceTypePreservingResultOperatorBase, IQueryAnnotation
  {
    public IQuerySource QuerySource { get; set; }
    public QueryModel QueryModel { get; set; }
    
    public override ResultOperatorBase Clone(CloneContext cloneContext)
      => new WithSqlTweaksResultOperator();
      
    public override StreamedSequence ExecuteInMemory<T>(StreamedSequence input) => input;
    
    public override void TransformExpressions(Func<Expression, Expression> transformation)
    {
    }
  }
}