Support conflicting partition key predicates in Cosmos queries

src/EFCore.Cosmos/Properties/CosmosStrings.resx

@@ -1,4 +1,4 @@
-<?xml version="1.0" encoding="utf-8"?>
+<?xml version="1.0" encoding="utf-8"?>
 <root>
   <!-- 
     Microsoft ResX Schema 
@@ -424,4 +424,7 @@
   <data name="WithPartitionKeyNotConstantOrParameter" xml:space="preserve">
     <value>'WithPartitionKey' only accepts simple constant or parameter arguments. See https://aka.ms/efdocs-cosmos-partition-keys for more information.</value>
   </data>
+  <data name="WithPartitionKeyConflictingPartitionKeyPredicate" xml:space="preserve">
+    <value>The partition key value specified via 'WithPartitionKey' conflicts with a partition key predicate in the query. Remove the predicate or ensure both specify the same value.</value>
+  </data>
 </root>

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

@@ -1,6 +1,7 @@
 // Licensed to the .NET Foundation under one or more agreements.
 // The .NET Foundation licenses this file to you under the MIT license.
 
+using Microsoft.EntityFrameworkCore.Cosmos.Internal;
 using Microsoft.EntityFrameworkCore.Cosmos.Metadata.Internal;
 
 namespace Microsoft.EntityFrameworkCore.Cosmos.Query.Internal;
@@ -55,6 +56,22 @@ public virtual Expression ExtractPartitionKeysAndId(
 
         _rootAlias = rootSource.Alias;
 
+        Expression UnwrapShaperForReadItem(Expression shaper)
+        {
+            if (shaper is UnaryExpression { NodeType: ExpressionType.Convert } convert
+                && convert.Type == typeof(object))
+            {
+                shaper = convert.Operand;
+            }
+
+            while (shaper is IncludeExpression { EntityExpression: var nested })
+            {
+                shaper = nested;
+            }
+
+            return shaper;
+        }
+
         // We're going to be looking for equality comparisons on the JSON id definition properties and the partition key properties of the
         // entity type; build a dictionary where the properties are the keys, and where the values are expressions that will get populated
         // from the tree (either constants or parameters).
@@ -88,6 +105,15 @@ public virtual Expression ExtractPartitionKeysAndId(
         var allIdPropertiesSpecified =
             _jsonIdPropertyValues.Values.All(p => p is not null) && _jsonIdPropertyValues.Count > 0;
 
+        // WithPartitionKey will clear _partitionKeyPropertyValues during the lift pass below; snapshot predicate partition key
+        // comparisons first so we can validate conflicts against WithPartitionKey when using ReadItem (see #38238).
+        var hadWithPartitionKey = queryCompilationContext.PartitionKeyPropertyValues.Count > 0;
+        Dictionary<IProperty, (Expression? ValueExpression, Expression? OriginalExpression)>? predicatePartitionKeySnapshot = null;
+        if (hadWithPartitionKey)
+        {
+            predicatePartitionKeySnapshot = new Dictionary<IProperty, (Expression?, Expression?)>(_partitionKeyPropertyValues);
+        }
+
         // First, go over the partition key properties and lift them from the predicate to the query compilation context, as possible.
         // We do this only as long as all partition key values are provided; the moment there's a gap we stop (so if PK1 and PK3 are
         // provided but not PK2, only PK1 will be lifted out).
@@ -111,9 +137,7 @@ public virtual Expression ExtractPartitionKeysAndId(
             }
         }
 
-        // Now, attempt to also transform the query to ReadItem form; this is only possible if all JSON ID properties were compared in the
-        // predicate, and *all* partition key values are specified(in the predicate or via WithPartitionKey)
-        if (_isPredicateCompatibleWithReadItem
+        var willUseReadItemOptimization = _isPredicateCompatibleWithReadItem
             && allIdPropertiesSpecified
             && queryCompilationContext.PartitionKeyPropertyValues.Count == partitionKeyProperties.Count
             && select is
@@ -123,8 +147,20 @@ public virtual Expression ExtractPartitionKeysAndId(
             }
             // We only transform to ReadItem if the entire document (i.e. root entity type) is being projected out.
             // Using ReadItem even when a projection is present is tracked by #34163.
-            && Unwrap(shapedQuery.ShaperExpression) is StructuralTypeShaperExpression { StructuralType: var projectedStructuralType }
-            && projectedStructuralType == _entityType)
+            && UnwrapShaperForReadItem(shapedQuery.ShaperExpression) is StructuralTypeShaperExpression { StructuralType: var projectedStructuralType }
+            && projectedStructuralType == _entityType;
+
+        if (hadWithPartitionKey && predicatePartitionKeySnapshot is not null && willUseReadItemOptimization)
+        {
+            ValidateNoWithPartitionKeyConflict(
+                queryCompilationContext,
+                partitionKeyProperties,
+                predicatePartitionKeySnapshot);
+        }
+
+        // Now, attempt to also transform the query to ReadItem form; this is only possible if all JSON ID properties were compared in the
+        // predicate, and *all* partition key values are specified(in the predicate or via WithPartitionKey)
+        if (willUseReadItemOptimization)
         {
             return shapedQuery.UpdateQueryExpression(select.WithReadItemInfo(new ReadItemInfo(_jsonIdPropertyValues!)));
         }
@@ -152,21 +188,87 @@ public virtual Expression ExtractPartitionKeysAndId(
         }
 
         return shapedQuery;
+    }
 
-        Expression Unwrap(Expression shaper)
+    private void ValidateNoWithPartitionKeyConflict(
+        CosmosQueryCompilationContext queryCompilationContext,
+        IReadOnlyList<IProperty> partitionKeyProperties,
+        IReadOnlyDictionary<IProperty, (Expression? ValueExpression, Expression? OriginalExpression)>
+            predicatePartitionKeyPropertyValues)
+    {
+        // WithPartitionKey(...) already populated partition key values in the query compilation context.
+        // If the predicate also contains partition key comparisons, we must validate that they match; otherwise, a ReadItem
+        // optimization would execute with the WithPartitionKey partition key and ignore the conflicting predicate.
+        for (var i = 0; i < partitionKeyProperties.Count; i++)
         {
-            if (shaper is UnaryExpression { NodeType: ExpressionType.Convert } convert
-                && convert.Type == typeof(object))
+            var property = partitionKeyProperties[i];
+            var predicateValueExpression = predicatePartitionKeyPropertyValues[property].ValueExpression;
+            if (predicateValueExpression is null)
             {
-                shaper = convert.Operand;
+                continue;
             }
 
-            while (shaper is IncludeExpression { EntityExpression: var nested })
+            if (queryCompilationContext.PartitionKeyPropertyValues.Count <= i)
             {
-                shaper = nested;
+                // Shouldn't happen: WithPartitionKey doesn't specify enough PK components; let the existing missing-PK flow handle it.
+                break;
             }
 
-            return shaper;
+            var withPartitionKeyValueExpression = queryCompilationContext.PartitionKeyPropertyValues[i];
+
+            if (!PartitionKeySqlValuesMatch(
+                    predicateValueExpression,
+                    withPartitionKeyValueExpression,
+                    property))
+            {
+                throw new InvalidOperationException(CosmosStrings.WithPartitionKeyConflictingPartitionKeyPredicate);
+            }
+        }
+    }
+
+    private static bool PartitionKeySqlValuesMatch(Expression left, Expression right, IProperty partitionKeyProperty)
+    {
+        if (ReferenceEquals(left, right))
+        {
+            return true;
+        }
+
+        left = UnwrapPartitionKeySqlExpression(left);
+        right = UnwrapPartitionKeySqlExpression(right);
+
+        if (ReferenceEquals(left, right))
+        {
+            return true;
+        }
+
+        if (left is SqlExpression sqlLeft && right is SqlExpression sqlRight && sqlLeft.Equals(sqlRight))
+        {
+            return true;
+        }
+
+        var comparer = partitionKeyProperty.GetValueComparer();
+
+        switch (left, right)
+        {
+            case (SqlConstantExpression { Value: var leftValue }, SqlConstantExpression { Value: var rightValue }):
+                return comparer.Equals(leftValue, rightValue);
+
+            case (SqlParameterExpression { Name: var leftName }, SqlParameterExpression { Name: var rightName }):
+                return leftName == rightName;
+
+            default:
+                return false;
+        }
+
+        static Expression UnwrapPartitionKeySqlExpression(Expression expression)
+        {
+            while (expression is SqlUnaryExpression { OperatorType: ExpressionType.Convert or ExpressionType.ConvertChecked } unary
+                && unary.Operand is SqlExpression operand)
+            {
+                expression = operand;
+            }
+
+            return expression;
         }
     }
 
@@ -320,7 +422,8 @@ void ProcessPropertyComparison(string propertyName, SqlExpression propertyValue,
                 // call. Note that this is always considered a compatible comparison for ReadItem.
                 if (propertyName == property.GetJsonPropertyName()
                     && _partitionKeyPropertyValues.TryGetValue(property, out var previousValues)
-                    && (previousValues.ValueExpression is null || previousValues.Equals(propertyValue)))
+                    && (previousValues.ValueExpression is null
+                        || PartitionKeySqlValuesMatch(previousValues.ValueExpression, propertyValue, property)))
                 {
                     _partitionKeyPropertyValues[property] = (ValueExpression: propertyValue, OriginalExpression: originalExpression);
                     return;

test/EFCore.Cosmos.FunctionalTests/Query/ReadItemPartitionKeyQueryDiscriminatorInIdTest.cs

@@ -1,4 +1,4 @@
-// Licensed to the .NET Foundation under one or more agreements.
+// 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.Query;
@@ -442,6 +442,40 @@ FROM root c
 """);
     }
 
+    public override async Task ReadItem_with_WithPartitionKey_and_conflicting_partition_key_predicate_throws()
+    {
+        // Not ReadItem when the discriminator is stored in the JSON id: the conflicting partition key predicate remains in SQL
+        // and the query returns no results (see #38238; validation only applies to the ReadItem optimization path).
+        await AssertQuery(
+            async: true,
+            ss => ss.Set<SinglePartitionKeyEntity>().WithPartitionKey("PK1")
+                .Where(e => e.Id == Guid.Parse("B29BCED8-E1E5-420E-82D7-1C7A51703D34") && e.PartitionKey == "PK2"),
+            ss => ss.Set<SinglePartitionKeyEntity>().Where(e => e.PartitionKey == "PK1")
+                .Where(e => e.Id == Guid.Parse("B29BCED8-E1E5-420E-82D7-1C7A51703D34") && e.PartitionKey == "PK2"),
+            assertEmpty: true);
+
+        AssertSql(
+            """
+SELECT VALUE c
+FROM root c
+WHERE (c["$type"] IN ("SinglePartitionKeyEntity", "DerivedSinglePartitionKeyEntity") AND ((c["Id"] = "b29bced8-e1e5-420e-82d7-1c7a51703d34") AND (c["PartitionKey"] = "PK2")))
+""");
+    }
+
+    public override async Task ReadItem_with_WithPartitionKey_and_partition_key_predicate_same_parameter_does_not_throw()
+    {
+        await base.ReadItem_with_WithPartitionKey_and_partition_key_predicate_same_parameter_does_not_throw();
+
+        AssertSql(
+            """
+@partitionKey='PK1'
+
+SELECT VALUE c
+FROM root c
+WHERE (c["$type"] IN ("SinglePartitionKeyEntity", "DerivedSinglePartitionKeyEntity") AND ((c["Id"] = "b29bced8-e1e5-420e-82d7-1c7a51703d34") AND (c["PartitionKey"] = @partitionKey)))
+""");
+    }
+
     public override async Task ReadItem_with_WithPartitionKey_with_only_partition_key()
     {
         await base.ReadItem_with_WithPartitionKey_with_only_partition_key();

test/EFCore.Cosmos.FunctionalTests/Query/ReadItemPartitionKeyQueryNoDiscriminatorInIdTest.cs

@@ -1,6 +1,8 @@
-// Licensed to the .NET Foundation under one or more agreements.
+// Licensed to the .NET Foundation under one or more agreements.
 // The .NET Foundation licenses this file to you under the MIT license.
 
+using Microsoft.EntityFrameworkCore.Cosmos.Internal;
+
 namespace Microsoft.EntityFrameworkCore.Query;
 
 public class ReadItemPartitionKeyQueryNoDiscriminatorInIdTest
@@ -359,6 +361,25 @@ public override async Task ReadItem_with_WithPartitionKey()
         AssertSql("""ReadItem(["PK1"], b29bced8-e1e5-420e-82d7-1c7a51703d34)""");
     }
 
+    public override async Task ReadItem_with_WithPartitionKey_and_conflicting_partition_key_predicate_throws()
+    {
+        var exception = await Assert.ThrowsAsync<InvalidOperationException>(() => AssertQuery(
+            async: true,
+            ss => ss.Set<SinglePartitionKeyEntity>().WithPartitionKey("PK1")
+                .Where(e => e.Id == Guid.Parse("B29BCED8-E1E5-420E-82D7-1C7A51703D34") && e.PartitionKey == "PK2")));
+
+        Assert.Equal(CosmosStrings.WithPartitionKeyConflictingPartitionKeyPredicate, exception.Message);
+
+        AssertSql();
+    }
+
+    public override async Task ReadItem_with_WithPartitionKey_and_partition_key_predicate_same_parameter_does_not_throw()
+    {
+        await base.ReadItem_with_WithPartitionKey_and_partition_key_predicate_same_parameter_does_not_throw();
+
+        AssertSql("""ReadItem(["PK1"], b29bced8-e1e5-420e-82d7-1c7a51703d34)""");
+    }
+
     public override async Task ReadItem_with_WithPartitionKey_with_only_partition_key()
     {
         await base.ReadItem_with_WithPartitionKey_with_only_partition_key();

test/EFCore.Cosmos.FunctionalTests/Query/ReadItemPartitionKeyQueryRootDiscriminatorInIdTest.cs

@@ -1,6 +1,8 @@
-// Licensed to the .NET Foundation under one or more agreements.
+// Licensed to the .NET Foundation under one or more agreements.
 // The .NET Foundation licenses this file to you under the MIT license.
 
+using Microsoft.EntityFrameworkCore.Cosmos.Internal;
+
 namespace Microsoft.EntityFrameworkCore.Query;
 
 public class ReadItemPartitionKeyQueryRootDiscriminatorInIdTest
@@ -352,6 +354,25 @@ public override async Task ReadItem_with_WithPartitionKey()
         AssertSql("""ReadItem(["PK1"], SinglePartitionKeyEntity|b29bced8-e1e5-420e-82d7-1c7a51703d34)""");
     }
 
+    public override async Task ReadItem_with_WithPartitionKey_and_conflicting_partition_key_predicate_throws()
+    {
+        var exception = await Assert.ThrowsAsync<InvalidOperationException>(() => AssertQuery(
+            async: true,
+            ss => ss.Set<SinglePartitionKeyEntity>().WithPartitionKey("PK1")
+                .Where(e => e.Id == Guid.Parse("B29BCED8-E1E5-420E-82D7-1C7A51703D34") && e.PartitionKey == "PK2")));
+
+        Assert.Equal(CosmosStrings.WithPartitionKeyConflictingPartitionKeyPredicate, exception.Message);
+
+        AssertSql();
+    }
+
+    public override async Task ReadItem_with_WithPartitionKey_and_partition_key_predicate_same_parameter_does_not_throw()
+    {
+        await base.ReadItem_with_WithPartitionKey_and_partition_key_predicate_same_parameter_does_not_throw();
+
+        AssertSql("""ReadItem(["PK1"], SinglePartitionKeyEntity|b29bced8-e1e5-420e-82d7-1c7a51703d34)""");
+    }
+
     public override async Task ReadItem_with_WithPartitionKey_with_only_partition_key()
     {
         await base.ReadItem_with_WithPartitionKey_with_only_partition_key();

test/EFCore.Cosmos.FunctionalTests/Query/ReadItemPartitionKeyQueryTest.cs

@@ -1,6 +1,8 @@
 // Licensed to the .NET Foundation under one or more agreements.
 // The .NET Foundation licenses this file to you under the MIT license.
 
+using Microsoft.EntityFrameworkCore.Cosmos.Internal;
+
 namespace Microsoft.EntityFrameworkCore.Query;
 
 public class ReadItemPartitionKeyQueryTest : ReadItemPartitionKeyQueryTestBase<
@@ -351,6 +353,25 @@ public override async Task ReadItem_with_WithPartitionKey()
         AssertSql("""ReadItem(["PK1"], b29bced8-e1e5-420e-82d7-1c7a51703d34)""");
     }
 
+    public override async Task ReadItem_with_WithPartitionKey_and_conflicting_partition_key_predicate_throws()
+    {
+        var exception = await Assert.ThrowsAsync<InvalidOperationException>(() => AssertQuery(
+            async: true,
+            ss => ss.Set<SinglePartitionKeyEntity>().WithPartitionKey("PK1")
+                .Where(e => e.Id == Guid.Parse("B29BCED8-E1E5-420E-82D7-1C7A51703D34") && e.PartitionKey == "PK2")));
+
+        Assert.Equal(CosmosStrings.WithPartitionKeyConflictingPartitionKeyPredicate, exception.Message);
+
+        AssertSql();
+    }
+
+    public override async Task ReadItem_with_WithPartitionKey_and_partition_key_predicate_same_parameter_does_not_throw()
+    {
+        await base.ReadItem_with_WithPartitionKey_and_partition_key_predicate_same_parameter_does_not_throw();
+
+        AssertSql("""ReadItem(["PK1"], b29bced8-e1e5-420e-82d7-1c7a51703d34)""");
+    }
+
     public override async Task ReadItem_with_WithPartitionKey_with_only_partition_key()
     {
         await base.ReadItem_with_WithPartitionKey_with_only_partition_key();