Cosmos: Add API to configure partition keys and use them in the update pipeline.

src/EFCore.Cosmos/Extensions/CosmosEntityTypeBuilderExtensions.cs

@@ -1,8 +1,12 @@
 // 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;
+using System.Linq.Expressions;
+using System.Reflection;
 using JetBrains.Annotations;
 using Microsoft.EntityFrameworkCore.Cosmos.Metadata.Internal;
+using Microsoft.EntityFrameworkCore.Internal;
 using Microsoft.EntityFrameworkCore.Metadata.Builders;
 using Microsoft.EntityFrameworkCore.Utilities;
 
@@ -160,5 +164,96 @@ public static bool ForCosmosCanSetProperty(
 
             return entityTypeBuilder.CanSetAnnotation(CosmosAnnotationNames.PropertyName, name, fromDataAnnotation);
         }
+
+        /// <summary>
+        ///     Configures the property that is used to store the partition key.
+        /// </summary>
+        /// <param name="entityTypeBuilder"> The builder for the entity type being configured. </param>
+        /// <param name="name"> The name of the partition key property. </param>
+        /// <returns> The same builder instance so that multiple calls can be chained. </returns>
+        public static EntityTypeBuilder ForCosmosHasPartitionKey(
+            [NotNull] this EntityTypeBuilder entityTypeBuilder,
+            [CanBeNull] string name)
+        {
+            entityTypeBuilder.Metadata.SetCosmosPartitionKeyPropertyName(name);
+
+            return entityTypeBuilder;
+        }
+
+        /// <summary>
+        ///     Configures the property that is used to store the partition key.
+        /// </summary>
+        /// <param name="entityTypeBuilder"> The builder for the entity type being configured. </param>
+        /// <param name="name"> The name of the partition key property. </param>
+        /// <returns> The same builder instance so that multiple calls can be chained. </returns>
+        public static EntityTypeBuilder<TEntity> ForCosmosHasPartitionKey<TEntity>(
+            [NotNull] this EntityTypeBuilder<TEntity> entityTypeBuilder,
+            [CanBeNull] string name)
+            where TEntity : class
+        {
+            entityTypeBuilder.Metadata.SetCosmosPartitionKeyPropertyName(name);
+
+            return entityTypeBuilder;
+        }
+
+        /// <summary>
+        ///     Configures the property that is used to store the partition key.
+        /// </summary>
+        /// <param name="entityTypeBuilder"> The builder for the entity type being configured. </param>
+        /// <param name="propertyExpression"> The  partition key property. </param>
+        /// <returns> The same builder instance so that multiple calls can be chained. </returns>
+        public static EntityTypeBuilder<TEntity> ForCosmosHasPartitionKey<TEntity, TProperty>(
+            [NotNull] this EntityTypeBuilder<TEntity> entityTypeBuilder,
+            [NotNull] Expression<Func<TEntity, TProperty>> propertyExpression)
+            where TEntity : class
+        {
+            Check.NotNull(propertyExpression, nameof(propertyExpression));
+
+            entityTypeBuilder.Metadata.SetCosmosPartitionKeyPropertyName(propertyExpression.GetPropertyAccess().GetSimpleMemberName());
+
+            return entityTypeBuilder;
+        }
+
+        /// <summary>
+        ///     Configures the property that is used to store the partition key.
+        /// </summary>
+        /// <param name="entityTypeBuilder"> The builder for the entity type being configured. </param>
+        /// <param name="name"> The name of the partition key property. </param>
+        /// <param name="fromDataAnnotation"> Indicates whether the configuration was specified using a data annotation. </param>
+        /// <returns>
+        ///     The same builder instance if the configuration was applied,
+        ///     <c>null</c> otherwise.
+        /// </returns>
+        public static IConventionEntityTypeBuilder ForCosmosHasPartitionKey(
+            [NotNull] this IConventionEntityTypeBuilder entityTypeBuilder,
+            [CanBeNull] string name,
+            bool fromDataAnnotation = false)
+        {
+            if (!entityTypeBuilder.ForCosmosCanSetPartitionKey(name, fromDataAnnotation))
+            {
+                return null;
+            }
+
+            entityTypeBuilder.Metadata.SetCosmosPartitionKeyPropertyName(name, fromDataAnnotation);
+
+            return entityTypeBuilder;
+        }
+
+        /// <summary>
+        ///     Returns a value indicating whether the property that is used to store the partition key can be set
+        ///     from the current configuration source
+        /// </summary>
+        /// <param name="entityTypeBuilder"> The builder for the entity type being configured. </param>
+        /// <param name="name"> The name of the partition key property. </param>
+        /// <param name="fromDataAnnotation"> Indicates whether the configuration was specified using a data annotation. </param>
+        /// <returns> <c>true</c> if the configuration can be applied. </returns>
+        public static bool ForCosmosCanSetPartitionKey(
+            [NotNull] this IConventionEntityTypeBuilder entityTypeBuilder, [CanBeNull] string name, bool fromDataAnnotation = false)
+        {
+            Check.NotNull(entityTypeBuilder, nameof(entityTypeBuilder));
+            Check.NullButNotEmpty(name, nameof(name));
+
+            return entityTypeBuilder.CanSetAnnotation(CosmosAnnotationNames.PartitionKeyName, name, fromDataAnnotation);
+        }
     }
 }

src/EFCore.Cosmos/Extensions/CosmosEntityTypeExtensions.cs

@@ -3,6 +3,7 @@
 
 using JetBrains.Annotations;
 using Microsoft.EntityFrameworkCore.Cosmos.Metadata.Internal;
+using Microsoft.EntityFrameworkCore.Cosmos.Storage.Internal;
 using Microsoft.EntityFrameworkCore.Metadata;
 using Microsoft.EntityFrameworkCore.Utilities;
 
@@ -104,5 +105,61 @@ public static void SetCosmosContainingPropertyName(
         public static ConfigurationSource? GetCosmosContainingPropertyNameConfigurationSource([NotNull] this IConventionEntityType entityType)
             => entityType.FindAnnotation(CosmosAnnotationNames.PropertyName)
                 ?.GetConfigurationSource();
+
+        /// <summary>
+        ///     Returns the name of the property that is used to store the partition key.
+        /// </summary>
+        /// <param name="entityType"> The entity type to get the partition key property name for. </param>
+        /// <returns> The name of the partition key property. </returns>
+        public static string GetCosmosPartitionKeyPropertyName([NotNull] this IEntityType entityType) =>
+            entityType[CosmosAnnotationNames.PartitionKeyName] as string;
+
+        /// <summary>
+        ///     Sets the name of the property that is used to store the partition key key.
+        /// </summary>
+        /// <param name="entityType"> The entity type to set the partition key property name for. </param>
+        /// <param name="name"> The name to set. </param>
+        public static void SetCosmosPartitionKeyPropertyName([NotNull] this IMutableEntityType entityType, [CanBeNull] string name)
+            => entityType.SetOrRemoveAnnotation(
+                CosmosAnnotationNames.PartitionKeyName,
+                Check.NullButNotEmpty(name, nameof(name)));
+
+        /// <summary>
+        ///     Sets the name of the property that is used to store the partition key.
+        /// </summary>
+        /// <param name="entityType"> The entity type to set the partition key property name for. </param>
+        /// <param name="name"> The name to set. </param>
+        /// <param name="fromDataAnnotation"> Indicates whether the configuration was specified using a data annotation. </param>
+        public static void SetCosmosPartitionKeyPropertyName(
+            [NotNull] this IConventionEntityType entityType, [CanBeNull] string name, bool fromDataAnnotation = false)
+            => entityType.SetOrRemoveAnnotation(
+                CosmosAnnotationNames.PartitionKeyName,
+                Check.NullButNotEmpty(name, nameof(name)),
+                fromDataAnnotation);
+
+        /// <summary>
+        ///     Gets the <see cref="ConfigurationSource" /> for the property that is used to store the partition key.
+        /// </summary>
+        /// <param name="entityType"> The entity type to find configuration source for. </param>
+        /// <returns> The <see cref="ConfigurationSource" /> for the partition key property. </returns>
+        public static ConfigurationSource? GetCosmosPartitionKeyPropertyNameConfigurationSource([NotNull] this IConventionEntityType entityType)
+            => entityType.FindAnnotation(CosmosAnnotationNames.PartitionKeyName)
+                ?.GetConfigurationSource();
+
+        /// <summary>
+        ///     Returns the store name of the property that is used to store the partition key.
+        /// </summary>
+        /// <param name="entityType"> The entity type to get the partition key property name for. </param>
+        /// <returns> The name of the partition key property. </returns>
+        public static string GetCosmosPartitionKeyStoreName([NotNull] this IEntityType entityType)
+        {
+            var name = entityType.GetCosmosPartitionKeyPropertyName();
+            if (name != null)
+            {
+                return entityType.FindProperty(name).GetCosmosPropertyName();
+            }
+
+            return CosmosClientWrapper.DefaultPartitionKey;
+        }
     }
 }

src/EFCore.Cosmos/Extensions/CosmosModelBuilderExtensions.cs

@@ -49,7 +49,7 @@ public static IConventionModelBuilder ForCosmosHasDefaultContainerName(
             [CanBeNull] string name,
             bool fromDataAnnotation = false)
         {
-            if (modelBuilder.ForCosmosCanSetDefaultContainerName(name, fromDataAnnotation))
+            if (!modelBuilder.ForCosmosCanSetDefaultContainerName(name, fromDataAnnotation))
             {
                 return null;
             }

src/EFCore.Cosmos/Extensions/CosmosServiceCollectionExtensions.cs

@@ -33,6 +33,7 @@ public static IServiceCollection AddEntityFrameworkCosmos([NotNull] this IServic
                 .TryAdd<IExecutionStrategyFactory, CosmosExecutionStrategyFactory>()
                 .TryAdd<IDbContextTransactionManager, CosmosTransactionManager>()
                 .TryAdd<IModelCustomizer, CosmosModelCustomizer>()
+                .TryAdd<IModelValidator, CosmosModelValidator>()
                 .TryAdd<IProviderConventionSetBuilder, CosmosConventionSetBuilder>()
                 .TryAdd<IDatabaseCreator, CosmosDatabaseCreator>()
                 .TryAdd<IQueryContextFactory, CosmosQueryContextFactory>()

src/EFCore.Cosmos/Infrastructure/CosmosDbContextOptionsBuilder.cs

@@ -1,4 +1,4 @@
-// Copyright (c) .NET Foundation. All rights reserved.
+// 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;

src/EFCore.Cosmos/Infrastructure/CosmosModelCustomizer.cs

@@ -1,4 +1,4 @@
-// Copyright (c) .NET Foundation. All rights reserved.
+// 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 Microsoft.EntityFrameworkCore.Infrastructure;

src/EFCore.Cosmos/Infrastructure/CosmosModelValidator.cs

@@ -0,0 +1,173 @@
+// 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;
+using System.Collections.Generic;
+using System.Linq;
+using JetBrains.Annotations;
+using Microsoft.EntityFrameworkCore.Cosmos.Internal;
+using Microsoft.EntityFrameworkCore.Diagnostics;
+using Microsoft.EntityFrameworkCore.Infrastructure;
+using Microsoft.EntityFrameworkCore.Metadata;
+
+namespace Microsoft.EntityFrameworkCore.Cosmos.Infrastructure
+{
+    public class CosmosModelValidator : ModelValidator
+    {
+        public CosmosModelValidator([NotNull] ModelValidatorDependencies dependencies)
+            : base(dependencies)
+        {
+        }
+
+        /// <summary>
+        ///     Validates a model, throwing an exception if any errors are found.
+        /// </summary>
+        /// <param name="model"> The model to validate. </param>
+        /// <param name="logger"> The logger to use. </param>
+        public override void Validate(IModel model, IDiagnosticsLogger<DbLoggerCategory.Model.Validation> logger)
+        {
+            base.Validate(model, logger);
+
+            ValidateSharedContainerCompatibility(model, logger);
+        }
+
+        /// <summary>
+        ///     Validates the mapping/configuration of shared containers in the model.
+        /// </summary>
+        /// <param name="model"> The model to validate. </param>
+        /// <param name="logger"> The logger to use. </param>
+        protected virtual void ValidateSharedContainerCompatibility(
+            [NotNull] IModel model,
+            [NotNull] IDiagnosticsLogger<DbLoggerCategory.Model.Validation> logger)
+        {
+            var containers = new Dictionary<string, List<IEntityType>>();
+            foreach (var entityType in model.GetEntityTypes().Where(et => et.FindPrimaryKey() != null))
+            {
+                var containerName = entityType.GetCosmosContainerName();
+
+                if (!containers.TryGetValue(containerName, out var mappedTypes))
+                {
+                    mappedTypes = new List<IEntityType>();
+                    containers[containerName] = mappedTypes;
+                }
+
+                mappedTypes.Add(entityType);
+            }
+
+            foreach (var containerMapping in containers)
+            {
+                var mappedTypes = containerMapping.Value;
+                var containerName = containerMapping.Key;
+                ValidateSharedContainerCompatibility(mappedTypes, containerName, logger);
+            }
+        }
+
+        /// <summary>
+        ///     Validates the compatibility of entity types sharing a given container.
+        /// </summary>
+        /// <param name="mappedTypes"> The mapped entity types. </param>
+        /// <param name="containerName"> The container name. </param>
+        /// <param name="logger"> The logger to use. </param>
+        protected virtual void ValidateSharedContainerCompatibility(
+            [NotNull] IReadOnlyList<IEntityType> mappedTypes,
+            [NotNull] string containerName,
+            [NotNull] IDiagnosticsLogger<DbLoggerCategory.Model.Validation> logger)
+        {
+            if (mappedTypes.Count == 1)
+            {
+                var entityType = mappedTypes[0];
+                var partitionKeyPropertyName = entityType.GetCosmosPartitionKeyPropertyName();
+                if (partitionKeyPropertyName != null)
+                {
+                    var nextPartitionKeyProperty = entityType.FindProperty(partitionKeyPropertyName);
+                    if (nextPartitionKeyProperty == null)
+                    {
+                        throw new InvalidOperationException(
+                            CosmosStrings.PartitionKeyMissingProperty(entityType.DisplayName(), partitionKeyPropertyName));
+                    }
+                }
+                return;
+            }
+
+            var discriminatorValues = new Dictionary<object, IEntityType>();
+            IProperty partitionKey = null;
+            IEntityType firstEntityType = null;
+            foreach (var entityType in mappedTypes)
+            {
+                var partitionKeyPropertyName = entityType.GetCosmosPartitionKeyPropertyName();
+                if (partitionKeyPropertyName != null)
+                {
+                    var nextPartitionKeyProperty = entityType.FindProperty(partitionKeyPropertyName);
+                    if (nextPartitionKeyProperty == null)
+                    {
+                        throw new InvalidOperationException(
+                            CosmosStrings.PartitionKeyMissingProperty(entityType.DisplayName(), partitionKeyPropertyName));
+                    }
+
+                    if (partitionKey == null)
+                    {
+                        if (firstEntityType != null)
+                        {
+                            throw new InvalidOperationException(CosmosStrings.NoPartitionKey(firstEntityType.DisplayName(), containerName));
+                        }
+                        partitionKey = nextPartitionKeyProperty;
+                    }
+                    else if (partitionKey.GetCosmosPropertyName() != nextPartitionKeyProperty.GetCosmosPropertyName())
+                    {
+                        throw new InvalidOperationException(
+                            CosmosStrings.PartitionKeyStoreNameMismatch(
+                                partitionKey.Name, firstEntityType.DisplayName(), partitionKey.GetCosmosPropertyName(),
+                                nextPartitionKeyProperty.Name, entityType.DisplayName(), nextPartitionKeyProperty.GetCosmosPropertyName()));
+                    }
+                    else if ((partitionKey.FindMapping().Converter?.ProviderClrType ?? partitionKey.ClrType)
+                      != (nextPartitionKeyProperty.FindMapping().Converter?.ProviderClrType ?? nextPartitionKeyProperty.ClrType))
+                    {
+                        throw new InvalidOperationException(
+                            CosmosStrings.PartitionKeyStoreTypeMismatch(
+                                partitionKey.Name,
+                                firstEntityType.DisplayName(),
+                                (partitionKey.FindMapping().Converter?.ProviderClrType ?? partitionKey.ClrType).ShortDisplayName(),
+                                nextPartitionKeyProperty.Name,
+                                entityType.DisplayName(),
+                                (nextPartitionKeyProperty.FindMapping().Converter?.ProviderClrType ?? nextPartitionKeyProperty.ClrType)
+                                    .ShortDisplayName()));
+                    }
+                }
+                else if (partitionKey != null)
+                {
+                    throw new InvalidOperationException(CosmosStrings.NoPartitionKey(entityType.DisplayName(), containerName));
+                }
+
+                if (firstEntityType == null)
+                {
+                    firstEntityType = entityType;
+                }
+
+                if (entityType.ClrType?.IsInstantiable() == true)
+                {
+                    if (entityType.GetDiscriminatorProperty() == null)
+                    {
+                        throw new InvalidOperationException(
+                            CosmosStrings.NoDiscriminatorProperty(entityType.DisplayName(), containerName));
+                    }
+
+                    var discriminatorValue = entityType.GetDiscriminatorValue();
+                    if (discriminatorValue == null)
+                    {
+                        throw new InvalidOperationException(
+                            CosmosStrings.NoDiscriminatorValue(entityType.DisplayName(), containerName));
+                    }
+
+                    if (discriminatorValues.TryGetValue(discriminatorValue, out var duplicateEntityType))
+                    {
+                        throw new InvalidOperationException(
+                            CosmosStrings.DuplicateDiscriminatorValue(
+                                entityType.DisplayName(), discriminatorValue, duplicateEntityType.DisplayName(), containerName));
+                    }
+
+                    discriminatorValues[discriminatorValue] = entityType;
+                }
+            }
+        }
+    }
+}

src/EFCore.Cosmos/Metadata/Internal/CosmosAnnotationNames.cs

@@ -14,5 +14,6 @@ public static class CosmosAnnotationNames
         public const string Prefix = "Cosmos:";
         public const string ContainerName = Prefix + "ContainerName";
         public const string PropertyName = Prefix + "PropertyName";
+        public const string PartitionKeyName = Prefix + "PartitionKeyName";
     }
 }