Add Cosmos db manual session token management

src/EFCore.Cosmos/Extensions/CosmosDatabaseFacadeExtensions.cs

@@ -3,6 +3,7 @@
 
 using Microsoft.EntityFrameworkCore.Cosmos.Infrastructure.Internal;
 using Microsoft.EntityFrameworkCore.Cosmos.Internal;
+using Microsoft.EntityFrameworkCore.Cosmos.Storage;
 using Microsoft.EntityFrameworkCore.Cosmos.Storage.Internal;
 
 // ReSharper disable once CheckNamespace
@@ -25,6 +26,58 @@ public static class CosmosDatabaseFacadeExtensions
     public static CosmosClient GetCosmosClient(this DatabaseFacade databaseFacade)
         => GetService<ISingletonCosmosClientWrapper>(databaseFacade).Client;
 
+    /// <summary>
+    ///     Gets the composite session token for the default container for this <see cref="DbContext" />.
+    /// </summary>
+    /// <remarks>Use this when using only 1 container in the same <see cref="DbContext"/></remarks>
+    /// <param name="databaseFacade">The <see cref="DatabaseFacade" /> for the context.</param>
+    /// <returns>The session token dictionary.</returns>
+    public static string? GetSessionToken(this DatabaseFacade databaseFacade)
+        => GetSessionTokenStorage(databaseFacade).GetSessionToken();
+
+    /// <summary>
+    ///     Gets a dictionary that contains the composite session token per container for this <see cref="DbContext" />.
+    /// </summary>
+    /// <remarks>Use this when using multiple containers in the same <see cref="DbContext"/></remarks>
+    /// <param name="databaseFacade">The <see cref="DatabaseFacade" /> for the context.</param>
+    /// <returns>The session token dictionary.</returns>
+    public static IReadOnlyDictionary<string, string?> GetSessionTokens(this DatabaseFacade databaseFacade)
+        => GetSessionTokenStorage(databaseFacade);
+
+    /// <summary>
+    /// Appends the composite session token for the default container for this <see cref="DbContext" />.
+    /// </summary>
+    /// <remarks>Use this when using only 1 container in the same <see cref="DbContext"/></remarks>
+    /// <param name="databaseFacade">The <see cref="DatabaseFacade" /> for the context.</param>
+    /// <param name="sessionToken">The session token to append.</param>
+    public static void AppendSessionToken(this DatabaseFacade databaseFacade, string sessionToken)
+        => GetSessionTokenStorage(databaseFacade).AppendSessionToken(sessionToken);
+
+    /// <summary>
+    ///     Appends the composite session token per container for this <see cref="DbContext" />.
+    /// </summary>
+    /// <remarks>Use this when using multiple containers in the same <see cref="DbContext"/></remarks>
+    /// <param name="databaseFacade">The <see cref="DatabaseFacade" /> for the context.</param>
+    /// <param name="sessionTokens">The session tokens to append per container.</param>
+    public static void AppendSessionTokens(this DatabaseFacade databaseFacade, IReadOnlyDictionary<string, string> sessionTokens)
+        => GetSessionTokenStorage(databaseFacade).AppendSessionTokens(sessionTokens);
+
+    private static SessionTokenStorage GetSessionTokenStorage(DatabaseFacade databaseFacade)
+    {
+        var db = GetService<IDatabase>(databaseFacade);
+        if (db is not CosmosDatabaseWrapper dbWrapper)
+        {
+            throw new InvalidOperationException(CosmosStrings.CosmosNotInUse);
+        }
+
+        if (dbWrapper.SessionTokenStorage is not SessionTokenStorage sts)
+        {
+            throw new InvalidOperationException(CosmosStrings.EnableManualSessionTokenManagement);
+        }
+
+        return sts;
+    }
+
     private static TService GetService<TService>(IInfrastructure<IServiceProvider> databaseFacade)
         where TService : class
     {

src/EFCore.Cosmos/Extensions/CosmosModelExtensions.cs

@@ -23,6 +23,14 @@ public static class CosmosModelExtensions
     public static string? GetDefaultContainer(this IReadOnlyModel model)
         => (string?)model[CosmosAnnotationNames.ContainerName];
 
+    /// <summary>
+    ///     Returns the all container names used in the model.
+    /// </summary>
+    /// <param name="model">The model.</param>
+    /// <returns>A set of the names of the containers used in the model.</returns>
+    public static HashSet<string> GetContainerNames(this IReadOnlyModel model)
+        => (HashSet<string>)model[CosmosAnnotationNames.ContainerNames]!;
+
     /// <summary>
     ///     Sets the default container name.
     /// </summary>

src/EFCore.Cosmos/Extensions/CosmosServiceCollectionExtensions.cs

@@ -97,6 +97,7 @@ public static IServiceCollection AddEntityFrameworkCosmos(this IServiceCollectio
             .TryAdd<LoggingDefinitions, CosmosLoggingDefinitions>()
             .TryAdd<IDatabaseProvider, DatabaseProvider<CosmosOptionsExtension>>()
             .TryAdd<IDatabase, CosmosDatabaseWrapper>()
+            .TryAdd<IResettableService, CosmosDatabaseWrapper>(sp => (CosmosDatabaseWrapper)sp.GetRequiredService<IDatabase>())
             .TryAdd<IExecutionStrategyFactory, CosmosExecutionStrategyFactory>()
             .TryAdd<IDbContextTransactionManager, CosmosTransactionManager>()
             .TryAdd<IModelValidator, CosmosModelValidator>()

src/EFCore.Cosmos/Infrastructure/CosmosDbContextOptionsBuilder.cs

@@ -211,6 +211,23 @@ public virtual CosmosDbContextOptionsBuilder MaxRequestsPerTcpConnection(int req
     public virtual CosmosDbContextOptionsBuilder ContentResponseOnWriteEnabled(bool enabled = true)
         => WithOption(e => e.ContentResponseOnWriteEnabled(Check.NotNull(enabled)));
 
+
+    /// <summary>
+    ///     Sets the boolean to track and manage session tokens for requests made to Cosmos DB
+    ///     and being able to access them via the <see cref="CosmosDatabaseFacadeExtensions.GetSessionTokens(DatabaseFacade)"/> and <see cref="CosmosDatabaseFacadeExtensions.AppendSessionTokens(DatabaseFacade, IReadOnlyDictionary{string, string})"/> methods.
+    ///     This is only relevant when your application needs to manage session tokens manually.
+    ///     For example: If you're using a round-robin load balancer that doesn't maintain session affinity between requests.
+    ///     Enabling manual session token management can break session consistency when not handled properly.
+    ///     See <see href="https://learn.microsoft.com/en-us/azure/cosmos-db/nosql/how-to-manage-consistency?tabs=portal%2Cdotnetv2%2Capi-async#utilize-session-tokens">Utilize session tokens</see> for more details.
+    /// </summary>
+    /// <remarks>
+    ///     See <see href="https://aka.ms/efcore-docs-dbcontext-options">Using DbContextOptions</see>, and
+    ///     <see href="https://aka.ms/efcore-docs-cosmos">Accessing Azure Cosmos DB with EF Core</see> for more information and examples.
+    /// </remarks>
+    /// <param name="enabled"><see langword="true" /> to track and manually manage session tokens in EF.</param>
+    public virtual CosmosDbContextOptionsBuilder ManualSessionTokenManagementEnabled(bool enabled = true)
+        => WithOption(e => e.ManualSessionTokenManagementEnabled(enabled));
+
     /// <summary>
     ///     Sets an option by cloning the extension used to store the settings. This ensures the builder
     ///     does not modify options that are already in use elsewhere.

src/EFCore.Cosmos/Infrastructure/Internal/CosmosDbOptionExtension.cs

@@ -36,6 +36,7 @@ public class CosmosOptionsExtension : IDbContextOptionsExtension
     private bool? _enableContentResponseOnWrite;
     private DbContextOptionsExtensionInfo? _info;
     private Func<HttpClient>? _httpClientFactory;
+    private bool _enableManualSessionTokenManagement;
 
     /// <summary>
     ///     This is an internal API that supports the Entity Framework Core infrastructure and not subject to
@@ -73,6 +74,7 @@ protected CosmosOptionsExtension(CosmosOptionsExtension copyFrom)
         _maxTcpConnectionsPerEndpoint = copyFrom._maxTcpConnectionsPerEndpoint;
         _maxRequestsPerTcpConnection = copyFrom._maxRequestsPerTcpConnection;
         _httpClientFactory = copyFrom._httpClientFactory;
+        _enableManualSessionTokenManagement = copyFrom._enableManualSessionTokenManagement;
     }
 
     /// <summary>
@@ -564,6 +566,30 @@ public virtual CosmosOptionsExtension WithHttpClientFactory(Func<HttpClient>? ht
         return clone;
     }
 
+    /// <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 bool EnableManualSessionTokenManagement
+        => _enableManualSessionTokenManagement;
+
+    /// <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 CosmosOptionsExtension ManualSessionTokenManagementEnabled(bool enabled)
+    {
+        var clone = Clone();
+
+        clone._enableManualSessionTokenManagement = enabled;
+
+        return clone;
+    }
+
     /// <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
@@ -632,6 +658,7 @@ public override int GetServiceProviderHashCode()
                 hashCode.Add(Extension._maxTcpConnectionsPerEndpoint);
                 hashCode.Add(Extension._maxRequestsPerTcpConnection);
                 hashCode.Add(Extension._httpClientFactory);
+                hashCode.Add(Extension._enableManualSessionTokenManagement);
 
                 _serviceProviderHash = hashCode.ToHashCode();
             }
@@ -656,7 +683,8 @@ public override bool ShouldUseSameServiceProvider(DbContextOptionsExtensionInfo
                 && Extension._gatewayModeMaxConnectionLimit == otherInfo.Extension._gatewayModeMaxConnectionLimit
                 && Extension._maxTcpConnectionsPerEndpoint == otherInfo.Extension._maxTcpConnectionsPerEndpoint
                 && Extension._maxRequestsPerTcpConnection == otherInfo.Extension._maxRequestsPerTcpConnection
-                && Extension._httpClientFactory == otherInfo.Extension._httpClientFactory;
+                && Extension._httpClientFactory == otherInfo.Extension._httpClientFactory
+                && Extension._enableManualSessionTokenManagement == otherInfo.Extension._enableManualSessionTokenManagement;
 
         public override void PopulateDebugInfo(IDictionary<string, string> debugInfo)
         {

src/EFCore.Cosmos/Infrastructure/Internal/CosmosSingletonOptions.cs

@@ -151,6 +151,14 @@ public class CosmosSingletonOptions : ICosmosSingletonOptions
     /// </summary>
     public virtual Func<HttpClient>? HttpClientFactory { get; private 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 bool EnableManualSessionTokenManagement { get; private 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
@@ -178,6 +186,7 @@ public virtual void Initialize(IDbContextOptions options)
             MaxTcpConnectionsPerEndpoint = cosmosOptions.MaxTcpConnectionsPerEndpoint;
             MaxRequestsPerTcpConnection = cosmosOptions.MaxRequestsPerTcpConnection;
             HttpClientFactory = cosmosOptions.HttpClientFactory;
+            EnableManualSessionTokenManagement = cosmosOptions.EnableManualSessionTokenManagement;
         }
     }
 
@@ -208,6 +217,7 @@ public virtual void Validate(IDbContextOptions options)
                 || MaxTcpConnectionsPerEndpoint != cosmosOptions.MaxTcpConnectionsPerEndpoint
                 || MaxRequestsPerTcpConnection != cosmosOptions.MaxRequestsPerTcpConnection
                 || HttpClientFactory != cosmosOptions.HttpClientFactory
+                || EnableManualSessionTokenManagement != cosmosOptions.EnableManualSessionTokenManagement
             ))
         {
             throw new InvalidOperationException(

src/EFCore.Cosmos/Infrastructure/Internal/ICosmosSingletonOptions.cs

@@ -155,4 +155,12 @@ public interface ICosmosSingletonOptions : ISingletonOptions
     ///     doing so can result in application failures when updating to a new Entity Framework Core release.
     /// </summary>
     Func<HttpClient>? HttpClientFactory { get; }
+
+    /// <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>
+    bool EnableManualSessionTokenManagement { get; }
 }

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

@@ -44,8 +44,19 @@ protected override void ProcessModelAnnotations(
         {
             annotations.Remove(CosmosAnnotationNames.Throughput);
         }
+
+        // @TODO: Is this the right place for this?
+        annotations.Add(CosmosAnnotationNames.ContainerNames, GetContainerNames(model));
     }
 
+    private HashSet<string> GetContainerNames(IModel model)
+        => model.GetEntityTypes()
+            .Where(et => et.FindPrimaryKey() != null)
+            .Select(et => et.GetContainer())
+            .Where(container => container != null)
+            .Distinct()!
+            .ToHashSet()!;
+
     /// <summary>
     ///     Updates the entity type annotations that will be set on the read-only object.
     /// </summary>

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

@@ -27,6 +27,14 @@ public static class CosmosAnnotationNames
     /// </summary>
     public const string ContainerName = Prefix + "ContainerName";
 
+    /// <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 const string ContainerNames = Prefix + "ContainerNames"; // @TODO: is this the right way?
+
     /// <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

src/EFCore.Cosmos/Properties/CosmosStrings.resx

@@ -147,6 +147,10 @@
   <data name="ContainerContainingPropertyConflict" xml:space="preserve">
     <value>The entity type '{entityType}' is mapped to the container '{container}' but it is also configured as being contained in property '{property}'.</value>
   </data>
+  <data name="ContainerNameDoesNotExist" xml:space="preserve">
+    <value>The container with the name '{containerName}' does not exist.</value>
+    <comment>string</comment>
+  </data>
   <data name="ContainerNotOnRoot" xml:space="preserve">
     <value>An Azure Cosmos DB container name is defined on entity type '{entityType}', which inherits from '{baseEntityType}'. Container names must be defined on the root entity type of a hierarchy.</value>
   </data>
@@ -171,6 +175,9 @@
   <data name="ElementWithValueConverter" xml:space="preserve">
     <value>The property '{propertyType} {structuralType}.{property}' has element type '{elementType}', which requires a value converter. Elements types requiring value converters are not currently supported with the Azure Cosmos DB database provider.</value>
   </data>
+  <data name="EnableManualSessionTokenManagement" xml:space="preserve">
+    <value>Enable manual session token management using CosmosDbContextOptionsBuilder.ManualSessionTokenManagementEnabled to use this method.</value>
+  </data>
   <data name="ETagNonStringStoreType" xml:space="preserve">
     <value>The type of the etag property '{property}' on '{entityType}' is '{propertyType}'. All etag properties must be strings or have a string value converter.</value>
   </data>
@@ -350,6 +357,9 @@
   <data name="SaveChangesAutoTransactionBehaviorAlwaysTriggerAtomicity" xml:space="preserve">
     <value>When using AutoTransactionBehavior.Always with the Cosmos DB provider, only 1 entity can be saved at a time when using pre- or post- triggers to ensure atomicity.</value>
   </data>
+  <data name="SessionTokenCanNotBeWhiteSpace" xml:space="preserve">
+    <value>Session token can not be white space.</value>
+  </data>
   <data name="SingleFirstOrDefaultNotSupportedOnNonNullableQueries" xml:space="preserve">
     <value>SingleOrDefault and FirstOrDefault cannot be used Cosmos SQL does not allow Offset without Limit. Consider specifying a 'Take' operation on the query.</value>
   </data>