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;
@@ -88,6 +89,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 even when the query does not use ReadItem.
+        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,6 +121,14 @@ public virtual Expression ExtractPartitionKeysAndId(
             }
         }
 
+        if (hadWithPartitionKey && predicatePartitionKeySnapshot is not null)
+        {
+            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 (_isPredicateCompatibleWithReadItem
@@ -170,6 +188,88 @@ 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++)
+        {
+            var property = partitionKeyProperties[i];
+            var predicateValueExpression = predicatePartitionKeyPropertyValues[property].ValueExpression;
+            if (predicateValueExpression is null)
+            {
+                continue;
+            }
+
+            if (queryCompilationContext.PartitionKeyPropertyValues.Count <= i)
+            {
+                // Shouldn't happen: WithPartitionKey doesn't specify enough PK components; let the existing missing-PK flow handle it.
+                break;
+            }
+
+            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;
+        }
+    }
+
     /// <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
@@ -320,7 +420,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,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 ReadItemPartitionKeyQueryDiscriminatorInIdTest
@@ -442,6 +444,32 @@ FROM root c
 """);
     }
 
+    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(
+            """
+@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();

test/EFCore.Cosmos.FunctionalTests/Query/ReadItemPartitionKeyQueryTestBase.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 abstract class ReadItemPartitionKeyQueryTestBase<TFixture> : QueryTestBase<TFixture>
@@ -285,6 +287,28 @@ public virtual Task ReadItem_with_WithPartitionKey()
             ss => ss.Set<SinglePartitionKeyEntity>().Where(e => e.PartitionKey == "PK1")
                 .Where(e => e.Id == Guid.Parse("B29BCED8-E1E5-420E-82D7-1C7A51703D34")));
 
+    [ConditionalFact] // Issue #38238
+    public virtual 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);
+    }
+
+    [ConditionalFact] // Issue #38238
+    public virtual async Task ReadItem_with_WithPartitionKey_and_partition_key_predicate_same_parameter_does_not_throw()
+    {
+        var partitionKey = "PK1";
+
+        await AssertQuery(
+            async: true,
+            ss => ss.Set<SinglePartitionKeyEntity>().WithPartitionKey(partitionKey)
+                .Where(e => e.Id == Guid.Parse("B29BCED8-E1E5-420E-82D7-1C7A51703D34") && e.PartitionKey == partitionKey));
+    }
+
     [ConditionalFact]
     public virtual Task ReadItem_with_WithPartitionKey_with_only_partition_key()
         => AssertQuery(