Add Cosmos db manual session token management

src/EFCore.Cosmos/Extensions/CosmosDatabaseFacadeExtensions.cs

@@ -25,6 +25,79 @@ 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 for the default container in the context, or <see langword="null"/> if none present.</returns>
+    public static string? GetSessionToken(this DatabaseFacade databaseFacade)
+        => GetSessionTokenStorage(databaseFacade).GetDefaultContainerTrackedToken();
+
+    /// <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).GetTrackedTokens();
+
+    /// <summary>
+    ///     Sets 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 set.</param>
+    public static void UseSessionToken(this DatabaseFacade databaseFacade, string sessionToken)
+        => GetSessionTokenStorage(databaseFacade).SetDefaultContainerSessionToken(sessionToken);
+
+    /// <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).AppendDefaultContainerSessionToken(sessionToken);
+
+    /// <summary>
+    ///     Sets the composite sessions token per container for this <see cref="DbContext" /> with the tokens specified in <paramref name="sessionTokens"/>.
+    /// </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 set per container.</param>
+    public static void UseSessionTokens(this DatabaseFacade databaseFacade, IReadOnlyDictionary<string, string?> sessionTokens)
+    {
+        var sessionTokenStorage = GetSessionTokenStorage(databaseFacade, sessionTokens);
+
+        sessionTokenStorage.SetSessionTokens(sessionTokens);
+    }
+
+    /// <summary>
+    ///     Appends the composite sessions token per container for this <see cref="DbContext" /> with the tokens specified in <paramref name="sessionTokens"/>.
+    /// </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)
+    {
+        var sessionTokenStorage = GetSessionTokenStorage(databaseFacade, (IReadOnlyDictionary<string, string?>)sessionTokens);
+
+        sessionTokenStorage.AppendSessionTokens(sessionTokens);
+    }
+
+    private static ISessionTokenStorage GetSessionTokenStorage(DatabaseFacade databaseFacade, IReadOnlyDictionary<string, string?>? sessionTokens = null)
+    {
+        var db = GetService<IDatabase>(databaseFacade);
+        if (db is not CosmosDatabaseWrapper dbWrapper)
+        {
+            throw new InvalidOperationException(CosmosStrings.CosmosNotInUse);
+        }
+
+        return dbWrapper.SessionTokenStorage;
+    }
+
     private static TService GetService<TService>(IInfrastructure<IServiceProvider> databaseFacade)
         where TService : class
     {

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>()
@@ -121,7 +122,8 @@ public static IServiceCollection AddEntityFrameworkCosmos(this IServiceCollectio
                 .TryAddScoped<ISqlExpressionFactory, SqlExpressionFactory>()
                 .TryAddScoped<IMemberTranslatorProvider, CosmosMemberTranslatorProvider>()
                 .TryAddScoped<IMethodCallTranslatorProvider, CosmosMethodCallTranslatorProvider>()
-                .TryAddScoped<ICosmosClientWrapper, CosmosClientWrapper>());
+                .TryAddScoped<ICosmosClientWrapper, CosmosClientWrapper>()
+                .TryAddSingleton<ISessionTokenStorageFactory, SessionTokenStorageFactory>());
 
         builder.TryAddCoreServices();
 

src/EFCore.Cosmos/Infrastructure/CosmosDbContextOptionsBuilder.cs

@@ -3,6 +3,7 @@
 
 using System.ComponentModel;
 using System.Net;
+using Microsoft.EntityFrameworkCore.Cosmos.Infrastructure;
 using Microsoft.EntityFrameworkCore.Cosmos.Infrastructure.Internal;
 
 namespace Microsoft.EntityFrameworkCore.Infrastructure;
@@ -211,6 +212,23 @@ public virtual CosmosDbContextOptionsBuilder MaxRequestsPerTcpConnection(int req
     public virtual CosmosDbContextOptionsBuilder ContentResponseOnWriteEnabled(bool enabled = true)
         => WithOption(e => e.ContentResponseOnWriteEnabled(Check.NotNull(enabled)));
 
+
+    /// <summary>
+    ///     Sets the <see cref="Cosmos.Infrastructure.SessionTokenManagementMode"/> to use.
+    ///     By default, <see cref="SessionTokenManagementMode.FullyAutomatic"/> will be used.
+    ///     Any other mode 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.
+    ///     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="mode">The <see cref="Cosmos.Infrastructure.SessionTokenManagementMode"/> to use.</param>
+    public virtual CosmosDbContextOptionsBuilder SessionTokenManagementMode(SessionTokenManagementMode mode)
+        => WithOption(e => e.WithSessionTokenManagementMode(mode));
+
     /// <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 SessionTokenManagementMode _sessionTokenManagementMode = SessionTokenManagementMode.FullyAutomatic;
 
     /// <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;
+        _sessionTokenManagementMode = copyFrom._sessionTokenManagementMode;
     }
 
     /// <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 SessionTokenManagementMode SessionTokenManagementMode
+        => _sessionTokenManagementMode;
+
+    /// <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 WithSessionTokenManagementMode(SessionTokenManagementMode mode)
+    {
+        var clone = Clone();
+
+        clone._sessionTokenManagementMode = mode;
+
+        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._sessionTokenManagementMode);
 
                 _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._sessionTokenManagementMode == otherInfo.Extension._sessionTokenManagementMode;
 
         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 SessionTokenManagementMode SessionTokenManagementMode { 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;
+            SessionTokenManagementMode = cosmosOptions.SessionTokenManagementMode;
         }
     }
 
@@ -208,6 +217,7 @@ public virtual void Validate(IDbContextOptions options)
                 || MaxTcpConnectionsPerEndpoint != cosmosOptions.MaxTcpConnectionsPerEndpoint
                 || MaxRequestsPerTcpConnection != cosmosOptions.MaxRequestsPerTcpConnection
                 || HttpClientFactory != cosmosOptions.HttpClientFactory
+                || SessionTokenManagementMode != cosmosOptions.SessionTokenManagementMode
             ))
         {
             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>
+    SessionTokenManagementMode SessionTokenManagementMode { get; }
 }

src/EFCore.Cosmos/Infrastructure/SessionTokenManagementMode.cs

@@ -0,0 +1,42 @@
+// Licensed to the .NET Foundation under one or more agreements.
+// The .NET Foundation licenses this file to you under the MIT license.
+
+namespace Microsoft.EntityFrameworkCore.Cosmos.Infrastructure;
+
+/// <summary>
+///     Defines the behaviour of EF regarding the management of Cosmos DB session tokens.
+/// </summary>
+/// <remarks>
+///     See <see href="https://docs.azure.cn/en-us/cosmos-db/consistency-levels#session-consistency">Consistency level choices</see> for more info.
+/// </remarks>
+public enum SessionTokenManagementMode
+{
+    /// <summary>
+    ///     The default mode.
+    ///     Uses the underlying Cosmos DB SDK automatic session token management.
+    ///     EF will not track or parse session tokens returned from Cosmos DB. <see cref="CosmosDatabaseFacadeExtensions.UseSessionTokens(DatabaseFacade, IReadOnlyDictionary{string, string?})"/> and <see cref="CosmosDatabaseFacadeExtensions.GetSessionTokens(DatabaseFacade)"/> methods will throw when invoked.
+    ///     Use this mode when every request for the same user will land on the same instance of your app.
+    ///     This means you either have 1 application instance, or maintain session affinity between requests.
+    ///     Otherwhise, use of one of the other modes is required to guarantee session consistency between requests.
+    /// </summary>
+    FullyAutomatic,
+
+    /// <summary>
+    ///     Allows the usage of <see cref="CosmosDatabaseFacadeExtensions.UseSessionTokens(DatabaseFacade, IReadOnlyDictionary{string, string?})"/> to overwrite the default Cosmos DB SDK automatic session token management by use of the <see cref="CosmosDatabaseFacadeExtensions.UseSessionTokens(DatabaseFacade, IReadOnlyDictionary{string, string?})"/> method on a <see cref="DbContext.Database"/> instance.
+    ///     If <see cref="CosmosDatabaseFacadeExtensions.UseSessionTokens(DatabaseFacade, IReadOnlyDictionary{string, string?})"/> has not been invoked for an container, the default Cosmos DB SDK automatic session token management will be used.
+    ///     EF will track and parse session tokens returned from Cosmos DB, which can be retrieved via <see cref="CosmosDatabaseFacadeExtensions.GetSessionTokens(DatabaseFacade)"/>.
+    /// </summary>
+    SemiAutomatic,
+
+    /// <summary>
+    ///     Fully overwrites the Cosmos DB SDK automatic session token management, and only uses session tokens specified via <see cref="CosmosDatabaseFacadeExtensions.UseSessionTokens(DatabaseFacade, IReadOnlyDictionary{string, string?})"/>.
+    ///     If <see cref="CosmosDatabaseFacadeExtensions.UseSessionTokens(DatabaseFacade, IReadOnlyDictionary{string, string?})"/> has not been invoked for an container, no session token will be used.
+    ///     EF will track and parse session tokens returned from Cosmos DB, which can be retrieved via <see cref="CosmosDatabaseFacadeExtensions.GetSessionTokens(DatabaseFacade)"/>.
+    /// </summary>
+    Manual,
+
+    /// <summary>
+    ///     Same as <see cref="Manual"/>, but will throw an exception if <see cref="CosmosDatabaseFacadeExtensions.UseSessionTokens(DatabaseFacade, IReadOnlyDictionary{string, string?})"/> was not invoked before executong a read.
+    /// </summary>
+    EnforcedManual
+}