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,27 @@ protected override ShapedQueryExpression TranslateWhere(ShapedQueryExpression so
             Check.NotNull(source, nameof(source));
             Check.NotNull(predicate, nameof(predicate));
 
+            if (source.ShaperExpression is EntityShaperExpression entityShaperExpression)
+            {
+                var predicateVisitor = new PredicateVisitor(entityShaperExpression.EntityType);
+
+                var modifiedPredicateBody = predicateVisitor.ExtractPartitionKeyFromPredicate(predicate.Body);
+
+                if (predicateVisitor.PartitionKeyValueExpression != null
+                    && predicateVisitor.PartitionKeyProperty != null)
+                {
+                    ((SelectExpression)source.QueryExpression).SetPartitionKeyProperty(
+                        predicateVisitor.PartitionKeyProperty, predicateVisitor.PartitionKeyValueExpression);
+
+                    if (modifiedPredicateBody is null)
+                    {
+                        return source;
+                    }
+
+                    predicate = Expression.Lambda(modifiedPredicateBody, predicate.Parameters);
+                }
+            }
+
             var translation = TranslateLambdaExpression(source, predicate);
             if (translation != null)
             {
@@ -1072,5 +1093,180 @@ private ShapedQueryExpression AggregateResultShaper(
 
             return source.UpdateShaperExpression(shaper);
         }
+
+        private sealed class PredicateVisitor : ExpressionVisitor
+        {
+            private readonly IEntityType _entityType;
+
+            internal IProperty PartitionKeyProperty { get; private set; }
+            internal Expression PartitionKeyValueExpression { get; private set; }
+
+            internal PredicateVisitor(IEntityType entityType)
+            {
+                _entityType = entityType;
+            }
+
+            internal Expression ExtractPartitionKeyFromPredicate(Expression expression)
+            {
+                if (expression is BinaryExpression binaryExpression)
+                {
+                    if (binaryExpression.NodeType == ExpressionType.Equal
+                        && TryGetPartitionKeyPropertyAndValuExpression(binaryExpression.Left, binaryExpression.Right,
+                            out var property,
+                            out var valueExpression))
+                    {
+                        PartitionKeyProperty = property;
+                        PartitionKeyValueExpression = valueExpression;
+                        return null;
+                    }
+
+                    return VisitBinary(binaryExpression);
+                }
+
+                return expression;
+            }
+
+            protected override Expression VisitBinary(BinaryExpression binaryExpression)
+            {
+                if (binaryExpression.Left is BinaryExpression leftBinaryExpression
+                    && binaryExpression.Right is BinaryExpression rightBinaryExpression)
+                {
+                    if (binaryExpression.NodeType == ExpressionType.AndAlso)
+                    {
+                        if (binaryExpression.Left.NodeType == ExpressionType.Equal
+                            && TryGetPartitionKeyPropertyAndValuExpression(
+                                leftBinaryExpression.Left, leftBinaryExpression.Right,
+                                out var leftProperty,
+                                out var leftValueExpression))
+                        {
+                            PartitionKeyProperty = leftProperty;
+                            PartitionKeyValueExpression = leftValueExpression;
+                            return binaryExpression.Right;
+                        }
+
+                        if (binaryExpression.Right.NodeType == ExpressionType.Equal
+                            && TryGetPartitionKeyPropertyAndValuExpression(
+                                rightBinaryExpression.Left, rightBinaryExpression.Right,
+                                out var rightProperty,
+                                out var rightValueExpression))
+                        {
+                            PartitionKeyProperty = rightProperty;
+                            PartitionKeyValueExpression = rightValueExpression;
+                            return binaryExpression.Left;
+                        }
+                    }
+
+                    return Expression.MakeBinary(
+                        binaryExpression.NodeType, VisitBinary(leftBinaryExpression), VisitBinary(rightBinaryExpression));
+                }
+
+                return base.VisitBinary(binaryExpression);
+            }
+
+            private bool TryGetPartitionKeyPropertyAndValuExpression(Expression leftBinaryExpression, Expression rightBinaryExpression,
+                out IProperty partitionKeyProperty,
+                out Expression paritionKeyPropertyValueExpression)
+            {
+                partitionKeyProperty = null;
+                paritionKeyPropertyValueExpression = null;
+
+                if (TryGetPartitionKeyValueConstantExpression(leftBinaryExpression, rightBinaryExpression,
+                    out var leftConstantProperty,
+                    out var leftConstantExpression))
+                {
+                    partitionKeyProperty = leftConstantProperty;
+                    paritionKeyPropertyValueExpression = leftConstantExpression;
+                    return true;
+                }
+
+                if (TryGetPartitionKeyValueConstantExpression(rightBinaryExpression, leftConstantExpression,
+                    out var rightConstantProperty,
+                    out var rightConstantExpression))
+                {
+                    partitionKeyProperty = rightConstantProperty;
+                    paritionKeyPropertyValueExpression = rightConstantExpression;
+                    return true;
+                }
+
+                if (TryGetPartitionKeyValueParameterExpression(leftConstantExpression, rightBinaryExpression,
+                    out var leftParameterProperty,
+                    out var leftParamenterValueExpression))
+                {
+                    partitionKeyProperty = leftParameterProperty;
+                    paritionKeyPropertyValueExpression = leftParamenterValueExpression;
+                    return true;
+                }
+
+                if (TryGetPartitionKeyValueParameterExpression(rightBinaryExpression, leftConstantExpression,
+                    out var rightParameterProperty,
+                    out var rightParamenterValueExpression))
+                {
+                    partitionKeyProperty = rightParameterProperty;
+                    paritionKeyPropertyValueExpression = rightParamenterValueExpression;
+                    return true;
+                }
+
+                return false;
+            }
+
+            private bool TryGetPartitionKeyValueConstantExpression(
+                Expression leftExpression, Expression rightExpression,
+                out IProperty property,
+                out Expression constantExpression)
+            {
+                property = null;
+                constantExpression = null;
+
+                if (leftExpression is MemberExpression memberExpression
+                    && (rightExpression is ConstantExpression rightConstantExpression
+                        && memberExpression.Member.GetSimpleMemberName() == _entityType.GetPartitionKeyPropertyName()))
+                {
+                    property = _entityType.FindProperty(memberExpression.Member.GetSimpleMemberName());
+
+                    if (property is null)
+                    {
+                        return false;
+                    }
+
+                    constantExpression = rightConstantExpression;
+                    return true;
+                }
+
+                return false;
+            }
+
+            private bool TryGetPartitionKeyValueParameterExpression(
+                Expression leftExpression, Expression rightExpression,
+                out IProperty property,
+                out Expression parameterValueExpression)
+            {
+                property = null;
+                parameterValueExpression = null;
+
+                if (leftExpression is MethodCallExpression methodCallExpression
+                    && rightExpression is ParameterExpression parameterExpression)
+                {
+                    property = GetPropertyFromMethodCall(methodCallExpression);
+
+                    if (property is null)
+                    {
+                        return false;
+                    }
+
+                    parameterValueExpression = parameterExpression;
+                    return true;
+                }
+
+                return false;
+            }
+
+            private IProperty GetPropertyFromMethodCall(MethodCallExpression methodCallExpression) =>
+                methodCallExpression.TryGetEFPropertyArguments(out _, out var parameterName)
+                    ? _entityType.FindProperty(parameterName)
+                    : methodCallExpression.TryGetIndexerArguments(_entityType.Model, out _, out var indexerName)
+                        ? _entityType.FindProperty(indexerName)
+                        : default;
+
+        }
     }
 }

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);
+
+                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([CanBeNull]IProperty partitionKeyProperty, [CanBeNull]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]CosmosQueryContext cosmosQueryContext)
+        {
+            return _partitionKeyProperty != null && _paritionKeyValueExpression is ConstantExpression constantExpression
+                ? GetString(_partitionKeyProperty, constantExpression.Value)
+                : _partitionKeyProperty != null
+                    && _paritionKeyValueExpression is ParameterExpression parameterExpression
+                    && cosmosQueryContext.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