Cosmos: add partition key to the PK

Add model building tests for Cosmos and fix found bugs.

Fixes #20509

src/EFCore.Cosmos/Internal/CosmosModelValidator.cs

@@ -5,9 +5,12 @@
 using System.Collections.Generic;
 using System.Linq;
 using JetBrains.Annotations;
+using Microsoft.EntityFrameworkCore.Cosmos.Metadata.Internal;
 using Microsoft.EntityFrameworkCore.Diagnostics;
 using Microsoft.EntityFrameworkCore.Infrastructure;
 using Microsoft.EntityFrameworkCore.Metadata;
+using Microsoft.EntityFrameworkCore.Metadata.Conventions;
+using Microsoft.EntityFrameworkCore.Utilities;
 
 namespace Microsoft.EntityFrameworkCore.Cosmos.Internal
 {
@@ -40,6 +43,7 @@ public override void Validate(IModel model, IDiagnosticsLogger<DbLoggerCategory.
         {
             base.Validate(model, logger);
 
+            ValidateKeys(model, logger);
             ValidateSharedContainerCompatibility(model, logger);
             ValidateOnlyETagConcurrencyToken(model, logger);
         }
@@ -96,25 +100,11 @@ protected virtual void ValidateSharedContainerCompatibility(
             IEntityType firstEntityType = null;
             foreach (var entityType in mappedTypes)
             {
+                Check.DebugAssert(entityType.IsDocumentRoot(), "Only document roots expected here.");
                 var partitionKeyPropertyName = entityType.GetPartitionKeyPropertyName();
                 if (partitionKeyPropertyName != null)
                 {
                     var nextPartitionKeyProperty = entityType.FindProperty(partitionKeyPropertyName);
-                    if (nextPartitionKeyProperty == null)
-                    {
-                        throw new InvalidOperationException(
-                            CosmosStrings.PartitionKeyMissingProperty(entityType.DisplayName(), partitionKeyPropertyName));
-                    }
-
-                    var keyType = nextPartitionKeyProperty.GetTypeMapping().Converter?.ProviderClrType
-                                  ?? nextPartitionKeyProperty.ClrType;
-                    if (keyType != typeof(string))
-                    {
-                        throw new InvalidOperationException(
-                            CosmosStrings.PartitionKeyNonStringStoreType(
-                                partitionKeyPropertyName, entityType.DisplayName(), keyType.ShortDisplayName()));
-                    }
-
                     if (partitionKey == null)
                     {
                         if (firstEntityType != null)
@@ -209,5 +199,72 @@ protected virtual void ValidateOnlyETagConcurrencyToken(
                 }
             }
         }
+
+        /// <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>
+        protected virtual void ValidateKeys(
+            [NotNull] IModel model,
+            [NotNull] IDiagnosticsLogger<DbLoggerCategory.Model.Validation> logger)
+        {
+            foreach (var entityType in model.GetEntityTypes())
+            {
+                var primaryKey = entityType.FindPrimaryKey();
+                if (primaryKey == null
+                    || !entityType.IsDocumentRoot())
+                {
+                    continue;
+                }
+
+                var idProperty = entityType.GetProperties().FirstOrDefault(p => p.GetJsonPropertyName() == StoreKeyConvention.IdPropertyName);
+                if (idProperty == null)
+                {
+                    throw new InvalidOperationException(CosmosStrings.NoIdProperty(entityType.DisplayName()));
+                }
+
+                var idType = idProperty.GetTypeMapping().Converter?.ProviderClrType
+                    ?? idProperty.ClrType;
+                if (idType != typeof(string))
+                {
+                    throw new InvalidOperationException(
+                        CosmosStrings.IdNonStringStoreType(
+                            idProperty.Name, entityType.DisplayName(), idType.ShortDisplayName()));
+                }
+
+                if (!idProperty.IsKey())
+                {
+                    throw new InvalidOperationException(CosmosStrings.NoIdKey(entityType.DisplayName(), idProperty.Name));
+                }
+
+                var partitionKeyPropertyName = entityType.GetPartitionKeyPropertyName();
+                if (partitionKeyPropertyName != null)
+                {
+                    var partitionKey = entityType.FindProperty(partitionKeyPropertyName);
+                    if (partitionKey == null)
+                    {
+                        throw new InvalidOperationException(
+                            CosmosStrings.PartitionKeyMissingProperty(entityType.DisplayName(), partitionKeyPropertyName));
+                    }
+
+                    var partitionKeyType = partitionKey.GetTypeMapping().Converter?.ProviderClrType
+                        ?? partitionKey.ClrType;
+                    if (partitionKeyType != typeof(string))
+                    {
+                        throw new InvalidOperationException(
+                            CosmosStrings.PartitionKeyNonStringStoreType(
+                                partitionKeyPropertyName, entityType.DisplayName(), partitionKeyType.ShortDisplayName()));
+                    }
+
+                    if (!partitionKey.GetContainingKeys().Any(k => k.Properties.Contains(idProperty)))
+                    {
+                        throw new InvalidOperationException(CosmosStrings.NoPartitionKeyKey(
+                            entityType.DisplayName(), partitionKeyPropertyName, idProperty.Name));
+                    }
+                }
+            }
+        }
     }
 }

src/EFCore.Cosmos/Metadata/Conventions/CosmosDiscriminatorConvention.cs

@@ -42,7 +42,6 @@ public virtual void ProcessEntityTypeAdded(
 
             var entityType = entityTypeBuilder.Metadata;
             if (entityType.BaseType == null
-                && !entityType.GetDerivedTypes().Any()
                 && entityType.IsDocumentRoot())
             {
                 entityTypeBuilder.HasDiscriminator(typeof(string))
@@ -122,7 +121,7 @@ public override void ProcessEntityTypeBaseTypeChanged(
             }
             else
             {
-                discriminator = newBaseType.Builder?.HasDiscriminator(typeof(string));
+                discriminator = newBaseType.GetRootType().Builder?.HasDiscriminator(typeof(string));
 
                 if (newBaseType.BaseType == null)
                 {

src/EFCore.Cosmos/Metadata/Conventions/CosmosKeyDiscoveryConvention.cs

@@ -0,0 +1,75 @@
+// Copyright (c) .NET Foundation. All rights reserved.
+// Licensed under the Apache License, Version 2.0. See License.txt in the project root for license information.
+
+using System.Collections.Generic;
+using JetBrains.Annotations;
+using Microsoft.EntityFrameworkCore.Cosmos.Metadata.Internal;
+using Microsoft.EntityFrameworkCore.Metadata.Builders;
+using Microsoft.EntityFrameworkCore.Metadata.Conventions.Infrastructure;
+using Microsoft.EntityFrameworkCore.Utilities;
+
+namespace Microsoft.EntityFrameworkCore.Metadata.Conventions
+{
+    /// <summary>
+    ///     A convention that finds primary key property for the entity type based on the names
+    ///     and adds the partition key to it if present.
+    /// </summary>
+    public class CosmosKeyDiscoveryConvention :
+        KeyDiscoveryConvention,
+        IEntityTypeAnnotationChangedConvention
+    {
+        /// <summary>
+        ///     Creates a new instance of <see cref="KeyDiscoveryConvention" />.
+        /// </summary>
+        /// <param name="dependencies"> Parameter object containing dependencies for this convention. </param>
+        public CosmosKeyDiscoveryConvention([NotNull] ProviderConventionSetBuilderDependencies dependencies)
+            : base(dependencies)
+        {
+        }
+
+        /// <summary>
+        ///     Called after an annotation is changed on an entity type.
+        /// </summary>
+        /// <param name="entityTypeBuilder"> The builder for the entity type. </param>
+        /// <param name="name"> The annotation name. </param>
+        /// <param name="annotation"> The new annotation. </param>
+        /// <param name="oldAnnotation"> The old annotation.  </param>
+        /// <param name="context"> Additional information associated with convention execution. </param>
+        public virtual void ProcessEntityTypeAnnotationChanged(
+            IConventionEntityTypeBuilder entityTypeBuilder,
+            string name,
+            IConventionAnnotation annotation,
+            IConventionAnnotation oldAnnotation,
+            IConventionContext<IConventionAnnotation> context)
+        {
+            Check.NotNull(entityTypeBuilder, nameof(entityTypeBuilder));
+            Check.NotEmpty(name, nameof(name));
+            Check.NotNull(context, nameof(context));
+
+            if (name == CosmosAnnotationNames.PartitionKeyName)
+            {
+                TryConfigurePrimaryKey(entityTypeBuilder);
+            }
+        }
+
+        /// <inheritdoc />
+        protected override void ProcessKeyProperties(IList<IConventionProperty> keyProperties, IConventionEntityType entityType)
+        {
+            if (keyProperties.Count == 0)
+            {
+                return;
+            }
+
+            var partitionKey = entityType.GetPartitionKeyPropertyName();
+            if (partitionKey != null)
+            {
+                var partitionKeyProperty = entityType.FindProperty(partitionKey);
+                if (partitionKeyProperty != null
+                    && !keyProperties.Contains(partitionKeyProperty))
+                {
+                    keyProperties.Add(partitionKeyProperty);
+                }
+            }
+        }
+    }
+}

src/EFCore.Cosmos/Metadata/Conventions/Internal/CosmosConventionSetBuilder.cs

@@ -41,23 +41,39 @@ public override ConventionSet CreateConventionSet()
 
             conventionSet.ModelFinalizingConventions.Add(new ETagPropertyConvention());
 
-            var discriminatorConvention = new CosmosDiscriminatorConvention(Dependencies);
             var storeKeyConvention = new StoreKeyConvention(Dependencies);
+            var discriminatorConvention = new CosmosDiscriminatorConvention(Dependencies);
+            var keyDiscoveryConvention = new CosmosKeyDiscoveryConvention(Dependencies);
             conventionSet.EntityTypeAddedConventions.Add(storeKeyConvention);
             conventionSet.EntityTypeAddedConventions.Add(discriminatorConvention);
+            ReplaceConvention(conventionSet.EntityTypeAddedConventions, (KeyDiscoveryConvention)keyDiscoveryConvention);
 
             ReplaceConvention(conventionSet.EntityTypeRemovedConventions, (DiscriminatorConvention)discriminatorConvention);
 
             conventionSet.EntityTypeBaseTypeChangedConventions.Add(storeKeyConvention);
             ReplaceConvention(conventionSet.EntityTypeBaseTypeChangedConventions, (DiscriminatorConvention)discriminatorConvention);
+            ReplaceConvention(conventionSet.EntityTypeBaseTypeChangedConventions, (KeyDiscoveryConvention)keyDiscoveryConvention);
+
+            ReplaceConvention(conventionSet.PropertyAddedConventions, (KeyDiscoveryConvention)keyDiscoveryConvention);
+
+            ReplaceConvention(conventionSet.KeyRemovedConventions, (KeyDiscoveryConvention)keyDiscoveryConvention);
+
+            ReplaceConvention(conventionSet.ForeignKeyAddedConventions, (KeyDiscoveryConvention)keyDiscoveryConvention);
 
             conventionSet.ForeignKeyRemovedConventions.Add(discriminatorConvention);
             conventionSet.ForeignKeyRemovedConventions.Add(storeKeyConvention);
+            ReplaceConvention(conventionSet.ForeignKeyRemovedConventions, (KeyDiscoveryConvention)keyDiscoveryConvention);
+
+            ReplaceConvention(conventionSet.ForeignKeyPropertiesChangedConventions, (KeyDiscoveryConvention)keyDiscoveryConvention);
+
+            ReplaceConvention(conventionSet.ForeignKeyUniquenessChangedConventions, (KeyDiscoveryConvention)keyDiscoveryConvention);
 
             conventionSet.ForeignKeyOwnershipChangedConventions.Add(discriminatorConvention);
             conventionSet.ForeignKeyOwnershipChangedConventions.Add(storeKeyConvention);
+            ReplaceConvention(conventionSet.ForeignKeyOwnershipChangedConventions, (KeyDiscoveryConvention)keyDiscoveryConvention);
 
             conventionSet.EntityTypeAnnotationChangedConventions.Add(storeKeyConvention);
+            conventionSet.EntityTypeAnnotationChangedConventions.Add(keyDiscoveryConvention);
 
             return conventionSet;
         }

src/EFCore.Cosmos/Metadata/Conventions/StoreKeyConvention.cs

@@ -84,9 +84,17 @@ private static void Process(IConventionEntityTypeBuilder entityTypeBuilder)
                     if (partitionKey != null)
                     {
                         var partitionKeyProperty = entityType.FindProperty(partitionKey);
-                        if (partitionKeyProperty != null)
+                        if (partitionKeyProperty == null)
                         {
-                            newKey = entityTypeBuilder.HasKey(new[] { idProperty, partitionKeyProperty })?.Metadata;
+                            newKey = entityTypeBuilder.HasKey(new[] { idProperty })?.Metadata;
+                        }
+                        else
+                        {
+                            if (entityType.FindKey(new[] { partitionKeyProperty, idProperty }) == null)
+                            {
+                                newKey = entityTypeBuilder.HasKey(new[] { idProperty, partitionKeyProperty })?.Metadata;
+                            }
+                            entityTypeBuilder.HasNoKey(new[] { idProperty });
                         }
                     }
                     else
@@ -115,8 +123,8 @@ private static void Process(IConventionEntityTypeBuilder entityTypeBuilder)
                 && !entityType.IsKeyless)
             {
                 var jObjectProperty = entityTypeBuilder.Property(typeof(JObject), JObjectPropertyName);
-                jObjectProperty.ToJsonProperty("");
-                jObjectProperty.ValueGenerated(ValueGenerated.OnAddOrUpdate);
+                jObjectProperty?.ToJsonProperty("");
+                jObjectProperty?.ValueGenerated(ValueGenerated.OnAddOrUpdate);
             }
             else
             {