Query expression interceptor

Part of #626
Fixes #19748

src/EFCore/Diagnostics/CoreLoggerExtensions.cs

@@ -448,10 +448,13 @@ private static string QueryCanceled(EventDefinitionBase definition, EventData pa
     ///     Logs for the <see cref="CoreEventId.QueryCompilationStarting" /> event.
     /// </summary>
     /// <param name="diagnostics">The diagnostics logger to use.</param>
+    /// <param name="context">The current <see cref="DbContext" />, or null if not known.</param>
     /// <param name="expressionPrinter">Used to create a human-readable representation of the expression tree.</param>
     /// <param name="queryExpression">The query expression tree.</param>
-    public static void QueryCompilationStarting(
+    /// <returns>The query expression and event data.</returns>
+    public static (Expression Query, QueryExpressionEventData? EventData) QueryCompilationStarting(
         this IDiagnosticsLogger<DbLoggerCategory.Query> diagnostics,
+        DbContext? context,
         ExpressionPrinter expressionPrinter,
         Expression queryExpression)
     {
@@ -462,16 +465,25 @@ public static void QueryCompilationStarting(
             definition.Log(diagnostics, Environment.NewLine, expressionPrinter.Print(queryExpression));
         }
 
-        if (diagnostics.NeedsEventData(definition, out var diagnosticSourceEnabled, out var simpleLogEnabled))
+        if (diagnostics.NeedsEventData<IQueryExpressionInterceptor>(
+                definition, out var interceptor, out var diagnosticSourceEnabled, out var simpleLogEnabled))
         {
             var eventData = new QueryExpressionEventData(
                 definition,
                 QueryCompilationStarting,
+                context,
                 queryExpression,
                 expressionPrinter);
 
             diagnostics.DispatchEventData(definition, eventData, diagnosticSourceEnabled, simpleLogEnabled);
+
+            if (interceptor != null)
+            {
+                return (interceptor.ProcessingQuery(queryExpression, eventData), eventData);
+            }
         }
+
+        return (queryExpression, null);
     }
 
     private static string QueryCompilationStarting(EventDefinitionBase definition, EventData payload)
@@ -643,10 +655,12 @@ private static string NavigationBaseIncludeIgnored(EventDefinitionBase definitio
     ///     Logs for the <see cref="CoreEventId.QueryExecutionPlanned" /> event.
     /// </summary>
     /// <param name="diagnostics">The diagnostics logger to use.</param>
+    /// <param name="context">The current <see cref="DbContext" />, or null if not known.</param>
     /// <param name="expressionPrinter">Used to create a human-readable representation of the expression tree.</param>
     /// <param name="queryExecutorExpression">The query expression tree.</param>
     public static void QueryExecutionPlanned(
         this IDiagnosticsLogger<DbLoggerCategory.Query> diagnostics,
+        DbContext? context,
         ExpressionPrinter expressionPrinter,
         Expression queryExecutorExpression)
     {
@@ -662,6 +676,7 @@ public static void QueryExecutionPlanned(
             var eventData = new QueryExpressionEventData(
                 definition,
                 QueryExecutionPlanned,
+                context,
                 queryExecutorExpression,
                 expressionPrinter);
 

src/EFCore/Diagnostics/IQueryExpressionInterceptor.cs

@@ -0,0 +1,63 @@
+// Licensed to the .NET Foundation under one or more agreements.
+// The .NET Foundation licenses this file to you under the MIT license.
+
+namespace Microsoft.EntityFrameworkCore.Diagnostics;
+
+/// <summary>
+///     Allows interception of query expression trees and resulting compiled delegates.
+/// </summary>
+/// <remarks>
+///     <para>
+///         Use <see cref="DbContextOptionsBuilder.AddInterceptors(Microsoft.EntityFrameworkCore.Diagnostics.IInterceptor[])" />
+///         to register application interceptors.
+///     </para>
+///     <para>
+///         Extensions can also register interceptors in the internal service provider.
+///         If both injected and application interceptors are found, then the injected interceptors are run in the
+///         order that they are resolved from the service provider, and then the application interceptors are run last.
+///     </para>
+///     <para>
+///         See <see href="https://aka.ms/efcore-docs-interceptors">EF Core interceptors</see> for more information and examples.
+///     </para>
+/// </remarks>
+public interface IQueryExpressionInterceptor : IInterceptor
+{
+    /// <summary>
+    ///     Called with the LINQ expression tree for a query before it is compiled.
+    /// </summary>
+    /// <param name="queryExpression">The query expression.</param>
+    /// <param name="eventData">Contextual information about the query environment.</param>
+    /// <returns>The query expression tree to continue with, which may have been changed by the interceptor.</returns>
+    Expression ProcessingQuery(
+        Expression queryExpression,
+        QueryExpressionEventData eventData)
+        => queryExpression;
+
+    /// <summary>
+    ///     Called when EF is about to compile the query delegate that will be used to execute the query.
+    /// </summary>
+    /// <param name="queryExpression">The query expression.</param>
+    /// <param name="queryExecutorExpression">The expression that will be compiled into the execution delegate.</param>
+    /// <param name="eventData">Contextual information about the query environment.</param>
+    /// <typeparam name="TResult">The return type of the execution delegate.</typeparam>
+    /// <returns>The expression that will be compiled into the execution delegate, which may have been changed by the interceptor.</returns>
+    Expression<Func<QueryContext, TResult>> CompilingQuery<TResult>(
+        Expression queryExpression,
+        Expression<Func<QueryContext, TResult>> queryExecutorExpression,
+        QueryExpressionEventData eventData)
+        => queryExecutorExpression;
+
+    /// <summary>
+    ///     Called when EF is about to compile the query delegate that will be used to execute the query.
+    /// </summary>
+    /// <param name="queryExpression">The query expression.</param>
+    /// <param name="eventData">Contextual information about the query environment.</param>
+    /// <param name="queryExecutor">The delegate that will be used to execute the query.</param>
+    /// <typeparam name="TResult">The return type of the execution delegate.</typeparam>
+    /// <returns>The delegate that will be used to execute the query, which may have been changed by the interceptor.</returns>
+    Func<QueryContext, TResult> CompiledQuery<TResult>(
+        Expression queryExpression,
+        QueryExpressionEventData eventData,
+        Func<QueryContext, TResult> queryExecutor)
+        => queryExecutor;
+}

src/EFCore/Diagnostics/Internal/QueryExpressionInterceptorAggregator.cs

@@ -0,0 +1,65 @@
+// Licensed to the .NET Foundation under one or more agreements.
+// The .NET Foundation licenses this file to you under the MIT license.
+
+namespace Microsoft.EntityFrameworkCore.Diagnostics.Internal;
+
+/// <summary>
+///     This is an internal API that supports the Entity Framework Core infrastructure and not subject to
+///     the same compatibility standards as public APIs. It may be changed or removed without notice in
+///     any release. You should only use it directly in your code with extreme caution and knowing that
+///     doing so can result in application failures when updating to a new Entity Framework Core release.
+/// </summary>
+public class QueryExpressionInterceptorAggregator : InterceptorAggregator<IQueryExpressionInterceptor>
+{
+    /// <inheritdoc />
+    protected override IQueryExpressionInterceptor CreateChain(IEnumerable<IQueryExpressionInterceptor> interceptors)
+        => new CompositeQueryExpressionInterceptor(interceptors);
+
+    private sealed class CompositeQueryExpressionInterceptor : IQueryExpressionInterceptor
+    {
+        private readonly IQueryExpressionInterceptor[] _interceptors;
+
+        public CompositeQueryExpressionInterceptor(IEnumerable<IQueryExpressionInterceptor> interceptors)
+        {
+            _interceptors = interceptors.ToArray();
+        }
+
+        public Expression ProcessingQuery(
+            Expression queryExpression,
+            QueryExpressionEventData eventData)
+        {
+            for (var i = 0; i < _interceptors.Length; i++)
+            {
+                queryExpression = _interceptors[i].ProcessingQuery(queryExpression, eventData);
+            }
+
+            return queryExpression;
+        }
+
+        public Expression<Func<QueryContext, TResult>> CompilingQuery<TResult>(
+            Expression queryExpression,
+            Expression<Func<QueryContext, TResult>> queryExecutorExpression,
+            QueryExpressionEventData eventData)
+        {
+            for (var i = 0; i < _interceptors.Length; i++)
+            {
+                queryExecutorExpression = _interceptors[i].CompilingQuery(queryExpression, queryExecutorExpression, eventData);
+            }
+
+            return queryExecutorExpression;
+        }
+
+        public Func<QueryContext, TResult> CompiledQuery<TResult>(
+            Expression queryExpression,
+            QueryExpressionEventData eventData,
+            Func<QueryContext, TResult> queryExecutor)
+        {
+            for (var i = 0; i < _interceptors.Length; i++)
+            {
+                queryExecutor = _interceptors[i].CompiledQuery(queryExpression, eventData, queryExecutor);
+            }
+
+            return queryExecutor;
+        }
+    }
+}

src/EFCore/Diagnostics/QueryExpressionEventData.cs

@@ -10,21 +10,23 @@ namespace Microsoft.EntityFrameworkCore.Diagnostics;
 /// <remarks>
 ///     See <see href="https://aka.ms/efcore-docs-diagnostics">Logging, events, and diagnostics</see> for more information and examples.
 /// </remarks>
-public class QueryExpressionEventData : EventData
+public class QueryExpressionEventData : DbContextEventData
 {
     /// <summary>
     ///     Constructs the event payload.
     /// </summary>
     /// <param name="eventDefinition">The event definition.</param>
     /// <param name="messageGenerator">A delegate that generates a log message for this event.</param>
+    /// <param name="context">The current <see cref="DbContext" />, or null if not known.</param>
     /// <param name="queryExpression">The <see cref="Expression" />.</param>
     /// <param name="expressionPrinter">An <see cref="ExpressionPrinter" /> that can be used to render the <see cref="Expression" />.</param>
     public QueryExpressionEventData(
         EventDefinitionBase eventDefinition,
         Func<EventDefinitionBase, EventData, string> messageGenerator,
+        DbContext? context,
         Expression queryExpression,
         ExpressionPrinter expressionPrinter)
-        : base(eventDefinition, messageGenerator)
+        : base(eventDefinition, messageGenerator, context)
     {
         Expression = queryExpression;
         ExpressionPrinter = expressionPrinter;

src/EFCore/Infrastructure/EntityFrameworkServicesBuilder.cs

@@ -273,6 +273,7 @@ public virtual EntityFrameworkServicesBuilder TryAddCoreServices()
         TryAdd<IInterceptors, Interceptors>();
         TryAdd<IInterceptorAggregator, SaveChangesInterceptorAggregator>();
         TryAdd<IInterceptorAggregator, IdentityResolutionInterceptorAggregator>();
+        TryAdd<IInterceptorAggregator, QueryExpressionInterceptorAggregator>();
         TryAdd<ILoggingOptions, LoggingOptions>();
         TryAdd<ICoreSingletonOptions, CoreSingletonOptions>();
         TryAdd<ISingletonOptions, ILoggingOptions>(p => p.GetRequiredService<ILoggingOptions>());

src/EFCore/Query/QueryCompilationContext.cs

@@ -55,6 +55,7 @@ public class QueryCompilationContext
     private readonly IQueryableMethodTranslatingExpressionVisitorFactory _queryableMethodTranslatingExpressionVisitorFactory;
     private readonly IQueryTranslationPostprocessorFactory _queryTranslationPostprocessorFactory;
     private readonly IShapedQueryCompilingExpressionVisitorFactory _shapedQueryCompilingExpressionVisitorFactory;
+    private readonly IQueryExpressionInterceptor? _queryExpressionInterceptor;
 
     private readonly ExpressionPrinter _expressionPrinter;
 
@@ -84,6 +85,7 @@ public QueryCompilationContext(
         _shapedQueryCompilingExpressionVisitorFactory = dependencies.ShapedQueryCompilingExpressionVisitorFactory;
 
         _expressionPrinter = new ExpressionPrinter();
+        _queryExpressionInterceptor = dependencies.Interceptors.Aggregate<IQueryExpressionInterceptor>();
     }
 
     /// <summary>
@@ -156,7 +158,8 @@ public virtual void AddTag(string tag)
     /// <returns>Returns <see cref="Func{QueryContext, TResult}" /> which can be invoked to get results of this query.</returns>
     public virtual Func<QueryContext, TResult> CreateQueryExecutor<TResult>(Expression query)
     {
-        Logger.QueryCompilationStarting(_expressionPrinter, query);
+        var queryAndEventData = Logger.QueryCompilationStarting(Dependencies.Context, _expressionPrinter, query);
+        query = queryAndEventData.Query;
 
         query = _queryTranslationPreprocessorFactory.Create(this).Process(query);
         // Convert EntityQueryable to ShapedQueryExpression
@@ -175,13 +178,24 @@ public virtual Func<QueryContext, TResult> CreateQueryExecutor<TResult>(Expressi
             query,
             QueryContextParameter);
 
+        if (_queryExpressionInterceptor != null)
+        {
+            queryExecutorExpression = _queryExpressionInterceptor.CompilingQuery(
+                queryAndEventData.Query, queryExecutorExpression, queryAndEventData.EventData!);
+        }
+
         try
         {
-            return queryExecutorExpression.Compile();
+            var queryExecutor = queryExecutorExpression.Compile();
+
+            return _queryExpressionInterceptor != null
+                ? _queryExpressionInterceptor.CompiledQuery(
+                    queryAndEventData.Query, queryAndEventData.EventData!, queryExecutor)
+                : queryExecutor;
         }
         finally
         {
-            Logger.QueryExecutionPlanned(_expressionPrinter, queryExecutorExpression);
+            Logger.QueryExecutionPlanned(Dependencies.Context, _expressionPrinter, queryExecutorExpression);
         }
     }
 

src/EFCore/Query/QueryCompilationContextDependencies.cs

@@ -56,7 +56,8 @@ public QueryCompilationContextDependencies(
         IExecutionStrategy executionStrategy,
         ICurrentDbContext currentContext,
         IDbContextOptions contextOptions,
-        IDiagnosticsLogger<DbLoggerCategory.Query> logger)
+        IDiagnosticsLogger<DbLoggerCategory.Query> logger,
+        IInterceptors interceptors)
     {
         _currentContext = currentContext;
         Model = model;
@@ -67,8 +68,15 @@ public QueryCompilationContextDependencies(
         IsRetryingExecutionStrategy = executionStrategy.RetriesOnFailure;
         ContextOptions = contextOptions;
         Logger = logger;
+        Interceptors = interceptors;
     }
 
+    /// <summary>
+    ///     The current context.
+    /// </summary>
+    public DbContext Context
+        => _currentContext.Context;
+
     /// <summary>
     ///     The CLR type of DbContext.
     /// </summary>
@@ -120,4 +128,9 @@ public QueryTrackingBehavior QueryTrackingBehavior
     ///     The logger.
     /// </summary>
     public IDiagnosticsLogger<DbLoggerCategory.Query> Logger { get; init; }
+
+    /// <summary>
+    ///     Registered interceptors.
+    /// </summary>
+    public IInterceptors Interceptors { get; }
 }

test/EFCore.Cosmos.FunctionalTests/QueryExpressionInterceptionWithDiagnosticsCosmosTest.cs

@@ -0,0 +1,31 @@
+// Licensed to the .NET Foundation under one or more agreements.
+// The .NET Foundation licenses this file to you under the MIT license.
+
+namespace Microsoft.EntityFrameworkCore.Cosmos;
+
+public class QueryExpressionInterceptionWithDiagnosticsCosmosTest
+    : QueryExpressionInterceptionTestBase,
+        IClassFixture<QueryExpressionInterceptionWithDiagnosticsCosmosTest.InterceptionCosmosFixture>
+{
+    public QueryExpressionInterceptionWithDiagnosticsCosmosTest(InterceptionCosmosFixture fixture)
+        : base(fixture)
+    {
+    }
+
+    public class InterceptionCosmosFixture : InterceptionFixtureBase
+    {
+        protected override ITestStoreFactory TestStoreFactory
+            => CosmosTestStoreFactory.Instance;
+
+        protected override IServiceCollection InjectInterceptors(
+            IServiceCollection serviceCollection,
+            IEnumerable<IInterceptor> injectedInterceptors)
+            => base.InjectInterceptors(serviceCollection.AddEntityFrameworkCosmos(), injectedInterceptors);
+
+        protected override string StoreName
+            => "QueryExpressionInterceptionWithDiagnostics";
+
+        protected override bool ShouldSubscribeToDiagnosticListener
+            => true;
+    }
+}