Cosmos: Detect partition key filters in queries

src/EFCore.Cosmos/Properties/CosmosStrings.resx

@@ -177,4 +177,7 @@
   <data name="ResourceIdMissing" xml:space="preserve">
     <value>A ReadItem query was detected, but the id value is missing and cannot be generated.</value>
   </data>
+  <data name="PartitionKeyMismatch" xml:space="preserve">
+    <value>Partition key specified in the WithPartitionKey call '{paritionKey1}' and the partition key specified in the Where predicate '{paritionKey2}' must be identical. Remove one of them .</value>
+  </data>
 </root>

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

@@ -5,12 +5,12 @@
 using System.Collections.Generic;
 using System.Linq;
 using System.Linq.Expressions;
+using System.Reflection;
 using JetBrains.Annotations;
 using Microsoft.EntityFrameworkCore.Diagnostics;
 using Microsoft.EntityFrameworkCore.Infrastructure;
 using Microsoft.EntityFrameworkCore.Metadata;
 using Microsoft.EntityFrameworkCore.Metadata.Conventions;
-using Microsoft.EntityFrameworkCore.Metadata.Internal;
 using Microsoft.EntityFrameworkCore.Query;
 using Microsoft.EntityFrameworkCore.Storage;
 using Microsoft.EntityFrameworkCore.Utilities;
@@ -135,27 +135,27 @@ public override Expression Visit(Expression expression)
                             bool ProcessJoinCondition(
                                 Expression joinCondition, ICollection<IProperty> properties, ICollection<string> paramNames)
                             {
-                                if (joinCondition is BinaryExpression binaryExpression)
+                                if (joinCondition is BinaryExpression joinBinaryExpression)
                                 {
-                                    switch (binaryExpression.NodeType)
+                                    switch (joinBinaryExpression.NodeType)
                                     {
                                         case ExpressionType.AndAlso:
-                                            return ProcessJoinCondition(binaryExpression.Left, properties, paramNames)
-                                                && ProcessJoinCondition(binaryExpression.Right, properties, paramNames);
+                                            return ProcessJoinCondition(joinBinaryExpression.Left, properties, paramNames)
+                                                && ProcessJoinCondition(joinBinaryExpression.Right, properties, paramNames);
 
                                         case ExpressionType.Equal:
-                                            if (binaryExpression.Left is MethodCallExpression methodCallExpr
-                                                && binaryExpression.Right is ParameterExpression parameterExpr)
+                                            if (joinBinaryExpression.Left is MethodCallExpression equalMethodCallExpression
+                                                && joinBinaryExpression.Right is ParameterExpression equalParameterExpresion)
                                             {
-                                                if (methodCallExpr.TryGetEFPropertyArguments(out _, out var propertyName))
+                                                if (equalMethodCallExpression.TryGetEFPropertyArguments(out _, out var propertyName))
                                                 {
                                                     var property = entityType.FindProperty(propertyName);
                                                     if (property == null)
                                                     {
                                                         return false;
                                                     }
                                                     properties.Add(property);
-                                                    paramNames.Add(parameterExpr.Name);
+                                                    paramNames.Add(equalParameterExpresion.Name);
                                                     return true;
                                                 }
                                             }
@@ -1003,6 +1003,23 @@ protected override ShapedQueryExpression TranslateWhere(ShapedQueryExpression so
             Check.NotNull(source, nameof(source));
             Check.NotNull(predicate, nameof(predicate));
 
+            if (source.ShaperExpression is EntityShaperExpression entityShaperExpression)
+            {
+                var (predicateExpression, partitionKeyProperty, partitionKeyValueExpression) = ExtractPartitionKey(predicate.Body, entityShaperExpression.EntityType);
+
+                if (partitionKeyProperty != null && partitionKeyValueExpression != null)
+                {
+                    ((SelectExpression)source.QueryExpression).SetPartitionKeyProperty(partitionKeyProperty, partitionKeyValueExpression);
+
+                    if (predicateExpression is null)
+                    {
+                        return source;
+                    }
+
+                    predicate = Expression.Lambda(predicateExpression, predicate.Parameters);
+                }
+            }
+
             var translation = TranslateLambdaExpression(source, predicate);
             if (translation != null)
             {
@@ -1012,6 +1029,128 @@ protected override ShapedQueryExpression TranslateWhere(ShapedQueryExpression so
             }
 
             return null;
+
+            static (Expression predicateExpression, IProperty partitionKeyProperty, Expression partitionKeyValueExpression)
+                ExtractPartitionKey(Expression expression, IEntityType entityType)
+            {
+                if (expression is BinaryExpression binaryExpression)
+                {
+                    if (binaryExpression.Left is BinaryExpression leftBinaryExpression
+                        && binaryExpression.Right is BinaryExpression rightBinaryExpression)
+                    {
+                        switch (binaryExpression.NodeType)
+                        {
+                            case ExpressionType.AndAlso when leftBinaryExpression.NodeType == ExpressionType.Equal
+                                && TryGetPartitionKeyPropertyAndValuExpression(
+                                    leftBinaryExpression.Left, leftBinaryExpression.Right, entityType,
+                                    out var leftProperty,
+                                    out var leftValueExpression):
+
+                                return (rightBinaryExpression, leftProperty, leftValueExpression);
+
+                            case ExpressionType.AndAlso when rightBinaryExpression.NodeType == ExpressionType.Equal
+                                && TryGetPartitionKeyPropertyAndValuExpression(
+                                    rightBinaryExpression.Left, rightBinaryExpression.Right, entityType,
+                                    out var rightProperty,
+                                    out var rightValueExpression):
+
+                                return (leftBinaryExpression, rightProperty, rightValueExpression);
+
+                            case ExpressionType.Equal when TryGetPartitionKeyPropertyAndValuExpression(
+                                binaryExpression.Left, binaryExpression.Right, entityType,
+                                out var property,
+                                out var valueExpression):
+
+                                return (null, property, valueExpression);
+                        }
+
+                        var left = ExtractPartitionKey(leftBinaryExpression, entityType);
+                        var right = ExtractPartitionKey(rightBinaryExpression, entityType);
+
+                        return (
+                            Expression.MakeBinary(binaryExpression.NodeType, left.predicateExpression, right.predicateExpression),
+                            left.partitionKeyProperty ?? right.partitionKeyProperty,
+                            left.partitionKeyValueExpression ?? right.partitionKeyValueExpression);
+                    }
+                    return (binaryExpression, null, null);
+                }
+
+                return (expression, null, null);
+            }
+
+            static bool TryGetPartitionKeyPropertyAndValuExpression(Expression leftBinaryExpression, Expression rightBinaryExpression, IEntityType entityType,
+                out IProperty partitionKeyProperty,
+                out Expression paritionKeyPropertyValueExpression)
+            {
+                partitionKeyProperty = null;
+                paritionKeyPropertyValueExpression = null;
+
+                if (TryGetPropertyAndValueExpression(leftBinaryExpression, rightBinaryExpression, entityType,
+                    out var leftSidedProperty,
+                    out var leftSidedExpression))
+                {
+                    partitionKeyProperty = leftSidedProperty;
+                    paritionKeyPropertyValueExpression = leftSidedExpression;
+                    return true;
+                }
+
+                if (TryGetPropertyAndValueExpression(rightBinaryExpression, leftSidedExpression, entityType,
+                    out var rightSidedProperty,
+                    out var rightSidedExpression))
+                {
+                    partitionKeyProperty = rightSidedProperty;
+                    paritionKeyPropertyValueExpression = rightSidedExpression;
+                    return true;
+                }
+
+                return false;
+            }
+
+            static bool TryGetPropertyAndValueExpression(Expression leftExpression, Expression rightExpression, IEntityType entityType,
+                out IProperty property,
+                out Expression expression)
+            {
+                property = null;
+                expression = null;
+
+                if (leftExpression is MemberExpression memberExpression
+                    && rightExpression is ConstantExpression constantExpression
+                    && memberExpression.Member.GetSimpleMemberName() == entityType.GetPartitionKeyPropertyName())
+                {
+                    property = entityType.FindProperty(memberExpression.Member.GetSimpleMemberName());
+
+                    if (property is null)
+                    {
+                        return false;
+                    }
+
+                    expression = constantExpression;
+                    return true;
+                }
+
+                if (leftExpression is MethodCallExpression methodCallExpression
+                    && rightExpression is ParameterExpression parameterExpression)
+                {
+                    property = GetPropertyFromMethodCall(methodCallExpression, entityType);
+
+                    if (property is null)
+                    {
+                        return false;
+                    }
+
+                    expression = parameterExpression;
+                    return true;
+                }
+
+                return false;
+            }
+
+            static IProperty GetPropertyFromMethodCall(MethodCallExpression methodCallExpression, IEntityType entityType) =>
+                methodCallExpression.TryGetEFPropertyArguments(out _, out var parameterName)
+                    ? entityType.FindProperty(parameterName)
+                    : methodCallExpression.TryGetIndexerArguments(entityType.Model, out _, out var indexerName)
+                        ? entityType.FindProperty(indexerName)
+                        : default;
         }
 
         private SqlExpression TranslateExpression(Expression expression)

src/EFCore.Cosmos/Query/Internal/CosmosShapedQueryCompilingExpressionVisitor.QueryingEnumerable.cs

@@ -7,6 +7,7 @@
 using System.Text;
 using System.Threading;
 using System.Threading.Tasks;
+using Microsoft.EntityFrameworkCore.Cosmos.Internal;
 using Microsoft.EntityFrameworkCore.Cosmos.Storage.Internal;
 using Microsoft.EntityFrameworkCore.Diagnostics;
 using Microsoft.EntityFrameworkCore.Query;
@@ -50,9 +51,17 @@ public QueryingEnumerable(
                 _shaper = shaper;
                 _contextType = contextType;
                 _logger = logger;
-                _partitionKey = partitionKeyFromExtension;
-            }
 
+                var partitionKey = selectExpression.GetPartitionKey(cosmosQueryContext.ParameterValues);
+
+                if (partitionKey != null && partitionKeyFromExtension != null && partitionKeyFromExtension != partitionKey)
+                {
+                    throw new InvalidOperationException(CosmosStrings.PartitionKeyMismatch(partitionKeyFromExtension, partitionKey));
+                }
+
+                _partitionKey = partitionKey ?? partitionKeyFromExtension;
+            }
+
             public IAsyncEnumerator<T> GetAsyncEnumerator(CancellationToken cancellationToken = default)
                 => new AsyncEnumerator(this, cancellationToken);
 

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

@@ -28,6 +28,9 @@ public class SelectExpression : Expression
         private readonly List<ProjectionExpression> _projection = new List<ProjectionExpression>();
         private readonly List<OrderingExpression> _orderings = new List<OrderingExpression>();
 
+        private IProperty _partitionKeyProperty;
+        private Expression _paritionKeyValueExpression;
+
         /// <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
@@ -130,6 +133,45 @@ public SelectExpression(
         public virtual Expression GetMappedProjection([NotNull] ProjectionMember projectionMember)
             => _projectionMapping[projectionMember];
 
+
+        /// <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 virtual void SetPartitionKeyProperty([NotNull]IProperty partitionKeyProperty, [NotNull]Expression expression)
+        {
+            _partitionKeyProperty = partitionKeyProperty;
+            _paritionKeyValueExpression = 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
+        ///     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 virtual string GetPartitionKey([NotNull]IReadOnlyDictionary<string, object> parameterValues)
+        {
+            return _partitionKeyProperty != null && _paritionKeyValueExpression is ConstantExpression constantExpression
+                ? GetString(_partitionKeyProperty, constantExpression.Value)
+                : _partitionKeyProperty != null
+                    && _paritionKeyValueExpression is ParameterExpression parameterExpression
+                    && parameterValues.TryGetValue(parameterExpression.Name, out var value)
+                        ? GetString(_partitionKeyProperty, value)
+                        : null;
+
+            static string GetString(IProperty property, object value)
+            {
+                var converter = property.GetTypeMapping().Converter;
+
+                return converter is null
+                    ? (string)value
+                    : (string)converter.ConvertToProvider(value);
+            }
+        }
+
         /// <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