Properly implement support for Cosmos hierarchical partition keys

src/EFCore.Cosmos/Extensions/CosmosQueryableExtensions.cs

@@ -18,58 +18,118 @@ namespace Microsoft.EntityFrameworkCore;
 /// </remarks>
 public static class CosmosQueryableExtensions
 {
-    internal static readonly MethodInfo WithPartitionKeyMethodInfo
+    internal static readonly MethodInfo WithPartitionKeyMethodInfo1
+        = typeof(CosmosQueryableExtensions).GetTypeInfo()
+            .GetDeclaredMethods(nameof(WithPartitionKey))
+            .Single(mi => mi.GetParameters().Length == 2);
+
+    internal static readonly MethodInfo WithPartitionKeyMethodInfo2
         = typeof(CosmosQueryableExtensions).GetTypeInfo()
             .GetDeclaredMethods(nameof(WithPartitionKey))
             .Single(mi => mi.GetParameters().Length == 3);
 
+    internal static readonly MethodInfo WithPartitionKeyMethodInfo3
+        = typeof(CosmosQueryableExtensions).GetTypeInfo()
+            .GetDeclaredMethods(nameof(WithPartitionKey))
+            .Single(mi => mi.GetParameters().Length == 4);
+
     /// <summary>
-    ///     Specify the partition key value for partition used for the query. Required when using
-    ///     a resource token that provides permission based on a partition key for authentication.
+    ///     Specify the partition key for partition used for the query.
+    ///     Required when using a resource token that provides permission based on a partition key for authentication,
     /// </summary>
     /// <remarks>
     ///     See <see href="https://aka.ms/efcore-docs-query">Querying data with EF Core</see>, and
     ///     <see href="https://aka.ms/efcore-docs-cosmos">Accessing Azure Cosmos DB with EF Core</see> for more information and examples.
     /// </remarks>
     /// <typeparam name="TEntity">The type of entity being queried.</typeparam>
     /// <param name="source">The source query.</param>
-    /// <param name="partitionKey">The partition key value.</param>
+    /// <param name="partitionKeyValue">The partition key value.</param>
     /// <returns>A new query with the set partition key.</returns>
-    public static IQueryable<TEntity> WithPartitionKey<TEntity>(this IQueryable<TEntity> source, string partitionKey)
+    public static IQueryable<TEntity> WithPartitionKey<TEntity>(this IQueryable<TEntity> source, object partitionKeyValue)
         where TEntity : class
-        => WithPartitionKey(source, partitionKey, []);
+    {
+        Check.NotNull(partitionKeyValue, nameof(partitionKeyValue));
+
+        return
+            source.Provider is EntityQueryProvider
+                ? source.Provider.CreateQuery<TEntity>(
+                    Expression.Call(
+                        instance: null,
+                        method: WithPartitionKeyMethodInfo1.MakeGenericMethod(typeof(TEntity)),
+                        source.Expression,
+                        Expression.Constant(partitionKeyValue, typeof(object))))
+                : source;
+    }
 
     /// <summary>
-    ///     Specify the partition key for partition used for the query. Required when using
-    ///     a resource token that provides permission based on a partition key for authentication,
+    ///     Specify the partition key for partition used for the query.
+    ///     Required when using a resource token that provides permission based on a partition key for authentication,
     /// </summary>
     /// <remarks>
     ///     See <see href="https://aka.ms/efcore-docs-query">Querying data with EF Core</see>, and
     ///     <see href="https://aka.ms/efcore-docs-cosmos">Accessing Azure Cosmos DB with EF Core</see> for more information and examples.
     /// </remarks>
     /// <typeparam name="TEntity">The type of entity being queried.</typeparam>
     /// <param name="source">The source query.</param>
-    /// <param name="partitionKeyValue">The partition key value.</param>
-    /// <param name="additionalPartitionKeyValues">Additional values for hierarchical partitions.</param>
+    /// <param name="partitionKeyValue1">The first value in a hierarchical partition key.</param>
+    /// <param name="partitionKeyValue2">The second value in a hierarchical partition key.</param>
     /// <returns>A new query with the set partition key.</returns>
     public static IQueryable<TEntity> WithPartitionKey<TEntity>(
         this IQueryable<TEntity> source,
-        object partitionKeyValue,
-        params object[] additionalPartitionKeyValues)
+        object partitionKeyValue1,
+        object partitionKeyValue2)
         where TEntity : class
     {
-        Check.NotNull(partitionKeyValue, nameof(partitionKeyValue));
-        Check.HasNoNulls(additionalPartitionKeyValues, nameof(additionalPartitionKeyValues));
+        Check.NotNull(partitionKeyValue1, nameof(partitionKeyValue1));
+        Check.NotNull(partitionKeyValue2, nameof(partitionKeyValue2));
+
+        return
+            source.Provider is EntityQueryProvider
+                ? source.Provider.CreateQuery<TEntity>(
+                    Expression.Call(
+                        instance: null,
+                        method: WithPartitionKeyMethodInfo2.MakeGenericMethod(typeof(TEntity)),
+                        source.Expression,
+                        Expression.Constant(partitionKeyValue1, typeof(object)),
+                        Expression.Constant(partitionKeyValue2, typeof(object))))
+                : source;
+    }
+
+    /// <summary>
+    ///     Specify the partition key for partition used for the query.
+    ///     Required when using a resource token that provides permission based on a partition key for authentication,
+    /// </summary>
+    /// <remarks>
+    ///     See <see href="https://aka.ms/efcore-docs-query">Querying data with EF Core</see>, and
+    ///     <see href="https://aka.ms/efcore-docs-cosmos">Accessing Azure Cosmos DB with EF Core</see> for more information and examples.
+    /// </remarks>
+    /// <typeparam name="TEntity">The type of entity being queried.</typeparam>
+    /// <param name="source">The source query.</param>
+    /// <param name="partitionKeyValue1">The first value in a hierarchical partition key.</param>
+    /// <param name="partitionKeyValue2">The second value in a hierarchical partition key.</param>
+    /// <param name="partitionKeyValue3">The third value in a hierarchical partition key.</param>
+    /// <returns>A new query with the set partition key.</returns>
+    public static IQueryable<TEntity> WithPartitionKey<TEntity>(
+        this IQueryable<TEntity> source,
+        object partitionKeyValue1,
+        object partitionKeyValue2,
+        object partitionKeyValue3)
+        where TEntity : class
+    {
+        Check.NotNull(partitionKeyValue1, nameof(partitionKeyValue1));
+        Check.NotNull(partitionKeyValue2, nameof(partitionKeyValue2));
+        Check.NotNull(partitionKeyValue3, nameof(partitionKeyValue3));
 
         return
             source.Provider is EntityQueryProvider
                 ? source.Provider.CreateQuery<TEntity>(
                     Expression.Call(
                         instance: null,
-                        method: WithPartitionKeyMethodInfo.MakeGenericMethod(typeof(TEntity)),
+                        method: WithPartitionKeyMethodInfo3.MakeGenericMethod(typeof(TEntity)),
                         source.Expression,
-                        Expression.Constant(partitionKeyValue, typeof(object)),
-                        Expression.Constant(additionalPartitionKeyValues, typeof(object[]))))
+                        Expression.Constant(partitionKeyValue1, typeof(object)),
+                        Expression.Constant(partitionKeyValue2, typeof(object)),
+                        Expression.Constant(partitionKeyValue3, typeof(object))))
                 : source;
     }
 

src/EFCore.Cosmos/Properties/CosmosStrings.resx

@@ -174,9 +174,6 @@
   <data name="IdNonStringStoreType" xml:space="preserve">
     <value>The type of the '{idProperty}' property on '{entityType}' is '{propertyType}'. All 'id' properties must be strings or have a string value converter.</value>
   </data>
-  <data name="IncorrectPartitionKeyNumber" xml:space="preserve">
-    <value>{actual} partition key values were provided, but the entity type '{entityType}' has {expected} partition key values defined.</value>
-  </data>
   <data name="IndexesExist" xml:space="preserve">
     <value>The entity type '{entityType}' has an index defined over properties '{properties}'. The Azure Cosmos DB provider for EF Core currently does not support index definitions.</value>
   </data>

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

@@ -172,23 +172,15 @@ protected override Expression VisitMethodCall(MethodCallExpression methodCallExp
 
             var innerQueryable = Visit(methodCallExpression.Arguments[0]);
 
-            var firstValue = _sqlTranslator.Translate(methodCallExpression.Arguments[1], applyDefaultTypeMapping: false);
-            if (firstValue is not SqlConstantExpression and not SqlParameterExpression)
+            for (var i = 1; i < methodCallExpression.Arguments.Count; i++)
             {
-                throw new InvalidOperationException(CosmosStrings.WithPartitionKeyNotConstantOrParameter);
-            }
-
-            _queryCompilationContext.PartitionKeyPropertyValues.Add(firstValue);
-
-            if (methodCallExpression.Arguments.Count == 3)
-            {
-                var remainingValuesArray = _sqlTranslator.Translate(methodCallExpression.Arguments[2], applyDefaultTypeMapping: false);
-                if (remainingValuesArray is not SqlParameterExpression)
+                var value = _sqlTranslator.Translate(methodCallExpression.Arguments[i], applyDefaultTypeMapping: false);
+                if (value is not SqlConstantExpression and not SqlParameterExpression)
                 {
                     throw new InvalidOperationException(CosmosStrings.WithPartitionKeyNotConstantOrParameter);
                 }
 
-                _queryCompilationContext.PartitionKeyPropertyValues.Add(remainingValuesArray);
+                _queryCompilationContext.PartitionKeyPropertyValues.Add(value);
             }
 
             return innerQueryable;