Query: Allow using STET overload for Set in query filter/defining query

- Update with correct query root when replacing with runtime entity type in query filter

Resolves #24601

src/EFCore.Cosmos/Query/Internal/FromSqlQueryRootExpression.cs

@@ -3,6 +3,7 @@
 
 using System;
 using System.Linq.Expressions;
+using Microsoft.EntityFrameworkCore.Diagnostics;
 using Microsoft.EntityFrameworkCore.Metadata;
 using Microsoft.EntityFrameworkCore.Query;
 using Microsoft.EntityFrameworkCore.Utilities;
@@ -81,6 +82,18 @@ public FromSqlQueryRootExpression(
         public override Expression DetachQueryProvider()
             => new FromSqlQueryRootExpression(EntityType, Sql, Argument);
 
+        /// <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 override QueryRootExpression UpdateEntityType(IEntityType entityType)
+            => entityType.ClrType != EntityType.ClrType
+                || entityType.Name != EntityType.Name
+                ? throw new InvalidOperationException(CoreStrings.QueryRootDifferentEntityType(entityType.DisplayName()))
+                : new FromSqlQueryRootExpression(entityType, Sql, Argument);
+
         /// <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

src/EFCore.Relational/Query/Internal/FromSqlQueryRootExpression.cs

@@ -3,6 +3,7 @@
 
 using System;
 using System.Linq.Expressions;
+using Microsoft.EntityFrameworkCore.Diagnostics;
 using Microsoft.EntityFrameworkCore.Metadata;
 using Microsoft.EntityFrameworkCore.Utilities;
 
@@ -80,6 +81,18 @@ public FromSqlQueryRootExpression(
         public override Expression DetachQueryProvider()
             => new FromSqlQueryRootExpression(EntityType, Sql, Argument);
 
+        /// <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 override QueryRootExpression UpdateEntityType(IEntityType entityType)
+            => entityType.ClrType != EntityType.ClrType
+                || entityType.Name != EntityType.Name
+                ? throw new InvalidOperationException(CoreStrings.QueryRootDifferentEntityType(entityType.DisplayName()))
+                : new FromSqlQueryRootExpression(entityType, Sql, Argument);
+
         /// <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

src/EFCore.Relational/Query/Internal/TableValuedFunctionQueryRootExpression.cs

@@ -5,6 +5,7 @@
 using System.Collections.Generic;
 using System.Linq;
 using System.Linq.Expressions;
+using Microsoft.EntityFrameworkCore.Diagnostics;
 using Microsoft.EntityFrameworkCore.Metadata;
 using Microsoft.EntityFrameworkCore.Utilities;
 
@@ -76,6 +77,18 @@ protected override Expression VisitChildren(ExpressionVisitor visitor)
                 : this;
         }
 
+        /// <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 override QueryRootExpression UpdateEntityType(IEntityType entityType)
+            => entityType.ClrType != EntityType.ClrType
+                || entityType.Name != EntityType.Name
+                ? throw new InvalidOperationException(CoreStrings.QueryRootDifferentEntityType(entityType.DisplayName()))
+                : new TableValuedFunctionQueryRootExpression(entityType, Function, Arguments);
+
         /// <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

src/EFCore.SqlServer/Query/Internal/TemporalAllQueryRootExpression.cs

@@ -1,7 +1,9 @@
 // Licensed to the .NET Foundation under one or more agreements.
 // The .NET Foundation licenses this file to you under the MIT license.
 
+using System;
 using System.Linq.Expressions;
+using Microsoft.EntityFrameworkCore.Diagnostics;
 using Microsoft.EntityFrameworkCore.Metadata;
 using Microsoft.EntityFrameworkCore.Query;
 
@@ -46,6 +48,18 @@ public TemporalAllQueryRootExpression(IAsyncQueryProvider queryProvider, IEntity
         public override Expression DetachQueryProvider()
             => new TemporalAllQueryRootExpression(EntityType);
 
+        /// <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 override QueryRootExpression UpdateEntityType(IEntityType entityType)
+            => entityType.ClrType != EntityType.ClrType
+                || entityType.Name != EntityType.Name
+                ? throw new InvalidOperationException(CoreStrings.QueryRootDifferentEntityType(entityType.DisplayName()))
+                : new TemporalAllQueryRootExpression(entityType);
+
         /// <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

src/EFCore.SqlServer/Query/Internal/TemporalAsOfQueryRootExpression.cs

@@ -3,6 +3,7 @@
 
 using System;
 using System.Linq.Expressions;
+using Microsoft.EntityFrameworkCore.Diagnostics;
 using Microsoft.EntityFrameworkCore.Metadata;
 using Microsoft.EntityFrameworkCore.Query;
 
@@ -58,6 +59,18 @@ public TemporalAsOfQueryRootExpression(
         public override Expression DetachQueryProvider()
             => new TemporalAsOfQueryRootExpression(EntityType, PointInTime);
 
+        /// <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 override QueryRootExpression UpdateEntityType(IEntityType entityType)
+            => entityType.ClrType != EntityType.ClrType
+                || entityType.Name != EntityType.Name
+                ? throw new InvalidOperationException(CoreStrings.QueryRootDifferentEntityType(entityType.DisplayName()))
+                : new TemporalAsOfQueryRootExpression(entityType, PointInTime);
+
         /// <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

src/EFCore.SqlServer/Query/Internal/TemporalBetweenQueryRootExpression.cs

@@ -3,6 +3,7 @@
 
 using System;
 using System.Linq.Expressions;
+using Microsoft.EntityFrameworkCore.Diagnostics;
 using Microsoft.EntityFrameworkCore.Metadata;
 using Microsoft.EntityFrameworkCore.Query;
 
@@ -54,6 +55,18 @@ public TemporalBetweenQueryRootExpression(
         public override Expression DetachQueryProvider()
             => new TemporalBetweenQueryRootExpression(EntityType, From, To);
 
+        /// <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 override QueryRootExpression UpdateEntityType(IEntityType entityType)
+            => entityType.ClrType != EntityType.ClrType
+                || entityType.Name != EntityType.Name
+                ? throw new InvalidOperationException(CoreStrings.QueryRootDifferentEntityType(entityType.DisplayName()))
+                : new TemporalBetweenQueryRootExpression(entityType, From, To);
+
         /// <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

src/EFCore.SqlServer/Query/Internal/TemporalContainedInQueryRootExpression.cs

@@ -3,6 +3,7 @@
 
 using System;
 using System.Linq.Expressions;
+using Microsoft.EntityFrameworkCore.Diagnostics;
 using Microsoft.EntityFrameworkCore.Metadata;
 using Microsoft.EntityFrameworkCore.Query;
 
@@ -54,6 +55,18 @@ public TemporalContainedInQueryRootExpression(
         public override Expression DetachQueryProvider()
             => new TemporalContainedInQueryRootExpression(EntityType, From, To);
 
+        /// <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 override QueryRootExpression UpdateEntityType(IEntityType entityType)
+            => entityType.ClrType != EntityType.ClrType
+                || entityType.Name != EntityType.Name
+                ? throw new InvalidOperationException(CoreStrings.QueryRootDifferentEntityType(entityType.DisplayName()))
+                : new TemporalContainedInQueryRootExpression(entityType, From, To);
+
         /// <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

src/EFCore.SqlServer/Query/Internal/TemporalFromToQueryRootExpression.cs

@@ -3,6 +3,7 @@
 
 using System;
 using System.Linq.Expressions;
+using Microsoft.EntityFrameworkCore.Diagnostics;
 using Microsoft.EntityFrameworkCore.Metadata;
 using Microsoft.EntityFrameworkCore.Query;
 
@@ -54,6 +55,18 @@ public TemporalFromToQueryRootExpression(
         public override Expression DetachQueryProvider()
             => new TemporalFromToQueryRootExpression(EntityType, From, To);
 
+        /// <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 override QueryRootExpression UpdateEntityType(IEntityType entityType)
+            => entityType.ClrType != EntityType.ClrType
+                || entityType.Name != EntityType.Name
+                ? throw new InvalidOperationException(CoreStrings.QueryRootDifferentEntityType(entityType.DisplayName()))
+                : new TemporalFromToQueryRootExpression(entityType, From, To);
+
         /// <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

src/EFCore/Metadata/Conventions/QueryFilterRewritingConvention.cs

@@ -3,8 +3,11 @@
 
 using System;
 using System.Linq.Expressions;
+using Microsoft.EntityFrameworkCore.Diagnostics;
+using Microsoft.EntityFrameworkCore.Infrastructure;
 using Microsoft.EntityFrameworkCore.Metadata.Builders;
 using Microsoft.EntityFrameworkCore.Metadata.Conventions.Infrastructure;
+using Microsoft.EntityFrameworkCore.Metadata.Internal;
 using Microsoft.EntityFrameworkCore.Query;
 using Microsoft.EntityFrameworkCore.Utilities;
 
@@ -93,7 +96,8 @@ protected override Expression VisitMember(MemberExpression memberExpression)
                     && memberExpression.Type.GetGenericTypeDefinition() == typeof(DbSet<>)
                     && _model != null)
                 {
-                    return new QueryRootExpression(FindEntityType(memberExpression.Type)!);
+                    var entityClrType = memberExpression.Type.GetGenericArguments()[0];
+                    return new QueryRootExpression(FindEntityType(entityClrType)!);
                 }
 
                 return base.VisitMember(memberExpression);
@@ -111,14 +115,62 @@ protected override Expression VisitMethodCall(MethodCallExpression methodCallExp
                     && methodCallExpression.Type.GetGenericTypeDefinition() == typeof(DbSet<>)
                     && _model != null)
                 {
-                    return new QueryRootExpression(FindEntityType(methodCallExpression.Type)!);
+                    IEntityType? entityType;
+                    var entityClrType = methodCallExpression.Type.GetGenericArguments()[0];
+                    if (methodCallExpression.Arguments.Count == 1)
+                    {
+                        // STET Set method
+                        var entityTypeName = methodCallExpression.Arguments[0].GetConstantValue<string>();
+                        entityType = (IEntityType?)_model.FindEntityType(entityTypeName);
+                    }
+                    else
+                    {
+                        entityType = FindEntityType(entityClrType);
+                    }
+
+                    if (entityType == null)
+                    {
+                        if (_model.IsShared(entityClrType))
+                        {
+                            throw new InvalidOperationException(CoreStrings.InvalidSetSharedType(entityClrType.ShortDisplayName()));
+                        }
+
+                        var findSameTypeName = ((IModel)_model).FindSameTypeNameWithDifferentNamespace(entityClrType);
+                        //if the same name exists in your entity types we will show you the full namespace of the type
+                        if (!string.IsNullOrEmpty(findSameTypeName))
+                        {
+                            throw new InvalidOperationException(CoreStrings.InvalidSetSameTypeWithDifferentNamespace(entityClrType.DisplayName(), findSameTypeName));
+                        }
+                        else
+                        {
+                            throw new InvalidOperationException(CoreStrings.InvalidSetType(entityClrType.ShortDisplayName()));
+                        }
+                    }
+
+                    if (entityType.IsOwned())
+                    {
+                        var message = CoreStrings.InvalidSetTypeOwned(
+                            entityType.DisplayName(), entityType.FindOwnership()!.PrincipalEntityType.DisplayName());
+
+                        throw new InvalidOperationException(message);
+                    }
+
+                    if (entityType.ClrType != entityClrType)
+                    {
+                        var message = CoreStrings.DbSetIncorrectGenericType(
+                            entityType.ShortName(), entityType.ClrType.ShortDisplayName(), entityClrType.ShortDisplayName());
+
+                        throw new InvalidOperationException(message);
+                    }
+
+                    return new QueryRootExpression(entityType);
                 }
 
                 return base.VisitMethodCall(methodCallExpression);
             }
 
-            private IEntityType? FindEntityType(Type dbSetType)
-                => ((IModel)_model!).FindRuntimeEntityType(dbSetType.GetGenericArguments()[0]);
+            private IEntityType? FindEntityType(Type entityClrType)
+                => ((IModel)_model!).FindRuntimeEntityType(entityClrType);
         }
     }
 }

src/EFCore/Metadata/Conventions/RuntimeModelConvention.cs

@@ -613,7 +613,7 @@ public Expression Rewrite(Expression expression)
             /// <inheritdoc />
             protected override Expression VisitExtension(Expression extensionExpression)
                 => extensionExpression is QueryRootExpression queryRootExpression
-                    ? new QueryRootExpression(_model.FindEntityType(queryRootExpression.EntityType.Name)!)
+                    ? queryRootExpression.UpdateEntityType(_model.FindEntityType(queryRootExpression.EntityType.Name)!)
                     : base.VisitExtension(extensionExpression);
         }
     }