Set partition key on join entity type by convention.

Fixes #23491

src/EFCore.Cosmos/Extensions/CosmosEntityTypeExtensions.cs

@@ -155,6 +155,58 @@ public static void SetPartitionKeyPropertyName(
             => entityType.FindAnnotation(CosmosAnnotationNames.PartitionKeyName)
                 ?.GetConfigurationSource();
 
+        /// <summary>
+        ///     Returns the property that is used to store the partition key.
+        /// </summary>
+        /// <param name="entityType"> The entity type to get the partition key property for. </param>
+        /// <returns> The name of the partition key property. </returns>
+        public static IReadOnlyProperty? GetPartitionKeyProperty(this IReadOnlyEntityType entityType)
+        {
+            var partitionKeyPropertyName = entityType.GetPartitionKeyPropertyName();
+            return partitionKeyPropertyName == null
+                ? null
+                : entityType.FindProperty(partitionKeyPropertyName);
+        }
+
+        /// <summary>
+        ///     Returns the property that is used to store the partition key.
+        /// </summary>
+        /// <param name="entityType"> The entity type to get the partition key property for. </param>
+        /// <returns> The name of the partition key property. </returns>
+        public static IMutableProperty? GetPartitionKeyProperty(this IMutableEntityType entityType)
+        {
+            var partitionKeyPropertyName = entityType.GetPartitionKeyPropertyName();
+            return partitionKeyPropertyName == null
+                ? null
+                : entityType.FindProperty(partitionKeyPropertyName);
+        }
+
+        /// <summary>
+        ///     Returns the property that is used to store the partition key.
+        /// </summary>
+        /// <param name="entityType"> The entity type to get the partition key property for. </param>
+        /// <returns> The name of the partition key property. </returns>
+        public static IConventionProperty? GetPartitionKeyProperty(this IConventionEntityType entityType)
+        {
+            var partitionKeyPropertyName = entityType.GetPartitionKeyPropertyName();
+            return partitionKeyPropertyName == null
+                ? null
+                : entityType.FindProperty(partitionKeyPropertyName);
+        }
+
+        /// <summary>
+        ///     Returns the property that is used to store the partition key.
+        /// </summary>
+        /// <param name="entityType"> The entity type to get the partition key property for. </param>
+        /// <returns> The name of the partition key property. </returns>
+        public static IProperty? GetPartitionKeyProperty(this IEntityType entityType)
+        {
+            var partitionKeyPropertyName = entityType.GetPartitionKeyPropertyName();
+            return partitionKeyPropertyName == null
+                ? null
+                : entityType.FindProperty(partitionKeyPropertyName);
+        }
+
         /// <summary>
         ///     Returns the name of the property that is used to store the ETag.
         /// </summary>

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

@@ -1,7 +1,6 @@
 // Licensed to the .NET Foundation under one or more agreements.
 // The .NET Foundation licenses this file to you under the MIT license.
 
-using System.Collections.Generic;
 using Microsoft.EntityFrameworkCore.Cosmos.Metadata.Internal;
 using Microsoft.EntityFrameworkCore.Metadata.Builders;
 using Microsoft.EntityFrameworkCore.Metadata.Conventions.Infrastructure;
@@ -18,7 +17,7 @@ public class CosmosKeyDiscoveryConvention :
         IEntityTypeAnnotationChangedConvention
     {
         /// <summary>
-        ///     Creates a new instance of <see cref="KeyDiscoveryConvention" />.
+        ///     Creates a new instance of <see cref="CosmosKeyDiscoveryConvention" />.
         /// </summary>
         /// <param name="dependencies"> Parameter object containing dependencies for this convention. </param>
         public CosmosKeyDiscoveryConvention(ProviderConventionSetBuilderDependencies dependencies)

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

@@ -0,0 +1,195 @@
+// Licensed to the .NET Foundation under one or more agreements.
+// The .NET Foundation licenses this file to you under the MIT license.
+
+using Microsoft.EntityFrameworkCore.Cosmos.Metadata.Internal;
+using Microsoft.EntityFrameworkCore.Metadata.Builders;
+using Microsoft.EntityFrameworkCore.Metadata.Conventions.Infrastructure;
+using Microsoft.EntityFrameworkCore.Metadata.Internal;
+using Microsoft.EntityFrameworkCore.Utilities;
+
+namespace Microsoft.EntityFrameworkCore.Metadata.Conventions
+{
+    /// <summary>
+    ///     A convention that creates a join entity type for a many-to-many relationship
+    ///     and adds a partition key to it if the related types share one.
+    /// </summary>
+    public class CosmosManyToManyJoinEntityTypeConvention :
+        ManyToManyJoinEntityTypeConvention,
+        IEntityTypeAnnotationChangedConvention
+    {
+        /// <summary>
+        ///     Creates a new instance of <see cref="CosmosManyToManyJoinEntityTypeConvention" />.
+        /// </summary>
+        /// <param name="dependencies"> Parameter object containing dependencies for this convention. </param>
+        public CosmosManyToManyJoinEntityTypeConvention(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
+                || name == CosmosAnnotationNames.ContainerName)
+            {
+                foreach (var skipNavigation in entityTypeBuilder.Metadata.GetSkipNavigations())
+                {
+                    ProcessJoinPartitionKey(skipNavigation);
+                }
+            }
+        }
+
+        /// <inheritdoc />
+        public override void ProcessSkipNavigationForeignKeyChanged(
+            IConventionSkipNavigationBuilder skipNavigationBuilder,
+            IConventionForeignKey? foreignKey,
+            IConventionForeignKey? oldForeignKey,
+            IConventionContext<IConventionForeignKey> context)
+        {
+            base.ProcessSkipNavigationForeignKeyChanged(skipNavigationBuilder, foreignKey, oldForeignKey, context);
+
+            if (oldForeignKey != null)
+            {
+                ProcessJoinPartitionKey(skipNavigationBuilder.Metadata);
+            }
+        }
+
+        /// <inheritdoc />
+        protected override void CreateJoinEntityType(string joinEntityTypeName, IConventionSkipNavigation skipNavigation)
+        {
+            if (ShouldSharePartitionKey(skipNavigation))
+            {
+                var model = skipNavigation.DeclaringEntityType.Model;
+                var joinEntityTypeBuilder = model.Builder.SharedTypeEntity(joinEntityTypeName, typeof(Dictionary<string, object>))!;
+                ConfigurePartitionKeyJoinEntityType(skipNavigation, joinEntityTypeBuilder);
+            }
+            else
+            {
+                base.CreateJoinEntityType(joinEntityTypeName, skipNavigation);
+            }
+        }
+
+        private void ConfigurePartitionKeyJoinEntityType(IConventionSkipNavigation skipNavigation, IConventionEntityTypeBuilder joinEntityTypeBuilder)
+        {
+            var principalPartitionKey = skipNavigation.DeclaringEntityType.GetPartitionKeyProperty()!;
+            var partitionKey = joinEntityTypeBuilder.Property(principalPartitionKey.ClrType, principalPartitionKey.Name)!.Metadata;
+            joinEntityTypeBuilder.HasPartitionKey(partitionKey.Name);
+
+            CreateSkipNavigationForeignKey(skipNavigation, joinEntityTypeBuilder, partitionKey);
+            CreateSkipNavigationForeignKey(skipNavigation.Inverse!, joinEntityTypeBuilder, partitionKey);
+        }
+
+        private IConventionForeignKey CreateSkipNavigationForeignKey(
+            IConventionSkipNavigation skipNavigation,
+            IConventionEntityTypeBuilder joinEntityTypeBuilder,
+            IConventionProperty partitionKeyProperty)
+        {
+            if (skipNavigation.ForeignKey != null
+                && !skipNavigation.Builder.CanSetForeignKey(null))
+            {
+                return skipNavigation.ForeignKey;
+            }
+
+            var principalKey = skipNavigation.DeclaringEntityType.FindPrimaryKey();
+            if (principalKey == null
+                || principalKey.Properties.All(p => p.Name != partitionKeyProperty.Name))
+            {
+                return CreateSkipNavigationForeignKey(skipNavigation, joinEntityTypeBuilder);
+            }
+
+            if (skipNavigation.ForeignKey?.Properties.Contains(partitionKeyProperty) == true)
+            {
+                return skipNavigation.ForeignKey;
+            }
+
+            var dependentProperties = new IConventionProperty[principalKey.Properties.Count];
+            for (var i = 0; i < principalKey.Properties.Count; i++)
+            {
+                var principalProperty = principalKey.Properties[i];
+                if (principalProperty.Name == partitionKeyProperty.Name)
+                {
+                    dependentProperties[i] = partitionKeyProperty;
+                }
+                else
+                {
+                    dependentProperties[i] = joinEntityTypeBuilder.CreateUniqueProperty(
+                        principalProperty.ClrType, principalProperty.Name, required: true)!.Metadata;
+                }
+            }
+
+            var foreignKey = joinEntityTypeBuilder.HasRelationship(skipNavigation.DeclaringEntityType, dependentProperties, principalKey)!
+                .IsUnique(false)!
+                .Metadata;
+
+            skipNavigation.Builder.HasForeignKey(foreignKey);
+
+            return foreignKey;
+        }
+
+        private void ProcessJoinPartitionKey(IConventionSkipNavigation skipNavigation)
+        {
+            var inverseSkipNavigation = skipNavigation.Inverse;
+            if (skipNavigation.JoinEntityType != null
+                    && skipNavigation.IsCollection
+                    && inverseSkipNavigation != null
+                    && inverseSkipNavigation.IsCollection
+                    && inverseSkipNavigation.JoinEntityType == skipNavigation.JoinEntityType)
+            {
+                var joinEntityType = skipNavigation.JoinEntityType;
+                var joinEntityTypeBuilder = joinEntityType.Builder;
+                if (ShouldSharePartitionKey(skipNavigation))
+                {
+                    var principalPartitionKey = skipNavigation.DeclaringEntityType.GetPartitionKeyProperty()!;
+                    var partitionKey = joinEntityType.GetPartitionKeyProperty();
+                    if ((partitionKey != null
+                        && (!joinEntityTypeBuilder.CanSetPartitionKey(principalPartitionKey.Name)
+                            || (skipNavigation.ForeignKey!.Properties.Contains(partitionKey)
+                                && inverseSkipNavigation.ForeignKey!.Properties.Contains(partitionKey))))
+                        || !skipNavigation.Builder.CanSetForeignKey(null)
+                        || !inverseSkipNavigation.Builder.CanSetForeignKey(null))
+                    {
+                        return;
+                    }
+
+                    ConfigurePartitionKeyJoinEntityType(skipNavigation, joinEntityTypeBuilder);
+                }
+                else
+                {
+                    var partitionKey = joinEntityType.GetPartitionKeyProperty();
+                    if (partitionKey != null
+                        && joinEntityTypeBuilder.HasPartitionKey(null) != null
+                        && ((skipNavigation.ForeignKey!.Properties.Contains(partitionKey)
+                            && skipNavigation.Builder.CanSetForeignKey(null))
+                            || (inverseSkipNavigation.ForeignKey!.Properties.Contains(partitionKey)
+                                && inverseSkipNavigation.Builder.CanSetForeignKey(null))))
+                    {
+                        CreateSkipNavigationForeignKey(skipNavigation, joinEntityTypeBuilder);
+                        CreateSkipNavigationForeignKey(inverseSkipNavigation, joinEntityTypeBuilder);
+                    }
+                }
+            }
+        }
+
+        private bool ShouldSharePartitionKey(IConventionSkipNavigation skipNavigation)
+            => skipNavigation.DeclaringEntityType.GetContainer() == skipNavigation.TargetEntityType.GetContainer()
+                && skipNavigation.DeclaringEntityType.GetPartitionKeyPropertyName() != null
+                && skipNavigation.Inverse?.DeclaringEntityType.GetPartitionKeyPropertyName()
+                    == skipNavigation.DeclaringEntityType.GetPartitionKeyPropertyName();
+    }
+}

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

@@ -100,9 +100,19 @@ public override ConventionSet CreateConventionSet()
 
             ReplaceConvention(conventionSet.NavigationRemovedConventions, relationshipDiscoveryConvention);
 
+            ManyToManyJoinEntityTypeConvention manyToManyJoinEntityTypeConvention = new CosmosManyToManyJoinEntityTypeConvention(Dependencies);
+            ReplaceConvention(conventionSet.SkipNavigationAddedConventions, manyToManyJoinEntityTypeConvention);
+
+            ReplaceConvention(conventionSet.SkipNavigationRemovedConventions, manyToManyJoinEntityTypeConvention);
+
+            ReplaceConvention(conventionSet.SkipNavigationInverseChangedConventions, manyToManyJoinEntityTypeConvention);
+
+            ReplaceConvention(conventionSet.SkipNavigationForeignKeyChangedConventions, manyToManyJoinEntityTypeConvention);
+
             conventionSet.EntityTypeAnnotationChangedConventions.Add(discriminatorConvention);
             conventionSet.EntityTypeAnnotationChangedConventions.Add(storeKeyConvention);
             conventionSet.EntityTypeAnnotationChangedConventions.Add((CosmosKeyDiscoveryConvention)keyDiscoveryConvention);
+            conventionSet.EntityTypeAnnotationChangedConventions.Add((CosmosManyToManyJoinEntityTypeConvention)manyToManyJoinEntityTypeConvention);
 
             ReplaceConvention(conventionSet.PropertyAddedConventions, keyDiscoveryConvention);