Cosmos: Redo ReadItem and partition key management

src/EFCore.Cosmos/Extensions/CosmosQueryableExtensions.cs

@@ -35,9 +35,7 @@ internal static readonly MethodInfo WithPartitionKeyMethodInfo
     /// <param name="source">The source query.</param>
     /// <param name="partitionKey">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,
-        [NotParameterized] string partitionKey)
+    public static IQueryable<TEntity> WithPartitionKey<TEntity>(this IQueryable<TEntity> source, string partitionKey)
         where TEntity : class
         => WithPartitionKey(source, partitionKey, []);
 
@@ -56,8 +54,8 @@ public static IQueryable<TEntity> WithPartitionKey<TEntity>(
     /// <returns>A new query with the set partition key.</returns>
     public static IQueryable<TEntity> WithPartitionKey<TEntity>(
         this IQueryable<TEntity> source,
-        [NotParameterized] object partitionKeyValue,
-        [NotParameterized] params object[] additionalPartitionKeyValues)
+        object partitionKeyValue,
+        params object[] additionalPartitionKeyValues)
         where TEntity : class
     {
         Check.NotNull(partitionKeyValue, nameof(partitionKeyValue));

src/EFCore.Cosmos/Extensions/Internal/PartitionKeyBuilderExtensions.cs

@@ -35,36 +35,37 @@ public static PartitionKeyBuilder Add(this PartitionKeyBuilder builder, object?
         else
         {
             var expectedType = (converter?.ProviderClrType ?? property?.ClrType)?.UnwrapNullableType();
-            if (value is string stringValue)
+            switch (value)
             {
-                if (expectedType != null && expectedType != typeof(string))
-                {
-                    CheckType(typeof(string));
-                }
+                case string stringValue:
+                    if (expectedType != null && expectedType != typeof(string))
+                    {
+                        CheckType(typeof(string));
+                    }
 
-                builder.Add(stringValue);
-            }
-            else if (value is bool boolValue)
-            {
-                if (expectedType != null && expectedType != typeof(bool))
-                {
-                    CheckType(typeof(bool));
-                }
+                    builder.Add(stringValue);
+                    break;
 
-                builder.Add(boolValue);
-            }
-            else if (value.GetType().IsNumeric())
-            {
-                if (expectedType != null && !expectedType.IsNumeric())
-                {
-                    CheckType(value.GetType());
-                }
+                case bool boolValue:
+                    if (expectedType != null && expectedType != typeof(bool))
+                    {
+                        CheckType(typeof(bool));
+                    }
 
-                builder.Add(Convert.ToDouble(value));
-            }
-            else
-            {
-                throw new InvalidOperationException(CosmosStrings.PartitionKeyBadValue(value.GetType()));
+                    builder.Add(boolValue);
+                    break;
+
+                case var _ when value.GetType().IsNumeric():
+                    if (expectedType != null && !expectedType.IsNumeric())
+                    {
+                        CheckType(value.GetType());
+                    }
+
+                    builder.Add(Convert.ToDouble(value));
+                    break;
+
+                default:
+                    throw new InvalidOperationException(CosmosStrings.PartitionKeyBadValue(value.GetType()));
             }
 
             void CheckType(Type actualType)

src/EFCore.Cosmos/Properties/CosmosStrings.resx

@@ -153,6 +153,9 @@
   <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>
@@ -213,8 +216,8 @@
   <data name="MissingOrderingInSelectExpression" xml:space="preserve">
     <value>'Reverse' could not be translated to the server because there is no ordering on the server side.</value>
   </data>
-  <data name="MultipleContainersReferencedInQuery" xml:space="preserve">
-    <value>Cosmos container '{container1}' is referenced by the query, but '{container2}' is already being referenced. A query can only reference a single Cosmos container.</value>
+  <data name="MultipleRootEntityTypesReferencedInQuery" xml:space="preserve">
+    <value>Root entity type '{entityType1}' is referenced by the query, but '{entityType2}' is already being referenced. A query can only reference a single root entity type.</value>
   </data>
   <data name="NavigationPropertyIsNotAnEmbeddedEntity" xml:space="preserve">
     <value>Navigation '{entityType}.{navigationName}' doesn't point to an embedded entity.</value>
@@ -282,9 +285,6 @@
   <data name="PartitionKeyBadValueType" xml:space="preserve">
     <value>The partition key value supplied for '{propertyType}' property '{entityType}.{property}' is of type '{valueType}'. Partition key values must be of a type assignable to the property.</value>
   </data>
-  <data name="PartitionKeyMismatch" xml:space="preserve">
-    <value>The partition key specified in the 'WithPartitionKey' call '{partitionKey1}' and the partition key specified in the 'Where' predicate '{partitionKey2}' must be identical to return any results. Remove one of them.</value>
-  </data>
   <data name="PartitionKeyMissing" xml:space="preserve">
     <value>Unable to execute a 'ReadItem' query since the partition key value is missing. Consider using the 'WithPartitionKey' method on the query to specify partition key to use.</value>
   </data>
@@ -327,7 +327,13 @@
   <data name="VisitChildrenMustBeOverridden" xml:space="preserve">
     <value>'VisitChildren' must be overridden in the class deriving from 'SqlExpression'.</value>
   </data>
+  <data name="WithPartitionKeyAlreadyCalled" xml:space="preserve">
+    <value>'WithPartitionKey' can only be called once in a query. See https://aka.ms/efdocs-cosmos-partition-keys for more information.</value>
+  </data>
   <data name="WithPartitionKeyBadNode" xml:space="preserve">
-    <value>'WithPartitionKeyMethodInfo' can only be called on a entity query root. See https://aka.ms/efdocs-cosmos-partition-keys for more information.</value>
+    <value>'WithPartitionKey' can only be called on a entity query root. See https://aka.ms/efdocs-cosmos-partition-keys for more information.</value>
+  </data>
+  <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>
 </root>

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

@@ -13,23 +13,23 @@ public class CosmosQueryCompilationContext(QueryCompilationContextDependencies d
     : QueryCompilationContext(dependencies, async)
 {
     /// <summary>
-    ///     The name of the Cosmos container against which this query will be executed.
+    ///     The root entity type being queried.
     /// </summary>
     /// <remarks>
     ///     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.
     /// </remarks>
-    public virtual string? CosmosContainer { get; internal set; }
+    public virtual IEntityType? RootEntityType { get; internal set; }
 
     /// <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 PartitionKey? PartitionKeyValueFromExtension { get; internal set; }
+    public virtual List<Expression> PartitionKeyPropertyValues { get; internal set; } = new();
 
     /// <summary>
     ///     A manager for aliases, capable of generate uniquified source aliases.

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

@@ -294,7 +294,7 @@ protected override Expression VisitSelect(SelectExpression selectExpression)
             // Both methods produce the exact same results; we usually prefer the 1st, but in some cases we use the 2nd.
             else if ((projection.Count > 1
                          // Cosmos does not support "AS Value" projections, specifically for the alias "Value"
-                         || projection is [{ Alias: var alias }] && alias.Equals("value", StringComparison.OrdinalIgnoreCase))
+                         || projection is [{ Alias: string alias }] && alias.Equals("value", StringComparison.OrdinalIgnoreCase))
                      && projection.Any(p => !string.IsNullOrEmpty(p.Alias) && p.Alias != p.Name)
                      && !projection.Any(p => p.Expression is SqlFunctionExpression)) // Aggregates are not allowed
             {

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

@@ -27,13 +27,14 @@ public override Expression Process(Expression query)
 
         if (query is ShapedQueryExpression { QueryExpression: SelectExpression selectExpression })
         {
-            // Cosmos does not have nested select expression so this should be safe.
             selectExpression.ApplyProjection();
         }
 
         var afterValueConverterCompensation = new CosmosValueConverterCompensatingExpressionVisitor(sqlExpressionFactory).Visit(query);
         var afterAliases = queryCompilationContext.AliasManager.PostprocessAliases(afterValueConverterCompensation);
+        var afterExtraction = new CosmosReadItemAndPartitionKeysExtractor().ExtractPartitionKeysAndId(
+            queryCompilationContext, sqlExpressionFactory, afterAliases);
 
-        return afterAliases;
+        return afterExtraction;
     }
 }

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

@@ -14,20 +14,6 @@ public class CosmosQueryTranslationPreprocessor(
         CosmosQueryCompilationContext cosmosQueryCompilationContext)
     : QueryTranslationPreprocessor(dependencies, cosmosQueryCompilationContext)
 {
-    /// <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 override Expression NormalizeQueryableMethod(Expression query)
-    {
-        query = new CosmosQueryMetadataExtractingExpressionVisitor(cosmosQueryCompilationContext).Visit(query);
-        query = base.NormalizeQueryableMethod(query);
-
-        return query;
-    }
-
     /// <inheritdoc />
     protected override Expression ProcessQueryRoots(Expression expression)
         => new CosmosQueryRootProcessor(Dependencies, QueryCompilationContext).Visit(expression);