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,22 @@ public static class CosmosDatabaseFacadeExtensions
     public static CosmosClient GetCosmosClient(this DatabaseFacade databaseFacade)
         => GetService<ISingletonCosmosClientWrapper>(databaseFacade).Client;
 
+    /// <summary>
+    ///     Gets <see cref="ISessionTokenStorage"/> used to manage the session tokens for this <see cref="DbContext" />.
+    /// </summary>
+    /// <param name="databaseFacade">The <see cref="DatabaseFacade" /> for the context.</param>
+    /// <returns>The <see cref="ISessionTokenStorage"/>.</returns>
+    public static ISessionTokenStorage GetSessionTokens(this DatabaseFacade databaseFacade)
+    {
+        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/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/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>
@@ -350,6 +354,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>

src/EFCore.Cosmos/Query/Internal/CosmosQueryCompilationContext.cs

@@ -1,6 +1,8 @@
 // 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.Storage;
+
 namespace Microsoft.EntityFrameworkCore.Cosmos.Query.Internal;
 
 /// <summary>
@@ -9,9 +11,19 @@ namespace Microsoft.EntityFrameworkCore.Cosmos.Query.Internal;
 ///     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 class CosmosQueryCompilationContext(QueryCompilationContextDependencies dependencies, bool async)
-    : QueryCompilationContext(dependencies, async)
+public class CosmosQueryCompilationContext : QueryCompilationContext
 {
+    /// <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 CosmosQueryCompilationContext(QueryCompilationContextDependencies dependencies, ISessionTokenStorage sessionTokenStorage, bool async) : base(dependencies, async)
+    {
+        SessionTokenStorage = sessionTokenStorage;
+    }
+
     /// <summary>
     ///     The root entity type being queried.
     /// </summary>
@@ -41,4 +53,12 @@ public class CosmosQueryCompilationContext(QueryCompilationContextDependencies d
     ///     doing so can result in application failures when updating to a new Entity Framework Core release.
     /// </remarks>
     public virtual CosmosAliasManager AliasManager { get; } = new();
+
+    /// <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 ISessionTokenStorage SessionTokenStorage { get; }
 }

src/EFCore.Cosmos/Query/Internal/CosmosQueryCompilationContextFactory.cs

@@ -1,6 +1,8 @@
 // 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.Storage;
+
 namespace Microsoft.EntityFrameworkCore.Cosmos.Query.Internal;
 
 /// <summary>
@@ -11,26 +13,41 @@ namespace Microsoft.EntityFrameworkCore.Cosmos.Query.Internal;
 /// </summary>
 public class CosmosQueryCompilationContextFactory : IQueryCompilationContextFactory
 {
+    private ISessionTokenStorage? _sessionTokenStorage;
+
     /// <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 CosmosQueryCompilationContextFactory(QueryCompilationContextDependencies dependencies)
-        => Dependencies = dependencies;
+    {
+        Dependencies = dependencies;
+    }
 
     /// <summary>
-    ///     Dependencies for this service.
+    ///     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 QueryCompilationContextDependencies Dependencies { 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>
+    protected virtual ISessionTokenStorage SessionTokenStorage => _sessionTokenStorage ??= Dependencies.Context.Database.GetSessionTokens();
+
     /// <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 QueryCompilationContext Create(bool async)
-        => new CosmosQueryCompilationContext(Dependencies, async);
+        => new CosmosQueryCompilationContext(Dependencies, SessionTokenStorage, async);
 }

src/EFCore.Cosmos/Query/Internal/CosmosQueryContext.cs

@@ -1,6 +1,7 @@
 // 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.Storage;
 using Microsoft.EntityFrameworkCore.Cosmos.Storage.Internal;
 
 namespace Microsoft.EntityFrameworkCore.Cosmos.Query.Internal;

src/EFCore.Cosmos/Query/Internal/CosmosShapedQueryCompilingExpressionVisitor.PagingQueryingEnumerable.cs

@@ -4,6 +4,7 @@
 #nullable disable
 
 using Microsoft.EntityFrameworkCore.Cosmos.Diagnostics.Internal;
+using Microsoft.EntityFrameworkCore.Cosmos.Storage;
 using Microsoft.EntityFrameworkCore.Cosmos.Storage.Internal;
 using Newtonsoft.Json.Linq;
 
@@ -34,6 +35,7 @@ private sealed class PagingQueryingEnumerable<T> : IAsyncEnumerable<CosmosPage<T
         private readonly string _maxItemCountParameterName;
         private readonly string _continuationTokenParameterName;
         private readonly string _responseContinuationTokenLimitInKbParameterName;
+        private readonly ISessionTokenStorage _sessionTokenStorage;
 
         public PagingQueryingEnumerable(
             CosmosQueryContext cosmosQueryContext,
@@ -48,7 +50,8 @@ public PagingQueryingEnumerable(
             bool threadSafetyChecksEnabled,
             string maxItemCountParameterName,
             string continuationTokenParameterName,
-            string responseContinuationTokenLimitInKbParameterName)
+            string responseContinuationTokenLimitInKbParameterName,
+            ISessionTokenStorage sessionTokenStorage)
         {
             _cosmosQueryContext = cosmosQueryContext;
             _sqlExpressionFactory = sqlExpressionFactory;
@@ -63,7 +66,7 @@ public PagingQueryingEnumerable(
             _maxItemCountParameterName = maxItemCountParameterName;
             _continuationTokenParameterName = continuationTokenParameterName;
             _responseContinuationTokenLimitInKbParameterName = responseContinuationTokenLimitInKbParameterName;
-
+            _sessionTokenStorage = sessionTokenStorage;
             _cosmosContainer = rootEntityType.GetContainer()
                 ?? throw new UnreachableException("Root entity type without a Cosmos container.");
             _cosmosPartitionKey = GeneratePartitionKey(
@@ -92,10 +95,10 @@ private sealed class AsyncEnumerator : IAsyncEnumerator<CosmosPage<T>>
             private readonly IDiagnosticsLogger<DbLoggerCategory.Query> _queryLogger;
             private readonly IDiagnosticsLogger<DbLoggerCategory.Database.Command> _commandLogger;
             private readonly bool _standAloneStateManager;
+            private readonly ISessionTokenStorage _sessionTokenStorage;
             private readonly CancellationToken _cancellationToken;
             private readonly IConcurrencyDetector _concurrencyDetector;
             private readonly IExceptionDetector _exceptionDetector;
-
             private bool _hasExecuted;
             private bool _isDisposed;
 
@@ -111,6 +114,7 @@ public AsyncEnumerator(PagingQueryingEnumerable<T> queryingEnumerable, Cancellat
                 _commandLogger = queryingEnumerable._commandLogger;
                 _standAloneStateManager = queryingEnumerable._standAloneStateManager;
                 _exceptionDetector = _cosmosQueryContext.ExceptionDetector;
+                _sessionTokenStorage = queryingEnumerable._sessionTokenStorage;
                 _cancellationToken = cancellationToken;
 
                 _concurrencyDetector = queryingEnumerable._threadSafetyChecksEnabled
@@ -155,6 +159,8 @@ public async ValueTask<bool> MoveNextAsync()
                         queryRequestOptions.PartitionKey = _cosmosPartitionKey;
                     }
 
+                    queryRequestOptions.SessionToken = _sessionTokenStorage.GetSessionToken(_cosmosContainer);
+
                     var cosmosClient = _cosmosQueryContext.CosmosClient;
                     _commandLogger.ExecutingSqlQuery(_cosmosContainer, _cosmosPartitionKey, sqlQuery);
                     _cosmosQueryContext.InitializeStateManager(_standAloneStateManager);
@@ -165,7 +171,7 @@ public async ValueTask<bool> MoveNextAsync()
                     {
                         queryRequestOptions.MaxItemCount = maxItemCount;
                         using var feedIterator = cosmosClient.CreateQuery(
-                            _cosmosContainer, sqlQuery, continuationToken, queryRequestOptions);
+                            _cosmosContainer, sqlQuery, _sessionTokenStorage, continuationToken, queryRequestOptions);
 
                         using var responseMessage = await feedIterator.ReadNextAsync(_cancellationToken).ConfigureAwait(false);