From 2853accb7fa7a39afa1289b83ce6058422aca8f4 Mon Sep 17 00:00:00 2001 From: JoasE <32096708+JoasE@users.noreply.github.com> Date: Sun, 19 Oct 2025 17:29:09 +0200 Subject: [PATCH 01/37] WIP --- .../CosmosDatabaseFacadeExtensions.cs | 17 + .../Extensions/CosmosModelExtensions.cs | 8 + .../CosmosRuntimeModelConvention.cs | 10 + .../Internal/CosmosAnnotationNames.cs | 8 + .../Properties/CosmosStrings.Designer.cs | 8 + .../Properties/CosmosStrings.resx | 4 + .../Internal/CosmosQueryCompilationContext.cs | 35 +- ...ressionVisitor.PagingQueryingEnumerable.cs | 14 +- ...ingExpressionVisitor.QueryingEnumerable.cs | 13 +- ...ssionVisitor.ReadItemQueryingEnumerable.cs | 12 +- ...osShapedQueryCompilingExpressionVisitor.cs | 10 +- .../Storage/Internal/CosmosClientWrapper.cs | 186 +++--- .../Storage/Internal/CosmosDatabaseWrapper.cs | 75 ++- .../CosmosTransactionalBatchResult.cs | 25 +- .../Storage/Internal/CosmosWriteResult.cs | 56 ++ .../Storage/Internal/ICosmosClientWrapper.cs | 23 +- .../Storage/Internal/VectorSessionToken.cs | 8 + .../Storage/SessionTokenStorage.cs | 533 ++++++++++++++++++ .../CosmosSessionTokensTest.cs | 387 +++++++++++++ 19 files changed, 1316 insertions(+), 116 deletions(-) create mode 100644 src/EFCore.Cosmos/Storage/Internal/CosmosWriteResult.cs create mode 100644 src/EFCore.Cosmos/Storage/Internal/VectorSessionToken.cs create mode 100644 src/EFCore.Cosmos/Storage/SessionTokenStorage.cs create mode 100644 test/EFCore.Cosmos.FunctionalTests/CosmosSessionTokensTest.cs diff --git a/src/EFCore.Cosmos/Extensions/CosmosDatabaseFacadeExtensions.cs b/src/EFCore.Cosmos/Extensions/CosmosDatabaseFacadeExtensions.cs index f5c5d179b0b..5d1ba8e9ded 100644 --- a/src/EFCore.Cosmos/Extensions/CosmosDatabaseFacadeExtensions.cs +++ b/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(databaseFacade).Client; + /// + /// Gets used to manage the session tokens for this . + /// + /// The for the context. + /// The Gets . + public static SessionTokenStorage GetSessionTokens(this DatabaseFacade databaseFacade) + { + var db = GetService(databaseFacade); + if (db is not CosmosDatabaseWrapper dbWrapper) + { + throw new InvalidOperationException(CosmosStrings.CosmosNotInUse); + } + + return dbWrapper.SessionTokenStorage!; + } + private static TService GetService(IInfrastructure databaseFacade) where TService : class { diff --git a/src/EFCore.Cosmos/Extensions/CosmosModelExtensions.cs b/src/EFCore.Cosmos/Extensions/CosmosModelExtensions.cs index be6f8f79366..7d52fc47071 100644 --- a/src/EFCore.Cosmos/Extensions/CosmosModelExtensions.cs +++ b/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]; + /// + /// Returns the all container names used in the model. + /// + /// The model. + /// A set of the names of the containers used in the model. + public static HashSet GetContainerNames(this IReadOnlyModel model) + => (HashSet)model[CosmosAnnotationNames.ContainerNames]!; + /// /// Sets the default container name. /// diff --git a/src/EFCore.Cosmos/Metadata/Conventions/CosmosRuntimeModelConvention.cs b/src/EFCore.Cosmos/Metadata/Conventions/CosmosRuntimeModelConvention.cs index 4ae848ac437..f86192daaae 100644 --- a/src/EFCore.Cosmos/Metadata/Conventions/CosmosRuntimeModelConvention.cs +++ b/src/EFCore.Cosmos/Metadata/Conventions/CosmosRuntimeModelConvention.cs @@ -44,8 +44,18 @@ protected override void ProcessModelAnnotations( { annotations.Remove(CosmosAnnotationNames.Throughput); } + + annotations.Add(CosmosAnnotationNames.ContainerNames, GetContainerNames(model)); } + private HashSet GetContainerNames(IModel model) + => model.GetEntityTypes() + .Where(et => et.FindPrimaryKey() != null) + .Select(et => et.GetContainer()) + .Where(container => container != null) + .Distinct()! + .ToHashSet()!; + /// /// Updates the entity type annotations that will be set on the read-only object. /// diff --git a/src/EFCore.Cosmos/Metadata/Internal/CosmosAnnotationNames.cs b/src/EFCore.Cosmos/Metadata/Internal/CosmosAnnotationNames.cs index 5aac3bede02..85f259dfa12 100644 --- a/src/EFCore.Cosmos/Metadata/Internal/CosmosAnnotationNames.cs +++ b/src/EFCore.Cosmos/Metadata/Internal/CosmosAnnotationNames.cs @@ -27,6 +27,14 @@ public static class CosmosAnnotationNames /// public const string ContainerName = Prefix + "ContainerName"; + /// + /// 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. + /// + public const string ContainerNames = Prefix + "ContainerNames"; // @TODO: is this the right way? + /// /// 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 diff --git a/src/EFCore.Cosmos/Properties/CosmosStrings.Designer.cs b/src/EFCore.Cosmos/Properties/CosmosStrings.Designer.cs index f52af8503fa..6a8bf326061 100644 --- a/src/EFCore.Cosmos/Properties/CosmosStrings.Designer.cs +++ b/src/EFCore.Cosmos/Properties/CosmosStrings.Designer.cs @@ -97,6 +97,14 @@ public static string ContainerContainingPropertyConflict(object? entityType, obj GetString("ContainerContainingPropertyConflict", nameof(entityType), nameof(container), nameof(property)), entityType, container, property); + /// + /// The container with the name '{containerName}' does not exist. + /// + public static string ContainerNameDoesNotExist(object? containerName) + => string.Format( + GetString("ContainerNameDoesNotExist", nameof(containerName)), + containerName); + /// /// 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. /// diff --git a/src/EFCore.Cosmos/Properties/CosmosStrings.resx b/src/EFCore.Cosmos/Properties/CosmosStrings.resx index fd4e48bbad7..9c0e0259e41 100644 --- a/src/EFCore.Cosmos/Properties/CosmosStrings.resx +++ b/src/EFCore.Cosmos/Properties/CosmosStrings.resx @@ -147,6 +147,10 @@ The entity type '{entityType}' is mapped to the container '{container}' but it is also configured as being contained in property '{property}'. + + The container with the name '{containerName}' does not exist. + string + 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. diff --git a/src/EFCore.Cosmos/Query/Internal/CosmosQueryCompilationContext.cs b/src/EFCore.Cosmos/Query/Internal/CosmosQueryCompilationContext.cs index bc804f1f869..67360709119 100644 --- a/src/EFCore.Cosmos/Query/Internal/CosmosQueryCompilationContext.cs +++ b/src/EFCore.Cosmos/Query/Internal/CosmosQueryCompilationContext.cs @@ -9,9 +9,18 @@ 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. /// -public class CosmosQueryCompilationContext(QueryCompilationContextDependencies dependencies, bool async) - : QueryCompilationContext(dependencies, async) +public class CosmosQueryCompilationContext : QueryCompilationContext { + /// + /// 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. + /// + public CosmosQueryCompilationContext(QueryCompilationContextDependencies dependencies, bool async) : base(dependencies, async) + { + } + /// /// The root entity type being queried. /// @@ -41,4 +50,26 @@ public class CosmosQueryCompilationContext(QueryCompilationContextDependencies d /// doing so can result in application failures when updating to a new Entity Framework Core release. /// public virtual CosmosAliasManager AliasManager { get; } = new(); + + /// + /// 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. + /// + public string? SessionToken { get; internal set; } + + /// + /// 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. + /// + public string? GetSessionToken() + { + Debug.Assert(RootEntityType != null); + var container = RootEntityType.GetContainer(); + Debug.Assert(container != null); + return SessionToken ?? Dependencies.Context.Database.GetSessionTokens().GetSessionToken(container); + } } diff --git a/src/EFCore.Cosmos/Query/Internal/CosmosShapedQueryCompilingExpressionVisitor.PagingQueryingEnumerable.cs b/src/EFCore.Cosmos/Query/Internal/CosmosShapedQueryCompilingExpressionVisitor.PagingQueryingEnumerable.cs index 2766b35a3c9..41f7daddab5 100644 --- a/src/EFCore.Cosmos/Query/Internal/CosmosShapedQueryCompilingExpressionVisitor.PagingQueryingEnumerable.cs +++ b/src/EFCore.Cosmos/Query/Internal/CosmosShapedQueryCompilingExpressionVisitor.PagingQueryingEnumerable.cs @@ -34,6 +34,7 @@ private sealed class PagingQueryingEnumerable : IAsyncEnumerable> private readonly Type _contextType; private readonly string _cosmosContainer; private readonly PartitionKey _cosmosPartitionKey; + private readonly string _sessionToken; private readonly IDiagnosticsLogger _queryLogger; private readonly IDiagnosticsLogger _commandLogger; private readonly bool _standAloneStateManager; private readonly CancellationToken _cancellationToken; private readonly IConcurrencyDetector _concurrencyDetector; private readonly IExceptionDetector _exceptionDetector; - private bool _hasExecuted; private bool _isDisposed; @@ -107,6 +109,7 @@ public AsyncEnumerator(PagingQueryingEnumerable queryingEnumerable, Cancellat _contextType = queryingEnumerable._contextType; _cosmosContainer = queryingEnumerable._cosmosContainer; _cosmosPartitionKey = queryingEnumerable._cosmosPartitionKey; + _sessionToken = queryingEnumerable._sessionToken; _queryLogger = queryingEnumerable._queryLogger; _commandLogger = queryingEnumerable._commandLogger; _standAloneStateManager = queryingEnumerable._standAloneStateManager; @@ -155,6 +158,11 @@ public async ValueTask MoveNextAsync() queryRequestOptions.PartitionKey = _cosmosPartitionKey; } + if (_sessionToken is not null) + { + queryRequestOptions.SessionToken = _sessionToken; + } + var cosmosClient = _cosmosQueryContext.CosmosClient; _commandLogger.ExecutingSqlQuery(_cosmosContainer, _cosmosPartitionKey, sqlQuery); _cosmosQueryContext.InitializeStateManager(_standAloneStateManager); diff --git a/src/EFCore.Cosmos/Query/Internal/CosmosShapedQueryCompilingExpressionVisitor.QueryingEnumerable.cs b/src/EFCore.Cosmos/Query/Internal/CosmosShapedQueryCompilingExpressionVisitor.QueryingEnumerable.cs index ced9ff54222..1d6be1fe8bf 100644 --- a/src/EFCore.Cosmos/Query/Internal/CosmosShapedQueryCompilingExpressionVisitor.QueryingEnumerable.cs +++ b/src/EFCore.Cosmos/Query/Internal/CosmosShapedQueryCompilingExpressionVisitor.QueryingEnumerable.cs @@ -27,6 +27,7 @@ private sealed class QueryingEnumerable : IEnumerable, IAsyncEnumerable private readonly IQuerySqlGeneratorFactory _querySqlGeneratorFactory; private readonly Type _contextType; private readonly string _cosmosContainer; + private readonly string _sessionToken; private readonly PartitionKey _cosmosPartitionKey; private readonly IDiagnosticsLogger _queryLogger; private readonly bool _standAloneStateManager; @@ -42,7 +43,8 @@ public QueryingEnumerable( IEntityType rootEntityType, List partitionKeyPropertyValues, bool standAloneStateManager, - bool threadSafetyChecksEnabled) + bool threadSafetyChecksEnabled, + string sessionToken) { _cosmosQueryContext = cosmosQueryContext; _sqlExpressionFactory = sqlExpressionFactory; @@ -53,6 +55,7 @@ public QueryingEnumerable( _queryLogger = cosmosQueryContext.QueryLogger; _standAloneStateManager = standAloneStateManager; _threadSafetyChecksEnabled = threadSafetyChecksEnabled; + _sessionToken = sessionToken; _cosmosContainer = rootEntityType.GetContainer() ?? throw new UnreachableException("Root entity type without a Cosmos container."); @@ -106,6 +109,7 @@ private sealed class Enumerator : IEnumerator private readonly Func _shaper; private readonly Type _contextType; private readonly string _cosmosContainer; + private readonly string _sessionToken; private readonly PartitionKey _cosmosPartitionKey; private readonly IDiagnosticsLogger _queryLogger; private readonly bool _standAloneStateManager; @@ -121,6 +125,7 @@ public Enumerator(QueryingEnumerable queryingEnumerable) _shaper = queryingEnumerable._shaper; _contextType = queryingEnumerable._contextType; _cosmosContainer = queryingEnumerable._cosmosContainer; + _sessionToken = queryingEnumerable._sessionToken; _cosmosPartitionKey = queryingEnumerable._cosmosPartitionKey; _queryLogger = queryingEnumerable._queryLogger; _standAloneStateManager = queryingEnumerable._standAloneStateManager; @@ -149,7 +154,7 @@ public bool MoveNext() EntityFrameworkMetricsData.ReportQueryExecuting(); _enumerator = _cosmosQueryContext.CosmosClient - .ExecuteSqlQuery(_cosmosContainer, _cosmosPartitionKey, sqlQuery) + .ExecuteSqlQuery(_cosmosContainer, _cosmosPartitionKey, sqlQuery, _sessionToken) .GetEnumerator(); _cosmosQueryContext.InitializeStateManager(_standAloneStateManager); } @@ -196,6 +201,7 @@ private sealed class AsyncEnumerator : IAsyncEnumerator private readonly Type _contextType; private readonly string _cosmosContainer; private readonly PartitionKey _cosmosPartitionKey; + private readonly string _sessionToken; private readonly IDiagnosticsLogger _queryLogger; private readonly bool _standAloneStateManager; private readonly CancellationToken _cancellationToken; @@ -212,6 +218,7 @@ public AsyncEnumerator(QueryingEnumerable queryingEnumerable, CancellationTok _contextType = queryingEnumerable._contextType; _cosmosContainer = queryingEnumerable._cosmosContainer; _cosmosPartitionKey = queryingEnumerable._cosmosPartitionKey; + _sessionToken = queryingEnumerable._sessionToken; _queryLogger = queryingEnumerable._queryLogger; _standAloneStateManager = queryingEnumerable._standAloneStateManager; _exceptionDetector = _cosmosQueryContext.ExceptionDetector; @@ -237,7 +244,7 @@ public async ValueTask MoveNextAsync() EntityFrameworkMetricsData.ReportQueryExecuting(); _enumerator = _cosmosQueryContext.CosmosClient - .ExecuteSqlQueryAsync(_cosmosContainer, _cosmosPartitionKey, sqlQuery) + .ExecuteSqlQueryAsync(_cosmosContainer, _cosmosPartitionKey, sqlQuery, _sessionToken) .GetAsyncEnumerator(_cancellationToken); _cosmosQueryContext.InitializeStateManager(_standAloneStateManager); } diff --git a/src/EFCore.Cosmos/Query/Internal/CosmosShapedQueryCompilingExpressionVisitor.ReadItemQueryingEnumerable.cs b/src/EFCore.Cosmos/Query/Internal/CosmosShapedQueryCompilingExpressionVisitor.ReadItemQueryingEnumerable.cs index 68843902984..79e45ae3dea 100644 --- a/src/EFCore.Cosmos/Query/Internal/CosmosShapedQueryCompilingExpressionVisitor.ReadItemQueryingEnumerable.cs +++ b/src/EFCore.Cosmos/Query/Internal/CosmosShapedQueryCompilingExpressionVisitor.ReadItemQueryingEnumerable.cs @@ -25,6 +25,7 @@ private sealed class ReadItemQueryingEnumerable : IEnumerable, IAsyncEnume private readonly string _cosmosContainer; private readonly ReadItemInfo _readItemInfo; private readonly PartitionKey _cosmosPartitionKey; + private readonly string _sessionToken; private readonly Func _shaper; private readonly Type _contextType; private readonly IDiagnosticsLogger _queryLogger; @@ -39,7 +40,8 @@ public ReadItemQueryingEnumerable( Func shaper, Type contextType, bool standAloneStateManager, - bool threadSafetyChecksEnabled) + bool threadSafetyChecksEnabled, + string sessionToken) { _cosmosQueryContext = cosmosQueryContext; _rootEntityType = rootEntityType; @@ -49,7 +51,7 @@ public ReadItemQueryingEnumerable( _queryLogger = _cosmosQueryContext.QueryLogger; _standAloneStateManager = standAloneStateManager; _threadSafetyChecksEnabled = threadSafetyChecksEnabled; - + _sessionToken = sessionToken; _cosmosContainer = rootEntityType.GetContainer() ?? throw new UnreachableException("Root entity type without a Cosmos container."); _cosmosPartitionKey = GeneratePartitionKey( @@ -105,6 +107,7 @@ private sealed class Enumerator : IEnumerator, IAsyncEnumerator private readonly CosmosQueryContext _cosmosQueryContext; private readonly string _cosmosContainer; private readonly PartitionKey _cosmosPartitionKey; + private readonly string _sessionToken; private readonly Func _shaper; private readonly Type _contextType; private readonly IDiagnosticsLogger _queryLogger; @@ -122,6 +125,7 @@ public Enumerator(ReadItemQueryingEnumerable readItemEnumerable, Cancellation _cosmosQueryContext = readItemEnumerable._cosmosQueryContext; _cosmosContainer = readItemEnumerable._cosmosContainer; _cosmosPartitionKey = readItemEnumerable._cosmosPartitionKey; + _sessionToken = readItemEnumerable._sessionToken; _shaper = readItemEnumerable._shaper; _contextType = readItemEnumerable._contextType; _queryLogger = readItemEnumerable._queryLogger; @@ -161,7 +165,8 @@ public bool MoveNext() _item = _cosmosQueryContext.CosmosClient.ExecuteReadItem( _cosmosContainer, _cosmosPartitionKey, - resourceId); + resourceId, + _sessionToken); return ShapeResult(); } @@ -202,6 +207,7 @@ public async ValueTask MoveNextAsync() _cosmosContainer, _cosmosPartitionKey, resourceId, + _sessionToken, _cancellationToken) .ConfigureAwait(false); diff --git a/src/EFCore.Cosmos/Query/Internal/CosmosShapedQueryCompilingExpressionVisitor.cs b/src/EFCore.Cosmos/Query/Internal/CosmosShapedQueryCompilingExpressionVisitor.cs index f79539eb0b7..d7c462389a0 100644 --- a/src/EFCore.Cosmos/Query/Internal/CosmosShapedQueryCompilingExpressionVisitor.cs +++ b/src/EFCore.Cosmos/Query/Internal/CosmosShapedQueryCompilingExpressionVisitor.cs @@ -83,6 +83,7 @@ protected override Expression VisitShapedQuery(ShapedQueryExpression shapedQuery var threadSafetyConstant = Constant(_threadSafetyChecksEnabled); var standAloneStateManagerConstant = Constant( QueryCompilationContext.QueryTrackingBehavior == QueryTrackingBehavior.NoTrackingWithIdentityResolution); + var sessionTokenConstant = Constant(cosmosQueryCompilationContext.GetSessionToken()); Check.DebugAssert(!paging || selectExpression.ReadItemInfo is null, "ReadItem is being with paging, impossible."); @@ -97,7 +98,8 @@ protected override Expression VisitShapedQuery(ShapedQueryExpression shapedQuery shaperConstant, contextTypeConstant, standAloneStateManagerConstant, - threadSafetyConstant), + threadSafetyConstant, + sessionTokenConstant), _ when paging => New( typeof(PagingQueryingEnumerable<>).MakeGenericType(shaperLambda.ReturnType).GetConstructors()[0], @@ -113,7 +115,8 @@ protected override Expression VisitShapedQuery(ShapedQueryExpression shapedQuery threadSafetyConstant, Constant(maxItemCount.Name), Constant(continuationToken.Name), - Constant(responseContinuationTokenLimitInKb.Name)), + Constant(responseContinuationTokenLimitInKb.Name), + sessionTokenConstant), _ => New( typeof(QueryingEnumerable<>).MakeGenericType(shaperLambda.ReturnType).GetConstructors()[0], cosmosQueryContextConstant, @@ -125,7 +128,8 @@ protected override Expression VisitShapedQuery(ShapedQueryExpression shapedQuery rootEntityTypeConstant, Constant(cosmosQueryCompilationContext.PartitionKeyPropertyValues), standAloneStateManagerConstant, - threadSafetyConstant) + threadSafetyConstant, + sessionTokenConstant) }; } diff --git a/src/EFCore.Cosmos/Storage/Internal/CosmosClientWrapper.cs b/src/EFCore.Cosmos/Storage/Internal/CosmosClientWrapper.cs index 3f1b6f62643..03654ac6e11 100644 --- a/src/EFCore.Cosmos/Storage/Internal/CosmosClientWrapper.cs +++ b/src/EFCore.Cosmos/Storage/Internal/CosmosClientWrapper.cs @@ -385,7 +385,7 @@ private static string GetPathFromRoot(IReadOnlyEntityType entityType) /// 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. /// - public virtual bool CreateItem( + public virtual CosmosWriteResult CreateItem( string containerId, JToken document, IUpdateEntry entry) @@ -395,7 +395,7 @@ public virtual bool CreateItem( return _executionStrategy.Execute((containerId, document, entry, this), CreateItemOnce, null); } - private static bool CreateItemOnce( + private static CosmosWriteResult CreateItemOnce( DbContext context, (string ContainerId, JToken Document, IUpdateEntry Entry, CosmosClientWrapper Wrapper) parameters) => CreateItemOnceAsync(context, parameters).GetAwaiter().GetResult(); @@ -406,14 +406,14 @@ private static bool CreateItemOnce( /// 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. /// - public virtual Task CreateItemAsync( + public virtual Task CreateItemAsync( string containerId, JToken document, IUpdateEntry updateEntry, CancellationToken cancellationToken = default) => _executionStrategy.ExecuteAsync((containerId, document, updateEntry, this), CreateItemOnceAsync, null, cancellationToken); - private static async Task CreateItemOnceAsync( + private static async Task CreateItemOnceAsync( DbContext _, (string ContainerId, JToken Document, IUpdateEntry Entry, CosmosClientWrapper Wrapper) parameters, CancellationToken cancellationToken = default) @@ -458,7 +458,7 @@ private static async Task CreateItemOnceAsync( ProcessResponse(response, entry); - return response.StatusCode == HttpStatusCode.Created; + return CosmosWriteResult.Success(response.Headers.Session); } /// @@ -467,7 +467,7 @@ private static async Task CreateItemOnceAsync( /// 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. /// - public virtual bool ReplaceItem( + public virtual CosmosWriteResult ReplaceItem( string collectionId, string documentId, JObject document, @@ -478,7 +478,7 @@ public virtual bool ReplaceItem( return _executionStrategy.Execute((collectionId, documentId, document, entry, this), ReplaceItemOnce, null); } - private static bool ReplaceItemOnce( + private static CosmosWriteResult ReplaceItemOnce( DbContext context, (string ContainerId, string ItemId, JObject Document, IUpdateEntry Entry, CosmosClientWrapper Wrapper) parameters) => ReplaceItemOnceAsync(context, parameters).GetAwaiter().GetResult(); @@ -489,7 +489,7 @@ private static bool ReplaceItemOnce( /// 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. /// - public virtual Task ReplaceItemAsync( + public virtual Task ReplaceItemAsync( string collectionId, string documentId, JObject document, @@ -498,7 +498,7 @@ public virtual Task ReplaceItemAsync( => _executionStrategy.ExecuteAsync( (collectionId, documentId, document, updateEntry, this), ReplaceItemOnceAsync, null, cancellationToken); - private static async Task ReplaceItemOnceAsync( + private static async Task ReplaceItemOnceAsync( DbContext _, (string ContainerId, string ResourceId, JObject Document, IUpdateEntry Entry, CosmosClientWrapper Wrapper) parameters, CancellationToken cancellationToken = default) @@ -544,7 +544,7 @@ private static async Task ReplaceItemOnceAsync( ProcessResponse(response, entry); - return response.StatusCode == HttpStatusCode.OK; + return CosmosWriteResult.Success(response.Headers.Session); } /// @@ -553,7 +553,7 @@ private static async Task ReplaceItemOnceAsync( /// 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. /// - public virtual bool DeleteItem( + public virtual CosmosWriteResult DeleteItem( string containerId, string documentId, IUpdateEntry entry) @@ -563,7 +563,7 @@ public virtual bool DeleteItem( return _executionStrategy.Execute((containerId, documentId, entry, this), DeleteItemOnce, null); } - private static bool DeleteItemOnce( + private static CosmosWriteResult DeleteItemOnce( DbContext context, (string ContainerId, string DocumentId, IUpdateEntry Entry, CosmosClientWrapper Wrapper) parameters) => DeleteItemOnceAsync(context, parameters).GetAwaiter().GetResult(); @@ -574,13 +574,60 @@ private static bool DeleteItemOnce( /// 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. /// - public virtual Task DeleteItemAsync( + public virtual Task DeleteItemAsync( string containerId, string documentId, IUpdateEntry entry, CancellationToken cancellationToken = default) => _executionStrategy.ExecuteAsync((containerId, documentId, entry, this), DeleteItemOnceAsync, null, cancellationToken); + private static async Task DeleteItemOnceAsync( + DbContext? _, + (string ContainerId, string ResourceId, IUpdateEntry Entry, CosmosClientWrapper Wrapper) parameters, + CancellationToken cancellationToken = default) + { + var entry = parameters.Entry; + var wrapper = parameters.Wrapper; + var items = wrapper.Client.GetDatabase(wrapper._databaseId).GetContainer(parameters.ContainerId); + + var itemRequestOptions = CreateItemRequestOptions(entry, wrapper._enableContentResponseOnWrite); + var partitionKeyValue = ExtractPartitionKeyValue(entry); + var preTriggers = GetTriggers(entry, TriggerType.Pre, TriggerOperation.Delete); + var postTriggers = GetTriggers(entry, TriggerType.Post, TriggerOperation.Delete); + if (preTriggers != null || postTriggers != null) + { + itemRequestOptions ??= new ItemRequestOptions(); + if (preTriggers != null) + { + itemRequestOptions.PreTriggers = preTriggers; + } + + if (postTriggers != null) + { + itemRequestOptions.PostTriggers = postTriggers; + } + } + + using var response = await items.DeleteItemStreamAsync( + parameters.ResourceId, + partitionKeyValue, + itemRequestOptions, + cancellationToken: cancellationToken) + .ConfigureAwait(false); + + wrapper._commandLogger.ExecutedDeleteItem( + response.Diagnostics.GetClientElapsedTime(), + response.Headers.RequestCharge, + response.Headers.ActivityId, + parameters.ResourceId, + parameters.ContainerId, + partitionKeyValue); + + ProcessResponse(response, entry); + + return CosmosWriteResult.Success(response.Headers.Session); + } + /// /// 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 @@ -650,7 +697,7 @@ private static async Task ExecuteBatchOnceAsync( .ToList(); var exception = new CosmosException(response.ErrorMessage, errorCode, 0, response.ActivityId, response.RequestCharge); - return new CosmosTransactionalBatchResult(errorEntries, exception); + return CosmosTransactionalBatchResult.Failure(errorEntries, exception); } wrapper._commandLogger.ExecutedTransactionalBatch( @@ -663,54 +710,7 @@ private static async Task ExecuteBatchOnceAsync( ProcessResponse(response, batch.Entries); - return CosmosTransactionalBatchResult.Success; - } - - private static async Task DeleteItemOnceAsync( - DbContext? _, - (string ContainerId, string ResourceId, IUpdateEntry Entry, CosmosClientWrapper Wrapper) parameters, - CancellationToken cancellationToken = default) - { - var entry = parameters.Entry; - var wrapper = parameters.Wrapper; - var items = wrapper.Client.GetDatabase(wrapper._databaseId).GetContainer(parameters.ContainerId); - - var itemRequestOptions = CreateItemRequestOptions(entry, wrapper._enableContentResponseOnWrite); - var partitionKeyValue = ExtractPartitionKeyValue(entry); - var preTriggers = GetTriggers(entry, TriggerType.Pre, TriggerOperation.Delete); - var postTriggers = GetTriggers(entry, TriggerType.Post, TriggerOperation.Delete); - if (preTriggers != null || postTriggers != null) - { - itemRequestOptions ??= new ItemRequestOptions(); - if (preTriggers != null) - { - itemRequestOptions.PreTriggers = preTriggers; - } - - if (postTriggers != null) - { - itemRequestOptions.PostTriggers = postTriggers; - } - } - - using var response = await items.DeleteItemStreamAsync( - parameters.ResourceId, - partitionKeyValue, - itemRequestOptions, - cancellationToken: cancellationToken) - .ConfigureAwait(false); - - wrapper._commandLogger.ExecutedDeleteItem( - response.Diagnostics.GetClientElapsedTime(), - response.Headers.RequestCharge, - response.Headers.ActivityId, - parameters.ResourceId, - parameters.ContainerId, - partitionKeyValue); - - ProcessResponse(response, entry); - - return response.StatusCode == HttpStatusCode.NoContent; + return CosmosTransactionalBatchResult.Success(response.Headers.Session); } private static ItemRequestOptions? CreateItemRequestOptions(IUpdateEntry entry, bool? enableContentResponseOnWrite) @@ -803,13 +803,14 @@ private static void ProcessResponse(IUpdateEntry entry, string eTag, Stream? con public virtual IEnumerable ExecuteSqlQuery( string containerId, PartitionKey partitionKeyValue, - CosmosSqlQuery query) + CosmosSqlQuery query, + string? sessionToken) { _databaseLogger.SyncNotSupported(); _commandLogger.ExecutingSqlQuery(containerId, partitionKeyValue, query); - return new DocumentEnumerable(this, containerId, partitionKeyValue, query); + return new DocumentEnumerable(this, containerId, partitionKeyValue, query, sessionToken); } /// @@ -821,11 +822,12 @@ public virtual IEnumerable ExecuteSqlQuery( public virtual IAsyncEnumerable ExecuteSqlQueryAsync( string containerId, PartitionKey partitionKeyValue, - CosmosSqlQuery query) + CosmosSqlQuery query, + string? sessionToken) { _commandLogger.ExecutingSqlQuery(containerId, partitionKeyValue, query); - return new DocumentAsyncEnumerable(this, containerId, partitionKeyValue, query); + return new DocumentAsyncEnumerable(this, containerId, partitionKeyValue, query, sessionToken); } /// @@ -837,13 +839,14 @@ public virtual IAsyncEnumerable ExecuteSqlQueryAsync( public virtual JObject? ExecuteReadItem( string containerId, PartitionKey partitionKeyValue, - string resourceId) + string resourceId, + string? sessionToken) { _databaseLogger.SyncNotSupported(); _commandLogger.ExecutingReadItem(containerId, partitionKeyValue, resourceId); - var response = _executionStrategy.Execute((containerId, partitionKeyValue, resourceId, this), CreateSingleItemQuery, null); + var response = _executionStrategy.Execute((containerId, partitionKeyValue, resourceId, sessionToken, this), CreateSingleItemQuery, null); _commandLogger.ExecutedReadItem( response.Diagnostics.GetClientElapsedTime(), @@ -866,12 +869,13 @@ public virtual IAsyncEnumerable ExecuteSqlQueryAsync( string containerId, PartitionKey partitionKeyValue, string resourceId, + string? sessionToken, CancellationToken cancellationToken = default) { _commandLogger.ExecutingReadItem(containerId, partitionKeyValue, resourceId); var response = await _executionStrategy.ExecuteAsync( - (containerId, partitionKeyValue, resourceId, this), + (containerId, partitionKeyValue, resourceId, sessionToken, this), CreateSingleItemQueryAsync, null, cancellationToken) @@ -890,32 +894,44 @@ public virtual IAsyncEnumerable ExecuteSqlQueryAsync( private static ResponseMessage CreateSingleItemQuery( DbContext? context, - (string ContainerId, PartitionKey PartitionKeyValue, string ResourceId, CosmosClientWrapper Wrapper) parameters) + (string ContainerId, PartitionKey PartitionKeyValue, string ResourceId, string? SessionToken, CosmosClientWrapper Wrapper) parameters) => CreateSingleItemQueryAsync(context, parameters).GetAwaiter().GetResult(); private static Task CreateSingleItemQueryAsync( DbContext? _, - (string ContainerId, PartitionKey PartitionKeyValue, string ResourceId, CosmosClientWrapper Wrapper) parameters, + (string ContainerId, PartitionKey PartitionKeyValue, string ResourceId, string? SessionToken, CosmosClientWrapper Wrapper) parameters, CancellationToken cancellationToken = default) { - var (containerId, partitionKeyValue, resourceId, wrapper) = parameters; + var (containerId, partitionKeyValue, resourceId, sessionToken, wrapper) = parameters; var container = wrapper.Client.GetDatabase(wrapper._databaseId).GetContainer(containerId); + ItemRequestOptions? itemRequestOptions = null; + if (sessionToken != null) + { + itemRequestOptions = new ItemRequestOptions { SessionToken = sessionToken }; + } + return container.ReadItemStreamAsync( resourceId, partitionKeyValue, + itemRequestOptions, cancellationToken: cancellationToken); } private static JObject? JObjectFromReadItemResponseMessage(ResponseMessage responseMessage) { - if (responseMessage.StatusCode == HttpStatusCode.NotFound) + const int resourceNotFoundSubStatusCode = 0; + + try { + responseMessage.EnsureSuccessStatusCode(); + } + catch (CosmosException ex) when (ex.StatusCode == HttpStatusCode.NotFound && ex.SubStatusCode == resourceNotFoundSubStatusCode) + { + // @TODO: Is there a test for this return null? return null; } - responseMessage.EnsureSuccessStatusCode(); - var responseStream = responseMessage.Content; using var reader = new StreamReader(responseStream); using var jsonReader = new JsonTextReader(reader); @@ -973,9 +989,11 @@ private sealed class DocumentEnumerable( CosmosClientWrapper cosmosClient, string containerId, PartitionKey partitionKeyValue, - CosmosSqlQuery cosmosSqlQuery) + CosmosSqlQuery cosmosSqlQuery, + string? sessionToken) : IEnumerable { + private readonly string? _sessionToken = sessionToken; private readonly CosmosClientWrapper _cosmosClient = cosmosClient; private readonly string _containerId = containerId; private readonly PartitionKey _partitionKeyValue = partitionKeyValue; @@ -989,6 +1007,7 @@ IEnumerator IEnumerable.GetEnumerator() private sealed class Enumerator(DocumentEnumerable documentEnumerable) : IEnumerator { + private readonly string? _sessionToken = documentEnumerable._sessionToken; private readonly CosmosClientWrapper _cosmosClientWrapper = documentEnumerable._cosmosClient; private readonly string _containerId = documentEnumerable._containerId; private readonly PartitionKey _partitionKeyValue = documentEnumerable._partitionKeyValue; @@ -1019,6 +1038,11 @@ public bool MoveNext() queryRequestOptions.PartitionKey = _partitionKeyValue; } + if (_sessionToken is not null) + { + queryRequestOptions.SessionToken = _sessionToken; + } + _query = _cosmosClientWrapper.CreateQuery( _containerId, _cosmosSqlQuery, continuationToken: null, queryRequestOptions); } @@ -1080,12 +1104,14 @@ private sealed class DocumentAsyncEnumerable( CosmosClientWrapper cosmosClient, string containerId, PartitionKey partitionKeyValue, - CosmosSqlQuery cosmosSqlQuery) + CosmosSqlQuery cosmosSqlQuery, + string? sessionToken) : IAsyncEnumerable { private readonly CosmosClientWrapper _cosmosClient = cosmosClient; private readonly string _containerId = containerId; private readonly PartitionKey _partitionKeyValue = partitionKeyValue; + private readonly string? _sessionToken = sessionToken; private readonly CosmosSqlQuery _cosmosSqlQuery = cosmosSqlQuery; public IAsyncEnumerator GetAsyncEnumerator(CancellationToken cancellationToken = default) @@ -1097,6 +1123,7 @@ private sealed class AsyncEnumerator(DocumentAsyncEnumerable documentEnumerable, private readonly CosmosClientWrapper _cosmosClientWrapper = documentEnumerable._cosmosClient; private readonly string _containerId = documentEnumerable._containerId; private readonly PartitionKey _partitionKeyValue = documentEnumerable._partitionKeyValue; + private readonly string? _sessionToken = documentEnumerable._sessionToken; private readonly CosmosSqlQuery _cosmosSqlQuery = documentEnumerable._cosmosSqlQuery; private JToken? _current; @@ -1123,6 +1150,11 @@ public async ValueTask MoveNextAsync() queryRequestOptions.PartitionKey = _partitionKeyValue; } + if (_sessionToken is not null) + { + queryRequestOptions.SessionToken = _sessionToken; + } + _query = _cosmosClientWrapper.CreateQuery( _containerId, _cosmosSqlQuery, continuationToken: null, queryRequestOptions); } diff --git a/src/EFCore.Cosmos/Storage/Internal/CosmosDatabaseWrapper.cs b/src/EFCore.Cosmos/Storage/Internal/CosmosDatabaseWrapper.cs index a8608b8e48b..0a50b49ce81 100644 --- a/src/EFCore.Cosmos/Storage/Internal/CosmosDatabaseWrapper.cs +++ b/src/EFCore.Cosmos/Storage/Internal/CosmosDatabaseWrapper.cs @@ -44,12 +44,22 @@ public CosmosDatabaseWrapper( _currentDbContext = currentDbContext; _cosmosClient = cosmosClient; + SessionTokenStorage = new(_currentDbContext.Context); + if (loggingOptions.IsSensitiveDataLoggingEnabled) { _sensitiveLoggingEnabled = true; } } + /// + /// 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. + /// + public SessionTokenStorage SessionTokenStorage { get; } + /// /// 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 @@ -91,7 +101,14 @@ public override int SaveChanges(IList entries) try { var response = _cosmosClient.ExecuteTransactionalBatch(transaction); - if (!response.IsSuccess) + if (response.IsSuccess) + { + if (!string.IsNullOrWhiteSpace(response.SessionToken)) + { + SessionTokenStorage.AppendSessionToken(batch.Key.ContainerId, response.SessionToken); + } + } + else { var exception = WrapUpdateException(response.Exception, response.ErroredEntries); if (exception is not DbUpdateConcurrencyException @@ -159,7 +176,14 @@ public override async Task SaveChangesAsync( try { var response = await _cosmosClient.ExecuteTransactionalBatchAsync(transaction, cancellationToken).ConfigureAwait(false); - if (!response.IsSuccess) + if (response.IsSuccess) + { + if (!string.IsNullOrWhiteSpace(response.SessionToken)) + { + SessionTokenStorage.AppendSessionToken(batch.Key.ContainerId, response.SessionToken); + } + } + else { var exception = WrapUpdateException(response.Exception, response.ErroredEntries); if (exception is not DbUpdateConcurrencyException @@ -184,6 +208,16 @@ public override async Task SaveChangesAsync( return rowsAffected; } + private void ProcessSessionToken(string containerId, string? sessionToken) + { + if (string.IsNullOrEmpty(sessionToken)) + { + return; + } + + + } + private SaveGroups CreateSaveGroups(IList entries) { var count = entries.Count; @@ -518,14 +552,19 @@ private bool Save(CosmosUpdateEntry updateEntry) { return updateEntry.Operation switch { - CosmosCudOperation.Create => _cosmosClient.CreateItem( - updateEntry.CollectionId, updateEntry.Document!, updateEntry.Entry), - CosmosCudOperation.Update => _cosmosClient.ReplaceItem( + CosmosCudOperation.Create => ProcessWriteResult(updateEntry, _cosmosClient.CreateItem( + updateEntry.CollectionId, + updateEntry.Document!, + updateEntry.Entry)), + CosmosCudOperation.Update => ProcessWriteResult(updateEntry, _cosmosClient.ReplaceItem( updateEntry.CollectionId, updateEntry.DocumentSource.GetId(updateEntry.Entry.SharedIdentityEntry ?? updateEntry.Entry), updateEntry.Document!, - updateEntry.Entry), - CosmosCudOperation.Delete => _cosmosClient.DeleteItem(updateEntry.CollectionId, updateEntry.DocumentSource.GetId(updateEntry.Entry), updateEntry.Entry), + updateEntry.Entry)), + CosmosCudOperation.Delete => ProcessWriteResult(updateEntry, _cosmosClient.DeleteItem( + updateEntry.CollectionId, + updateEntry.DocumentSource.GetId(updateEntry.Entry), + updateEntry.Entry)), _ => throw new UnreachableException(), }; } @@ -551,22 +590,22 @@ private async Task SaveAsync(CosmosUpdateEntry updateEntry, CancellationTo { return updateEntry.Operation switch { - CosmosCudOperation.Create => await _cosmosClient.CreateItemAsync( + CosmosCudOperation.Create => ProcessWriteResult(updateEntry, await _cosmosClient.CreateItemAsync( updateEntry.CollectionId, updateEntry.Document!, updateEntry.Entry, - cancellationToken).ConfigureAwait(false), - CosmosCudOperation.Update => await _cosmosClient.ReplaceItemAsync( + cancellationToken).ConfigureAwait(false)), + CosmosCudOperation.Update => ProcessWriteResult(updateEntry, await _cosmosClient.ReplaceItemAsync( updateEntry.CollectionId, updateEntry.DocumentSource.GetId(updateEntry.Entry.SharedIdentityEntry ?? updateEntry.Entry), updateEntry.Document!, updateEntry.Entry, - cancellationToken).ConfigureAwait(false), - CosmosCudOperation.Delete => await _cosmosClient.DeleteItemAsync( + cancellationToken).ConfigureAwait(false)), + CosmosCudOperation.Delete => ProcessWriteResult(updateEntry, await _cosmosClient.DeleteItemAsync( updateEntry.CollectionId, updateEntry.DocumentSource.GetId(updateEntry.Entry), updateEntry.Entry, - cancellationToken).ConfigureAwait(false), + cancellationToken).ConfigureAwait(false)), _ => throw new UnreachableException(), }; } @@ -587,6 +626,16 @@ private async Task SaveAsync(CosmosUpdateEntry updateEntry, CancellationTo } } + private bool ProcessWriteResult(CosmosUpdateEntry updateEntry, CosmosWriteResult result) + { + if (!string.IsNullOrWhiteSpace(result.SessionToken)) + { + SessionTokenStorage.AppendSessionToken(updateEntry.CollectionId, result.SessionToken); + } + + return result.IsSuccess; + } + /// /// 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 diff --git a/src/EFCore.Cosmos/Storage/Internal/CosmosTransactionalBatchResult.cs b/src/EFCore.Cosmos/Storage/Internal/CosmosTransactionalBatchResult.cs index d882ff8da1c..df6a535f71b 100644 --- a/src/EFCore.Cosmos/Storage/Internal/CosmosTransactionalBatchResult.cs +++ b/src/EFCore.Cosmos/Storage/Internal/CosmosTransactionalBatchResult.cs @@ -19,11 +19,22 @@ public sealed class CosmosTransactionalBatchResult /// 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. /// - public static CosmosTransactionalBatchResult Success { get; } = new CosmosTransactionalBatchResult(); + public static CosmosTransactionalBatchResult Success(string? sessionToken) + => new CosmosTransactionalBatchResult(sessionToken); - private CosmosTransactionalBatchResult() + /// + /// 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. + /// + public static CosmosTransactionalBatchResult Failure(IReadOnlyList entries, CosmosException exception) + => new CosmosTransactionalBatchResult(entries, exception); + + private CosmosTransactionalBatchResult(string? sessionToken) { IsSuccess = true; + SessionToken = sessionToken; } /// @@ -32,7 +43,7 @@ private CosmosTransactionalBatchResult() /// 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. /// - public CosmosTransactionalBatchResult(IReadOnlyList entries, CosmosException exception) + private CosmosTransactionalBatchResult(IReadOnlyList entries, CosmosException exception) { IsSuccess = false; ErroredEntries = entries; @@ -48,6 +59,14 @@ public CosmosTransactionalBatchResult(IReadOnlyList entries, Cosmo [MemberNotNullWhen(false, nameof(ErroredEntries), nameof(Exception))] public bool IsSuccess { get; } + /// + /// 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. + /// + public string? SessionToken { get; } + /// /// 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 diff --git a/src/EFCore.Cosmos/Storage/Internal/CosmosWriteResult.cs b/src/EFCore.Cosmos/Storage/Internal/CosmosWriteResult.cs new file mode 100644 index 00000000000..5f8e122ae21 --- /dev/null +++ b/src/EFCore.Cosmos/Storage/Internal/CosmosWriteResult.cs @@ -0,0 +1,56 @@ +// 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.Storage.Internal; + +/// +/// 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. +/// +public class CosmosWriteResult +{ + /// + /// 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. + /// + public static CosmosWriteResult Failure { get; } = new CosmosWriteResult(); + + /// + /// 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. + /// + public static CosmosWriteResult Success(string? sessionToken) + => new CosmosWriteResult(sessionToken); + + private CosmosWriteResult() + { + } + + private CosmosWriteResult(string? sessionToken) + { + IsSuccess = true; + SessionToken = sessionToken; + } + + /// + /// 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. + /// + public bool IsSuccess { get; } = true; + + /// + /// 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. + /// + public string? SessionToken { get; } +} diff --git a/src/EFCore.Cosmos/Storage/Internal/ICosmosClientWrapper.cs b/src/EFCore.Cosmos/Storage/Internal/ICosmosClientWrapper.cs index 3c3d545ec90..b0e09dba277 100644 --- a/src/EFCore.Cosmos/Storage/Internal/ICosmosClientWrapper.cs +++ b/src/EFCore.Cosmos/Storage/Internal/ICosmosClientWrapper.cs @@ -67,7 +67,7 @@ public interface ICosmosClientWrapper /// 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. /// - bool CreateItem(string containerId, JToken document, IUpdateEntry entry); + CosmosWriteResult CreateItem(string containerId, JToken document, IUpdateEntry entry); /// /// This is an internal API that supports the Entity Framework Core infrastructure and not subject to @@ -75,7 +75,7 @@ public interface ICosmosClientWrapper /// 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. /// - bool ReplaceItem( + CosmosWriteResult ReplaceItem( string collectionId, string documentId, JObject document, @@ -87,7 +87,7 @@ bool ReplaceItem( /// 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. /// - bool DeleteItem( + CosmosWriteResult DeleteItem( string containerId, string documentId, IUpdateEntry entry); @@ -98,7 +98,7 @@ bool DeleteItem( /// 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. /// - Task CreateItemAsync( + Task CreateItemAsync( string containerId, JToken document, IUpdateEntry updateEntry, @@ -110,7 +110,7 @@ Task CreateItemAsync( /// 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. /// - Task ReplaceItemAsync( + Task ReplaceItemAsync( string collectionId, string documentId, JObject document, @@ -123,7 +123,8 @@ Task ReplaceItemAsync( /// 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. /// - Task DeleteItemAsync( + // @TODO: We also need to send session token on writes to keep the same chain or not? + Task DeleteItemAsync( string containerId, string documentId, IUpdateEntry entry, @@ -150,7 +151,8 @@ FeedIterator CreateQuery( JObject? ExecuteReadItem( string containerId, PartitionKey partitionKeyValue, - string resourceId); + string resourceId, + string? sessionToken); /// /// This is an internal API that supports the Entity Framework Core infrastructure and not subject to @@ -162,6 +164,7 @@ FeedIterator CreateQuery( string containerId, PartitionKey partitionKeyValue, string resourceId, + string? sessionToken, CancellationToken cancellationToken = default); /// @@ -173,7 +176,8 @@ FeedIterator CreateQuery( IEnumerable ExecuteSqlQuery( string containerId, PartitionKey partitionKeyValue, - CosmosSqlQuery query); + CosmosSqlQuery query, + string? sessionToken); /// /// This is an internal API that supports the Entity Framework Core infrastructure and not subject to @@ -184,7 +188,8 @@ IEnumerable ExecuteSqlQuery( IAsyncEnumerable ExecuteSqlQueryAsync( string containerId, PartitionKey partitionKeyValue, - CosmosSqlQuery query); + CosmosSqlQuery query, + string? sessionToken); /// /// This is an internal API that supports the Entity Framework Core infrastructure and not subject to diff --git a/src/EFCore.Cosmos/Storage/Internal/VectorSessionToken.cs b/src/EFCore.Cosmos/Storage/Internal/VectorSessionToken.cs new file mode 100644 index 00000000000..467c70590d2 --- /dev/null +++ b/src/EFCore.Cosmos/Storage/Internal/VectorSessionToken.cs @@ -0,0 +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 System.Diagnostics.CodeAnalysis; +using System.Globalization; + +namespace Microsoft.EntityFrameworkCore.Cosmos.Storage.Internal; + diff --git a/src/EFCore.Cosmos/Storage/SessionTokenStorage.cs b/src/EFCore.Cosmos/Storage/SessionTokenStorage.cs new file mode 100644 index 00000000000..bedf3b9e5a7 --- /dev/null +++ b/src/EFCore.Cosmos/Storage/SessionTokenStorage.cs @@ -0,0 +1,533 @@ +// Licensed to the .NET Foundation under one or more agreements. +// The .NET Foundation licenses this file to you under the MIT license. + +using System.Diagnostics.CodeAnalysis; +using System.Globalization; +using System.Runtime.CompilerServices; +using System.Runtime.InteropServices; +using Microsoft.EntityFrameworkCore.Cosmos.Internal; +using Microsoft.EntityFrameworkCore.Cosmos.Metadata.Internal; + +namespace Microsoft.EntityFrameworkCore.Cosmos.Storage; + +/// +/// Stores the session tokens for a DbContext. +/// +public sealed class SessionTokenStorage +{ + private readonly Dictionary _sessionTokens; + private readonly string _defaultContainerName; + + /// + /// 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. + /// + public SessionTokenStorage(DbContext dbContext) + { + var defaultContainerName = (string)dbContext.Model.GetAnnotation(CosmosAnnotationNames.ContainerName).Value!; + var containerNames = (HashSet)dbContext.Model.GetAnnotation(CosmosAnnotationNames.ContainerNames).Value!; + + _defaultContainerName = defaultContainerName; + _sessionTokens = containerNames.ToDictionary(containerName => containerName, _ => (string?)null); + } + + /// + /// Gets the composite session token for the default container. + /// + public string? GetSessionToken() + => GetSessionToken(_defaultContainerName); + + /// + /// Overwrites the composite session token for the default container + /// + public void OverwriteSessionToken(string? sessionToken) + => OverwriteSessionToken(_defaultContainerName, sessionToken); + + /// + /// Appends the session token specified to any composite already stored in the storage for the default container + /// + public void AppendSessionToken(string sessionToken) + => AppendSessionToken(_defaultContainerName, sessionToken); + + /// + /// Gets the composite session token for the specified container + /// + public string? GetSessionToken(string containerName) + { + ArgumentNullException.ThrowIfNullOrWhiteSpace(containerName, nameof(containerName)); + + if (!_sessionTokens.TryGetValue(containerName, out var value)) + { + throw new ArgumentException(CosmosStrings.ContainerNameDoesNotExist(containerName), nameof(containerName)); + } + + return value; + } + + /// + /// Appends the session token specified to any composite already stored in the storage for the container + /// + public void AppendSessionToken(string containerName, string sessionToken) + { + ArgumentNullException.ThrowIfNullOrWhiteSpace(containerName, nameof(containerName)); + ArgumentNullException.ThrowIfNullOrWhiteSpace(sessionToken, nameof(sessionToken)); + ref var value = ref CollectionsMarshal.GetValueRefOrNullRef(_sessionTokens, containerName); + + if (Unsafe.IsNullRef(ref value)) + { + throw new ArgumentException(CosmosStrings.ContainerNameDoesNotExist(containerName), nameof(containerName)); + } + + if (string.IsNullOrEmpty(value)) + { + value = sessionToken; + } + else + { + value += "," + sessionToken; + } + } + + /// + /// Overwrites the composite session token for the container + /// + public void OverwriteSessionToken(string containerName, string? sessionToken) + { + ArgumentNullException.ThrowIfNullOrWhiteSpace(containerName, nameof(containerName)); + if (sessionToken is not null && string.IsNullOrWhiteSpace(sessionToken)) + { + throw new ArgumentException("sessionToken cannot be whitespace.", sessionToken); + } + + ref var value = ref CollectionsMarshal.GetValueRefOrNullRef(_sessionTokens, containerName); + + if (Unsafe.IsNullRef(ref value)) + { + throw new ArgumentException(CosmosStrings.ContainerNameDoesNotExist(containerName), nameof(containerName)); + } + + value = sessionToken; + } + + + /// + private sealed class VectorSessionToken : IEquatable + { + private static readonly IReadOnlyDictionary DefaultLocalLsnByRegion = new Dictionary(0); + + private readonly string sessionToken; + + private readonly long version; + + private readonly long globalLsn; + + private readonly IReadOnlyDictionary localLsnByRegion; + + private static readonly bool isFalseProgressMergeDisabled = string.Equals(Environment.GetEnvironmentVariable("AZURE_COSMOS_SESSION_TOKEN_FALSE_PROGRESS_MERGE_DISABLED"), "true", StringComparison.OrdinalIgnoreCase); + + /// + /// 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. + /// + public long LSN => globalLsn; + + private VectorSessionToken(long version, long globalLsn, IReadOnlyDictionary localLsnByRegion, string? sessionToken = null) + { + this.version = version; + this.globalLsn = globalLsn; + this.localLsnByRegion = localLsnByRegion; + if (sessionToken != null) + { + this.sessionToken = sessionToken; + return; + } + + string? text = null; + if (localLsnByRegion.Any()) + { + text = string.Join("#", localLsnByRegion.Select((KeyValuePair kvp) => string.Format(CultureInfo.InvariantCulture, "{0}{1}{2}", kvp.Key, '=', kvp.Value))); + } + + if (string.IsNullOrEmpty(text)) + { + this.sessionToken = string.Format(CultureInfo.InvariantCulture, "{0}{1}{2}", this.version, "#", this.globalLsn); + return; + } + + this.sessionToken = string.Format(CultureInfo.InvariantCulture, "{0}{1}{2}{3}{4}", this.version, "#", this.globalLsn, "#", text); + } + + /// + /// 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. + /// + public VectorSessionToken(VectorSessionToken other, long globalLSN) + : this(other.version, globalLSN, other.localLsnByRegion) + { + } + + /// + /// 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. + /// + public static bool TryCreate(string sessionToken, out VectorSessionToken? parsedSessionToken) + { + parsedSessionToken = null; + if (TryParseSessionToken(sessionToken, out var num, out var num2, out var readOnlyDictionary)) + { + parsedSessionToken = new VectorSessionToken(num, num2, readOnlyDictionary, sessionToken); + return true; + } + + return false; + } + + /// + /// 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. + /// + public bool Equals(VectorSessionToken? obj) + { + if (!(obj is VectorSessionToken vectorSessionToken)) + { + return false; + } + + if (version == vectorSessionToken.version && globalLsn == vectorSessionToken.globalLsn) + { + return AreRegionProgressEqual(vectorSessionToken.localLsnByRegion); + } + + return false; + } + + /// + /// 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. + /// + public bool IsValid(VectorSessionToken otherSessionToken) + { + if (!(otherSessionToken is VectorSessionToken vectorSessionToken)) + { + throw new ArgumentNullException("otherSessionToken"); + } + + if (isFalseProgressMergeDisabled) + { + if (vectorSessionToken.version < version || vectorSessionToken.globalLsn < globalLsn) + { + return false; + } + } + else if (vectorSessionToken.version < version || (vectorSessionToken.version == version && vectorSessionToken.globalLsn < globalLsn)) + { + return false; + } + + if (vectorSessionToken.version == version && vectorSessionToken.localLsnByRegion.Count != localLsnByRegion.Count) + { + throw new InvalidOperationException("string.Format(CultureInfo.InvariantCulture, RMResources.InvalidRegionsInSessionToken, sessionToken, vectorSessionToken.sessionToken)"); + } + + foreach (KeyValuePair item in vectorSessionToken.localLsnByRegion) + { + uint key = item.Key; + long value = item.Value; + long value2 = -1L; + if (!localLsnByRegion.TryGetValue(key, out value2)) + { + if (version == vectorSessionToken.version) + { + throw new InvalidOperationException("string.Format(CultureInfo.InvariantCulture, RMResources.InvalidRegionsInSessionToken, sessionToken, vectorSessionToken.sessionToken)"); + } + } + else if (value < value2) + { + return false; + } + } + + return true; + } + + /// + /// 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. + /// + public VectorSessionToken Merge(VectorSessionToken obj) + { + if (!(obj is VectorSessionToken vectorSessionToken)) + { + throw new ArgumentNullException("obj"); + } + + if (version == vectorSessionToken.version && localLsnByRegion.Count != vectorSessionToken.localLsnByRegion.Count) + { + throw new InvalidOperationException("string.Format(CultureInfo.InvariantCulture, RMResources.InvalidRegionsInSessionToken, sessionToken, vectorSessionToken.sessionToken)"); + } + + if (version >= vectorSessionToken.version && globalLsn > vectorSessionToken.globalLsn) + { + if (AreAllLocalLsnByRegionsGreaterThanOrEqual(this, vectorSessionToken)) + { + return this; + } + } + else if (vectorSessionToken.version >= version && vectorSessionToken.globalLsn >= globalLsn && AreAllLocalLsnByRegionsGreaterThanOrEqual(vectorSessionToken, this)) + { + return vectorSessionToken; + } + + VectorSessionToken vectorSessionToken2; + VectorSessionToken vectorSessionToken3; + if (version < vectorSessionToken.version) + { + vectorSessionToken2 = this; + vectorSessionToken3 = vectorSessionToken; + } + else + { + vectorSessionToken2 = vectorSessionToken; + vectorSessionToken3 = this; + } + + Dictionary dictionary = new Dictionary(vectorSessionToken3.localLsnByRegion.Count); + foreach (KeyValuePair item in vectorSessionToken3.localLsnByRegion) + { + uint key = item.Key; + long value = item.Value; + long value2 = -1L; + if (vectorSessionToken2.localLsnByRegion.TryGetValue(key, out value2)) + { + dictionary[key] = Math.Max(value, value2); + continue; + } + + if (version == vectorSessionToken.version) + { + throw new InvalidOperationException("string.Format(CultureInfo.InvariantCulture, RMResources.InvalidRegionsInSessionToken, sessionToken, vectorSessionToken.sessionToken)"); + } + + dictionary[key] = value; + } + + long num = Math.Max(version, vectorSessionToken.version); + long num2 = ((version == vectorSessionToken.version || isFalseProgressMergeDisabled) ? Math.Max(globalLsn, vectorSessionToken.globalLsn) : vectorSessionToken3.globalLsn); + return new VectorSessionToken(num, num2, dictionary); + } + + /// + /// 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. + /// + public string ConvertToString() + { + return sessionToken; + } + + private bool AreRegionProgressEqual(IReadOnlyDictionary other) + { + if (localLsnByRegion.Count != other.Count) + { + return false; + } + + foreach (KeyValuePair item in localLsnByRegion) + { + uint key = item.Key; + long value = item.Value; + if (other.TryGetValue(key, out var value2) && value != value2) + { + return false; + } + } + + return true; + } + + private static bool AreAllLocalLsnByRegionsGreaterThanOrEqual(VectorSessionToken higherToken, VectorSessionToken lowerToken) + { + if (higherToken.localLsnByRegion.Count != lowerToken.localLsnByRegion.Count) + { + return false; + } + + if (!higherToken.localLsnByRegion.Any()) + { + return true; + } + + foreach (KeyValuePair item in higherToken.localLsnByRegion) + { + uint key = item.Key; + long value = item.Value; + if (lowerToken.localLsnByRegion.TryGetValue(key, out var value2)) + { + if (value2 > value) + { + return false; + } + + continue; + } + + return false; + } + + return true; + } + + private static bool TryParseSessionToken(string sessionToken, out long version, out long globalLsn, [NotNullWhen(true)] out IReadOnlyDictionary? localLsnByRegion) + { + version = 0L; + localLsnByRegion = null; + globalLsn = -1L; + if (string.IsNullOrEmpty(sessionToken)) + { + ////DefaultTrace.TraceWarning("Session token is empty"); + return false; + } + + int index = 0; + if (!TryParseLongSegment(sessionToken, ref index, out version)) + { + ////DefaultTrace.TraceWarning("Unexpected session token version number from token: " + sessionToken + " ."); + return false; + } + + if (index >= sessionToken.Length) + { + return false; + } + + if (!TryParseLongSegment(sessionToken, ref index, out globalLsn)) + { + //DefaultTrace.TraceWarning("Unexpected session token global lsn from token: " + sessionToken + " ."); + return false; + } + + if (index >= sessionToken.Length) + { + localLsnByRegion = DefaultLocalLsnByRegion; + return true; + } + + Dictionary dictionary = new Dictionary(); + while (index < sessionToken.Length) + { + if (!TryParseUintTillRegionProgressSeparator(sessionToken, ref index, out var value)) + { + //DefaultTrace.TraceWarning("Unexpected region progress segment in session token: " + sessionToken + "."); + return false; + } + + if (!TryParseLongSegment(sessionToken, ref index, out var value2)) + { + //DefaultTrace.TraceWarning("Unexpected local lsn for region id " + value.ToString(CultureInfo.InvariantCulture) + " for segment in session token: " + sessionToken + "."); + return false; + } + + dictionary[value] = value2; + } + + localLsnByRegion = dictionary; + return true; + } + + private static bool TryParseUintTillRegionProgressSeparator(string input, ref int index, out uint value) + { + value = 0u; + if (index >= input.Length) + { + return false; + } + + long num = 0L; + while (index < input.Length) + { + char c = input[index]; + if (c >= '0' && c <= '9') + { + num = num * 10 + (c - 48); + index++; + continue; + } + + if (c == '=') + { + index++; + break; + } + + return false; + } + + if (num > uint.MaxValue || num < 0) + { + return false; + } + + value = (uint)num; + return true; + } + + private static bool TryParseLongSegment(string input, ref int index, out long value) + { + value = 0L; + if (index >= input.Length) + { + return false; + } + + bool flag = false; + if (input[index] == '-') + { + index++; + flag = true; + } + + while (index < input.Length) + { + char c = input[index]; + if (c >= '0' && c <= '9') + { + value = value * 10 + (c - 48); + index++; + continue; + } + + if (c == '#') + { + index++; + break; + } + + return false; + } + + if (flag) + { + value *= -1L; + } + + return true; + } + } + +} diff --git a/test/EFCore.Cosmos.FunctionalTests/CosmosSessionTokensTest.cs b/test/EFCore.Cosmos.FunctionalTests/CosmosSessionTokensTest.cs new file mode 100644 index 00000000000..7a643d92724 --- /dev/null +++ b/test/EFCore.Cosmos.FunctionalTests/CosmosSessionTokensTest.cs @@ -0,0 +1,387 @@ +// 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.Internal; +using Microsoft.EntityFrameworkCore.Cosmos.Storage; + +namespace Microsoft.EntityFrameworkCore; + +public class CosmosSessionTokensTest(NonSharedFixture fixture) : NonSharedModelTestBase(fixture), IClassFixture +{ + protected override string StoreName + => "CosmosSessionTokensTest"; + + protected override ITestStoreFactory TestStoreFactory + => CosmosTestStoreFactory.Instance; + + [ConditionalFact] + public virtual async Task OverwriteSessionToken_ThrowsForNonExistentContainer() + { + var contextFactory = await InitializeAsync(); + + using var context = contextFactory.CreateContext(); + var sessionTokens = context.Database.GetSessionTokens(); + var exception = Assert.Throws(() => sessionTokens.OverwriteSessionToken("Not the store name", "0:-1#231")); + Assert.Equal(CosmosStrings.ContainerNameDoesNotExist("Not the store name") + " (Parameter 'containerName')", exception.Message); + } + + [ConditionalFact] + public virtual async Task AppendSessionToken_ThrowsForNonExistentContainer() + { + var contextFactory = await InitializeAsync(); + + using var context = contextFactory.CreateContext(); + var sessionTokens = context.Database.GetSessionTokens(); + var exception = Assert.Throws(() => sessionTokens.AppendSessionToken("Not the store name", "0:-1#231")); + Assert.Equal(CosmosStrings.ContainerNameDoesNotExist("Not the store name") + " (Parameter 'containerName')", exception.Message); + } + + [ConditionalFact] + public virtual async Task GetSessionToken_ThrowsForNonExistentContainer() + { + var contextFactory = await InitializeAsync(); + + using var context = contextFactory.CreateContext(); + var sessionTokens = context.Database.GetSessionTokens(); + var exception = Assert.Throws(() => sessionTokens.GetSessionToken("Not the store name")); + Assert.Equal(CosmosStrings.ContainerNameDoesNotExist("Not the store name") + " (Parameter 'containerName')", exception.Message); + } + + [ConditionalFact] + public virtual async Task AppendSessionToken_no_tokens_sets_token() + { + var contextFactory = await InitializeAsync(); + + using var context = contextFactory.CreateContext(); + + var sessionTokens = context.Database.GetSessionTokens(); + sessionTokens.AppendSessionToken("0:-1#231"); + var updatedToken = sessionTokens.GetSessionToken(); + + Assert.Equal("0:-1#231", updatedToken); + } + + [ConditionalTheory, InlineData(AutoTransactionBehavior.WhenNeeded), InlineData(AutoTransactionBehavior.Never), InlineData(AutoTransactionBehavior.Always)] + public virtual async Task AppendSessionToken(AutoTransactionBehavior autoTransactionBehavior) + { + var contextFactory = await InitializeAsync(); + + using var context = contextFactory.CreateContext(); + context.Database.AutoTransactionBehavior = autoTransactionBehavior; + context.Customers.Add(new Customer { Id = "1", PartitionKey = "1" }); + + await context.SaveChangesAsync(); + + var sessionTokens = context.Database.GetSessionTokens(); + var initialToken = sessionTokens.GetSessionToken(); + Assert.False(string.IsNullOrWhiteSpace(initialToken)); + + sessionTokens.AppendSessionToken("0:-1#231"); + + var updatedToken = sessionTokens.GetSessionToken(); + + Assert.Equal(initialToken + ",0:-1#231", updatedToken); + } + + [ConditionalTheory, InlineData(AutoTransactionBehavior.WhenNeeded), InlineData(AutoTransactionBehavior.Never), InlineData(AutoTransactionBehavior.Always)] + public virtual async Task OverwriteSessionToken(AutoTransactionBehavior autoTransactionBehavior) + { + var contextFactory = await InitializeAsync(); + + using var context = contextFactory.CreateContext(); + context.Database.AutoTransactionBehavior = autoTransactionBehavior; + context.Customers.Add(new Customer { Id = "1", PartitionKey = "1" }); + + await context.SaveChangesAsync(); + + var sessionTokens = context.Database.GetSessionTokens(); + var initialToken = sessionTokens.GetSessionToken(); + Assert.False(string.IsNullOrWhiteSpace(initialToken)); + + sessionTokens.OverwriteSessionToken("0:-1#231"); + + var updatedToken = sessionTokens.GetSessionToken(); + + Assert.Equal("0:-1#231", updatedToken); + } + + [ConditionalTheory, InlineData(AutoTransactionBehavior.WhenNeeded), InlineData(AutoTransactionBehavior.Never), InlineData(AutoTransactionBehavior.Always)] + public virtual async Task OverwriteSessionToken_null(AutoTransactionBehavior autoTransactionBehavior) + { + var contextFactory = await InitializeAsync(); + + using var context = contextFactory.CreateContext(); + context.Database.AutoTransactionBehavior = autoTransactionBehavior; + context.Customers.Add(new Customer { Id = "1", PartitionKey = "1" }); + + await context.SaveChangesAsync(); + + var sessionTokens = context.Database.GetSessionTokens(); + var initialToken = sessionTokens.GetSessionToken(); + Assert.False(string.IsNullOrWhiteSpace(initialToken)); + + sessionTokens.OverwriteSessionToken(null); + + var updatedToken = sessionTokens.GetSessionToken(); + + Assert.Null(updatedToken); + } + + [ConditionalFact] + public virtual async Task GetSessionToken_no_token_returns_null() + { + var contextFactory = await InitializeAsync(); + using var context = contextFactory.CreateContext(); + var sessionTokens = context.Database.GetSessionTokens(); + var sessionToken = sessionTokens.GetSessionToken(); + Assert.Null(sessionToken); + } + + [ConditionalFact] + // @TODO: Read item and select... + // @TODO: and sync.. + public virtual async Task AppendSessionToken_uses_session_token_list() + { + var contextFactory = await InitializeAsync(); + + using var context = contextFactory.CreateContext(); + context.Customers.Add(new Customer { Id = "1", PartitionKey = "1" }); + + await context.SaveChangesAsync(); + + var sessionTokens = context.Database.GetSessionTokens(); + var sessionToken = sessionTokens.GetSessionToken()!; + + // Only way we can test this is by setting a session token that will fail the request if used.. + // This will take a couple of seconds to fail + sessionTokens.OverwriteSessionToken(sessionToken.Substring(0, sessionToken.IndexOf('#') + 1) + int.MaxValue); + + var ex = await Assert.ThrowsAsync(() => context.Customers.ToListAsync()); + Assert.Contains("The read session is not available for the input session token.", ex.ResponseBody); + } + + [ConditionalFact] + // @TODO: and sync.. + public virtual async Task AppendSessionToken_uses_session_token_select_list() + { + var contextFactory = await InitializeAsync(); + + using var context = contextFactory.CreateContext(); + context.Customers.Add(new Customer { Id = "1", PartitionKey = "1" }); + + await context.SaveChangesAsync(); + + var sessionTokens = context.Database.GetSessionTokens(); + var sessionToken = sessionTokens.GetSessionToken()!; + + // Only way we can test this is by setting a session token that will fail the request if used.. + // This will take a couple of seconds to fail + sessionTokens.OverwriteSessionToken(sessionToken.Substring(0, sessionToken.IndexOf('#') + 1) + int.MaxValue); + + var ex = await Assert.ThrowsAsync(() => context.Customers.Select(x => new { x.Id, x.PartitionKey }).ToListAsync()); + Assert.Contains("The read session is not available for the input session token.", ex.ResponseBody); + } + + [ConditionalFact] + // @TODO: and sync.. + public virtual async Task AppendSessionToken_uses_session_token_read_item() + { + var contextFactory = await InitializeAsync(); + + using var context = contextFactory.CreateContext(); + context.Customers.Add(new Customer { Id = "1", PartitionKey = "1" }); + + await context.SaveChangesAsync(); + + var sessionTokens = context.Database.GetSessionTokens(); + var sessionToken = sessionTokens.GetSessionToken()!; + + // Only way we can test this is by setting a session token that will fail the request if used.. + // This will take a couple of seconds to fail + sessionTokens.OverwriteSessionToken(sessionToken.Substring(0, sessionToken.IndexOf('#') + 1) + int.MaxValue); + + var ex = await Assert.ThrowsAsync(() => context.Customers.FirstOrDefaultAsync(x => x.Id == "1" && x.PartitionKey == "1")); + Assert.Contains("The read session is not available for the input session token.", ex.ResponseBody); + } + + [ConditionalTheory, InlineData(AutoTransactionBehavior.WhenNeeded), InlineData(AutoTransactionBehavior.Never), InlineData(AutoTransactionBehavior.Always)] + public virtual async Task Add_sets_session_token(AutoTransactionBehavior autoTransactionBehavior) + { + var contextFactory = await InitializeAsync(); + + using var context = contextFactory.CreateContext(); + context.Database.AutoTransactionBehavior = autoTransactionBehavior; + context.Customers.Add(new Customer { Id = "1", PartitionKey = "1" }); + + await context.SaveChangesAsync(); + + var sessionTokens = context.Database.GetSessionTokens(); + var sessionToken = sessionTokens.GetSessionToken(); + Assert.False(string.IsNullOrWhiteSpace(sessionToken)); + } + + [ConditionalTheory, InlineData(AutoTransactionBehavior.WhenNeeded), InlineData(AutoTransactionBehavior.Never), InlineData(AutoTransactionBehavior.Always)] + public virtual async Task Delete_sets_session_token(AutoTransactionBehavior autoTransactionBehavior) + { + var contextFactory = await InitializeAsync(); + + using var context = contextFactory.CreateContext(); + context.Database.AutoTransactionBehavior = autoTransactionBehavior; + + var customer = new Customer { Id = "1", PartitionKey = "1" }; + context.Customers.Add(customer); + + await context.SaveChangesAsync(); + + var initialToken = context.Database.GetSessionTokens().GetSessionToken(); + + context.Remove(customer); + await context.SaveChangesAsync(); + + var sessionToken = context.Database.GetSessionTokens().GetSessionToken(); + Assert.False(string.IsNullOrWhiteSpace(sessionToken)); + Assert.StartsWith(initialToken + ",", sessionToken); + Assert.False(string.IsNullOrWhiteSpace(sessionToken)); + } + + //[ConditionalTheory, InlineData(AutoTransactionBehavior.WhenNeeded), InlineData(AutoTransactionBehavior.Never), InlineData(AutoTransactionBehavior.Always)] + //public virtual async Task Update_sets_session_token(AutoTransactionBehavior autoTransactionBehavior) + //{ + // var contextFactory = await InitializeAsync(); + + // using var context = contextFactory.CreateContext(); + // context.Database.AutoTransactionBehavior = autoTransactionBehavior; + + // var customer = new Customer { Id = "1", PartitionKey = "1" }; + // context.Customers.Add(customer); + + // await context.SaveChangesAsync(); + + // var initialToken = context.Database.GetSessionTokens().Single().Value; + + // customer.Name = "updated"; + // await context.SaveChangesAsync(); + + // var sessionTokens = context.Database.GetSessionTokens(); + // var key = new CosmosContainerPartitionScope("CosmosSessionTokenContext", new Azure.Cosmos.PartitionKey("1")); + // Assert.Equal(1, sessionTokens.Count); + // Assert.True(sessionTokens.ContainsKey(key)); + // Assert.NotEmpty(sessionTokens[key]); + // Assert.NotEqual(initialToken, sessionTokens[key]); + //} + + //[ConditionalFact] + //public virtual async Task Query_with_single_filter_uses_session_token() + //{ + // var contextFactory = await InitializeAsync(); + + // using var context = contextFactory.CreateContext(); + + // var customer = new Customer { Id = "1", PartitionKey = "1" }; + // context.Customers.Add(customer); + + // await context.SaveChangesAsync(); + + // await context.Customers.Where(x => x.PartitionKey == "1").ToListAsync(); + //} + + //[ConditionalFact] + //public virtual async Task Query_with_double_filter_uses_session_token() + //{ + // var contextFactory = await InitializeAsync(); + + // using var context = contextFactory.CreateContext(); + + // var customer = new Customer { Id = "1", PartitionKey = "1" }; + // context.Customers.Add(customer); + + // await context.SaveChangesAsync(); + + // await context.Customers.Where(x => x.PartitionKey == "1" || x.PartitionKey == "2").ToListAsync(); + //} + + + //[ConditionalFact] + //public virtual async Task Query_with_composite_partition_key_not_all_properties_does_not_use_sessiontoken() + //{ + // var contextFactory = await InitializeAsync(); + + // using var context = contextFactory.CreateContext(); + + // var customer = new CompositeCustomer { Id = "1", PartitionKey1 = "1", PartitionKey2 = "2" }; + // context.Add(customer); + + // await context.SaveChangesAsync(); + + // await context.CompositeCustomers.Where(x => x.PartitionKey1 == "1").ToListAsync(); + //} + + //[ConditionalFact] + //public virtual async Task Query_with_composite_partition_key_uses_session_token() + //{ + // var contextFactory = await InitializeAsync(); + + // using var context = contextFactory.CreateContext(); + + // var customer = new CompositeCustomer { Id = "1", PartitionKey1 = "1", PartitionKey2 = "2" }; + // context.Add(customer); + + // await context.SaveChangesAsync(); + + // await context.CompositeCustomers.Where(x => x.PartitionKey1 == "1" && x.PartitionKey2 == "2").ToListAsync(); + //} + + + public class CosmosSessionTokenContext(DbContextOptions options) : PoolableDbContext(options) + { + public DbSet Customers { get; set; } = null!; + // public DbSet CompositeCustomers { get; set; } = null!; + + + protected override void OnModelCreating(ModelBuilder builder) + { + builder.Entity( + b => + { + b.HasKey(c => c.Id); + b.Property(c => c.ETag).IsETagConcurrency(); + b.OwnsMany(x => x.Children); + b.HasPartitionKey(c => c.PartitionKey); + }); + + //builder.Entity( + // b => + // { + // b.HasKey(c => c.Id); + // b.HasPartitionKey(c => new { c.PartitionKey1, c.PartitionKey2 }); + // b.ToContainer("composite"); + // }); + } + } + + public class Customer + { + public string? Id { get; set; } + + public string? Name { get; set; } + + public string? ETag { get; set; } + + public string? PartitionKey { get; set; } + + public ICollection Children { get; } = new HashSet(); + } + + public class DummyChild + { + public string? Id { get; init; } + } + + //public class CompositeCustomer + //{ + // public string? Id { get; set; } + + // public string? PartitionKey1 { get; set; } + // public string? PartitionKey2 { get; set; } + //} +} From e374631a45f364e868b905a601293c48ef0ed1ea Mon Sep 17 00:00:00 2001 From: JoasE <32096708+JoasE@users.noreply.github.com> Date: Mon, 20 Oct 2025 10:59:01 +0200 Subject: [PATCH 02/37] Wip: merge session tokens --- .../Storage/Internal/VectorSessionToken.cs | 8 - .../Storage/SessionTokenStorage.cs | 101 ++++++-- .../CosmosSessionTokensTest.cs | 225 ++++++++---------- 3 files changed, 176 insertions(+), 158 deletions(-) delete mode 100644 src/EFCore.Cosmos/Storage/Internal/VectorSessionToken.cs diff --git a/src/EFCore.Cosmos/Storage/Internal/VectorSessionToken.cs b/src/EFCore.Cosmos/Storage/Internal/VectorSessionToken.cs deleted file mode 100644 index 467c70590d2..00000000000 --- a/src/EFCore.Cosmos/Storage/Internal/VectorSessionToken.cs +++ /dev/null @@ -1,8 +0,0 @@ -// Licensed to the .NET Foundation under one or more agreements. -// The .NET Foundation licenses this file to you under the MIT license. - -using System.Diagnostics.CodeAnalysis; -using System.Globalization; - -namespace Microsoft.EntityFrameworkCore.Cosmos.Storage.Internal; - diff --git a/src/EFCore.Cosmos/Storage/SessionTokenStorage.cs b/src/EFCore.Cosmos/Storage/SessionTokenStorage.cs index bedf3b9e5a7..4dc82caaf2a 100644 --- a/src/EFCore.Cosmos/Storage/SessionTokenStorage.cs +++ b/src/EFCore.Cosmos/Storage/SessionTokenStorage.cs @@ -15,7 +15,7 @@ namespace Microsoft.EntityFrameworkCore.Cosmos.Storage; /// public sealed class SessionTokenStorage { - private readonly Dictionary _sessionTokens; + private readonly Dictionary _sessionTokens; private readonly string _defaultContainerName; /// @@ -30,7 +30,7 @@ public SessionTokenStorage(DbContext dbContext) var containerNames = (HashSet)dbContext.Model.GetAnnotation(CosmosAnnotationNames.ContainerNames).Value!; _defaultContainerName = defaultContainerName; - _sessionTokens = containerNames.ToDictionary(containerName => containerName, _ => (string?)null); + _sessionTokens = containerNames.ToDictionary(containerName => containerName, _ => new CompositeSessionToken()); } /// @@ -42,11 +42,11 @@ public SessionTokenStorage(DbContext dbContext) /// /// Overwrites the composite session token for the default container /// - public void OverwriteSessionToken(string? sessionToken) - => OverwriteSessionToken(_defaultContainerName, sessionToken); + public void SetSessionToken(string? sessionToken) + => SetSessionToken(_defaultContainerName, sessionToken); /// - /// Appends the session token specified to any composite already stored in the storage for the default container + /// Appends or merges the session token specified to any composite already stored in the storage for the default container /// public void AppendSessionToken(string sessionToken) => AppendSessionToken(_defaultContainerName, sessionToken); @@ -63,37 +63,29 @@ public void AppendSessionToken(string sessionToken) throw new ArgumentException(CosmosStrings.ContainerNameDoesNotExist(containerName), nameof(containerName)); } - return value; + return value.ConvertToString(); } /// - /// Appends the session token specified to any composite already stored in the storage for the container + /// Appends or merges the session token specified to any composite already stored in the storage for the container /// public void AppendSessionToken(string containerName, string sessionToken) { ArgumentNullException.ThrowIfNullOrWhiteSpace(containerName, nameof(containerName)); ArgumentNullException.ThrowIfNullOrWhiteSpace(sessionToken, nameof(sessionToken)); - ref var value = ref CollectionsMarshal.GetValueRefOrNullRef(_sessionTokens, containerName); - if (Unsafe.IsNullRef(ref value)) + if (!_sessionTokens.TryGetValue(containerName, out var compositeSessionToken)) { throw new ArgumentException(CosmosStrings.ContainerNameDoesNotExist(containerName), nameof(containerName)); } - if (string.IsNullOrEmpty(value)) - { - value = sessionToken; - } - else - { - value += "," + sessionToken; - } + ParseAndMerge(compositeSessionToken, sessionToken); } /// /// Overwrites the composite session token for the container /// - public void OverwriteSessionToken(string containerName, string? sessionToken) + public void SetSessionToken(string containerName, string? sessionToken) { ArgumentNullException.ThrowIfNullOrWhiteSpace(containerName, nameof(containerName)); if (sessionToken is not null && string.IsNullOrWhiteSpace(sessionToken)) @@ -108,7 +100,69 @@ public void OverwriteSessionToken(string containerName, string? sessionToken) throw new ArgumentException(CosmosStrings.ContainerNameDoesNotExist(containerName), nameof(containerName)); } - value = sessionToken; + value = new CompositeSessionToken(); + + if (sessionToken is null) + { + return; + } + + ParseAndMerge(value, sessionToken); + } + + private void ParseAndMerge(CompositeSessionToken compositeSessionToken, string sessionToken) + { + var parts = sessionToken.Split(','); + foreach (var part in parts) + { + var index = part.IndexOf(':'); + if (index == -1) + { + throw new ArgumentException("CosmosStrings.InvalidSessionToken(sessionToken)", nameof(sessionToken)); + } + + var pkRangeId = sessionToken.Substring(0, index); + var vector = sessionToken.Substring(index + 1); + if (!VectorSessionToken.TryCreate(vector, out var vectorSessionToken)) + { + throw new ArgumentException("CosmosStrings.InvalidSessionToken(sessionToken)", nameof(sessionToken)); + } + + compositeSessionToken.Merge(pkRangeId, vectorSessionToken); + } + } + + private sealed class CompositeSessionToken + { + private string? _string; + private bool _isChanged; + public Dictionary Tokens { get; } = new(); + + public void Merge(string pkRangeId, VectorSessionToken token) + { + ref var existing = ref CollectionsMarshal.GetValueRefOrAddDefault(Tokens, pkRangeId, out var exists); + if (exists) + { + existing = existing!.Merge(token); + } + else + { + existing = token; + } + + _isChanged = true; + } + + public string? ConvertToString() + { + if (_isChanged) + { + _isChanged = false; + _string = string.Join(",", Tokens.Select(kvp => $"{kvp.Key}:{kvp.Value.ConvertToString()}")); + } + + return _string; + } } @@ -178,7 +232,7 @@ public VectorSessionToken(VectorSessionToken other, long globalLSN) /// 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. /// - public static bool TryCreate(string sessionToken, out VectorSessionToken? parsedSessionToken) + public static bool TryCreate(string sessionToken, [NotNullWhen(true)] out VectorSessionToken? parsedSessionToken) { parsedSessionToken = null; if (TryParseSessionToken(sessionToken, out var num, out var num2, out var readOnlyDictionary)) @@ -268,13 +322,8 @@ public bool IsValid(VectorSessionToken otherSessionToken) /// 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. /// - public VectorSessionToken Merge(VectorSessionToken obj) + public VectorSessionToken Merge(VectorSessionToken vectorSessionToken) { - if (!(obj is VectorSessionToken vectorSessionToken)) - { - throw new ArgumentNullException("obj"); - } - if (version == vectorSessionToken.version && localLsnByRegion.Count != vectorSessionToken.localLsnByRegion.Count) { throw new InvalidOperationException("string.Format(CultureInfo.InvariantCulture, RMResources.InvalidRegionsInSessionToken, sessionToken, vectorSessionToken.sessionToken)"); diff --git a/test/EFCore.Cosmos.FunctionalTests/CosmosSessionTokensTest.cs b/test/EFCore.Cosmos.FunctionalTests/CosmosSessionTokensTest.cs index 7a643d92724..8825232cc19 100644 --- a/test/EFCore.Cosmos.FunctionalTests/CosmosSessionTokensTest.cs +++ b/test/EFCore.Cosmos.FunctionalTests/CosmosSessionTokensTest.cs @@ -8,6 +8,7 @@ namespace Microsoft.EntityFrameworkCore; public class CosmosSessionTokensTest(NonSharedFixture fixture) : NonSharedModelTestBase(fixture), IClassFixture { + private const string OtherContainerName = "Other"; protected override string StoreName => "CosmosSessionTokensTest"; @@ -15,13 +16,13 @@ protected override ITestStoreFactory TestStoreFactory => CosmosTestStoreFactory.Instance; [ConditionalFact] - public virtual async Task OverwriteSessionToken_ThrowsForNonExistentContainer() + public virtual async Task SetSessionToken_ThrowsForNonExistentContainer() { var contextFactory = await InitializeAsync(); using var context = contextFactory.CreateContext(); var sessionTokens = context.Database.GetSessionTokens(); - var exception = Assert.Throws(() => sessionTokens.OverwriteSessionToken("Not the store name", "0:-1#231")); + var exception = Assert.Throws(() => sessionTokens.SetSessionToken("Not the store name", "0:-1#231")); Assert.Equal(CosmosStrings.ContainerNameDoesNotExist("Not the store name") + " (Parameter 'containerName')", exception.Message); } @@ -61,13 +62,12 @@ public virtual async Task AppendSessionToken_no_tokens_sets_token() Assert.Equal("0:-1#231", updatedToken); } - [ConditionalTheory, InlineData(AutoTransactionBehavior.WhenNeeded), InlineData(AutoTransactionBehavior.Never), InlineData(AutoTransactionBehavior.Always)] - public virtual async Task AppendSessionToken(AutoTransactionBehavior autoTransactionBehavior) + [ConditionalFact] + public virtual async Task AppendSessionToken_append_higher_lsn_same_pkrange_takes_higher_lsn() { var contextFactory = await InitializeAsync(); using var context = contextFactory.CreateContext(); - context.Database.AutoTransactionBehavior = autoTransactionBehavior; context.Customers.Add(new Customer { Id = "1", PartitionKey = "1" }); await context.SaveChangesAsync(); @@ -76,20 +76,20 @@ public virtual async Task AppendSessionToken(AutoTransactionBehavior autoTransac var initialToken = sessionTokens.GetSessionToken(); Assert.False(string.IsNullOrWhiteSpace(initialToken)); - sessionTokens.AppendSessionToken("0:-1#231"); + var newToken = initialToken.Substring(0, initialToken.IndexOf('#') + 1) + "999999"; + sessionTokens.AppendSessionToken(newToken); var updatedToken = sessionTokens.GetSessionToken(); - Assert.Equal(initialToken + ",0:-1#231", updatedToken); + Assert.Equal(newToken, updatedToken); } - [ConditionalTheory, InlineData(AutoTransactionBehavior.WhenNeeded), InlineData(AutoTransactionBehavior.Never), InlineData(AutoTransactionBehavior.Always)] - public virtual async Task OverwriteSessionToken(AutoTransactionBehavior autoTransactionBehavior) + [ConditionalFact] + public virtual async Task AppendSessionToken_append_lower_lsn_same_pkrange_takes_higher_lsn() { var contextFactory = await InitializeAsync(); using var context = contextFactory.CreateContext(); - context.Database.AutoTransactionBehavior = autoTransactionBehavior; context.Customers.Add(new Customer { Id = "1", PartitionKey = "1" }); await context.SaveChangesAsync(); @@ -98,15 +98,16 @@ public virtual async Task OverwriteSessionToken(AutoTransactionBehavior autoTran var initialToken = sessionTokens.GetSessionToken(); Assert.False(string.IsNullOrWhiteSpace(initialToken)); - sessionTokens.OverwriteSessionToken("0:-1#231"); + var newToken = initialToken.Substring(0, initialToken.IndexOf('#') + 1) + "1"; + sessionTokens.AppendSessionToken(newToken); var updatedToken = sessionTokens.GetSessionToken(); - Assert.Equal("0:-1#231", updatedToken); + Assert.Equal(initialToken, updatedToken); } [ConditionalTheory, InlineData(AutoTransactionBehavior.WhenNeeded), InlineData(AutoTransactionBehavior.Never), InlineData(AutoTransactionBehavior.Always)] - public virtual async Task OverwriteSessionToken_null(AutoTransactionBehavior autoTransactionBehavior) + public virtual async Task AppendSessionToken_different_pkrange_composites_tokens(AutoTransactionBehavior autoTransactionBehavior) { var contextFactory = await InitializeAsync(); @@ -120,7 +121,49 @@ public virtual async Task OverwriteSessionToken_null(AutoTransactionBehavior aut var initialToken = sessionTokens.GetSessionToken(); Assert.False(string.IsNullOrWhiteSpace(initialToken)); - sessionTokens.OverwriteSessionToken(null); + sessionTokens.AppendSessionToken("99:-1#999999"); + + var updatedToken = sessionTokens.GetSessionToken(); + + Assert.Equal(initialToken + ",99:-1#999999", updatedToken); + } + + [ConditionalFact] + public virtual async Task SetSessionToken_does_not_merge_session_token() + { + var contextFactory = await InitializeAsync(); + + using var context = contextFactory.CreateContext(); + context.Customers.Add(new Customer { Id = "1", PartitionKey = "1" }); + + await context.SaveChangesAsync(); + + var sessionTokens = context.Database.GetSessionTokens(); + var initialToken = sessionTokens.GetSessionToken(); + Assert.False(string.IsNullOrWhiteSpace(initialToken)); + + sessionTokens.SetSessionToken("0:-1#1"); + + var updatedToken = sessionTokens.GetSessionToken(); + + Assert.Equal("0:-1#1", updatedToken); + } + + [ConditionalFact] + public virtual async Task SetSessionToken_null_sets_session_token_null() + { + var contextFactory = await InitializeAsync(); + + using var context = contextFactory.CreateContext(); + context.Customers.Add(new Customer { Id = "1", PartitionKey = "1" }); + + await context.SaveChangesAsync(); + + var sessionTokens = context.Database.GetSessionTokens(); + var initialToken = sessionTokens.GetSessionToken(); + Assert.False(string.IsNullOrWhiteSpace(initialToken)); + + sessionTokens.SetSessionToken(null); var updatedToken = sessionTokens.GetSessionToken(); @@ -140,21 +183,18 @@ public virtual async Task GetSessionToken_no_token_returns_null() [ConditionalFact] // @TODO: Read item and select... // @TODO: and sync.. - public virtual async Task AppendSessionToken_uses_session_token_list() + public virtual async Task Query_uses_session_token() { var contextFactory = await InitializeAsync(); using var context = contextFactory.CreateContext(); - context.Customers.Add(new Customer { Id = "1", PartitionKey = "1" }); - - await context.SaveChangesAsync(); var sessionTokens = context.Database.GetSessionTokens(); var sessionToken = sessionTokens.GetSessionToken()!; // Only way we can test this is by setting a session token that will fail the request if used.. // This will take a couple of seconds to fail - sessionTokens.OverwriteSessionToken(sessionToken.Substring(0, sessionToken.IndexOf('#') + 1) + int.MaxValue); + sessionTokens.SetSessionToken(sessionToken.Substring(0, sessionToken.IndexOf('#') + 1) + int.MaxValue); var ex = await Assert.ThrowsAsync(() => context.Customers.ToListAsync()); Assert.Contains("The read session is not available for the input session token.", ex.ResponseBody); @@ -162,7 +202,7 @@ public virtual async Task AppendSessionToken_uses_session_token_list() [ConditionalFact] // @TODO: and sync.. - public virtual async Task AppendSessionToken_uses_session_token_select_list() + public virtual async Task Shaped_query_uses_session_token() { var contextFactory = await InitializeAsync(); @@ -173,10 +213,10 @@ public virtual async Task AppendSessionToken_uses_session_token_select_list() var sessionTokens = context.Database.GetSessionTokens(); var sessionToken = sessionTokens.GetSessionToken()!; - + // Only way we can test this is by setting a session token that will fail the request if used.. // This will take a couple of seconds to fail - sessionTokens.OverwriteSessionToken(sessionToken.Substring(0, sessionToken.IndexOf('#') + 1) + int.MaxValue); + sessionTokens.SetSessionToken(sessionToken.Substring(0, sessionToken.IndexOf('#') + 1) + int.MaxValue); var ex = await Assert.ThrowsAsync(() => context.Customers.Select(x => new { x.Id, x.PartitionKey }).ToListAsync()); Assert.Contains("The read session is not available for the input session token.", ex.ResponseBody); @@ -184,7 +224,7 @@ public virtual async Task AppendSessionToken_uses_session_token_select_list() [ConditionalFact] // @TODO: and sync.. - public virtual async Task AppendSessionToken_uses_session_token_read_item() + public virtual async Task Read_item_uses_session_token() { var contextFactory = await InitializeAsync(); @@ -198,13 +238,14 @@ public virtual async Task AppendSessionToken_uses_session_token_read_item() // Only way we can test this is by setting a session token that will fail the request if used.. // This will take a couple of seconds to fail - sessionTokens.OverwriteSessionToken(sessionToken.Substring(0, sessionToken.IndexOf('#') + 1) + int.MaxValue); + sessionTokens.SetSessionToken(sessionToken.Substring(0, sessionToken.IndexOf('#') + 1) + int.MaxValue); var ex = await Assert.ThrowsAsync(() => context.Customers.FirstOrDefaultAsync(x => x.Id == "1" && x.PartitionKey == "1")); Assert.Contains("The read session is not available for the input session token.", ex.ResponseBody); } [ConditionalTheory, InlineData(AutoTransactionBehavior.WhenNeeded), InlineData(AutoTransactionBehavior.Never), InlineData(AutoTransactionBehavior.Always)] + // @TODO: and sync.. public virtual async Task Add_sets_session_token(AutoTransactionBehavior autoTransactionBehavior) { var contextFactory = await InitializeAsync(); @@ -221,7 +262,8 @@ public virtual async Task Add_sets_session_token(AutoTransactionBehavior autoTra } [ConditionalTheory, InlineData(AutoTransactionBehavior.WhenNeeded), InlineData(AutoTransactionBehavior.Never), InlineData(AutoTransactionBehavior.Always)] - public virtual async Task Delete_sets_session_token(AutoTransactionBehavior autoTransactionBehavior) + // @TODO: and sync.. + public virtual async Task Delete_merges_session_token(AutoTransactionBehavior autoTransactionBehavior) { var contextFactory = await InitializeAsync(); @@ -233,110 +275,46 @@ public virtual async Task Delete_sets_session_token(AutoTransactionBehavior auto await context.SaveChangesAsync(); - var initialToken = context.Database.GetSessionTokens().GetSessionToken(); + var initialToken = context.Database.GetSessionTokens().GetSessionToken()!; context.Remove(customer); await context.SaveChangesAsync(); var sessionToken = context.Database.GetSessionTokens().GetSessionToken(); Assert.False(string.IsNullOrWhiteSpace(sessionToken)); - Assert.StartsWith(initialToken + ",", sessionToken); - Assert.False(string.IsNullOrWhiteSpace(sessionToken)); + Assert.NotEqual(sessionToken, initialToken); + Assert.StartsWith(initialToken.Substring(0, initialToken.IndexOf('#') + 1), sessionToken); } - //[ConditionalTheory, InlineData(AutoTransactionBehavior.WhenNeeded), InlineData(AutoTransactionBehavior.Never), InlineData(AutoTransactionBehavior.Always)] - //public virtual async Task Update_sets_session_token(AutoTransactionBehavior autoTransactionBehavior) - //{ - // var contextFactory = await InitializeAsync(); - - // using var context = contextFactory.CreateContext(); - // context.Database.AutoTransactionBehavior = autoTransactionBehavior; - - // var customer = new Customer { Id = "1", PartitionKey = "1" }; - // context.Customers.Add(customer); - - // await context.SaveChangesAsync(); - - // var initialToken = context.Database.GetSessionTokens().Single().Value; - - // customer.Name = "updated"; - // await context.SaveChangesAsync(); - - // var sessionTokens = context.Database.GetSessionTokens(); - // var key = new CosmosContainerPartitionScope("CosmosSessionTokenContext", new Azure.Cosmos.PartitionKey("1")); - // Assert.Equal(1, sessionTokens.Count); - // Assert.True(sessionTokens.ContainsKey(key)); - // Assert.NotEmpty(sessionTokens[key]); - // Assert.NotEqual(initialToken, sessionTokens[key]); - //} - - //[ConditionalFact] - //public virtual async Task Query_with_single_filter_uses_session_token() - //{ - // var contextFactory = await InitializeAsync(); - - // using var context = contextFactory.CreateContext(); - - // var customer = new Customer { Id = "1", PartitionKey = "1" }; - // context.Customers.Add(customer); - - // await context.SaveChangesAsync(); - - // await context.Customers.Where(x => x.PartitionKey == "1").ToListAsync(); - //} - - //[ConditionalFact] - //public virtual async Task Query_with_double_filter_uses_session_token() - //{ - // var contextFactory = await InitializeAsync(); - - // using var context = contextFactory.CreateContext(); - - // var customer = new Customer { Id = "1", PartitionKey = "1" }; - // context.Customers.Add(customer); - - // await context.SaveChangesAsync(); - - // await context.Customers.Where(x => x.PartitionKey == "1" || x.PartitionKey == "2").ToListAsync(); - //} - - - //[ConditionalFact] - //public virtual async Task Query_with_composite_partition_key_not_all_properties_does_not_use_sessiontoken() - //{ - // var contextFactory = await InitializeAsync(); - - // using var context = contextFactory.CreateContext(); - - // var customer = new CompositeCustomer { Id = "1", PartitionKey1 = "1", PartitionKey2 = "2" }; - // context.Add(customer); - - // await context.SaveChangesAsync(); - - // await context.CompositeCustomers.Where(x => x.PartitionKey1 == "1").ToListAsync(); - //} + [ConditionalTheory, InlineData(AutoTransactionBehavior.WhenNeeded), InlineData(AutoTransactionBehavior.Never), InlineData(AutoTransactionBehavior.Always)] + // @TODO: and sync.. + public virtual async Task Update_merges_session_token(AutoTransactionBehavior autoTransactionBehavior) + { + var contextFactory = await InitializeAsync(); - //[ConditionalFact] - //public virtual async Task Query_with_composite_partition_key_uses_session_token() - //{ - // var contextFactory = await InitializeAsync(); + using var context = contextFactory.CreateContext(); + context.Database.AutoTransactionBehavior = autoTransactionBehavior; - // using var context = contextFactory.CreateContext(); + var customer = new Customer { Id = "1", PartitionKey = "1" }; + context.Customers.Add(customer); - // var customer = new CompositeCustomer { Id = "1", PartitionKey1 = "1", PartitionKey2 = "2" }; - // context.Add(customer); + await context.SaveChangesAsync(); - // await context.SaveChangesAsync(); + var initialToken = context.Database.GetSessionTokens().GetSessionToken()!; - // await context.CompositeCustomers.Where(x => x.PartitionKey1 == "1" && x.PartitionKey2 == "2").ToListAsync(); - //} + customer.Name = "updated"; + await context.SaveChangesAsync(); + var sessionToken = context.Database.GetSessionTokens().GetSessionToken(); + Assert.False(string.IsNullOrWhiteSpace(sessionToken)); + Assert.NotEqual(initialToken, sessionToken); + Assert.StartsWith(initialToken.Substring(0, initialToken.IndexOf('#') + 1), sessionToken); + } public class CosmosSessionTokenContext(DbContextOptions options) : PoolableDbContext(options) { public DbSet Customers { get; set; } = null!; - // public DbSet CompositeCustomers { get; set; } = null!; - + public DbSet OtherContainerCustomers { get; set; } = null!; protected override void OnModelCreating(ModelBuilder builder) { @@ -349,13 +327,13 @@ protected override void OnModelCreating(ModelBuilder builder) b.HasPartitionKey(c => c.PartitionKey); }); - //builder.Entity( - // b => - // { - // b.HasKey(c => c.Id); - // b.HasPartitionKey(c => new { c.PartitionKey1, c.PartitionKey2 }); - // b.ToContainer("composite"); - // }); + builder.Entity( + b => + { + b.HasKey(c => c.Id); + b.HasPartitionKey(c => c.PartitionKey); + b.ToContainer(OtherContainerName); + }); } } @@ -377,11 +355,10 @@ public class DummyChild public string? Id { get; init; } } - //public class CompositeCustomer - //{ - // public string? Id { get; set; } + public class OtherContainerCustomer + { + public string? Id { get; set; } - // public string? PartitionKey1 { get; set; } - // public string? PartitionKey2 { get; set; } - //} + public string? PartitionKey { get; set; } + } } From f295ce9f6b14ebbd1f22b7031042cd28f62aa71b Mon Sep 17 00:00:00 2001 From: JoasE <32096708+JoasE@users.noreply.github.com> Date: Mon, 20 Oct 2025 11:54:27 +0200 Subject: [PATCH 03/37] WIP.. --- .../Storage/Internal/CosmosClientWrapper.cs | 36 ++++++++++++++----- .../CosmosSessionTokensTest.cs | 6 ++-- 2 files changed, 31 insertions(+), 11 deletions(-) diff --git a/src/EFCore.Cosmos/Storage/Internal/CosmosClientWrapper.cs b/src/EFCore.Cosmos/Storage/Internal/CosmosClientWrapper.cs index 03654ac6e11..b3c47e06307 100644 --- a/src/EFCore.Cosmos/Storage/Internal/CosmosClientWrapper.cs +++ b/src/EFCore.Cosmos/Storage/Internal/CosmosClientWrapper.cs @@ -687,6 +687,14 @@ private static async Task ExecuteBatchOnceAsync( using var response = await transactionalBatch.ExecuteAsync(cancellationToken).ConfigureAwait(false); + wrapper._commandLogger.ExecutedTransactionalBatch( + response.Diagnostics.GetClientElapsedTime(), + response.Headers.RequestCharge, + response.Headers.ActivityId, + batch.CollectionId, + batch.PartitionKeyValue, + "[ \"" + string.Join("\", \"", batch.Entries.Select(x => x.Id)) + "\" ]"); + if (!response.IsSuccessStatusCode) { var errorCode = response.StatusCode; @@ -700,14 +708,6 @@ private static async Task ExecuteBatchOnceAsync( return CosmosTransactionalBatchResult.Failure(errorEntries, exception); } - wrapper._commandLogger.ExecutedTransactionalBatch( - response.Diagnostics.GetClientElapsedTime(), - response.Headers.RequestCharge, - response.Headers.ActivityId, - batch.CollectionId, - batch.PartitionKeyValue, - "[ \"" + string.Join("\", \"", batch.Entries.Select(x => x.Id)) + "\" ]"); - ProcessResponse(response, batch.Entries); return CosmosTransactionalBatchResult.Success(response.Headers.Session); @@ -936,7 +936,14 @@ private static Task CreateSingleItemQueryAsync( using var reader = new StreamReader(responseStream); using var jsonReader = new JsonTextReader(reader); - return Serializer.Deserialize(jsonReader); + var jobject = Serializer.Deserialize(jsonReader); + + if (!string.IsNullOrWhiteSpace(responseMessage.Headers.Session)) + { + // @TODO: Set session token.. + } + + return jobject; } /// @@ -1068,6 +1075,11 @@ public bool MoveNext() _responseMessage.EnsureSuccessStatusCode(); + if (!string.IsNullOrWhiteSpace(_responseMessage.Headers.Session)) + { + // @TODO: set session token... + } + _responseMessageEnumerator = new ResponseMessageEnumerable(_responseMessage).GetEnumerator(); } @@ -1181,6 +1193,12 @@ public async ValueTask MoveNextAsync() _responseMessage.EnsureSuccessStatusCode(); + if (!string.IsNullOrWhiteSpace(_responseMessage.Headers.Session)) + { + // @TODO: set session token... + // if (appendSessionToken == higher.....) update _sessionToken?? + } + _responseMessageEnumerator = new ResponseMessageEnumerable(_responseMessage).GetAsyncEnumerator(cancellationToken); } diff --git a/test/EFCore.Cosmos.FunctionalTests/CosmosSessionTokensTest.cs b/test/EFCore.Cosmos.FunctionalTests/CosmosSessionTokensTest.cs index 8825232cc19..2d9c42ebc92 100644 --- a/test/EFCore.Cosmos.FunctionalTests/CosmosSessionTokensTest.cs +++ b/test/EFCore.Cosmos.FunctionalTests/CosmosSessionTokensTest.cs @@ -181,20 +181,22 @@ public virtual async Task GetSessionToken_no_token_returns_null() } [ConditionalFact] - // @TODO: Read item and select... // @TODO: and sync.. public virtual async Task Query_uses_session_token() { var contextFactory = await InitializeAsync(); using var context = contextFactory.CreateContext(); + context.Customers.Add(new Customer { Id = "1", PartitionKey = "1" }); + + await context.SaveChangesAsync(); var sessionTokens = context.Database.GetSessionTokens(); var sessionToken = sessionTokens.GetSessionToken()!; // Only way we can test this is by setting a session token that will fail the request if used.. // This will take a couple of seconds to fail - sessionTokens.SetSessionToken(sessionToken.Substring(0, sessionToken.IndexOf('#') + 1) + int.MaxValue); + // sessionTokens.SetSessionToken(sessionToken.Substring(0, sessionToken.IndexOf('#') + 1) + int.MaxValue); var ex = await Assert.ThrowsAsync(() => context.Customers.ToListAsync()); Assert.Contains("The read session is not available for the input session token.", ex.ResponseBody); From ad8d93066a89239ede29f983b3442e486cc3b510 Mon Sep 17 00:00:00 2001 From: JoasE <32096708+JoasE@users.noreply.github.com> Date: Mon, 20 Oct 2025 15:32:25 +0200 Subject: [PATCH 04/37] Make Client responsible --- .../CosmosDatabaseFacadeExtensions.cs | 6 +- .../Internal/CosmosQueryCompilationContext.cs | 14 +- .../Query/Internal/CosmosQueryContext.cs | 1 + ...ressionVisitor.PagingQueryingEnumerable.cs | 12 +- ...ingExpressionVisitor.QueryingEnumerable.cs | 24 +- ...ssionVisitor.ReadItemQueryingEnumerable.cs | 14 +- ...osShapedQueryCompilingExpressionVisitor.cs | 7 +- .../Storage/ISessionTokenStorage.cs | 49 ++++ .../Storage/Internal/CosmosClientWrapper.cs | 220 ++++++++++++------ .../Storage/Internal/CosmosDatabaseWrapper.cs | 76 ++---- .../Storage/Internal/ICosmosClientWrapper.cs | 30 ++- .../Internal/NullSessionTokenStorage.cs | 60 +++++ .../{ => Internal}/SessionTokenStorage.cs | 56 +++-- .../CosmosSessionTokensTest.cs | 182 ++++++++++++++- .../TestUtilities/CosmosTestStore.cs | 2 +- 15 files changed, 561 insertions(+), 192 deletions(-) create mode 100644 src/EFCore.Cosmos/Storage/ISessionTokenStorage.cs create mode 100644 src/EFCore.Cosmos/Storage/Internal/NullSessionTokenStorage.cs rename src/EFCore.Cosmos/Storage/{ => Internal}/SessionTokenStorage.cs (85%) diff --git a/src/EFCore.Cosmos/Extensions/CosmosDatabaseFacadeExtensions.cs b/src/EFCore.Cosmos/Extensions/CosmosDatabaseFacadeExtensions.cs index 5d1ba8e9ded..65080c57aa1 100644 --- a/src/EFCore.Cosmos/Extensions/CosmosDatabaseFacadeExtensions.cs +++ b/src/EFCore.Cosmos/Extensions/CosmosDatabaseFacadeExtensions.cs @@ -27,11 +27,11 @@ public static CosmosClient GetCosmosClient(this DatabaseFacade databaseFacade) => GetService(databaseFacade).Client; /// - /// Gets used to manage the session tokens for this . + /// Gets used to manage the session tokens for this . /// /// The for the context. - /// The Gets . - public static SessionTokenStorage GetSessionTokens(this DatabaseFacade databaseFacade) + /// The . + public static ISessionTokenStorage GetSessionTokens(this DatabaseFacade databaseFacade) { var db = GetService(databaseFacade); if (db is not CosmosDatabaseWrapper dbWrapper) diff --git a/src/EFCore.Cosmos/Query/Internal/CosmosQueryCompilationContext.cs b/src/EFCore.Cosmos/Query/Internal/CosmosQueryCompilationContext.cs index 67360709119..a106478358d 100644 --- a/src/EFCore.Cosmos/Query/Internal/CosmosQueryCompilationContext.cs +++ b/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; /// @@ -19,6 +21,8 @@ public class CosmosQueryCompilationContext : QueryCompilationContext /// public CosmosQueryCompilationContext(QueryCompilationContextDependencies dependencies, bool async) : base(dependencies, async) { + // @TODO: Faster way.. + SessionTokenStorage = Dependencies.Context.Database.GetSessionTokens(); } /// @@ -57,7 +61,7 @@ public CosmosQueryCompilationContext(QueryCompilationContextDependencies depende /// 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. /// - public string? SessionToken { get; internal set; } + public virtual string? SessionToken { get; internal set; } /// /// This is an internal API that supports the Entity Framework Core infrastructure and not subject to @@ -65,11 +69,5 @@ public CosmosQueryCompilationContext(QueryCompilationContextDependencies depende /// 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. /// - public string? GetSessionToken() - { - Debug.Assert(RootEntityType != null); - var container = RootEntityType.GetContainer(); - Debug.Assert(container != null); - return SessionToken ?? Dependencies.Context.Database.GetSessionTokens().GetSessionToken(container); - } + public virtual ISessionTokenStorage SessionTokenStorage { get; } } diff --git a/src/EFCore.Cosmos/Query/Internal/CosmosQueryContext.cs b/src/EFCore.Cosmos/Query/Internal/CosmosQueryContext.cs index 09f0723b959..842f2044047 100644 --- a/src/EFCore.Cosmos/Query/Internal/CosmosQueryContext.cs +++ b/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; diff --git a/src/EFCore.Cosmos/Query/Internal/CosmosShapedQueryCompilingExpressionVisitor.PagingQueryingEnumerable.cs b/src/EFCore.Cosmos/Query/Internal/CosmosShapedQueryCompilingExpressionVisitor.PagingQueryingEnumerable.cs index 41f7daddab5..855aa3eea60 100644 --- a/src/EFCore.Cosmos/Query/Internal/CosmosShapedQueryCompilingExpressionVisitor.PagingQueryingEnumerable.cs +++ b/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 : IAsyncEnumerable> private readonly Type _contextType; private readonly string _cosmosContainer; private readonly PartitionKey _cosmosPartitionKey; - private readonly string _sessionToken; private readonly IDiagnosticsLogger _queryLogger; private readonly IDiagnosticsLogger _commandLogger; private readonly bool _standAloneStateManager; + private readonly ISessionTokenStorage _sessionTokenStorage; + private readonly string _sessionToken; private readonly CancellationToken _cancellationToken; private readonly IConcurrencyDetector _concurrencyDetector; private readonly IExceptionDetector _exceptionDetector; @@ -109,11 +114,12 @@ public AsyncEnumerator(PagingQueryingEnumerable queryingEnumerable, Cancellat _contextType = queryingEnumerable._contextType; _cosmosContainer = queryingEnumerable._cosmosContainer; _cosmosPartitionKey = queryingEnumerable._cosmosPartitionKey; - _sessionToken = queryingEnumerable._sessionToken; _queryLogger = queryingEnumerable._queryLogger; _commandLogger = queryingEnumerable._commandLogger; _standAloneStateManager = queryingEnumerable._standAloneStateManager; _exceptionDetector = _cosmosQueryContext.ExceptionDetector; + _sessionTokenStorage = queryingEnumerable._sessionTokenStorage; + _sessionToken = queryingEnumerable._sessionToken; _cancellationToken = cancellationToken; _concurrencyDetector = queryingEnumerable._threadSafetyChecksEnabled @@ -173,7 +179,7 @@ public async ValueTask 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); diff --git a/src/EFCore.Cosmos/Query/Internal/CosmosShapedQueryCompilingExpressionVisitor.QueryingEnumerable.cs b/src/EFCore.Cosmos/Query/Internal/CosmosShapedQueryCompilingExpressionVisitor.QueryingEnumerable.cs index 1d6be1fe8bf..7143e8595cb 100644 --- a/src/EFCore.Cosmos/Query/Internal/CosmosShapedQueryCompilingExpressionVisitor.QueryingEnumerable.cs +++ b/src/EFCore.Cosmos/Query/Internal/CosmosShapedQueryCompilingExpressionVisitor.QueryingEnumerable.cs @@ -5,6 +5,7 @@ using System.Collections; using System.Text; +using Microsoft.EntityFrameworkCore.Cosmos.Storage; using Microsoft.EntityFrameworkCore.Cosmos.Storage.Internal; using Newtonsoft.Json.Linq; @@ -27,11 +28,12 @@ private sealed class QueryingEnumerable : IEnumerable, IAsyncEnumerable private readonly IQuerySqlGeneratorFactory _querySqlGeneratorFactory; private readonly Type _contextType; private readonly string _cosmosContainer; - private readonly string _sessionToken; private readonly PartitionKey _cosmosPartitionKey; private readonly IDiagnosticsLogger _queryLogger; private readonly bool _standAloneStateManager; private readonly bool _threadSafetyChecksEnabled; + private readonly ISessionTokenStorage _sessionTokenStorage; + private readonly string _sessionToken; public QueryingEnumerable( CosmosQueryContext cosmosQueryContext, @@ -44,6 +46,7 @@ public QueryingEnumerable( List partitionKeyPropertyValues, bool standAloneStateManager, bool threadSafetyChecksEnabled, + ISessionTokenStorage sessionTokenStorage, string sessionToken) { _cosmosQueryContext = cosmosQueryContext; @@ -55,6 +58,7 @@ public QueryingEnumerable( _queryLogger = cosmosQueryContext.QueryLogger; _standAloneStateManager = standAloneStateManager; _threadSafetyChecksEnabled = threadSafetyChecksEnabled; + _sessionTokenStorage = sessionTokenStorage; _sessionToken = sessionToken; _cosmosContainer = rootEntityType.GetContainer() @@ -109,12 +113,13 @@ private sealed class Enumerator : IEnumerator private readonly Func _shaper; private readonly Type _contextType; private readonly string _cosmosContainer; - private readonly string _sessionToken; private readonly PartitionKey _cosmosPartitionKey; private readonly IDiagnosticsLogger _queryLogger; private readonly bool _standAloneStateManager; private readonly IConcurrencyDetector _concurrencyDetector; private readonly IExceptionDetector _exceptionDetector; + private readonly string _sessionToken; + private readonly ISessionTokenStorage _sessionTokenStorage; private IEnumerator _enumerator; @@ -125,11 +130,12 @@ public Enumerator(QueryingEnumerable queryingEnumerable) _shaper = queryingEnumerable._shaper; _contextType = queryingEnumerable._contextType; _cosmosContainer = queryingEnumerable._cosmosContainer; - _sessionToken = queryingEnumerable._sessionToken; _cosmosPartitionKey = queryingEnumerable._cosmosPartitionKey; _queryLogger = queryingEnumerable._queryLogger; _standAloneStateManager = queryingEnumerable._standAloneStateManager; _exceptionDetector = _cosmosQueryContext.ExceptionDetector; + _sessionToken = queryingEnumerable._sessionToken; + _sessionTokenStorage = queryingEnumerable._sessionTokenStorage; _concurrencyDetector = queryingEnumerable._threadSafetyChecksEnabled ? _cosmosQueryContext.ConcurrencyDetector @@ -154,7 +160,7 @@ public bool MoveNext() EntityFrameworkMetricsData.ReportQueryExecuting(); _enumerator = _cosmosQueryContext.CosmosClient - .ExecuteSqlQuery(_cosmosContainer, _cosmosPartitionKey, sqlQuery, _sessionToken) + .ExecuteSqlQuery(_cosmosContainer, _cosmosPartitionKey, sqlQuery, _sessionTokenStorage, _sessionToken) .GetEnumerator(); _cosmosQueryContext.InitializeStateManager(_standAloneStateManager); } @@ -201,13 +207,15 @@ private sealed class AsyncEnumerator : IAsyncEnumerator private readonly Type _contextType; private readonly string _cosmosContainer; private readonly PartitionKey _cosmosPartitionKey; - private readonly string _sessionToken; private readonly IDiagnosticsLogger _queryLogger; private readonly bool _standAloneStateManager; + private readonly string _sessionToken; + private readonly ISessionTokenStorage _sessionTokenStorage; private readonly CancellationToken _cancellationToken; private readonly IConcurrencyDetector _concurrencyDetector; private readonly IExceptionDetector _exceptionDetector; + private IAsyncEnumerator _enumerator; public AsyncEnumerator(QueryingEnumerable queryingEnumerable, CancellationToken cancellationToken) @@ -218,10 +226,12 @@ public AsyncEnumerator(QueryingEnumerable queryingEnumerable, CancellationTok _contextType = queryingEnumerable._contextType; _cosmosContainer = queryingEnumerable._cosmosContainer; _cosmosPartitionKey = queryingEnumerable._cosmosPartitionKey; - _sessionToken = queryingEnumerable._sessionToken; _queryLogger = queryingEnumerable._queryLogger; _standAloneStateManager = queryingEnumerable._standAloneStateManager; _exceptionDetector = _cosmosQueryContext.ExceptionDetector; + _sessionToken = queryingEnumerable._sessionToken; + _sessionTokenStorage = queryingEnumerable._sessionTokenStorage; + _cancellationToken = cancellationToken; _concurrencyDetector = queryingEnumerable._threadSafetyChecksEnabled @@ -244,7 +254,7 @@ public async ValueTask MoveNextAsync() EntityFrameworkMetricsData.ReportQueryExecuting(); _enumerator = _cosmosQueryContext.CosmosClient - .ExecuteSqlQueryAsync(_cosmosContainer, _cosmosPartitionKey, sqlQuery, _sessionToken) + .ExecuteSqlQueryAsync(_cosmosContainer, _cosmosPartitionKey, sqlQuery, _sessionTokenStorage, _sessionToken) .GetAsyncEnumerator(_cancellationToken); _cosmosQueryContext.InitializeStateManager(_standAloneStateManager); } diff --git a/src/EFCore.Cosmos/Query/Internal/CosmosShapedQueryCompilingExpressionVisitor.ReadItemQueryingEnumerable.cs b/src/EFCore.Cosmos/Query/Internal/CosmosShapedQueryCompilingExpressionVisitor.ReadItemQueryingEnumerable.cs index 79e45ae3dea..9e6bcbfba33 100644 --- a/src/EFCore.Cosmos/Query/Internal/CosmosShapedQueryCompilingExpressionVisitor.ReadItemQueryingEnumerable.cs +++ b/src/EFCore.Cosmos/Query/Internal/CosmosShapedQueryCompilingExpressionVisitor.ReadItemQueryingEnumerable.cs @@ -6,6 +6,7 @@ using System.Collections; using Microsoft.EntityFrameworkCore.Cosmos.Internal; using Microsoft.EntityFrameworkCore.Cosmos.Metadata.Internal; +using Microsoft.EntityFrameworkCore.Cosmos.Storage; using Newtonsoft.Json.Linq; namespace Microsoft.EntityFrameworkCore.Cosmos.Query.Internal; @@ -25,12 +26,13 @@ private sealed class ReadItemQueryingEnumerable : IEnumerable, IAsyncEnume private readonly string _cosmosContainer; private readonly ReadItemInfo _readItemInfo; private readonly PartitionKey _cosmosPartitionKey; - private readonly string _sessionToken; private readonly Func _shaper; private readonly Type _contextType; private readonly IDiagnosticsLogger _queryLogger; private readonly bool _standAloneStateManager; private readonly bool _threadSafetyChecksEnabled; + private readonly ISessionTokenStorage _sessionTokenStorage; + private readonly string _sessionToken; public ReadItemQueryingEnumerable( CosmosQueryContext cosmosQueryContext, @@ -41,6 +43,7 @@ public ReadItemQueryingEnumerable( Type contextType, bool standAloneStateManager, bool threadSafetyChecksEnabled, + ISessionTokenStorage sessionTokenStorage, string sessionToken) { _cosmosQueryContext = cosmosQueryContext; @@ -51,6 +54,7 @@ public ReadItemQueryingEnumerable( _queryLogger = _cosmosQueryContext.QueryLogger; _standAloneStateManager = standAloneStateManager; _threadSafetyChecksEnabled = threadSafetyChecksEnabled; + _sessionTokenStorage = sessionTokenStorage; _sessionToken = sessionToken; _cosmosContainer = rootEntityType.GetContainer() ?? throw new UnreachableException("Root entity type without a Cosmos container."); @@ -107,7 +111,6 @@ private sealed class Enumerator : IEnumerator, IAsyncEnumerator private readonly CosmosQueryContext _cosmosQueryContext; private readonly string _cosmosContainer; private readonly PartitionKey _cosmosPartitionKey; - private readonly string _sessionToken; private readonly Func _shaper; private readonly Type _contextType; private readonly IDiagnosticsLogger _queryLogger; @@ -115,6 +118,8 @@ private sealed class Enumerator : IEnumerator, IAsyncEnumerator private readonly IConcurrencyDetector _concurrencyDetector; private readonly IExceptionDetector _exceptionDetector; private readonly ReadItemQueryingEnumerable _readItemEnumerable; + private readonly string _sessionToken; + private readonly ISessionTokenStorage _sessionTokenStorage; private readonly CancellationToken _cancellationToken; private JObject _item; @@ -125,13 +130,14 @@ public Enumerator(ReadItemQueryingEnumerable readItemEnumerable, Cancellation _cosmosQueryContext = readItemEnumerable._cosmosQueryContext; _cosmosContainer = readItemEnumerable._cosmosContainer; _cosmosPartitionKey = readItemEnumerable._cosmosPartitionKey; - _sessionToken = readItemEnumerable._sessionToken; _shaper = readItemEnumerable._shaper; _contextType = readItemEnumerable._contextType; _queryLogger = readItemEnumerable._queryLogger; _standAloneStateManager = readItemEnumerable._standAloneStateManager; _exceptionDetector = _cosmosQueryContext.ExceptionDetector; _readItemEnumerable = readItemEnumerable; + _sessionTokenStorage = readItemEnumerable._sessionTokenStorage; + _sessionToken = readItemEnumerable._sessionToken; _cancellationToken = cancellationToken; _concurrencyDetector = readItemEnumerable._threadSafetyChecksEnabled @@ -166,6 +172,7 @@ public bool MoveNext() _cosmosContainer, _cosmosPartitionKey, resourceId, + _sessionTokenStorage, _sessionToken); return ShapeResult(); @@ -207,6 +214,7 @@ public async ValueTask MoveNextAsync() _cosmosContainer, _cosmosPartitionKey, resourceId, + _sessionTokenStorage, _sessionToken, _cancellationToken) .ConfigureAwait(false); diff --git a/src/EFCore.Cosmos/Query/Internal/CosmosShapedQueryCompilingExpressionVisitor.cs b/src/EFCore.Cosmos/Query/Internal/CosmosShapedQueryCompilingExpressionVisitor.cs index d7c462389a0..9aeae2ebf5c 100644 --- a/src/EFCore.Cosmos/Query/Internal/CosmosShapedQueryCompilingExpressionVisitor.cs +++ b/src/EFCore.Cosmos/Query/Internal/CosmosShapedQueryCompilingExpressionVisitor.cs @@ -83,7 +83,9 @@ protected override Expression VisitShapedQuery(ShapedQueryExpression shapedQuery var threadSafetyConstant = Constant(_threadSafetyChecksEnabled); var standAloneStateManagerConstant = Constant( QueryCompilationContext.QueryTrackingBehavior == QueryTrackingBehavior.NoTrackingWithIdentityResolution); - var sessionTokenConstant = Constant(cosmosQueryCompilationContext.GetSessionToken()); + var sessionTokenStorageConstant = Constant(cosmosQueryCompilationContext.SessionTokenStorage); + var sessionToken = cosmosQueryCompilationContext.SessionToken ?? cosmosQueryCompilationContext.SessionTokenStorage.GetSessionToken(rootEntityType.GetContainer()!); + var sessionTokenConstant = Constant(sessionToken, typeof(string)); Check.DebugAssert(!paging || selectExpression.ReadItemInfo is null, "ReadItem is being with paging, impossible."); @@ -99,6 +101,7 @@ protected override Expression VisitShapedQuery(ShapedQueryExpression shapedQuery contextTypeConstant, standAloneStateManagerConstant, threadSafetyConstant, + sessionTokenStorageConstant, sessionTokenConstant), _ when paging => New( @@ -116,6 +119,7 @@ protected override Expression VisitShapedQuery(ShapedQueryExpression shapedQuery Constant(maxItemCount.Name), Constant(continuationToken.Name), Constant(responseContinuationTokenLimitInKb.Name), + sessionTokenStorageConstant, sessionTokenConstant), _ => New( @@ -129,6 +133,7 @@ protected override Expression VisitShapedQuery(ShapedQueryExpression shapedQuery Constant(cosmosQueryCompilationContext.PartitionKeyPropertyValues), standAloneStateManagerConstant, threadSafetyConstant, + sessionTokenStorageConstant, sessionTokenConstant) }; } diff --git a/src/EFCore.Cosmos/Storage/ISessionTokenStorage.cs b/src/EFCore.Cosmos/Storage/ISessionTokenStorage.cs new file mode 100644 index 00000000000..f4b6467fbb4 --- /dev/null +++ b/src/EFCore.Cosmos/Storage/ISessionTokenStorage.cs @@ -0,0 +1,49 @@ +// 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.Storage; + +/// +/// Defines methods for storing, retrieving, and managing session tokens associated with containers in a . +/// +public interface ISessionTokenStorage +{ + /// + /// Appends or merges the session token specified to any composite already stored in the storage for the default container. + /// + /// The session token to append or merge. + void AppendSessionToken(string sessionToken); + + /// + /// Appends or merges the session token specified to any composite already stored in the storage for the specified container. + /// + /// The name of the container to append the session token for. + /// The session token to append or merge. + void AppendSessionToken(string containerName, string sessionToken); + + /// + /// Gets the composite session token for the default container. + /// + /// The composite session token, or if none is stored. + string? GetSessionToken(); + + /// + /// Gets the composite session token for the specified container. + /// + /// The name of the container to get the session token for. + /// The composite session token, or if none is stored. + string? GetSessionToken(string containerName); + + /// + /// Overwrites the session token for the container. + /// + /// The name of the container to set the session token for. + /// The session token to set, or to clear any stored token. + void SetSessionToken(string containerName, string? sessionToken); + + /// + /// Overwrites the session token for the default container. + /// + /// The session token to set, or to clear any stored token. + void SetSessionToken(string? sessionToken); +} diff --git a/src/EFCore.Cosmos/Storage/Internal/CosmosClientWrapper.cs b/src/EFCore.Cosmos/Storage/Internal/CosmosClientWrapper.cs index b3c47e06307..f891fbed40b 100644 --- a/src/EFCore.Cosmos/Storage/Internal/CosmosClientWrapper.cs +++ b/src/EFCore.Cosmos/Storage/Internal/CosmosClientWrapper.cs @@ -385,19 +385,20 @@ private static string GetPathFromRoot(IReadOnlyEntityType entityType) /// 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. /// - public virtual CosmosWriteResult CreateItem( + public virtual bool CreateItem( string containerId, JToken document, - IUpdateEntry entry) + IUpdateEntry entry, + ISessionTokenStorage sessionTokenStorage) { _databaseLogger.SyncNotSupported(); - return _executionStrategy.Execute((containerId, document, entry, this), CreateItemOnce, null); + return _executionStrategy.Execute((containerId, document, entry, sessionTokenStorage, this), CreateItemOnce, null); } - private static CosmosWriteResult CreateItemOnce( + private static bool CreateItemOnce( DbContext context, - (string ContainerId, JToken Document, IUpdateEntry Entry, CosmosClientWrapper Wrapper) parameters) + (string ContainerId, JToken Document, IUpdateEntry Entry, ISessionTokenStorage SessionTokenStorage, CosmosClientWrapper Wrapper) parameters) => CreateItemOnceAsync(context, parameters).GetAwaiter().GetResult(); /// @@ -406,24 +407,27 @@ private static CosmosWriteResult CreateItemOnce( /// 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. /// - public virtual Task CreateItemAsync( + public virtual Task CreateItemAsync( string containerId, JToken document, IUpdateEntry updateEntry, + ISessionTokenStorage sessionTokenStorage, CancellationToken cancellationToken = default) - => _executionStrategy.ExecuteAsync((containerId, document, updateEntry, this), CreateItemOnceAsync, null, cancellationToken); + => _executionStrategy.ExecuteAsync((containerId, document, updateEntry, sessionTokenStorage, this), CreateItemOnceAsync, null, cancellationToken); - private static async Task CreateItemOnceAsync( + private static async Task CreateItemOnceAsync( DbContext _, - (string ContainerId, JToken Document, IUpdateEntry Entry, CosmosClientWrapper Wrapper) parameters, + (string ContainerId, JToken Document, IUpdateEntry Entry, ISessionTokenStorage SessionTokenStorage, CosmosClientWrapper Wrapper) parameters, CancellationToken cancellationToken = default) { using var stream = Serialize(parameters.Document); + var containerId = parameters.ContainerId; var entry = parameters.Entry; var wrapper = parameters.Wrapper; + var sessionTokenStorage = parameters.SessionTokenStorage; var container = wrapper.Client.GetDatabase(wrapper._databaseId).GetContainer(parameters.ContainerId); - var itemRequestOptions = CreateItemRequestOptions(entry, wrapper._enableContentResponseOnWrite); + var itemRequestOptions = CreateItemRequestOptions(entry, wrapper._enableContentResponseOnWrite, sessionTokenStorage.GetSessionToken(containerId)); var partitionKeyValue = ExtractPartitionKeyValue(entry); var preTriggers = GetTriggers(entry, TriggerType.Pre, TriggerOperation.Create); var postTriggers = GetTriggers(entry, TriggerType.Post, TriggerOperation.Create); @@ -453,12 +457,12 @@ private static async Task CreateItemOnceAsync( response.Headers.RequestCharge, response.Headers.ActivityId, parameters.Document["id"]!.ToString(), - parameters.ContainerId, + containerId, partitionKeyValue); - ProcessResponse(response, entry); + ProcessResponse(containerId, response, entry, sessionTokenStorage); - return CosmosWriteResult.Success(response.Headers.Session); + return response.StatusCode == HttpStatusCode.Created; } /// @@ -467,20 +471,21 @@ private static async Task CreateItemOnceAsync( /// 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. /// - public virtual CosmosWriteResult ReplaceItem( + public virtual bool ReplaceItem( string collectionId, string documentId, JObject document, - IUpdateEntry entry) + IUpdateEntry entry, + ISessionTokenStorage sessionTokenStorage) { _databaseLogger.SyncNotSupported(); - return _executionStrategy.Execute((collectionId, documentId, document, entry, this), ReplaceItemOnce, null); + return _executionStrategy.Execute((collectionId, documentId, document, entry, sessionTokenStorage, this), ReplaceItemOnce, null); } - private static CosmosWriteResult ReplaceItemOnce( + private static bool ReplaceItemOnce( DbContext context, - (string ContainerId, string ItemId, JObject Document, IUpdateEntry Entry, CosmosClientWrapper Wrapper) parameters) + (string ContainerId, string ItemId, JObject Document, IUpdateEntry Entry, ISessionTokenStorage SessionTokenStorage, CosmosClientWrapper Wrapper) parameters) => ReplaceItemOnceAsync(context, parameters).GetAwaiter().GetResult(); /// @@ -489,26 +494,29 @@ private static CosmosWriteResult ReplaceItemOnce( /// 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. /// - public virtual Task ReplaceItemAsync( + public virtual Task ReplaceItemAsync( string collectionId, string documentId, JObject document, IUpdateEntry updateEntry, + ISessionTokenStorage sessionTokenStorage, CancellationToken cancellationToken = default) => _executionStrategy.ExecuteAsync( - (collectionId, documentId, document, updateEntry, this), ReplaceItemOnceAsync, null, cancellationToken); + (collectionId, documentId, document, updateEntry, sessionTokenStorage, this), ReplaceItemOnceAsync, null, cancellationToken); - private static async Task ReplaceItemOnceAsync( + private static async Task ReplaceItemOnceAsync( DbContext _, - (string ContainerId, string ResourceId, JObject Document, IUpdateEntry Entry, CosmosClientWrapper Wrapper) parameters, + (string ContainerId, string ResourceId, JObject Document, IUpdateEntry Entry, ISessionTokenStorage SessionTokenStorage, CosmosClientWrapper Wrapper) parameters, CancellationToken cancellationToken = default) { using var stream = Serialize(parameters.Document); + var containerId = parameters.ContainerId; var entry = parameters.Entry; var wrapper = parameters.Wrapper; + var sessionTokenStorage = parameters.SessionTokenStorage; var container = wrapper.Client.GetDatabase(wrapper._databaseId).GetContainer(parameters.ContainerId); - var itemRequestOptions = CreateItemRequestOptions(entry, wrapper._enableContentResponseOnWrite); + var itemRequestOptions = CreateItemRequestOptions(entry, wrapper._enableContentResponseOnWrite, sessionTokenStorage.GetSessionToken(containerId)); var partitionKeyValue = ExtractPartitionKeyValue(entry); var preTriggers = GetTriggers(entry, TriggerType.Pre, TriggerOperation.Replace); var postTriggers = GetTriggers(entry, TriggerType.Post, TriggerOperation.Replace); @@ -539,12 +547,12 @@ private static async Task ReplaceItemOnceAsync( response.Headers.RequestCharge, response.Headers.ActivityId, parameters.ResourceId, - parameters.ContainerId, + containerId, partitionKeyValue); - ProcessResponse(response, entry); + ProcessResponse(containerId, response, entry, sessionTokenStorage); - return CosmosWriteResult.Success(response.Headers.Session); + return response.StatusCode == HttpStatusCode.OK; } /// @@ -553,19 +561,20 @@ private static async Task ReplaceItemOnceAsync( /// 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. /// - public virtual CosmosWriteResult DeleteItem( + public virtual bool DeleteItem( string containerId, string documentId, - IUpdateEntry entry) + IUpdateEntry entry, + ISessionTokenStorage sessionTokenStorage) { _databaseLogger.SyncNotSupported(); - return _executionStrategy.Execute((containerId, documentId, entry, this), DeleteItemOnce, null); + return _executionStrategy.Execute((containerId, documentId, entry, sessionTokenStorage, this), DeleteItemOnce, null); } - private static CosmosWriteResult DeleteItemOnce( + private static bool DeleteItemOnce( DbContext context, - (string ContainerId, string DocumentId, IUpdateEntry Entry, CosmosClientWrapper Wrapper) parameters) + (string ContainerId, string DocumentId, IUpdateEntry Entry, ISessionTokenStorage SessionTokenStorage, CosmosClientWrapper Wrapper) parameters) => DeleteItemOnceAsync(context, parameters).GetAwaiter().GetResult(); /// @@ -574,23 +583,26 @@ private static CosmosWriteResult DeleteItemOnce( /// 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. /// - public virtual Task DeleteItemAsync( + public virtual Task DeleteItemAsync( string containerId, string documentId, IUpdateEntry entry, + ISessionTokenStorage sessionTokenStorage, CancellationToken cancellationToken = default) - => _executionStrategy.ExecuteAsync((containerId, documentId, entry, this), DeleteItemOnceAsync, null, cancellationToken); + => _executionStrategy.ExecuteAsync((containerId, documentId, entry, sessionTokenStorage, this), DeleteItemOnceAsync, null, cancellationToken); - private static async Task DeleteItemOnceAsync( + private static async Task DeleteItemOnceAsync( DbContext? _, - (string ContainerId, string ResourceId, IUpdateEntry Entry, CosmosClientWrapper Wrapper) parameters, + (string ContainerId, string ResourceId, IUpdateEntry Entry, ISessionTokenStorage SessionTokenStorage, CosmosClientWrapper Wrapper) parameters, CancellationToken cancellationToken = default) { + var containerId = parameters.ContainerId; var entry = parameters.Entry; var wrapper = parameters.Wrapper; + var sessionTokenStorage = parameters.SessionTokenStorage; var items = wrapper.Client.GetDatabase(wrapper._databaseId).GetContainer(parameters.ContainerId); - var itemRequestOptions = CreateItemRequestOptions(entry, wrapper._enableContentResponseOnWrite); + var itemRequestOptions = CreateItemRequestOptions(entry, wrapper._enableContentResponseOnWrite, sessionTokenStorage.GetSessionToken(containerId)); var partitionKeyValue = ExtractPartitionKeyValue(entry); var preTriggers = GetTriggers(entry, TriggerType.Pre, TriggerOperation.Delete); var postTriggers = GetTriggers(entry, TriggerType.Post, TriggerOperation.Delete); @@ -620,12 +632,12 @@ private static async Task DeleteItemOnceAsync( response.Headers.RequestCharge, response.Headers.ActivityId, parameters.ResourceId, - parameters.ContainerId, + containerId, partitionKeyValue); - ProcessResponse(response, entry); + ProcessResponse(containerId, response, entry, sessionTokenStorage); - return CosmosWriteResult.Success(response.Headers.Session); + return response.StatusCode == HttpStatusCode.NoContent; } /// @@ -646,6 +658,7 @@ public virtual PartitionKey GetPartitionKeyValue(IUpdateEntry updateEntry) public virtual ICosmosTransactionalBatchWrapper CreateTransactionalBatch(string containerId, PartitionKey partitionKeyValue, bool checkSize) { var container = Client.GetDatabase(_databaseId).GetContainer(containerId); + var batch = container.CreateTransactionalBatch(partitionKeyValue); return new CosmosTransactionalBatchWrapper(batch, containerId, partitionKeyValue, checkSize, _enableContentResponseOnWrite); @@ -657,15 +670,15 @@ public virtual ICosmosTransactionalBatchWrapper CreateTransactionalBatch(string /// 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. /// - public virtual CosmosTransactionalBatchResult ExecuteTransactionalBatch(ICosmosTransactionalBatchWrapper batch) + public virtual CosmosTransactionalBatchResult ExecuteTransactionalBatch(ICosmosTransactionalBatchWrapper batch, ISessionTokenStorage sessionTokenStorage) { _databaseLogger.SyncNotSupported(); - return _executionStrategy.Execute((batch, this), ExecuteBatchOnce, null); + return _executionStrategy.Execute((batch, sessionTokenStorage, this), ExecuteBatchOnce, null); } private static CosmosTransactionalBatchResult ExecuteBatchOnce(DbContext _, - (ICosmosTransactionalBatchWrapper Batch, CosmosClientWrapper Wrapper) parameters) + (ICosmosTransactionalBatchWrapper Batch, ISessionTokenStorage SessionTokenStorage, CosmosClientWrapper Wrapper) parameters) => ExecuteBatchOnceAsync(_, parameters).GetAwaiter().GetResult(); /// @@ -674,18 +687,26 @@ private static CosmosTransactionalBatchResult ExecuteBatchOnce(DbContext _, /// 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. /// - public virtual Task ExecuteTransactionalBatchAsync(ICosmosTransactionalBatchWrapper batch, CancellationToken cancellationToken = default) - => _executionStrategy.ExecuteAsync((batch, this), ExecuteBatchOnceAsync, null, cancellationToken); + public virtual Task ExecuteTransactionalBatchAsync(ICosmosTransactionalBatchWrapper batch, ISessionTokenStorage sessionTokenStorage, CancellationToken cancellationToken = default) + => _executionStrategy.ExecuteAsync((batch, sessionTokenStorage, this), ExecuteBatchOnceAsync, null, cancellationToken); private static async Task ExecuteBatchOnceAsync(DbContext _, - (ICosmosTransactionalBatchWrapper Batch, CosmosClientWrapper Wrapper) parameters, + (ICosmosTransactionalBatchWrapper Batch, ISessionTokenStorage SessionTokenStorage, CosmosClientWrapper Wrapper) parameters, CancellationToken cancellationToken = default) { var batch = parameters.Batch; var transactionalBatch = batch.GetTransactionalBatch(); var wrapper = parameters.Wrapper; + var sessionTokenStorage = parameters.SessionTokenStorage; + + var options = new TransactionalBatchRequestOptions(); + var sessionToken = sessionTokenStorage.GetSessionToken(batch.CollectionId); + if (!string.IsNullOrWhiteSpace(sessionToken)) + { + options.SessionToken = sessionToken; + } - using var response = await transactionalBatch.ExecuteAsync(cancellationToken).ConfigureAwait(false); + using var response = await transactionalBatch.ExecuteAsync(options, cancellationToken).ConfigureAwait(false); wrapper._commandLogger.ExecutedTransactionalBatch( response.Diagnostics.GetClientElapsedTime(), @@ -708,18 +729,20 @@ private static async Task ExecuteBatchOnceAsync( return CosmosTransactionalBatchResult.Failure(errorEntries, exception); } - ProcessResponse(response, batch.Entries); + ProcessResponse(batch.CollectionId, response, batch.Entries, sessionTokenStorage); return CosmosTransactionalBatchResult.Success(response.Headers.Session); } - private static ItemRequestOptions? CreateItemRequestOptions(IUpdateEntry entry, bool? enableContentResponseOnWrite) + private static ItemRequestOptions? CreateItemRequestOptions(IUpdateEntry entry, bool? enableContentResponseOnWrite, string? sessionToken) { var helper = RequestOptionsHelper.Create(entry, enableContentResponseOnWrite); + sessionToken = string.IsNullOrWhiteSpace(sessionToken) ? null : sessionToken; + return helper == null ? null - : new ItemRequestOptions { IfMatchEtag = helper.IfMatchEtag, EnableContentResponseOnWrite = helper.EnableContentResponseOnWrite }; + : new ItemRequestOptions { IfMatchEtag = helper.IfMatchEtag, SessionToken = sessionToken, EnableContentResponseOnWrite = helper.EnableContentResponseOnWrite }; } private static IReadOnlyList? GetTriggers(IUpdateEntry entry, TriggerType type, TriggerOperation operation) @@ -755,14 +778,25 @@ private static PartitionKey ExtractPartitionKeyValue(IUpdateEntry entry) return builder.Build(); } - private static void ProcessResponse(ResponseMessage response, IUpdateEntry entry) + private static void ProcessResponse(string containerId, ResponseMessage response, IUpdateEntry entry, ISessionTokenStorage sessionTokenStorage) { response.EnsureSuccessStatusCode(); + + if (!string.IsNullOrWhiteSpace(response.Headers.Session)) + { + sessionTokenStorage.AppendSessionToken(containerId, response.Headers.Session); + } + ProcessResponse(entry, response.Headers.ETag, response.Content); } - private static void ProcessResponse(TransactionalBatchResponse batchResponse, IReadOnlyList entries) + private static void ProcessResponse(string containerId, TransactionalBatchResponse batchResponse, IReadOnlyList entries, ISessionTokenStorage sessionTokenStorage) { + if (!string.IsNullOrWhiteSpace(batchResponse.Headers.Session)) + { + sessionTokenStorage.AppendSessionToken(containerId, batchResponse.Headers.Session); + } + for (var i = 0; i < batchResponse.Count; i++) { var entry = entries[i]; @@ -804,13 +838,14 @@ public virtual IEnumerable ExecuteSqlQuery( string containerId, PartitionKey partitionKeyValue, CosmosSqlQuery query, + ISessionTokenStorage sessionTokenStorage, string? sessionToken) { _databaseLogger.SyncNotSupported(); _commandLogger.ExecutingSqlQuery(containerId, partitionKeyValue, query); - return new DocumentEnumerable(this, containerId, partitionKeyValue, query, sessionToken); + return new DocumentEnumerable(this, containerId, partitionKeyValue, query, sessionTokenStorage, sessionToken); } /// @@ -823,11 +858,12 @@ public virtual IAsyncEnumerable ExecuteSqlQueryAsync( string containerId, PartitionKey partitionKeyValue, CosmosSqlQuery query, + ISessionTokenStorage sessionTokenStorage, string? sessionToken) { _commandLogger.ExecutingSqlQuery(containerId, partitionKeyValue, query); - return new DocumentAsyncEnumerable(this, containerId, partitionKeyValue, query, sessionToken); + return new DocumentAsyncEnumerable(this, containerId, partitionKeyValue, query, sessionTokenStorage, sessionToken); } /// @@ -840,13 +876,14 @@ public virtual IAsyncEnumerable ExecuteSqlQueryAsync( string containerId, PartitionKey partitionKeyValue, string resourceId, + ISessionTokenStorage sessionTokenStorage, string? sessionToken) { _databaseLogger.SyncNotSupported(); _commandLogger.ExecutingReadItem(containerId, partitionKeyValue, resourceId); - var response = _executionStrategy.Execute((containerId, partitionKeyValue, resourceId, sessionToken, this), CreateSingleItemQuery, null); + var response = _executionStrategy.Execute((containerId, partitionKeyValue, resourceId, sessionTokenStorage, sessionToken, this), CreateSingleItemQuery, null); _commandLogger.ExecutedReadItem( response.Diagnostics.GetClientElapsedTime(), @@ -869,13 +906,14 @@ public virtual IAsyncEnumerable ExecuteSqlQueryAsync( string containerId, PartitionKey partitionKeyValue, string resourceId, + ISessionTokenStorage sessionTokenStorage, string? sessionToken, CancellationToken cancellationToken = default) { _commandLogger.ExecutingReadItem(containerId, partitionKeyValue, resourceId); var response = await _executionStrategy.ExecuteAsync( - (containerId, partitionKeyValue, resourceId, sessionToken, this), + (containerId, partitionKeyValue, resourceId, sessionTokenStorage, sessionToken, this), CreateSingleItemQueryAsync, null, cancellationToken) @@ -894,15 +932,15 @@ public virtual IAsyncEnumerable ExecuteSqlQueryAsync( private static ResponseMessage CreateSingleItemQuery( DbContext? context, - (string ContainerId, PartitionKey PartitionKeyValue, string ResourceId, string? SessionToken, CosmosClientWrapper Wrapper) parameters) + (string ContainerId, PartitionKey PartitionKeyValue, string ResourceId, ISessionTokenStorage SessionTokenStorage, string? SessionToken, CosmosClientWrapper Wrapper) parameters) => CreateSingleItemQueryAsync(context, parameters).GetAwaiter().GetResult(); - private static Task CreateSingleItemQueryAsync( + private static async Task CreateSingleItemQueryAsync( DbContext? _, - (string ContainerId, PartitionKey PartitionKeyValue, string ResourceId, string? SessionToken, CosmosClientWrapper Wrapper) parameters, + (string ContainerId, PartitionKey PartitionKeyValue, string ResourceId, ISessionTokenStorage SessionTokenStorage, string? SessionToken, CosmosClientWrapper Wrapper) parameters, CancellationToken cancellationToken = default) { - var (containerId, partitionKeyValue, resourceId, sessionToken, wrapper) = parameters; + var (containerId, partitionKeyValue, resourceId, sessionTokenStorage, sessionToken, wrapper) = parameters; var container = wrapper.Client.GetDatabase(wrapper._databaseId).GetContainer(containerId); ItemRequestOptions? itemRequestOptions = null; @@ -911,11 +949,18 @@ private static Task CreateSingleItemQueryAsync( itemRequestOptions = new ItemRequestOptions { SessionToken = sessionToken }; } - return container.ReadItemStreamAsync( + var response = await container.ReadItemStreamAsync( resourceId, partitionKeyValue, itemRequestOptions, - cancellationToken: cancellationToken); + cancellationToken: cancellationToken).ConfigureAwait(false); + + if (!string.IsNullOrWhiteSpace(response.Headers.Session)) + { + sessionTokenStorage.AppendSessionToken(containerId, response.Headers.Session); + } + + return response; } private static JObject? JObjectFromReadItemResponseMessage(ResponseMessage responseMessage) @@ -955,6 +1000,7 @@ private static Task CreateSingleItemQueryAsync( public virtual FeedIterator CreateQuery( string containerId, CosmosSqlQuery query, + ISessionTokenStorage sessionTokenStorage, string? continuationToken = null, QueryRequestOptions? queryRequestOptions = null) { @@ -966,7 +1012,7 @@ public virtual FeedIterator CreateQuery( queryDefinition, (current, parameter) => current.WithParameter(parameter.Name, parameter.Value)); - return container.GetItemQueryStreamIterator(queryDefinition, continuationToken, queryRequestOptions); + return new CosmosFeedIteratorWrapper(container.GetItemQueryStreamIterator(queryDefinition, continuationToken, queryRequestOptions), containerId, sessionTokenStorage); } [MethodImpl(MethodImplOptions.AggressiveInlining)] @@ -997,14 +1043,16 @@ private sealed class DocumentEnumerable( string containerId, PartitionKey partitionKeyValue, CosmosSqlQuery cosmosSqlQuery, + ISessionTokenStorage sessionTokenStorage, string? sessionToken) : IEnumerable { - private readonly string? _sessionToken = sessionToken; private readonly CosmosClientWrapper _cosmosClient = cosmosClient; private readonly string _containerId = containerId; private readonly PartitionKey _partitionKeyValue = partitionKeyValue; private readonly CosmosSqlQuery _cosmosSqlQuery = cosmosSqlQuery; + private readonly ISessionTokenStorage _sessionTokenStorage = sessionTokenStorage; + private readonly string? _sessionToken = sessionToken; public IEnumerator GetEnumerator() => new Enumerator(this); @@ -1014,11 +1062,12 @@ IEnumerator IEnumerable.GetEnumerator() private sealed class Enumerator(DocumentEnumerable documentEnumerable) : IEnumerator { - private readonly string? _sessionToken = documentEnumerable._sessionToken; private readonly CosmosClientWrapper _cosmosClientWrapper = documentEnumerable._cosmosClient; private readonly string _containerId = documentEnumerable._containerId; private readonly PartitionKey _partitionKeyValue = documentEnumerable._partitionKeyValue; private readonly CosmosSqlQuery _cosmosSqlQuery = documentEnumerable._cosmosSqlQuery; + private readonly ISessionTokenStorage _sessionTokenStorage = documentEnumerable._sessionTokenStorage; + private readonly string? _sessionToken = documentEnumerable._sessionToken; private JToken? _current; private ResponseMessage? _responseMessage; @@ -1051,7 +1100,7 @@ public bool MoveNext() } _query = _cosmosClientWrapper.CreateQuery( - _containerId, _cosmosSqlQuery, continuationToken: null, queryRequestOptions); + _containerId, _cosmosSqlQuery, _sessionTokenStorage, continuationToken: null, queryRequestOptions); } if (!_query.HasMoreResults) @@ -1075,11 +1124,6 @@ public bool MoveNext() _responseMessage.EnsureSuccessStatusCode(); - if (!string.IsNullOrWhiteSpace(_responseMessage.Headers.Session)) - { - // @TODO: set session token... - } - _responseMessageEnumerator = new ResponseMessageEnumerable(_responseMessage).GetEnumerator(); } @@ -1117,14 +1161,16 @@ private sealed class DocumentAsyncEnumerable( string containerId, PartitionKey partitionKeyValue, CosmosSqlQuery cosmosSqlQuery, + ISessionTokenStorage sessionTokenStorage, string? sessionToken) : IAsyncEnumerable { private readonly CosmosClientWrapper _cosmosClient = cosmosClient; private readonly string _containerId = containerId; private readonly PartitionKey _partitionKeyValue = partitionKeyValue; - private readonly string? _sessionToken = sessionToken; private readonly CosmosSqlQuery _cosmosSqlQuery = cosmosSqlQuery; + private readonly ISessionTokenStorage _sessionTokenStorage = sessionTokenStorage; + private readonly string? _sessionToken = sessionToken; public IAsyncEnumerator GetAsyncEnumerator(CancellationToken cancellationToken = default) => new AsyncEnumerator(this, cancellationToken); @@ -1135,8 +1181,9 @@ private sealed class AsyncEnumerator(DocumentAsyncEnumerable documentEnumerable, private readonly CosmosClientWrapper _cosmosClientWrapper = documentEnumerable._cosmosClient; private readonly string _containerId = documentEnumerable._containerId; private readonly PartitionKey _partitionKeyValue = documentEnumerable._partitionKeyValue; - private readonly string? _sessionToken = documentEnumerable._sessionToken; private readonly CosmosSqlQuery _cosmosSqlQuery = documentEnumerable._cosmosSqlQuery; + private readonly ISessionTokenStorage _sessionTokenStorage = documentEnumerable._sessionTokenStorage; + private readonly string? _sessionToken = documentEnumerable._sessionToken; private JToken? _current; private ResponseMessage? _responseMessage; @@ -1168,7 +1215,7 @@ public async ValueTask MoveNextAsync() } _query = _cosmosClientWrapper.CreateQuery( - _containerId, _cosmosSqlQuery, continuationToken: null, queryRequestOptions); + _containerId, _cosmosSqlQuery, _sessionTokenStorage, continuationToken: null, queryRequestOptions); } if (!_query.HasMoreResults) @@ -1334,4 +1381,31 @@ public async ValueTask DisposeAsync() } #endregion ResponseMessageEnumerable + + private class CosmosFeedIteratorWrapper : FeedIterator + { + private readonly FeedIterator _inner; + private readonly string _containerName; + private readonly ISessionTokenStorage _sessionTokenStorage; + + public CosmosFeedIteratorWrapper(FeedIterator inner, string containerName, ISessionTokenStorage sessionTokenStorage) + { + _inner = inner; + _containerName = containerName; + _sessionTokenStorage = sessionTokenStorage; + } + + public override bool HasMoreResults => _inner.HasMoreResults; + + public override async Task ReadNextAsync(CancellationToken cancellationToken = default) + { + var response = await _inner.ReadNextAsync(cancellationToken).ConfigureAwait(false); + if (!string.IsNullOrWhiteSpace(response.Headers.Session)) + { + _sessionTokenStorage.SetSessionToken(_containerName, response.Headers.Session); + } + return response; + } + } + } diff --git a/src/EFCore.Cosmos/Storage/Internal/CosmosDatabaseWrapper.cs b/src/EFCore.Cosmos/Storage/Internal/CosmosDatabaseWrapper.cs index 0a50b49ce81..8b147fcb0b7 100644 --- a/src/EFCore.Cosmos/Storage/Internal/CosmosDatabaseWrapper.cs +++ b/src/EFCore.Cosmos/Storage/Internal/CosmosDatabaseWrapper.cs @@ -44,7 +44,7 @@ public CosmosDatabaseWrapper( _currentDbContext = currentDbContext; _cosmosClient = cosmosClient; - SessionTokenStorage = new(_currentDbContext.Context); + SessionTokenStorage = new SessionTokenStorage(_currentDbContext.Context); if (loggingOptions.IsSensitiveDataLoggingEnabled) { @@ -58,7 +58,7 @@ public CosmosDatabaseWrapper( /// 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. /// - public SessionTokenStorage SessionTokenStorage { get; } + public ISessionTokenStorage SessionTokenStorage { get; } /// /// This is an internal API that supports the Entity Framework Core infrastructure and not subject to @@ -100,15 +100,8 @@ public override int SaveChanges(IList entries) { try { - var response = _cosmosClient.ExecuteTransactionalBatch(transaction); - if (response.IsSuccess) - { - if (!string.IsNullOrWhiteSpace(response.SessionToken)) - { - SessionTokenStorage.AppendSessionToken(batch.Key.ContainerId, response.SessionToken); - } - } - else + var response = _cosmosClient.ExecuteTransactionalBatch(transaction, SessionTokenStorage); + if (!response.IsSuccess) { var exception = WrapUpdateException(response.Exception, response.ErroredEntries); if (exception is not DbUpdateConcurrencyException @@ -175,15 +168,8 @@ public override async Task SaveChangesAsync( { try { - var response = await _cosmosClient.ExecuteTransactionalBatchAsync(transaction, cancellationToken).ConfigureAwait(false); - if (response.IsSuccess) - { - if (!string.IsNullOrWhiteSpace(response.SessionToken)) - { - SessionTokenStorage.AppendSessionToken(batch.Key.ContainerId, response.SessionToken); - } - } - else + var response = await _cosmosClient.ExecuteTransactionalBatchAsync(transaction, SessionTokenStorage, cancellationToken).ConfigureAwait(false); + if (!response.IsSuccess) { var exception = WrapUpdateException(response.Exception, response.ErroredEntries); if (exception is not DbUpdateConcurrencyException @@ -208,16 +194,6 @@ public override async Task SaveChangesAsync( return rowsAffected; } - private void ProcessSessionToken(string containerId, string? sessionToken) - { - if (string.IsNullOrEmpty(sessionToken)) - { - return; - } - - - } - private SaveGroups CreateSaveGroups(IList entries) { var count = entries.Count; @@ -552,19 +528,22 @@ private bool Save(CosmosUpdateEntry updateEntry) { return updateEntry.Operation switch { - CosmosCudOperation.Create => ProcessWriteResult(updateEntry, _cosmosClient.CreateItem( + CosmosCudOperation.Create => _cosmosClient.CreateItem( updateEntry.CollectionId, updateEntry.Document!, - updateEntry.Entry)), - CosmosCudOperation.Update => ProcessWriteResult(updateEntry, _cosmosClient.ReplaceItem( + updateEntry.Entry, + SessionTokenStorage), + CosmosCudOperation.Update => _cosmosClient.ReplaceItem( updateEntry.CollectionId, updateEntry.DocumentSource.GetId(updateEntry.Entry.SharedIdentityEntry ?? updateEntry.Entry), updateEntry.Document!, - updateEntry.Entry)), - CosmosCudOperation.Delete => ProcessWriteResult(updateEntry, _cosmosClient.DeleteItem( + updateEntry.Entry, + SessionTokenStorage), + CosmosCudOperation.Delete => _cosmosClient.DeleteItem( updateEntry.CollectionId, updateEntry.DocumentSource.GetId(updateEntry.Entry), - updateEntry.Entry)), + updateEntry.Entry, + SessionTokenStorage), _ => throw new UnreachableException(), }; } @@ -590,22 +569,25 @@ private async Task SaveAsync(CosmosUpdateEntry updateEntry, CancellationTo { return updateEntry.Operation switch { - CosmosCudOperation.Create => ProcessWriteResult(updateEntry, await _cosmosClient.CreateItemAsync( + CosmosCudOperation.Create => await _cosmosClient.CreateItemAsync( updateEntry.CollectionId, updateEntry.Document!, updateEntry.Entry, - cancellationToken).ConfigureAwait(false)), - CosmosCudOperation.Update => ProcessWriteResult(updateEntry, await _cosmosClient.ReplaceItemAsync( + SessionTokenStorage, + cancellationToken).ConfigureAwait(false), + CosmosCudOperation.Update => await _cosmosClient.ReplaceItemAsync( updateEntry.CollectionId, updateEntry.DocumentSource.GetId(updateEntry.Entry.SharedIdentityEntry ?? updateEntry.Entry), updateEntry.Document!, updateEntry.Entry, - cancellationToken).ConfigureAwait(false)), - CosmosCudOperation.Delete => ProcessWriteResult(updateEntry, await _cosmosClient.DeleteItemAsync( + SessionTokenStorage, + cancellationToken).ConfigureAwait(false), + CosmosCudOperation.Delete => await _cosmosClient.DeleteItemAsync( updateEntry.CollectionId, updateEntry.DocumentSource.GetId(updateEntry.Entry), updateEntry.Entry, - cancellationToken).ConfigureAwait(false)), + SessionTokenStorage, + cancellationToken).ConfigureAwait(false), _ => throw new UnreachableException(), }; } @@ -626,16 +608,6 @@ private async Task SaveAsync(CosmosUpdateEntry updateEntry, CancellationTo } } - private bool ProcessWriteResult(CosmosUpdateEntry updateEntry, CosmosWriteResult result) - { - if (!string.IsNullOrWhiteSpace(result.SessionToken)) - { - SessionTokenStorage.AppendSessionToken(updateEntry.CollectionId, result.SessionToken); - } - - return result.IsSuccess; - } - /// /// 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 diff --git a/src/EFCore.Cosmos/Storage/Internal/ICosmosClientWrapper.cs b/src/EFCore.Cosmos/Storage/Internal/ICosmosClientWrapper.cs index b0e09dba277..f9d0bb4bd16 100644 --- a/src/EFCore.Cosmos/Storage/Internal/ICosmosClientWrapper.cs +++ b/src/EFCore.Cosmos/Storage/Internal/ICosmosClientWrapper.cs @@ -67,7 +67,7 @@ public interface ICosmosClientWrapper /// 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. /// - CosmosWriteResult CreateItem(string containerId, JToken document, IUpdateEntry entry); + bool CreateItem(string containerId, JToken document, IUpdateEntry entry, ISessionTokenStorage sessionTokenStorage); /// /// This is an internal API that supports the Entity Framework Core infrastructure and not subject to @@ -75,11 +75,12 @@ public interface ICosmosClientWrapper /// 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. /// - CosmosWriteResult ReplaceItem( + bool ReplaceItem( string collectionId, string documentId, JObject document, - IUpdateEntry entry); + IUpdateEntry entry, + ISessionTokenStorage sessionTokenStorage); /// /// This is an internal API that supports the Entity Framework Core infrastructure and not subject to @@ -87,10 +88,11 @@ CosmosWriteResult ReplaceItem( /// 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. /// - CosmosWriteResult DeleteItem( + bool DeleteItem( string containerId, string documentId, - IUpdateEntry entry); + IUpdateEntry entry, + ISessionTokenStorage sessionTokenStorage); /// /// This is an internal API that supports the Entity Framework Core infrastructure and not subject to @@ -98,10 +100,11 @@ CosmosWriteResult DeleteItem( /// 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. /// - Task CreateItemAsync( + Task CreateItemAsync( string containerId, JToken document, IUpdateEntry updateEntry, + ISessionTokenStorage sessionTokenStorage, CancellationToken cancellationToken = default); /// @@ -110,11 +113,12 @@ Task CreateItemAsync( /// 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. /// - Task ReplaceItemAsync( + Task ReplaceItemAsync( string collectionId, string documentId, JObject document, IUpdateEntry updateEntry, + ISessionTokenStorage sessionTokenStorage, CancellationToken cancellationToken = default); /// @@ -124,10 +128,11 @@ Task ReplaceItemAsync( /// doing so can result in application failures when updating to a new Entity Framework Core release. /// // @TODO: We also need to send session token on writes to keep the same chain or not? - Task DeleteItemAsync( + Task DeleteItemAsync( string containerId, string documentId, IUpdateEntry entry, + ISessionTokenStorage sessionTokenStorage, CancellationToken cancellationToken = default); /// @@ -139,6 +144,7 @@ Task DeleteItemAsync( FeedIterator CreateQuery( string containerId, CosmosSqlQuery query, + ISessionTokenStorage sessionTokenStorage, string? continuationToken = null, QueryRequestOptions? queryRequestOptions = null); @@ -152,6 +158,7 @@ FeedIterator CreateQuery( string containerId, PartitionKey partitionKeyValue, string resourceId, + ISessionTokenStorage sessionTokenStorage, string? sessionToken); /// @@ -164,6 +171,7 @@ FeedIterator CreateQuery( string containerId, PartitionKey partitionKeyValue, string resourceId, + ISessionTokenStorage sessionTokenStorage, string? sessionToken, CancellationToken cancellationToken = default); @@ -177,6 +185,7 @@ IEnumerable ExecuteSqlQuery( string containerId, PartitionKey partitionKeyValue, CosmosSqlQuery query, + ISessionTokenStorage sessionTokenStorage, string? sessionToken); /// @@ -189,6 +198,7 @@ IAsyncEnumerable ExecuteSqlQueryAsync( string containerId, PartitionKey partitionKeyValue, CosmosSqlQuery query, + ISessionTokenStorage sessionTokenStorage, string? sessionToken); /// @@ -221,7 +231,7 @@ IAsyncEnumerable ExecuteSqlQueryAsync( /// 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. /// - Task ExecuteTransactionalBatchAsync(ICosmosTransactionalBatchWrapper batch, CancellationToken cancellationToken = default); + Task ExecuteTransactionalBatchAsync(ICosmosTransactionalBatchWrapper batch, ISessionTokenStorage sessionTokenStorage, CancellationToken cancellationToken = default); /// /// This is an internal API that supports the Entity Framework Core infrastructure and not subject to @@ -229,5 +239,5 @@ IAsyncEnumerable ExecuteSqlQueryAsync( /// 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. /// - CosmosTransactionalBatchResult ExecuteTransactionalBatch(ICosmosTransactionalBatchWrapper batch); + CosmosTransactionalBatchResult ExecuteTransactionalBatch(ICosmosTransactionalBatchWrapper batch, ISessionTokenStorage sessionTokenStorage); } diff --git a/src/EFCore.Cosmos/Storage/Internal/NullSessionTokenStorage.cs b/src/EFCore.Cosmos/Storage/Internal/NullSessionTokenStorage.cs new file mode 100644 index 00000000000..01891e78545 --- /dev/null +++ b/src/EFCore.Cosmos/Storage/Internal/NullSessionTokenStorage.cs @@ -0,0 +1,60 @@ +// 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.Storage.Internal; + +/// +/// 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. +/// +public sealed class NullSessionTokenStorage : ISessionTokenStorage +{ + /// + /// 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. + /// + public void AppendSessionToken(string sessionToken) { } + + /// + /// 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. + /// + public void AppendSessionToken(string containerName, string sessionToken) { } + + /// + /// 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. + /// + public string? GetSessionToken() => null; + /// + /// 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. + /// + public string? GetSessionToken(string containerName) => null; + + /// + /// 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. + /// + public void SetSessionToken(string containerName, string? sessionToken) { } + + /// + /// 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. + /// + public void SetSessionToken(string? sessionToken) { } +} diff --git a/src/EFCore.Cosmos/Storage/SessionTokenStorage.cs b/src/EFCore.Cosmos/Storage/Internal/SessionTokenStorage.cs similarity index 85% rename from src/EFCore.Cosmos/Storage/SessionTokenStorage.cs rename to src/EFCore.Cosmos/Storage/Internal/SessionTokenStorage.cs index 4dc82caaf2a..9dbe0b0fc18 100644 --- a/src/EFCore.Cosmos/Storage/SessionTokenStorage.cs +++ b/src/EFCore.Cosmos/Storage/Internal/SessionTokenStorage.cs @@ -8,14 +8,17 @@ using Microsoft.EntityFrameworkCore.Cosmos.Internal; using Microsoft.EntityFrameworkCore.Cosmos.Metadata.Internal; -namespace Microsoft.EntityFrameworkCore.Cosmos.Storage; +namespace Microsoft.EntityFrameworkCore.Cosmos.Storage.Internal; /// -/// Stores the session tokens for a DbContext. +/// 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. /// -public sealed class SessionTokenStorage +public sealed class SessionTokenStorage : ISessionTokenStorage { - private readonly Dictionary _sessionTokens; + private readonly Dictionary _containerSessionTokens; private readonly string _defaultContainerName; /// @@ -30,35 +33,47 @@ public SessionTokenStorage(DbContext dbContext) var containerNames = (HashSet)dbContext.Model.GetAnnotation(CosmosAnnotationNames.ContainerNames).Value!; _defaultContainerName = defaultContainerName; - _sessionTokens = containerNames.ToDictionary(containerName => containerName, _ => new CompositeSessionToken()); + _containerSessionTokens = containerNames.ToDictionary(containerName => containerName, _ => new CompositeSessionToken()); } /// - /// Gets the composite session token for the default container. + /// 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. /// public string? GetSessionToken() => GetSessionToken(_defaultContainerName); /// - /// Overwrites the composite session token for the default container + /// 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. /// public void SetSessionToken(string? sessionToken) => SetSessionToken(_defaultContainerName, sessionToken); /// - /// Appends or merges the session token specified to any composite already stored in the storage for the default container + /// 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. /// public void AppendSessionToken(string sessionToken) => AppendSessionToken(_defaultContainerName, sessionToken); /// - /// Gets the composite session token for the specified container + /// 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. /// public string? GetSessionToken(string containerName) { ArgumentNullException.ThrowIfNullOrWhiteSpace(containerName, nameof(containerName)); - if (!_sessionTokens.TryGetValue(containerName, out var value)) + if (!_containerSessionTokens.TryGetValue(containerName, out var value)) { throw new ArgumentException(CosmosStrings.ContainerNameDoesNotExist(containerName), nameof(containerName)); } @@ -67,14 +82,17 @@ public void AppendSessionToken(string sessionToken) } /// - /// Appends or merges the session token specified to any composite already stored in the storage for the container + /// 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. /// public void AppendSessionToken(string containerName, string sessionToken) { ArgumentNullException.ThrowIfNullOrWhiteSpace(containerName, nameof(containerName)); ArgumentNullException.ThrowIfNullOrWhiteSpace(sessionToken, nameof(sessionToken)); - if (!_sessionTokens.TryGetValue(containerName, out var compositeSessionToken)) + if (!_containerSessionTokens.TryGetValue(containerName, out var compositeSessionToken)) { throw new ArgumentException(CosmosStrings.ContainerNameDoesNotExist(containerName), nameof(containerName)); } @@ -83,7 +101,10 @@ public void AppendSessionToken(string containerName, string sessionToken) } /// - /// Overwrites the composite session token for the container + /// 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. /// public void SetSessionToken(string containerName, string? sessionToken) { @@ -93,7 +114,7 @@ public void SetSessionToken(string containerName, string? sessionToken) throw new ArgumentException("sessionToken cannot be whitespace.", sessionToken); } - ref var value = ref CollectionsMarshal.GetValueRefOrNullRef(_sessionTokens, containerName); + ref var value = ref CollectionsMarshal.GetValueRefOrNullRef(_containerSessionTokens, containerName); if (Unsafe.IsNullRef(ref value)) { @@ -166,7 +187,7 @@ public void Merge(string pkRangeId, VectorSessionToken token) } - /// + /// private sealed class VectorSessionToken : IEquatable { private static readonly IReadOnlyDictionary DefaultLocalLsnByRegion = new Dictionary(0); @@ -449,14 +470,12 @@ private static bool TryParseSessionToken(string sessionToken, out long version, globalLsn = -1L; if (string.IsNullOrEmpty(sessionToken)) { - ////DefaultTrace.TraceWarning("Session token is empty"); return false; } int index = 0; if (!TryParseLongSegment(sessionToken, ref index, out version)) { - ////DefaultTrace.TraceWarning("Unexpected session token version number from token: " + sessionToken + " ."); return false; } @@ -467,7 +486,6 @@ private static bool TryParseSessionToken(string sessionToken, out long version, if (!TryParseLongSegment(sessionToken, ref index, out globalLsn)) { - //DefaultTrace.TraceWarning("Unexpected session token global lsn from token: " + sessionToken + " ."); return false; } @@ -482,13 +500,11 @@ private static bool TryParseSessionToken(string sessionToken, out long version, { if (!TryParseUintTillRegionProgressSeparator(sessionToken, ref index, out var value)) { - //DefaultTrace.TraceWarning("Unexpected region progress segment in session token: " + sessionToken + "."); return false; } if (!TryParseLongSegment(sessionToken, ref index, out var value2)) { - //DefaultTrace.TraceWarning("Unexpected local lsn for region id " + value.ToString(CultureInfo.InvariantCulture) + " for segment in session token: " + sessionToken + "."); return false; } diff --git a/test/EFCore.Cosmos.FunctionalTests/CosmosSessionTokensTest.cs b/test/EFCore.Cosmos.FunctionalTests/CosmosSessionTokensTest.cs index 2d9c42ebc92..e679699ce9a 100644 --- a/test/EFCore.Cosmos.FunctionalTests/CosmosSessionTokensTest.cs +++ b/test/EFCore.Cosmos.FunctionalTests/CosmosSessionTokensTest.cs @@ -1,8 +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.Azure.Cosmos; using Microsoft.EntityFrameworkCore.Cosmos.Internal; -using Microsoft.EntityFrameworkCore.Cosmos.Storage; namespace Microsoft.EntityFrameworkCore; @@ -15,6 +15,8 @@ protected override string StoreName protected override ITestStoreFactory TestStoreFactory => CosmosTestStoreFactory.Instance; + // @TODO: Tests for other container? + [ConditionalFact] public virtual async Task SetSessionToken_ThrowsForNonExistentContainer() { @@ -22,8 +24,8 @@ public virtual async Task SetSessionToken_ThrowsForNonExistentContainer() using var context = contextFactory.CreateContext(); var sessionTokens = context.Database.GetSessionTokens(); - var exception = Assert.Throws(() => sessionTokens.SetSessionToken("Not the store name", "0:-1#231")); - Assert.Equal(CosmosStrings.ContainerNameDoesNotExist("Not the store name") + " (Parameter 'containerName')", exception.Message); + var exception = Assert.Throws(() => sessionTokens.SetSessionToken("Not the container name", "0:-1#231")); + Assert.Equal(CosmosStrings.ContainerNameDoesNotExist("Not the container name") + " (Parameter 'containerName')", exception.Message); } [ConditionalFact] @@ -33,8 +35,8 @@ public virtual async Task AppendSessionToken_ThrowsForNonExistentContainer() using var context = contextFactory.CreateContext(); var sessionTokens = context.Database.GetSessionTokens(); - var exception = Assert.Throws(() => sessionTokens.AppendSessionToken("Not the store name", "0:-1#231")); - Assert.Equal(CosmosStrings.ContainerNameDoesNotExist("Not the store name") + " (Parameter 'containerName')", exception.Message); + var exception = Assert.Throws(() => sessionTokens.AppendSessionToken("Not the container name", "0:-1#231")); + Assert.Equal(CosmosStrings.ContainerNameDoesNotExist("Not the container name") + " (Parameter 'containerName')", exception.Message); } [ConditionalFact] @@ -44,8 +46,8 @@ public virtual async Task GetSessionToken_ThrowsForNonExistentContainer() using var context = contextFactory.CreateContext(); var sessionTokens = context.Database.GetSessionTokens(); - var exception = Assert.Throws(() => sessionTokens.GetSessionToken("Not the store name")); - Assert.Equal(CosmosStrings.ContainerNameDoesNotExist("Not the store name") + " (Parameter 'containerName')", exception.Message); + var exception = Assert.Throws(() => sessionTokens.GetSessionToken("Not the container name")); + Assert.Equal(CosmosStrings.ContainerNameDoesNotExist("Not the container name") + " (Parameter 'containerName')", exception.Message); } [ConditionalFact] @@ -106,13 +108,12 @@ public virtual async Task AppendSessionToken_append_lower_lsn_same_pkrange_takes Assert.Equal(initialToken, updatedToken); } - [ConditionalTheory, InlineData(AutoTransactionBehavior.WhenNeeded), InlineData(AutoTransactionBehavior.Never), InlineData(AutoTransactionBehavior.Always)] - public virtual async Task AppendSessionToken_different_pkrange_composites_tokens(AutoTransactionBehavior autoTransactionBehavior) + [ConditionalFact] + public virtual async Task AppendSessionToken_different_pkrange_composites_tokens() { var contextFactory = await InitializeAsync(); using var context = contextFactory.CreateContext(); - context.Database.AutoTransactionBehavior = autoTransactionBehavior; context.Customers.Add(new Customer { Id = "1", PartitionKey = "1" }); await context.SaveChangesAsync(); @@ -196,12 +197,34 @@ public virtual async Task Query_uses_session_token() // Only way we can test this is by setting a session token that will fail the request if used.. // This will take a couple of seconds to fail - // sessionTokens.SetSessionToken(sessionToken.Substring(0, sessionToken.IndexOf('#') + 1) + int.MaxValue); + sessionTokens.SetSessionToken(sessionToken.Substring(0, sessionToken.IndexOf('#') + 1) + int.MaxValue); var ex = await Assert.ThrowsAsync(() => context.Customers.ToListAsync()); Assert.Contains("The read session is not available for the input session token.", ex.ResponseBody); } + [ConditionalFact] + // @TODO: and sync.. + public virtual async Task PagingQuery_uses_session_token() + { + var contextFactory = await InitializeAsync(); + + using var context = contextFactory.CreateContext(); + context.Customers.Add(new Customer { Id = "1", PartitionKey = "1" }); + + await context.SaveChangesAsync(); + + var sessionTokens = context.Database.GetSessionTokens(); + var sessionToken = sessionTokens.GetSessionToken()!; + + // Only way we can test this is by setting a session token that will fail the request if used.. + // This will take a couple of seconds to fail + sessionTokens.SetSessionToken(sessionToken.Substring(0, sessionToken.IndexOf('#') + 1) + int.MaxValue); + + var ex = await Assert.ThrowsAsync(() => context.Customers.ToPageAsync(1, null)); + Assert.Contains("The read session is not available for the input session token.", ex.ResponseBody); + } + [ConditionalFact] // @TODO: and sync.. public virtual async Task Shaped_query_uses_session_token() @@ -246,6 +269,71 @@ public virtual async Task Read_item_uses_session_token() Assert.Contains("The read session is not available for the input session token.", ex.ResponseBody); } + [ConditionalFact] + // @TODO: and sync.. + public virtual async Task Query_sets_session_token() + { + var contextFactory = await InitializeAsync(); + + using var context = contextFactory.CreateContext(); + await context.Customers.ToListAsync(); + + var sessionToken = context.Database.GetSessionTokens().GetSessionToken(); + Assert.True(!string.IsNullOrWhiteSpace(sessionToken)); + } + + [ConditionalFact] + // @TODO: and sync.. + public virtual async Task PagingQuery_sets_session_token() + { + var contextFactory = await InitializeAsync(); + + using var context = contextFactory.CreateContext(); + await context.Customers.ToPageAsync(1, null); + + var sessionToken = context.Database.GetSessionTokens().GetSessionToken(); + Assert.True(!string.IsNullOrWhiteSpace(sessionToken)); + } + + [ConditionalFact] + // @TODO: and sync.. + public virtual async Task Shaped_query_sets_session_token() + { + var contextFactory = await InitializeAsync(); + + using var context = contextFactory.CreateContext(); + await context.Customers.Select(x => new { x.Id, x.PartitionKey }).ToListAsync(); + + var sessionToken = context.Database.GetSessionTokens().GetSessionToken(); + Assert.True(!string.IsNullOrWhiteSpace(sessionToken)); + } + + [ConditionalFact] + // @TODO: and sync.. + public virtual async Task Read_item_sets_session_token() + { + var contextFactory = await InitializeAsync(); + + using var context = contextFactory.CreateContext(); + await context.Customers.FirstOrDefaultAsync(x => x.Id == "1" && x.PartitionKey == "1"); + + var sessionToken = context.Database.GetSessionTokens().GetSessionToken(); + Assert.True(!string.IsNullOrWhiteSpace(sessionToken)); + } + + [ConditionalFact] + // @TODO: and sync.. + public virtual async Task Read_item_enumerable_sets_session_token() + { + var contextFactory = await InitializeAsync(); + + using var context = contextFactory.CreateContext(); + await context.Customers.Where(x => x.Id == "1" && x.PartitionKey == "1").ToListAsync(); + + var sessionToken = context.Database.GetSessionTokens().GetSessionToken(); + Assert.True(!string.IsNullOrWhiteSpace(sessionToken)); + } + [ConditionalTheory, InlineData(AutoTransactionBehavior.WhenNeeded), InlineData(AutoTransactionBehavior.Never), InlineData(AutoTransactionBehavior.Always)] // @TODO: and sync.. public virtual async Task Add_sets_session_token(AutoTransactionBehavior autoTransactionBehavior) @@ -313,6 +401,78 @@ public virtual async Task Update_merges_session_token(AutoTransactionBehavior au Assert.StartsWith(initialToken.Substring(0, initialToken.IndexOf('#') + 1), sessionToken); } + [ConditionalTheory, InlineData(AutoTransactionBehavior.WhenNeeded), InlineData(AutoTransactionBehavior.Never), InlineData(AutoTransactionBehavior.Always)] + // @TODO: And sync.. + public virtual async Task Add_uses_session_token(AutoTransactionBehavior autoTransactionBehavior) + { + var contextFactory = await InitializeAsync(); + + using var context = contextFactory.CreateContext(); + context.Database.AutoTransactionBehavior = autoTransactionBehavior; + + var sessionTokens = context.Database.GetSessionTokens(); + var sessionToken = sessionTokens.GetSessionToken()!; + // Only way we can test this is by setting a session token that will fail the request if used.. + // Only way to do this for a write is to set an invalid session token.. + var internalDictionary = sessionTokens.GetType().GetField("_containerSessionTokens", BindingFlags.NonPublic | BindingFlags.Instance)!.GetValue(sessionTokens)!; + var internalComposite = internalDictionary.GetType().GetProperty("Item", BindingFlags.Public | BindingFlags.Instance)!.GetValue(internalDictionary, new object[] { nameof(CosmosSessionTokenContext) })!; + internalComposite.GetType().GetField("_string", BindingFlags.NonPublic | BindingFlags.Instance)!.SetValue(internalComposite, "invalidtoken"); + internalComposite.GetType().GetField("_isChanged", BindingFlags.NonPublic | BindingFlags.Instance)!.SetValue(internalComposite, false); + + context.Customers.Add(new Customer { Id = "1", PartitionKey = "1" }); + + var ex = await Assert.ThrowsAsync(() => context.SaveChangesAsync()); + Assert.Contains("The session token provided 'invalidtoken' is invalid.", ((CosmosException)ex.InnerException!).ResponseBody); + } + + [ConditionalTheory, InlineData(AutoTransactionBehavior.WhenNeeded), InlineData(AutoTransactionBehavior.Never), InlineData(AutoTransactionBehavior.Always)] + // @TODO: And sync.. + public virtual async Task Update_uses_session_token(AutoTransactionBehavior autoTransactionBehavior) + { + var contextFactory = await InitializeAsync(); + + using var context = contextFactory.CreateContext(); + context.Database.AutoTransactionBehavior = autoTransactionBehavior; + + var sessionTokens = context.Database.GetSessionTokens(); + var sessionToken = sessionTokens.GetSessionToken()!; + // Only way we can test this is by setting a session token that will fail the request if used.. + // Only way to do this for a write is to set an invalid session token.. + var internalDictionary = sessionTokens.GetType().GetField("_containerSessionTokens", BindingFlags.NonPublic | BindingFlags.Instance)!.GetValue(sessionTokens)!; + var internalComposite = internalDictionary.GetType().GetProperty("Item", BindingFlags.Public | BindingFlags.Instance)!.GetValue(internalDictionary, new object[] { nameof(CosmosSessionTokenContext) })!; + internalComposite.GetType().GetField("_string", BindingFlags.NonPublic | BindingFlags.Instance)!.SetValue(internalComposite, "invalidtoken"); + internalComposite.GetType().GetField("_isChanged", BindingFlags.NonPublic | BindingFlags.Instance)!.SetValue(internalComposite, false); + + context.Customers.Update(new Customer { Id = "1", PartitionKey = "1" }); + + var ex = await Assert.ThrowsAsync(() => context.SaveChangesAsync()); + Assert.Contains("The session token provided 'invalidtoken' is invalid.", ((CosmosException)ex.InnerException!).ResponseBody); + } + + [ConditionalTheory, InlineData(AutoTransactionBehavior.WhenNeeded), InlineData(AutoTransactionBehavior.Never), InlineData(AutoTransactionBehavior.Always)] + // @TODO: And sync.. + public virtual async Task Delete_uses_session_token(AutoTransactionBehavior autoTransactionBehavior) + { + var contextFactory = await InitializeAsync(); + + using var context = contextFactory.CreateContext(); + context.Database.AutoTransactionBehavior = autoTransactionBehavior; + + var sessionTokens = context.Database.GetSessionTokens(); + var sessionToken = sessionTokens.GetSessionToken()!; + // Only way we can test this is by setting a session token that will fail the request if used.. + // Only way to do this for a write is to set an invalid session token.. + var internalDictionary = sessionTokens.GetType().GetField("_containerSessionTokens", BindingFlags.NonPublic | BindingFlags.Instance)!.GetValue(sessionTokens)!; + var internalComposite = internalDictionary.GetType().GetProperty("Item", BindingFlags.Public | BindingFlags.Instance)!.GetValue(internalDictionary, new object[] { nameof(CosmosSessionTokenContext) })!; + internalComposite.GetType().GetField("_string", BindingFlags.NonPublic | BindingFlags.Instance)!.SetValue(internalComposite, "invalidtoken"); + internalComposite.GetType().GetField("_isChanged", BindingFlags.NonPublic | BindingFlags.Instance)!.SetValue(internalComposite, false); + + context.Customers.Remove(new Customer { Id = "1", PartitionKey = "1" }); + + var ex = await Assert.ThrowsAsync(() => context.SaveChangesAsync()); + Assert.Contains("The session token provided 'invalidtoken' is invalid.", ((CosmosException)ex.InnerException!).ResponseBody); + } + public class CosmosSessionTokenContext(DbContextOptions options) : PoolableDbContext(options) { public DbSet Customers { get; set; } = null!; diff --git a/test/EFCore.Cosmos.FunctionalTests/TestUtilities/CosmosTestStore.cs b/test/EFCore.Cosmos.FunctionalTests/TestUtilities/CosmosTestStore.cs index 346ba8887d0..7c4ff895950 100644 --- a/test/EFCore.Cosmos.FunctionalTests/TestUtilities/CosmosTestStore.cs +++ b/test/EFCore.Cosmos.FunctionalTests/TestUtilities/CosmosTestStore.cs @@ -241,7 +241,7 @@ private async Task CreateFromFile(DbContext context) document["$type"] = entityName; await cosmosClient.CreateItemAsync( - containerName!, document, new FakeUpdateEntry()).ConfigureAwait(false); + containerName!, document, new FakeUpdateEntry(), new NullSessionTokenStorage()).ConfigureAwait(false); } else if (reader.TokenType == JsonToken.EndObject) { From 3a3a51041adea63220dcc08e6253dceede4f2fdb Mon Sep 17 00:00:00 2001 From: JoasE <32096708+JoasE@users.noreply.github.com> Date: Mon, 20 Oct 2025 15:42:31 +0200 Subject: [PATCH 05/37] add todo --- src/EFCore.Cosmos/Storage/Internal/SessionTokenStorage.cs | 1 + 1 file changed, 1 insertion(+) diff --git a/src/EFCore.Cosmos/Storage/Internal/SessionTokenStorage.cs b/src/EFCore.Cosmos/Storage/Internal/SessionTokenStorage.cs index 9dbe0b0fc18..f87f93949e2 100644 --- a/src/EFCore.Cosmos/Storage/Internal/SessionTokenStorage.cs +++ b/src/EFCore.Cosmos/Storage/Internal/SessionTokenStorage.cs @@ -111,6 +111,7 @@ public void SetSessionToken(string containerName, string? sessionToken) ArgumentNullException.ThrowIfNullOrWhiteSpace(containerName, nameof(containerName)); if (sessionToken is not null && string.IsNullOrWhiteSpace(sessionToken)) { + // @TODO: Exception messages in this file. throw new ArgumentException("sessionToken cannot be whitespace.", sessionToken); } From 691aa791180985ec170c0bea1404acf729599856 Mon Sep 17 00:00:00 2001 From: JoasE <32096708+JoasE@users.noreply.github.com> Date: Mon, 20 Oct 2025 20:22:13 +0200 Subject: [PATCH 06/37] Cleanup and small improvements --- .../Internal/CosmosQueryCompilationContext.cs | 13 +---- .../CosmosQueryCompilationContextFactory.cs | 22 ++++++- ...ressionVisitor.PagingQueryingEnumerable.cs | 12 +--- ...ingExpressionVisitor.QueryingEnumerable.cs | 13 +---- ...ssionVisitor.ReadItemQueryingEnumerable.cs | 11 +--- ...osShapedQueryCompilingExpressionVisitor.cs | 11 +--- .../Storage/ISessionTokenStorage.cs | 1 + .../Storage/Internal/CosmosClientWrapper.cs | 58 ++++++------------- .../Storage/Internal/ICosmosClientWrapper.cs | 10 +--- 9 files changed, 52 insertions(+), 99 deletions(-) diff --git a/src/EFCore.Cosmos/Query/Internal/CosmosQueryCompilationContext.cs b/src/EFCore.Cosmos/Query/Internal/CosmosQueryCompilationContext.cs index a106478358d..1401fe31cc2 100644 --- a/src/EFCore.Cosmos/Query/Internal/CosmosQueryCompilationContext.cs +++ b/src/EFCore.Cosmos/Query/Internal/CosmosQueryCompilationContext.cs @@ -19,10 +19,9 @@ public class CosmosQueryCompilationContext : QueryCompilationContext /// 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. /// - public CosmosQueryCompilationContext(QueryCompilationContextDependencies dependencies, bool async) : base(dependencies, async) + public CosmosQueryCompilationContext(QueryCompilationContextDependencies dependencies, ISessionTokenStorage sessionTokenStorage, bool async) : base(dependencies, async) { - // @TODO: Faster way.. - SessionTokenStorage = Dependencies.Context.Database.GetSessionTokens(); + SessionTokenStorage = sessionTokenStorage; } /// @@ -55,14 +54,6 @@ public CosmosQueryCompilationContext(QueryCompilationContextDependencies depende /// public virtual CosmosAliasManager AliasManager { get; } = new(); - /// - /// 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. - /// - public virtual string? SessionToken { get; internal set; } - /// /// 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 diff --git a/src/EFCore.Cosmos/Query/Internal/CosmosQueryCompilationContextFactory.cs b/src/EFCore.Cosmos/Query/Internal/CosmosQueryCompilationContextFactory.cs index dc5b69cb0f4..ed0a7cb69c7 100644 --- a/src/EFCore.Cosmos/Query/Internal/CosmosQueryCompilationContextFactory.cs +++ b/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; /// @@ -18,13 +20,27 @@ public class CosmosQueryCompilationContextFactory : IQueryCompilationContextFact /// doing so can result in application failures when updating to a new Entity Framework Core release. /// public CosmosQueryCompilationContextFactory(QueryCompilationContextDependencies dependencies) - => Dependencies = dependencies; + { + Dependencies = dependencies; + SessionTokenStorage = dependencies.Context.Database.GetSessionTokens(); + } /// - /// 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. /// protected virtual QueryCompilationContextDependencies Dependencies { get; } + /// + /// 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. + /// + protected virtual ISessionTokenStorage SessionTokenStorage { get; } + /// /// 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 @@ -32,5 +48,5 @@ public CosmosQueryCompilationContextFactory(QueryCompilationContextDependencies /// doing so can result in application failures when updating to a new Entity Framework Core release. /// public virtual QueryCompilationContext Create(bool async) - => new CosmosQueryCompilationContext(Dependencies, async); + => new CosmosQueryCompilationContext(Dependencies, SessionTokenStorage, async); } diff --git a/src/EFCore.Cosmos/Query/Internal/CosmosShapedQueryCompilingExpressionVisitor.PagingQueryingEnumerable.cs b/src/EFCore.Cosmos/Query/Internal/CosmosShapedQueryCompilingExpressionVisitor.PagingQueryingEnumerable.cs index 855aa3eea60..ba5e8c1a13a 100644 --- a/src/EFCore.Cosmos/Query/Internal/CosmosShapedQueryCompilingExpressionVisitor.PagingQueryingEnumerable.cs +++ b/src/EFCore.Cosmos/Query/Internal/CosmosShapedQueryCompilingExpressionVisitor.PagingQueryingEnumerable.cs @@ -36,7 +36,6 @@ private sealed class PagingQueryingEnumerable : IAsyncEnumerable> private readonly IDiagnosticsLogger _commandLogger; private readonly bool _standAloneStateManager; private readonly ISessionTokenStorage _sessionTokenStorage; - private readonly string _sessionToken; private readonly CancellationToken _cancellationToken; private readonly IConcurrencyDetector _concurrencyDetector; private readonly IExceptionDetector _exceptionDetector; @@ -119,7 +115,6 @@ public AsyncEnumerator(PagingQueryingEnumerable queryingEnumerable, Cancellat _standAloneStateManager = queryingEnumerable._standAloneStateManager; _exceptionDetector = _cosmosQueryContext.ExceptionDetector; _sessionTokenStorage = queryingEnumerable._sessionTokenStorage; - _sessionToken = queryingEnumerable._sessionToken; _cancellationToken = cancellationToken; _concurrencyDetector = queryingEnumerable._threadSafetyChecksEnabled @@ -164,10 +159,7 @@ public async ValueTask MoveNextAsync() queryRequestOptions.PartitionKey = _cosmosPartitionKey; } - if (_sessionToken is not null) - { - queryRequestOptions.SessionToken = _sessionToken; - } + queryRequestOptions.SessionToken = _sessionTokenStorage.GetSessionToken(_cosmosContainer); var cosmosClient = _cosmosQueryContext.CosmosClient; _commandLogger.ExecutingSqlQuery(_cosmosContainer, _cosmosPartitionKey, sqlQuery); diff --git a/src/EFCore.Cosmos/Query/Internal/CosmosShapedQueryCompilingExpressionVisitor.QueryingEnumerable.cs b/src/EFCore.Cosmos/Query/Internal/CosmosShapedQueryCompilingExpressionVisitor.QueryingEnumerable.cs index 7143e8595cb..7cafa5b3805 100644 --- a/src/EFCore.Cosmos/Query/Internal/CosmosShapedQueryCompilingExpressionVisitor.QueryingEnumerable.cs +++ b/src/EFCore.Cosmos/Query/Internal/CosmosShapedQueryCompilingExpressionVisitor.QueryingEnumerable.cs @@ -33,7 +33,6 @@ private sealed class QueryingEnumerable : IEnumerable, IAsyncEnumerable private readonly bool _standAloneStateManager; private readonly bool _threadSafetyChecksEnabled; private readonly ISessionTokenStorage _sessionTokenStorage; - private readonly string _sessionToken; public QueryingEnumerable( CosmosQueryContext cosmosQueryContext, @@ -46,8 +45,7 @@ public QueryingEnumerable( List partitionKeyPropertyValues, bool standAloneStateManager, bool threadSafetyChecksEnabled, - ISessionTokenStorage sessionTokenStorage, - string sessionToken) + ISessionTokenStorage sessionTokenStorage) { _cosmosQueryContext = cosmosQueryContext; _sqlExpressionFactory = sqlExpressionFactory; @@ -59,7 +57,6 @@ public QueryingEnumerable( _standAloneStateManager = standAloneStateManager; _threadSafetyChecksEnabled = threadSafetyChecksEnabled; _sessionTokenStorage = sessionTokenStorage; - _sessionToken = sessionToken; _cosmosContainer = rootEntityType.GetContainer() ?? throw new UnreachableException("Root entity type without a Cosmos container."); @@ -118,7 +115,6 @@ private sealed class Enumerator : IEnumerator private readonly bool _standAloneStateManager; private readonly IConcurrencyDetector _concurrencyDetector; private readonly IExceptionDetector _exceptionDetector; - private readonly string _sessionToken; private readonly ISessionTokenStorage _sessionTokenStorage; private IEnumerator _enumerator; @@ -134,7 +130,6 @@ public Enumerator(QueryingEnumerable queryingEnumerable) _queryLogger = queryingEnumerable._queryLogger; _standAloneStateManager = queryingEnumerable._standAloneStateManager; _exceptionDetector = _cosmosQueryContext.ExceptionDetector; - _sessionToken = queryingEnumerable._sessionToken; _sessionTokenStorage = queryingEnumerable._sessionTokenStorage; _concurrencyDetector = queryingEnumerable._threadSafetyChecksEnabled @@ -160,7 +155,7 @@ public bool MoveNext() EntityFrameworkMetricsData.ReportQueryExecuting(); _enumerator = _cosmosQueryContext.CosmosClient - .ExecuteSqlQuery(_cosmosContainer, _cosmosPartitionKey, sqlQuery, _sessionTokenStorage, _sessionToken) + .ExecuteSqlQuery(_cosmosContainer, _cosmosPartitionKey, sqlQuery, _sessionTokenStorage) .GetEnumerator(); _cosmosQueryContext.InitializeStateManager(_standAloneStateManager); } @@ -209,7 +204,6 @@ private sealed class AsyncEnumerator : IAsyncEnumerator private readonly PartitionKey _cosmosPartitionKey; private readonly IDiagnosticsLogger _queryLogger; private readonly bool _standAloneStateManager; - private readonly string _sessionToken; private readonly ISessionTokenStorage _sessionTokenStorage; private readonly CancellationToken _cancellationToken; private readonly IConcurrencyDetector _concurrencyDetector; @@ -229,7 +223,6 @@ public AsyncEnumerator(QueryingEnumerable queryingEnumerable, CancellationTok _queryLogger = queryingEnumerable._queryLogger; _standAloneStateManager = queryingEnumerable._standAloneStateManager; _exceptionDetector = _cosmosQueryContext.ExceptionDetector; - _sessionToken = queryingEnumerable._sessionToken; _sessionTokenStorage = queryingEnumerable._sessionTokenStorage; _cancellationToken = cancellationToken; @@ -254,7 +247,7 @@ public async ValueTask MoveNextAsync() EntityFrameworkMetricsData.ReportQueryExecuting(); _enumerator = _cosmosQueryContext.CosmosClient - .ExecuteSqlQueryAsync(_cosmosContainer, _cosmosPartitionKey, sqlQuery, _sessionTokenStorage, _sessionToken) + .ExecuteSqlQueryAsync(_cosmosContainer, _cosmosPartitionKey, sqlQuery, _sessionTokenStorage) .GetAsyncEnumerator(_cancellationToken); _cosmosQueryContext.InitializeStateManager(_standAloneStateManager); } diff --git a/src/EFCore.Cosmos/Query/Internal/CosmosShapedQueryCompilingExpressionVisitor.ReadItemQueryingEnumerable.cs b/src/EFCore.Cosmos/Query/Internal/CosmosShapedQueryCompilingExpressionVisitor.ReadItemQueryingEnumerable.cs index 9e6bcbfba33..714cf3e3ab0 100644 --- a/src/EFCore.Cosmos/Query/Internal/CosmosShapedQueryCompilingExpressionVisitor.ReadItemQueryingEnumerable.cs +++ b/src/EFCore.Cosmos/Query/Internal/CosmosShapedQueryCompilingExpressionVisitor.ReadItemQueryingEnumerable.cs @@ -32,7 +32,6 @@ private sealed class ReadItemQueryingEnumerable : IEnumerable, IAsyncEnume private readonly bool _standAloneStateManager; private readonly bool _threadSafetyChecksEnabled; private readonly ISessionTokenStorage _sessionTokenStorage; - private readonly string _sessionToken; public ReadItemQueryingEnumerable( CosmosQueryContext cosmosQueryContext, @@ -43,8 +42,7 @@ public ReadItemQueryingEnumerable( Type contextType, bool standAloneStateManager, bool threadSafetyChecksEnabled, - ISessionTokenStorage sessionTokenStorage, - string sessionToken) + ISessionTokenStorage sessionTokenStorage) { _cosmosQueryContext = cosmosQueryContext; _rootEntityType = rootEntityType; @@ -55,7 +53,6 @@ public ReadItemQueryingEnumerable( _standAloneStateManager = standAloneStateManager; _threadSafetyChecksEnabled = threadSafetyChecksEnabled; _sessionTokenStorage = sessionTokenStorage; - _sessionToken = sessionToken; _cosmosContainer = rootEntityType.GetContainer() ?? throw new UnreachableException("Root entity type without a Cosmos container."); _cosmosPartitionKey = GeneratePartitionKey( @@ -118,7 +115,6 @@ private sealed class Enumerator : IEnumerator, IAsyncEnumerator private readonly IConcurrencyDetector _concurrencyDetector; private readonly IExceptionDetector _exceptionDetector; private readonly ReadItemQueryingEnumerable _readItemEnumerable; - private readonly string _sessionToken; private readonly ISessionTokenStorage _sessionTokenStorage; private readonly CancellationToken _cancellationToken; @@ -137,7 +133,6 @@ public Enumerator(ReadItemQueryingEnumerable readItemEnumerable, Cancellation _exceptionDetector = _cosmosQueryContext.ExceptionDetector; _readItemEnumerable = readItemEnumerable; _sessionTokenStorage = readItemEnumerable._sessionTokenStorage; - _sessionToken = readItemEnumerable._sessionToken; _cancellationToken = cancellationToken; _concurrencyDetector = readItemEnumerable._threadSafetyChecksEnabled @@ -172,8 +167,7 @@ public bool MoveNext() _cosmosContainer, _cosmosPartitionKey, resourceId, - _sessionTokenStorage, - _sessionToken); + _sessionTokenStorage); return ShapeResult(); } @@ -215,7 +209,6 @@ public async ValueTask MoveNextAsync() _cosmosPartitionKey, resourceId, _sessionTokenStorage, - _sessionToken, _cancellationToken) .ConfigureAwait(false); diff --git a/src/EFCore.Cosmos/Query/Internal/CosmosShapedQueryCompilingExpressionVisitor.cs b/src/EFCore.Cosmos/Query/Internal/CosmosShapedQueryCompilingExpressionVisitor.cs index 9aeae2ebf5c..a8e8b0cb5d4 100644 --- a/src/EFCore.Cosmos/Query/Internal/CosmosShapedQueryCompilingExpressionVisitor.cs +++ b/src/EFCore.Cosmos/Query/Internal/CosmosShapedQueryCompilingExpressionVisitor.cs @@ -84,8 +84,6 @@ protected override Expression VisitShapedQuery(ShapedQueryExpression shapedQuery var standAloneStateManagerConstant = Constant( QueryCompilationContext.QueryTrackingBehavior == QueryTrackingBehavior.NoTrackingWithIdentityResolution); var sessionTokenStorageConstant = Constant(cosmosQueryCompilationContext.SessionTokenStorage); - var sessionToken = cosmosQueryCompilationContext.SessionToken ?? cosmosQueryCompilationContext.SessionTokenStorage.GetSessionToken(rootEntityType.GetContainer()!); - var sessionTokenConstant = Constant(sessionToken, typeof(string)); Check.DebugAssert(!paging || selectExpression.ReadItemInfo is null, "ReadItem is being with paging, impossible."); @@ -101,8 +99,7 @@ protected override Expression VisitShapedQuery(ShapedQueryExpression shapedQuery contextTypeConstant, standAloneStateManagerConstant, threadSafetyConstant, - sessionTokenStorageConstant, - sessionTokenConstant), + sessionTokenStorageConstant), _ when paging => New( typeof(PagingQueryingEnumerable<>).MakeGenericType(shaperLambda.ReturnType).GetConstructors()[0], @@ -119,8 +116,7 @@ protected override Expression VisitShapedQuery(ShapedQueryExpression shapedQuery Constant(maxItemCount.Name), Constant(continuationToken.Name), Constant(responseContinuationTokenLimitInKb.Name), - sessionTokenStorageConstant, - sessionTokenConstant), + sessionTokenStorageConstant), _ => New( typeof(QueryingEnumerable<>).MakeGenericType(shaperLambda.ReturnType).GetConstructors()[0], cosmosQueryContextConstant, @@ -133,8 +129,7 @@ protected override Expression VisitShapedQuery(ShapedQueryExpression shapedQuery Constant(cosmosQueryCompilationContext.PartitionKeyPropertyValues), standAloneStateManagerConstant, threadSafetyConstant, - sessionTokenStorageConstant, - sessionTokenConstant) + sessionTokenStorageConstant) }; } diff --git a/src/EFCore.Cosmos/Storage/ISessionTokenStorage.cs b/src/EFCore.Cosmos/Storage/ISessionTokenStorage.cs index f4b6467fbb4..6fe753b76a1 100644 --- a/src/EFCore.Cosmos/Storage/ISessionTokenStorage.cs +++ b/src/EFCore.Cosmos/Storage/ISessionTokenStorage.cs @@ -3,6 +3,7 @@ namespace Microsoft.EntityFrameworkCore.Cosmos.Storage; +// @TODO: CosmosSession(Token)Context? /// /// Defines methods for storing, retrieving, and managing session tokens associated with containers in a . /// diff --git a/src/EFCore.Cosmos/Storage/Internal/CosmosClientWrapper.cs b/src/EFCore.Cosmos/Storage/Internal/CosmosClientWrapper.cs index f891fbed40b..5356ee4de78 100644 --- a/src/EFCore.Cosmos/Storage/Internal/CosmosClientWrapper.cs +++ b/src/EFCore.Cosmos/Storage/Internal/CosmosClientWrapper.cs @@ -838,14 +838,13 @@ public virtual IEnumerable ExecuteSqlQuery( string containerId, PartitionKey partitionKeyValue, CosmosSqlQuery query, - ISessionTokenStorage sessionTokenStorage, - string? sessionToken) + ISessionTokenStorage sessionTokenStorage) { _databaseLogger.SyncNotSupported(); _commandLogger.ExecutingSqlQuery(containerId, partitionKeyValue, query); - return new DocumentEnumerable(this, containerId, partitionKeyValue, query, sessionTokenStorage, sessionToken); + return new DocumentEnumerable(this, containerId, partitionKeyValue, query, sessionTokenStorage); } /// @@ -858,12 +857,11 @@ public virtual IAsyncEnumerable ExecuteSqlQueryAsync( string containerId, PartitionKey partitionKeyValue, CosmosSqlQuery query, - ISessionTokenStorage sessionTokenStorage, - string? sessionToken) + ISessionTokenStorage sessionTokenStorage) { _commandLogger.ExecutingSqlQuery(containerId, partitionKeyValue, query); - return new DocumentAsyncEnumerable(this, containerId, partitionKeyValue, query, sessionTokenStorage, sessionToken); + return new DocumentAsyncEnumerable(this, containerId, partitionKeyValue, query, sessionTokenStorage); } /// @@ -876,14 +874,13 @@ public virtual IAsyncEnumerable ExecuteSqlQueryAsync( string containerId, PartitionKey partitionKeyValue, string resourceId, - ISessionTokenStorage sessionTokenStorage, - string? sessionToken) + ISessionTokenStorage sessionTokenStorage) { _databaseLogger.SyncNotSupported(); _commandLogger.ExecutingReadItem(containerId, partitionKeyValue, resourceId); - var response = _executionStrategy.Execute((containerId, partitionKeyValue, resourceId, sessionTokenStorage, sessionToken, this), CreateSingleItemQuery, null); + var response = _executionStrategy.Execute((containerId, partitionKeyValue, resourceId, sessionTokenStorage, this), CreateSingleItemQuery, null); _commandLogger.ExecutedReadItem( response.Diagnostics.GetClientElapsedTime(), @@ -907,13 +904,12 @@ public virtual IAsyncEnumerable ExecuteSqlQueryAsync( PartitionKey partitionKeyValue, string resourceId, ISessionTokenStorage sessionTokenStorage, - string? sessionToken, CancellationToken cancellationToken = default) { _commandLogger.ExecutingReadItem(containerId, partitionKeyValue, resourceId); var response = await _executionStrategy.ExecuteAsync( - (containerId, partitionKeyValue, resourceId, sessionTokenStorage, sessionToken, this), + (containerId, partitionKeyValue, resourceId, sessionTokenStorage, this), CreateSingleItemQueryAsync, null, cancellationToken) @@ -932,22 +928,18 @@ public virtual IAsyncEnumerable ExecuteSqlQueryAsync( private static ResponseMessage CreateSingleItemQuery( DbContext? context, - (string ContainerId, PartitionKey PartitionKeyValue, string ResourceId, ISessionTokenStorage SessionTokenStorage, string? SessionToken, CosmosClientWrapper Wrapper) parameters) + (string ContainerId, PartitionKey PartitionKeyValue, string ResourceId, ISessionTokenStorage SessionTokenStorage, CosmosClientWrapper Wrapper) parameters) => CreateSingleItemQueryAsync(context, parameters).GetAwaiter().GetResult(); private static async Task CreateSingleItemQueryAsync( DbContext? _, - (string ContainerId, PartitionKey PartitionKeyValue, string ResourceId, ISessionTokenStorage SessionTokenStorage, string? SessionToken, CosmosClientWrapper Wrapper) parameters, + (string ContainerId, PartitionKey PartitionKeyValue, string ResourceId, ISessionTokenStorage SessionTokenStorage, CosmosClientWrapper Wrapper) parameters, CancellationToken cancellationToken = default) { - var (containerId, partitionKeyValue, resourceId, sessionTokenStorage, sessionToken, wrapper) = parameters; + var (containerId, partitionKeyValue, resourceId, sessionTokenStorage, wrapper) = parameters; var container = wrapper.Client.GetDatabase(wrapper._databaseId).GetContainer(containerId); - ItemRequestOptions? itemRequestOptions = null; - if (sessionToken != null) - { - itemRequestOptions = new ItemRequestOptions { SessionToken = sessionToken }; - } + var itemRequestOptions = new ItemRequestOptions { SessionToken = sessionTokenStorage.GetSessionToken(containerId) }; var response = await container.ReadItemStreamAsync( resourceId, @@ -1043,8 +1035,7 @@ private sealed class DocumentEnumerable( string containerId, PartitionKey partitionKeyValue, CosmosSqlQuery cosmosSqlQuery, - ISessionTokenStorage sessionTokenStorage, - string? sessionToken) + ISessionTokenStorage sessionTokenStorage) : IEnumerable { private readonly CosmosClientWrapper _cosmosClient = cosmosClient; @@ -1052,7 +1043,6 @@ private sealed class DocumentEnumerable( private readonly PartitionKey _partitionKeyValue = partitionKeyValue; private readonly CosmosSqlQuery _cosmosSqlQuery = cosmosSqlQuery; private readonly ISessionTokenStorage _sessionTokenStorage = sessionTokenStorage; - private readonly string? _sessionToken = sessionToken; public IEnumerator GetEnumerator() => new Enumerator(this); @@ -1067,7 +1057,6 @@ private sealed class Enumerator(DocumentEnumerable documentEnumerable) : IEnumer private readonly PartitionKey _partitionKeyValue = documentEnumerable._partitionKeyValue; private readonly CosmosSqlQuery _cosmosSqlQuery = documentEnumerable._cosmosSqlQuery; private readonly ISessionTokenStorage _sessionTokenStorage = documentEnumerable._sessionTokenStorage; - private readonly string? _sessionToken = documentEnumerable._sessionToken; private JToken? _current; private ResponseMessage? _responseMessage; @@ -1094,10 +1083,8 @@ public bool MoveNext() queryRequestOptions.PartitionKey = _partitionKeyValue; } - if (_sessionToken is not null) - { - queryRequestOptions.SessionToken = _sessionToken; - } + // @TODO: Or should this be inside CreateQuery... + queryRequestOptions.SessionToken = _sessionTokenStorage.GetSessionToken(_containerId); _query = _cosmosClientWrapper.CreateQuery( _containerId, _cosmosSqlQuery, _sessionTokenStorage, continuationToken: null, queryRequestOptions); @@ -1161,8 +1148,7 @@ private sealed class DocumentAsyncEnumerable( string containerId, PartitionKey partitionKeyValue, CosmosSqlQuery cosmosSqlQuery, - ISessionTokenStorage sessionTokenStorage, - string? sessionToken) + ISessionTokenStorage sessionTokenStorage) : IAsyncEnumerable { private readonly CosmosClientWrapper _cosmosClient = cosmosClient; @@ -1170,7 +1156,6 @@ private sealed class DocumentAsyncEnumerable( private readonly PartitionKey _partitionKeyValue = partitionKeyValue; private readonly CosmosSqlQuery _cosmosSqlQuery = cosmosSqlQuery; private readonly ISessionTokenStorage _sessionTokenStorage = sessionTokenStorage; - private readonly string? _sessionToken = sessionToken; public IAsyncEnumerator GetAsyncEnumerator(CancellationToken cancellationToken = default) => new AsyncEnumerator(this, cancellationToken); @@ -1183,7 +1168,6 @@ private sealed class AsyncEnumerator(DocumentAsyncEnumerable documentEnumerable, private readonly PartitionKey _partitionKeyValue = documentEnumerable._partitionKeyValue; private readonly CosmosSqlQuery _cosmosSqlQuery = documentEnumerable._cosmosSqlQuery; private readonly ISessionTokenStorage _sessionTokenStorage = documentEnumerable._sessionTokenStorage; - private readonly string? _sessionToken = documentEnumerable._sessionToken; private JToken? _current; private ResponseMessage? _responseMessage; @@ -1209,10 +1193,8 @@ public async ValueTask MoveNextAsync() queryRequestOptions.PartitionKey = _partitionKeyValue; } - if (_sessionToken is not null) - { - queryRequestOptions.SessionToken = _sessionToken; - } + // @TODO: Or should this be inside CreateQuery... + queryRequestOptions.SessionToken = _sessionTokenStorage.GetSessionToken(_containerId); _query = _cosmosClientWrapper.CreateQuery( _containerId, _cosmosSqlQuery, _sessionTokenStorage, continuationToken: null, queryRequestOptions); @@ -1240,12 +1222,6 @@ public async ValueTask MoveNextAsync() _responseMessage.EnsureSuccessStatusCode(); - if (!string.IsNullOrWhiteSpace(_responseMessage.Headers.Session)) - { - // @TODO: set session token... - // if (appendSessionToken == higher.....) update _sessionToken?? - } - _responseMessageEnumerator = new ResponseMessageEnumerable(_responseMessage).GetAsyncEnumerator(cancellationToken); } diff --git a/src/EFCore.Cosmos/Storage/Internal/ICosmosClientWrapper.cs b/src/EFCore.Cosmos/Storage/Internal/ICosmosClientWrapper.cs index f9d0bb4bd16..0337d55c197 100644 --- a/src/EFCore.Cosmos/Storage/Internal/ICosmosClientWrapper.cs +++ b/src/EFCore.Cosmos/Storage/Internal/ICosmosClientWrapper.cs @@ -158,8 +158,7 @@ FeedIterator CreateQuery( string containerId, PartitionKey partitionKeyValue, string resourceId, - ISessionTokenStorage sessionTokenStorage, - string? sessionToken); + ISessionTokenStorage sessionTokenStorage); /// /// This is an internal API that supports the Entity Framework Core infrastructure and not subject to @@ -172,7 +171,6 @@ FeedIterator CreateQuery( PartitionKey partitionKeyValue, string resourceId, ISessionTokenStorage sessionTokenStorage, - string? sessionToken, CancellationToken cancellationToken = default); /// @@ -185,8 +183,7 @@ IEnumerable ExecuteSqlQuery( string containerId, PartitionKey partitionKeyValue, CosmosSqlQuery query, - ISessionTokenStorage sessionTokenStorage, - string? sessionToken); + ISessionTokenStorage sessionTokenStorage); /// /// This is an internal API that supports the Entity Framework Core infrastructure and not subject to @@ -198,8 +195,7 @@ IAsyncEnumerable ExecuteSqlQueryAsync( string containerId, PartitionKey partitionKeyValue, CosmosSqlQuery query, - ISessionTokenStorage sessionTokenStorage, - string? sessionToken); + ISessionTokenStorage sessionTokenStorage); /// /// This is an internal API that supports the Entity Framework Core infrastructure and not subject to From c4a99772d7f06f4d5d9228ce1afe1b1cda7763b0 Mon Sep 17 00:00:00 2001 From: JoasE <32096708+JoasE@users.noreply.github.com> Date: Mon, 20 Oct 2025 21:12:01 +0200 Subject: [PATCH 07/37] Cleanup and small improvements --- .../CosmosQueryCompilationContextFactory.cs | 5 ++-- .../Storage/Internal/CosmosClientWrapper.cs | 14 ++++------- .../CosmosTransactionalBatchResult.cs | 25 +++---------------- .../Storage/Internal/SessionTokenStorage.cs | 14 +++++------ 4 files changed, 18 insertions(+), 40 deletions(-) diff --git a/src/EFCore.Cosmos/Query/Internal/CosmosQueryCompilationContextFactory.cs b/src/EFCore.Cosmos/Query/Internal/CosmosQueryCompilationContextFactory.cs index ed0a7cb69c7..99a937d82d6 100644 --- a/src/EFCore.Cosmos/Query/Internal/CosmosQueryCompilationContextFactory.cs +++ b/src/EFCore.Cosmos/Query/Internal/CosmosQueryCompilationContextFactory.cs @@ -13,6 +13,8 @@ namespace Microsoft.EntityFrameworkCore.Cosmos.Query.Internal; /// public class CosmosQueryCompilationContextFactory : IQueryCompilationContextFactory { + private ISessionTokenStorage? _sessionTokenStorage; + /// /// 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 @@ -22,7 +24,6 @@ public class CosmosQueryCompilationContextFactory : IQueryCompilationContextFact public CosmosQueryCompilationContextFactory(QueryCompilationContextDependencies dependencies) { Dependencies = dependencies; - SessionTokenStorage = dependencies.Context.Database.GetSessionTokens(); } /// @@ -39,7 +40,7 @@ public CosmosQueryCompilationContextFactory(QueryCompilationContextDependencies /// 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. /// - protected virtual ISessionTokenStorage SessionTokenStorage { get; } + protected virtual ISessionTokenStorage SessionTokenStorage => _sessionTokenStorage ??= Dependencies.Context.Database.GetSessionTokens(); /// /// This is an internal API that supports the Entity Framework Core infrastructure and not subject to diff --git a/src/EFCore.Cosmos/Storage/Internal/CosmosClientWrapper.cs b/src/EFCore.Cosmos/Storage/Internal/CosmosClientWrapper.cs index 5356ee4de78..280e7909588 100644 --- a/src/EFCore.Cosmos/Storage/Internal/CosmosClientWrapper.cs +++ b/src/EFCore.Cosmos/Storage/Internal/CosmosClientWrapper.cs @@ -699,12 +699,10 @@ private static async Task ExecuteBatchOnceAsync( var wrapper = parameters.Wrapper; var sessionTokenStorage = parameters.SessionTokenStorage; - var options = new TransactionalBatchRequestOptions(); - var sessionToken = sessionTokenStorage.GetSessionToken(batch.CollectionId); - if (!string.IsNullOrWhiteSpace(sessionToken)) + var options = new TransactionalBatchRequestOptions { - options.SessionToken = sessionToken; - } + SessionToken = sessionTokenStorage.GetSessionToken(batch.CollectionId) + }; using var response = await transactionalBatch.ExecuteAsync(options, cancellationToken).ConfigureAwait(false); @@ -726,20 +724,18 @@ private static async Task ExecuteBatchOnceAsync( .ToList(); var exception = new CosmosException(response.ErrorMessage, errorCode, 0, response.ActivityId, response.RequestCharge); - return CosmosTransactionalBatchResult.Failure(errorEntries, exception); + return new CosmosTransactionalBatchResult(errorEntries, exception); } ProcessResponse(batch.CollectionId, response, batch.Entries, sessionTokenStorage); - return CosmosTransactionalBatchResult.Success(response.Headers.Session); + return CosmosTransactionalBatchResult.Success; } private static ItemRequestOptions? CreateItemRequestOptions(IUpdateEntry entry, bool? enableContentResponseOnWrite, string? sessionToken) { var helper = RequestOptionsHelper.Create(entry, enableContentResponseOnWrite); - sessionToken = string.IsNullOrWhiteSpace(sessionToken) ? null : sessionToken; - return helper == null ? null : new ItemRequestOptions { IfMatchEtag = helper.IfMatchEtag, SessionToken = sessionToken, EnableContentResponseOnWrite = helper.EnableContentResponseOnWrite }; diff --git a/src/EFCore.Cosmos/Storage/Internal/CosmosTransactionalBatchResult.cs b/src/EFCore.Cosmos/Storage/Internal/CosmosTransactionalBatchResult.cs index df6a535f71b..d882ff8da1c 100644 --- a/src/EFCore.Cosmos/Storage/Internal/CosmosTransactionalBatchResult.cs +++ b/src/EFCore.Cosmos/Storage/Internal/CosmosTransactionalBatchResult.cs @@ -19,22 +19,11 @@ public sealed class CosmosTransactionalBatchResult /// 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. /// - public static CosmosTransactionalBatchResult Success(string? sessionToken) - => new CosmosTransactionalBatchResult(sessionToken); + public static CosmosTransactionalBatchResult Success { get; } = new CosmosTransactionalBatchResult(); - /// - /// 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. - /// - public static CosmosTransactionalBatchResult Failure(IReadOnlyList entries, CosmosException exception) - => new CosmosTransactionalBatchResult(entries, exception); - - private CosmosTransactionalBatchResult(string? sessionToken) + private CosmosTransactionalBatchResult() { IsSuccess = true; - SessionToken = sessionToken; } /// @@ -43,7 +32,7 @@ private CosmosTransactionalBatchResult(string? sessionToken) /// 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. /// - private CosmosTransactionalBatchResult(IReadOnlyList entries, CosmosException exception) + public CosmosTransactionalBatchResult(IReadOnlyList entries, CosmosException exception) { IsSuccess = false; ErroredEntries = entries; @@ -59,14 +48,6 @@ private CosmosTransactionalBatchResult(IReadOnlyList entries, Cosm [MemberNotNullWhen(false, nameof(ErroredEntries), nameof(Exception))] public bool IsSuccess { get; } - /// - /// 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. - /// - public string? SessionToken { get; } - /// /// 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 diff --git a/src/EFCore.Cosmos/Storage/Internal/SessionTokenStorage.cs b/src/EFCore.Cosmos/Storage/Internal/SessionTokenStorage.cs index f87f93949e2..fe42f30345b 100644 --- a/src/EFCore.Cosmos/Storage/Internal/SessionTokenStorage.cs +++ b/src/EFCore.Cosmos/Storage/Internal/SessionTokenStorage.cs @@ -16,7 +16,7 @@ namespace Microsoft.EntityFrameworkCore.Cosmos.Storage.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. /// -public sealed class SessionTokenStorage : ISessionTokenStorage +public class SessionTokenStorage : ISessionTokenStorage { private readonly Dictionary _containerSessionTokens; private readonly string _defaultContainerName; @@ -42,7 +42,7 @@ public SessionTokenStorage(DbContext dbContext) /// 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. /// - public string? GetSessionToken() + public virtual string? GetSessionToken() => GetSessionToken(_defaultContainerName); /// @@ -51,7 +51,7 @@ public SessionTokenStorage(DbContext dbContext) /// 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. /// - public void SetSessionToken(string? sessionToken) + public virtual void SetSessionToken(string? sessionToken) => SetSessionToken(_defaultContainerName, sessionToken); /// @@ -60,7 +60,7 @@ public void SetSessionToken(string? sessionToken) /// 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. /// - public void AppendSessionToken(string sessionToken) + public virtual void AppendSessionToken(string sessionToken) => AppendSessionToken(_defaultContainerName, sessionToken); /// @@ -69,7 +69,7 @@ public void AppendSessionToken(string sessionToken) /// 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. /// - public string? GetSessionToken(string containerName) + public virtual string? GetSessionToken(string containerName) { ArgumentNullException.ThrowIfNullOrWhiteSpace(containerName, nameof(containerName)); @@ -87,7 +87,7 @@ public void AppendSessionToken(string sessionToken) /// 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. /// - public void AppendSessionToken(string containerName, string sessionToken) + public virtual void AppendSessionToken(string containerName, string sessionToken) { ArgumentNullException.ThrowIfNullOrWhiteSpace(containerName, nameof(containerName)); ArgumentNullException.ThrowIfNullOrWhiteSpace(sessionToken, nameof(sessionToken)); @@ -106,7 +106,7 @@ public void AppendSessionToken(string containerName, string sessionToken) /// 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. /// - public void SetSessionToken(string containerName, string? sessionToken) + public virtual void SetSessionToken(string containerName, string? sessionToken) { ArgumentNullException.ThrowIfNullOrWhiteSpace(containerName, nameof(containerName)); if (sessionToken is not null && string.IsNullOrWhiteSpace(sessionToken)) From ac5ca96d2b3c2b7436e6c7927ccf7463f5fd1175 Mon Sep 17 00:00:00 2001 From: JoasE <32096708+JoasE@users.noreply.github.com> Date: Tue, 21 Oct 2025 15:19:56 +0200 Subject: [PATCH 08/37] Add option ManualSessionTokenManagementEnabled --- .../CosmosDatabaseFacadeExtensions.cs | 10 +++++-- .../CosmosDbContextOptionsBuilder.cs | 16 ++++++++++ .../Internal/CosmosDbOptionExtension.cs | 30 ++++++++++++++++++- .../Internal/CosmosSingletonOptions.cs | 10 +++++++ .../Internal/ICosmosSingletonOptions.cs | 8 +++++ .../CosmosQueryCompilationContextFactory.cs | 3 +- .../Storage/Internal/CosmosDatabaseWrapper.cs | 6 +++- .../CosmosSessionTokensTest.cs | 14 ++++++++- 8 files changed, 91 insertions(+), 6 deletions(-) diff --git a/src/EFCore.Cosmos/Extensions/CosmosDatabaseFacadeExtensions.cs b/src/EFCore.Cosmos/Extensions/CosmosDatabaseFacadeExtensions.cs index 65080c57aa1..7bbabdcdabf 100644 --- a/src/EFCore.Cosmos/Extensions/CosmosDatabaseFacadeExtensions.cs +++ b/src/EFCore.Cosmos/Extensions/CosmosDatabaseFacadeExtensions.cs @@ -27,7 +27,7 @@ public static CosmosClient GetCosmosClient(this DatabaseFacade databaseFacade) => GetService(databaseFacade).Client; /// - /// Gets used to manage the session tokens for this . + /// Gets the used to manage the session tokens for this . /// /// The for the context. /// The . @@ -39,7 +39,13 @@ public static ISessionTokenStorage GetSessionTokens(this DatabaseFacade database throw new InvalidOperationException(CosmosStrings.CosmosNotInUse); } - return dbWrapper.SessionTokenStorage!; + if (dbWrapper.SessionTokenStorage is not SessionTokenStorage) + { + // @TODO: string + throw new InvalidOperationException("CosmosStrings.EnableManualSessionTokenManagement"); + } + + return dbWrapper.SessionTokenStorage; } private static TService GetService(IInfrastructure databaseFacade) diff --git a/src/EFCore.Cosmos/Infrastructure/CosmosDbContextOptionsBuilder.cs b/src/EFCore.Cosmos/Infrastructure/CosmosDbContextOptionsBuilder.cs index e88fabf78da..2dd3ec58ff3 100644 --- a/src/EFCore.Cosmos/Infrastructure/CosmosDbContextOptionsBuilder.cs +++ b/src/EFCore.Cosmos/Infrastructure/CosmosDbContextOptionsBuilder.cs @@ -211,6 +211,22 @@ public virtual CosmosDbContextOptionsBuilder MaxRequestsPerTcpConnection(int req public virtual CosmosDbContextOptionsBuilder ContentResponseOnWriteEnabled(bool enabled = true) => WithOption(e => e.ContentResponseOnWriteEnabled(Check.NotNull(enabled))); + + /// + /// Sets the boolean to track and manage session tokens for requests made to Cosmos DB + /// and being able to access them via the method. + /// 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. + /// See Utilize session tokens for more details. + /// + /// + /// See Using DbContextOptions, and + /// Accessing Azure Cosmos DB with EF Core for more information and examples. + /// + /// to track and manually manage session tokens in EF. + public virtual CosmosDbContextOptionsBuilder ManualSessionTokenManagementEnabled(bool enabled = true) + => WithOption(e => e.ManualSessionTokenManagementEnabled(enabled)); + /// /// 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. diff --git a/src/EFCore.Cosmos/Infrastructure/Internal/CosmosDbOptionExtension.cs b/src/EFCore.Cosmos/Infrastructure/Internal/CosmosDbOptionExtension.cs index f2545174ac3..dbe9e67e273 100644 --- a/src/EFCore.Cosmos/Infrastructure/Internal/CosmosDbOptionExtension.cs +++ b/src/EFCore.Cosmos/Infrastructure/Internal/CosmosDbOptionExtension.cs @@ -36,6 +36,7 @@ public class CosmosOptionsExtension : IDbContextOptionsExtension private bool? _enableContentResponseOnWrite; private DbContextOptionsExtensionInfo? _info; private Func? _httpClientFactory; + private bool _enableManualSessionTokenManagement; /// /// 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; } /// @@ -564,6 +566,30 @@ public virtual CosmosOptionsExtension WithHttpClientFactory(Func? ht return clone; } + /// + /// 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. + /// + public virtual bool EnableManualSessionTokenManagement + => _enableManualSessionTokenManagement; + + /// + /// 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. + /// + public virtual CosmosOptionsExtension ManualSessionTokenManagementEnabled(bool enabled) + { + var clone = Clone(); + + clone._enableManualSessionTokenManagement = enabled; + + return clone; + } + /// /// 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 debugInfo) { diff --git a/src/EFCore.Cosmos/Infrastructure/Internal/CosmosSingletonOptions.cs b/src/EFCore.Cosmos/Infrastructure/Internal/CosmosSingletonOptions.cs index af29229cfa5..d5221c0eef3 100644 --- a/src/EFCore.Cosmos/Infrastructure/Internal/CosmosSingletonOptions.cs +++ b/src/EFCore.Cosmos/Infrastructure/Internal/CosmosSingletonOptions.cs @@ -151,6 +151,14 @@ public class CosmosSingletonOptions : ICosmosSingletonOptions /// public virtual Func? HttpClientFactory { get; private set; } + /// + /// 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. + /// + public virtual bool EnableManualSessionTokenManagement { get; private set; } + /// /// 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( diff --git a/src/EFCore.Cosmos/Infrastructure/Internal/ICosmosSingletonOptions.cs b/src/EFCore.Cosmos/Infrastructure/Internal/ICosmosSingletonOptions.cs index a26b79a82b3..8113da2b25f 100644 --- a/src/EFCore.Cosmos/Infrastructure/Internal/ICosmosSingletonOptions.cs +++ b/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. /// Func? HttpClientFactory { get; } + + /// + /// 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. + /// + bool EnableManualSessionTokenManagement { get; } } diff --git a/src/EFCore.Cosmos/Query/Internal/CosmosQueryCompilationContextFactory.cs b/src/EFCore.Cosmos/Query/Internal/CosmosQueryCompilationContextFactory.cs index 99a937d82d6..d409c850c97 100644 --- a/src/EFCore.Cosmos/Query/Internal/CosmosQueryCompilationContextFactory.cs +++ b/src/EFCore.Cosmos/Query/Internal/CosmosQueryCompilationContextFactory.cs @@ -2,6 +2,7 @@ // 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; @@ -40,7 +41,7 @@ public CosmosQueryCompilationContextFactory(QueryCompilationContextDependencies /// 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. /// - protected virtual ISessionTokenStorage SessionTokenStorage => _sessionTokenStorage ??= Dependencies.Context.Database.GetSessionTokens(); + protected virtual ISessionTokenStorage SessionTokenStorage => _sessionTokenStorage ??= ((CosmosDatabaseWrapper)Dependencies.Context.Database.GetService()).SessionTokenStorage; /// /// This is an internal API that supports the Entity Framework Core infrastructure and not subject to diff --git a/src/EFCore.Cosmos/Storage/Internal/CosmosDatabaseWrapper.cs b/src/EFCore.Cosmos/Storage/Internal/CosmosDatabaseWrapper.cs index 8b147fcb0b7..0955b427289 100644 --- a/src/EFCore.Cosmos/Storage/Internal/CosmosDatabaseWrapper.cs +++ b/src/EFCore.Cosmos/Storage/Internal/CosmosDatabaseWrapper.cs @@ -5,6 +5,7 @@ using System.Runtime.InteropServices; using Microsoft.EntityFrameworkCore.ChangeTracking.Internal; using Microsoft.EntityFrameworkCore.Cosmos.Diagnostics.Internal; +using Microsoft.EntityFrameworkCore.Cosmos.Infrastructure.Internal; using Microsoft.EntityFrameworkCore.Cosmos.Internal; using Microsoft.EntityFrameworkCore.Cosmos.Metadata.Internal; using Microsoft.EntityFrameworkCore.Cosmos.Update.Internal; @@ -38,13 +39,16 @@ public CosmosDatabaseWrapper( DatabaseDependencies dependencies, ICurrentDbContext currentDbContext, ICosmosClientWrapper cosmosClient, + ICosmosSingletonOptions cosmosSingletonOptions, ILoggingOptions loggingOptions) : base(dependencies) { _currentDbContext = currentDbContext; _cosmosClient = cosmosClient; - SessionTokenStorage = new SessionTokenStorage(_currentDbContext.Context); + SessionTokenStorage = cosmosSingletonOptions.EnableManualSessionTokenManagement ? + new SessionTokenStorage(_currentDbContext.Context) + : new NullSessionTokenStorage(); if (loggingOptions.IsSensitiveDataLoggingEnabled) { diff --git a/test/EFCore.Cosmos.FunctionalTests/CosmosSessionTokensTest.cs b/test/EFCore.Cosmos.FunctionalTests/CosmosSessionTokensTest.cs index e679699ce9a..af1e505372d 100644 --- a/test/EFCore.Cosmos.FunctionalTests/CosmosSessionTokensTest.cs +++ b/test/EFCore.Cosmos.FunctionalTests/CosmosSessionTokensTest.cs @@ -15,6 +15,19 @@ protected override string StoreName protected override ITestStoreFactory TestStoreFactory => CosmosTestStoreFactory.Instance; + protected override TestStore CreateTestStore() => CosmosTestStore.Create(StoreName, (c) => c.ManualSessionTokenManagementEnabled()); + + [ConditionalFact] + public virtual async Task GetSessionTokens_throws_if_not_enabled() + { + var contextFactory = await InitializeAsync(createTestStore: () => CosmosTestStore.Create(StoreName)); + + using var context = contextFactory.CreateContext(); + + var exception = Assert.Throws(() => context.Database.GetSessionTokens()); + Assert.Equal("CosmosStrings.EnableManualSessionTokenManagement", exception.Message); + } + // @TODO: Tests for other container? [ConditionalFact] @@ -411,7 +424,6 @@ public virtual async Task Add_uses_session_token(AutoTransactionBehavior autoTra context.Database.AutoTransactionBehavior = autoTransactionBehavior; var sessionTokens = context.Database.GetSessionTokens(); - var sessionToken = sessionTokens.GetSessionToken()!; // Only way we can test this is by setting a session token that will fail the request if used.. // Only way to do this for a write is to set an invalid session token.. var internalDictionary = sessionTokens.GetType().GetField("_containerSessionTokens", BindingFlags.NonPublic | BindingFlags.Instance)!.GetValue(sessionTokens)!; From 6f84317aac4779557cb8dfc83b77a296c45068b5 Mon Sep 17 00:00:00 2001 From: JoasE <32096708+JoasE@users.noreply.github.com> Date: Tue, 21 Oct 2025 16:29:00 +0200 Subject: [PATCH 09/37] fix no etag sessiontoken match --- .../Storage/Internal/CosmosClientWrapper.cs | 20 ++++++++++++------- 1 file changed, 13 insertions(+), 7 deletions(-) diff --git a/src/EFCore.Cosmos/Storage/Internal/CosmosClientWrapper.cs b/src/EFCore.Cosmos/Storage/Internal/CosmosClientWrapper.cs index 280e7909588..7c512b8b916 100644 --- a/src/EFCore.Cosmos/Storage/Internal/CosmosClientWrapper.cs +++ b/src/EFCore.Cosmos/Storage/Internal/CosmosClientWrapper.cs @@ -433,7 +433,6 @@ private static async Task CreateItemOnceAsync( var postTriggers = GetTriggers(entry, TriggerType.Post, TriggerOperation.Create); if (preTriggers != null || postTriggers != null) { - itemRequestOptions ??= new ItemRequestOptions(); if (preTriggers != null) { itemRequestOptions.PreTriggers = preTriggers; @@ -522,7 +521,6 @@ private static async Task ReplaceItemOnceAsync( var postTriggers = GetTriggers(entry, TriggerType.Post, TriggerOperation.Replace); if (preTriggers != null || postTriggers != null) { - itemRequestOptions ??= new ItemRequestOptions(); if (preTriggers != null) { itemRequestOptions.PreTriggers = preTriggers; @@ -608,7 +606,6 @@ private static async Task DeleteItemOnceAsync( var postTriggers = GetTriggers(entry, TriggerType.Post, TriggerOperation.Delete); if (preTriggers != null || postTriggers != null) { - itemRequestOptions ??= new ItemRequestOptions(); if (preTriggers != null) { itemRequestOptions.PreTriggers = preTriggers; @@ -732,13 +729,22 @@ private static async Task ExecuteBatchOnceAsync( return CosmosTransactionalBatchResult.Success; } - private static ItemRequestOptions? CreateItemRequestOptions(IUpdateEntry entry, bool? enableContentResponseOnWrite, string? sessionToken) + private static ItemRequestOptions CreateItemRequestOptions(IUpdateEntry entry, bool? enableContentResponseOnWrite, string? sessionToken) { var helper = RequestOptionsHelper.Create(entry, enableContentResponseOnWrite); - return helper == null - ? null - : new ItemRequestOptions { IfMatchEtag = helper.IfMatchEtag, SessionToken = sessionToken, EnableContentResponseOnWrite = helper.EnableContentResponseOnWrite }; + var itemRequestOptions = new ItemRequestOptions + { + SessionToken = sessionToken + }; + + if (helper != null) + { + itemRequestOptions.IfMatchEtag = helper.IfMatchEtag; + itemRequestOptions.EnableContentResponseOnWrite = helper.EnableContentResponseOnWrite; + } + + return itemRequestOptions; } private static IReadOnlyList? GetTriggers(IUpdateEntry entry, TriggerType type, TriggerOperation operation) From 64ab888a7be0f2d7f062d2443847c893405cdbf5 Mon Sep 17 00:00:00 2001 From: JoasE <32096708+JoasE@users.noreply.github.com> Date: Tue, 21 Oct 2025 16:38:40 +0200 Subject: [PATCH 10/37] Add some testcases --- .../CosmosSessionTokensTest.cs | 1055 ++++++++++++++--- 1 file changed, 898 insertions(+), 157 deletions(-) diff --git a/test/EFCore.Cosmos.FunctionalTests/CosmosSessionTokensTest.cs b/test/EFCore.Cosmos.FunctionalTests/CosmosSessionTokensTest.cs index af1e505372d..62914bb1a7d 100644 --- a/test/EFCore.Cosmos.FunctionalTests/CosmosSessionTokensTest.cs +++ b/test/EFCore.Cosmos.FunctionalTests/CosmosSessionTokensTest.cs @@ -15,6 +15,8 @@ protected override string StoreName protected override ITestStoreFactory TestStoreFactory => CosmosTestStoreFactory.Instance; + protected override DbContextOptionsBuilder AddOptions(DbContextOptionsBuilder builder) => base.AddOptions(builder).ConfigureWarnings(x => x.Ignore(CosmosEventId.SyncNotSupported)); + protected override TestStore CreateTestStore() => CosmosTestStore.Create(StoreName, (c) => c.ManualSessionTokenManagementEnabled()); [ConditionalFact] @@ -28,8 +30,6 @@ public virtual async Task GetSessionTokens_throws_if_not_enabled() Assert.Equal("CosmosStrings.EnableManualSessionTokenManagement", exception.Message); } - // @TODO: Tests for other container? - [ConditionalFact] public virtual async Task SetSessionToken_ThrowsForNonExistentContainer() { @@ -63,360 +63,1029 @@ public virtual async Task GetSessionToken_ThrowsForNonExistentContainer() Assert.Equal(CosmosStrings.ContainerNameDoesNotExist("Not the container name") + " (Parameter 'containerName')", exception.Message); } - [ConditionalFact] - public virtual async Task AppendSessionToken_no_tokens_sets_token() + + [ConditionalTheory, InlineData(true), InlineData(false)] + public virtual async Task AppendSessionToken_no_tokens_sets_token(bool defaultContainer) { var contextFactory = await InitializeAsync(); using var context = contextFactory.CreateContext(); var sessionTokens = context.Database.GetSessionTokens(); - sessionTokens.AppendSessionToken("0:-1#231"); - var updatedToken = sessionTokens.GetSessionToken(); + if (defaultContainer) + { + sessionTokens.AppendSessionToken("0:-1#231"); + } + else + { + sessionTokens.AppendSessionToken(OtherContainerName, "0:-1#231"); + } + + var updatedToken = defaultContainer ? sessionTokens.GetSessionToken() : sessionTokens.GetSessionToken(OtherContainerName); Assert.Equal("0:-1#231", updatedToken); } - [ConditionalFact] - public virtual async Task AppendSessionToken_append_higher_lsn_same_pkrange_takes_higher_lsn() + [ConditionalTheory, InlineData(true), InlineData(false)] + public virtual async Task AppendSessionToken_append_higher_lsn_same_pkrange_takes_higher_lsn(bool defaultContainer) { var contextFactory = await InitializeAsync(); - + using var context = contextFactory.CreateContext(); - context.Customers.Add(new Customer { Id = "1", PartitionKey = "1" }); - + if (defaultContainer) + { + context.Add(new Customer { Id = "1", PartitionKey = "1" }); + } + else + { + context.Add(new OtherContainerCustomer { Id = "1", PartitionKey = "1" }); + } await context.SaveChangesAsync(); var sessionTokens = context.Database.GetSessionTokens(); - var initialToken = sessionTokens.GetSessionToken(); + var initialToken = defaultContainer ? sessionTokens.GetSessionToken() : sessionTokens.GetSessionToken(OtherContainerName); Assert.False(string.IsNullOrWhiteSpace(initialToken)); var newToken = initialToken.Substring(0, initialToken.IndexOf('#') + 1) + "999999"; - sessionTokens.AppendSessionToken(newToken); + if (defaultContainer) + { + sessionTokens.AppendSessionToken(newToken); + } + else + { + sessionTokens.AppendSessionToken(OtherContainerName, newToken); + } - var updatedToken = sessionTokens.GetSessionToken(); + var updatedToken = defaultContainer ? sessionTokens.GetSessionToken() : sessionTokens.GetSessionToken(OtherContainerName); Assert.Equal(newToken, updatedToken); } - [ConditionalFact] - public virtual async Task AppendSessionToken_append_lower_lsn_same_pkrange_takes_higher_lsn() + [ConditionalTheory, InlineData(true), InlineData(false)] + public virtual async Task AppendSessionToken_append_lower_lsn_same_pkrange_takes_higher_lsn(bool defaultContainer) { var contextFactory = await InitializeAsync(); using var context = contextFactory.CreateContext(); - context.Customers.Add(new Customer { Id = "1", PartitionKey = "1" }); + if (defaultContainer) + { + context.Add(new Customer { Id = "1", PartitionKey = "1" }); + } + else + { + context.Add(new OtherContainerCustomer { Id = "1", PartitionKey = "1" }); + } await context.SaveChangesAsync(); var sessionTokens = context.Database.GetSessionTokens(); - var initialToken = sessionTokens.GetSessionToken(); + var initialToken = defaultContainer ? sessionTokens.GetSessionToken() : sessionTokens.GetSessionToken(OtherContainerName); Assert.False(string.IsNullOrWhiteSpace(initialToken)); var newToken = initialToken.Substring(0, initialToken.IndexOf('#') + 1) + "1"; - sessionTokens.AppendSessionToken(newToken); + if (defaultContainer) + { + sessionTokens.AppendSessionToken(newToken); + } + else + { + sessionTokens.AppendSessionToken(OtherContainerName, newToken); + } - var updatedToken = sessionTokens.GetSessionToken(); + var updatedToken = defaultContainer ? sessionTokens.GetSessionToken() : sessionTokens.GetSessionToken(OtherContainerName); Assert.Equal(initialToken, updatedToken); } - [ConditionalFact] - public virtual async Task AppendSessionToken_different_pkrange_composites_tokens() + [ConditionalTheory, InlineData(true), InlineData(false)] + public virtual async Task AppendSessionToken_different_pkrange_composites_tokens(bool defaultContainer) { var contextFactory = await InitializeAsync(); using var context = contextFactory.CreateContext(); - context.Customers.Add(new Customer { Id = "1", PartitionKey = "1" }); + if (defaultContainer) + { + context.Add(new Customer { Id = "1", PartitionKey = "1" }); + } + else + { + context.Add(new OtherContainerCustomer { Id = "1", PartitionKey = "1" }); + } await context.SaveChangesAsync(); var sessionTokens = context.Database.GetSessionTokens(); - var initialToken = sessionTokens.GetSessionToken(); + var initialToken = defaultContainer ? sessionTokens.GetSessionToken() : sessionTokens.GetSessionToken(OtherContainerName); Assert.False(string.IsNullOrWhiteSpace(initialToken)); - sessionTokens.AppendSessionToken("99:-1#999999"); + var newToken = "99:-1#999999"; + if (defaultContainer) + { + sessionTokens.AppendSessionToken(newToken); +} + else + { + sessionTokens.AppendSessionToken(OtherContainerName, newToken); + } - var updatedToken = sessionTokens.GetSessionToken(); + var updatedToken = defaultContainer ? sessionTokens.GetSessionToken() : sessionTokens.GetSessionToken(OtherContainerName); Assert.Equal(initialToken + ",99:-1#999999", updatedToken); } - [ConditionalFact] - public virtual async Task SetSessionToken_does_not_merge_session_token() + [ConditionalTheory, InlineData(true), InlineData(false)] + public virtual async Task SetSessionToken_does_not_merge_session_token(bool defaultContainer) { var contextFactory = await InitializeAsync(); using var context = contextFactory.CreateContext(); - context.Customers.Add(new Customer { Id = "1", PartitionKey = "1" }); + if (defaultContainer) + { + context.Add(new Customer { Id = "1", PartitionKey = "1" }); + } + else + { + context.Add(new OtherContainerCustomer { Id = "1", PartitionKey = "1" }); + } await context.SaveChangesAsync(); var sessionTokens = context.Database.GetSessionTokens(); - var initialToken = sessionTokens.GetSessionToken(); + var initialToken = defaultContainer ? sessionTokens.GetSessionToken() : sessionTokens.GetSessionToken(OtherContainerName); Assert.False(string.IsNullOrWhiteSpace(initialToken)); - sessionTokens.SetSessionToken("0:-1#1"); + var newToken = "0:-1#1"; + if (defaultContainer) + { + sessionTokens.SetSessionToken(newToken); + } + else + { + sessionTokens.SetSessionToken(OtherContainerName, newToken); + } - var updatedToken = sessionTokens.GetSessionToken(); + var updatedToken = defaultContainer ? sessionTokens.GetSessionToken() : sessionTokens.GetSessionToken(OtherContainerName); Assert.Equal("0:-1#1", updatedToken); } - [ConditionalFact] - public virtual async Task SetSessionToken_null_sets_session_token_null() + [ConditionalTheory, InlineData(true), InlineData(false)] + public virtual async Task SetSessionToken_null_sets_session_token_null(bool defaultContainer) { var contextFactory = await InitializeAsync(); using var context = contextFactory.CreateContext(); - context.Customers.Add(new Customer { Id = "1", PartitionKey = "1" }); + if (defaultContainer) + { + context.Add(new Customer { Id = "1", PartitionKey = "1" }); + } + else + { + context.Add(new OtherContainerCustomer { Id = "1", PartitionKey = "1" }); + } await context.SaveChangesAsync(); var sessionTokens = context.Database.GetSessionTokens(); - var initialToken = sessionTokens.GetSessionToken(); + var initialToken = defaultContainer ? sessionTokens.GetSessionToken() : sessionTokens.GetSessionToken(OtherContainerName); Assert.False(string.IsNullOrWhiteSpace(initialToken)); - sessionTokens.SetSessionToken(null); + if (defaultContainer) + { + sessionTokens.SetSessionToken(null); + } + else + { + sessionTokens.SetSessionToken(OtherContainerName, null); + } - var updatedToken = sessionTokens.GetSessionToken(); + var updatedToken = defaultContainer ? sessionTokens.GetSessionToken() : sessionTokens.GetSessionToken(OtherContainerName); Assert.Null(updatedToken); } - [ConditionalFact] - public virtual async Task GetSessionToken_no_token_returns_null() + [ConditionalTheory, InlineData(true), InlineData(false)] + public virtual async Task GetSessionToken_no_token_returns_null(bool defaultContainer) { var contextFactory = await InitializeAsync(); using var context = contextFactory.CreateContext(); var sessionTokens = context.Database.GetSessionTokens(); - var sessionToken = sessionTokens.GetSessionToken(); - Assert.Null(sessionToken); + var initialToken = defaultContainer ? sessionTokens.GetSessionToken() : sessionTokens.GetSessionToken(OtherContainerName); + Assert.Null(initialToken); } - [ConditionalFact] - // @TODO: and sync.. - public virtual async Task Query_uses_session_token() + [ConditionalTheory] + [InlineData(true, true)] + [InlineData(true, false)] + [InlineData(false, true)] + [InlineData(false, false)] + public virtual async Task Query_uses_session_token(bool async, bool defaultContainer) { var contextFactory = await InitializeAsync(); - using var context = contextFactory.CreateContext(); - context.Customers.Add(new Customer { Id = "1", PartitionKey = "1" }); - await context.SaveChangesAsync(); + if (defaultContainer) + { + context.Customers.Add(new Customer { Id = "1", PartitionKey = "1" }); + } + else + { + context.OtherContainerCustomers.Add(new OtherContainerCustomer { Id = "1", PartitionKey = "1" }); + } + + if (async) + { + await context.SaveChangesAsync(); + } + else + { + context.SaveChanges(); + } var sessionTokens = context.Database.GetSessionTokens(); - var sessionToken = sessionTokens.GetSessionToken()!; + + string sessionToken; + if (defaultContainer) + { + sessionToken = sessionTokens.GetSessionToken()!; + } + else + { + sessionToken = sessionTokens.GetSessionToken(OtherContainerName)!; + } // Only way we can test this is by setting a session token that will fail the request if used.. // This will take a couple of seconds to fail - sessionTokens.SetSessionToken(sessionToken.Substring(0, sessionToken.IndexOf('#') + 1) + int.MaxValue); + var newToken = sessionToken.Substring(0, sessionToken.IndexOf('#') + 1) + int.MaxValue; + + if (defaultContainer) + { + sessionTokens.SetSessionToken(newToken); + } + else + { + sessionTokens.SetSessionToken(OtherContainerName, newToken); + } + + CosmosException ex; + if (async) + { + if (defaultContainer) + { + ex = await Assert.ThrowsAsync(() => context.Customers.ToListAsync()); + } + else + { + ex = await Assert.ThrowsAsync(() => context.OtherContainerCustomers.ToListAsync()); + } + } + else + { + if (defaultContainer) + { + ex = Assert.Throws(() => context.Customers.ToList()); + } + else + { + ex = Assert.Throws(() => context.OtherContainerCustomers.ToList()); + } + } - var ex = await Assert.ThrowsAsync(() => context.Customers.ToListAsync()); Assert.Contains("The read session is not available for the input session token.", ex.ResponseBody); } - [ConditionalFact] - // @TODO: and sync.. - public virtual async Task PagingQuery_uses_session_token() + [ConditionalTheory] + [InlineData(true), InlineData(false)] + public virtual async Task PagingQuery_uses_session_token(bool defaultContainer) { var contextFactory = await InitializeAsync(); - using var context = contextFactory.CreateContext(); - context.Customers.Add(new Customer { Id = "1", PartitionKey = "1" }); + + if (defaultContainer) + { + context.Customers.Add(new Customer { Id = "1", PartitionKey = "1" }); + } + else + { + context.OtherContainerCustomers.Add(new OtherContainerCustomer { Id = "1", PartitionKey = "1" }); + } await context.SaveChangesAsync(); var sessionTokens = context.Database.GetSessionTokens(); - var sessionToken = sessionTokens.GetSessionToken()!; + + string sessionToken; + if (defaultContainer) + { + sessionToken = sessionTokens.GetSessionToken()!; + } + else + { + sessionToken = sessionTokens.GetSessionToken(OtherContainerName)!; + } // Only way we can test this is by setting a session token that will fail the request if used.. // This will take a couple of seconds to fail - sessionTokens.SetSessionToken(sessionToken.Substring(0, sessionToken.IndexOf('#') + 1) + int.MaxValue); + var newToken = sessionToken.Substring(0, sessionToken.IndexOf('#') + 1) + int.MaxValue; + + if (defaultContainer) + { + sessionTokens.SetSessionToken(newToken); + } + else + { + sessionTokens.SetSessionToken(OtherContainerName, newToken); + } + + CosmosException ex; + if (defaultContainer) + { + ex = await Assert.ThrowsAsync(() => context.Customers.ToPageAsync(1, null)); + } + else + { + ex = await Assert.ThrowsAsync(() => context.OtherContainerCustomers.ToPageAsync(1, null)); + } - var ex = await Assert.ThrowsAsync(() => context.Customers.ToPageAsync(1, null)); Assert.Contains("The read session is not available for the input session token.", ex.ResponseBody); } - [ConditionalFact] - // @TODO: and sync.. - public virtual async Task Shaped_query_uses_session_token() + [ConditionalTheory] + [InlineData(true, true)] + [InlineData(true, false)] + [InlineData(false, true)] + [InlineData(false, false)] + public virtual async Task Shaped_query_uses_session_token(bool async, bool defaultContainer) { var contextFactory = await InitializeAsync(); - using var context = contextFactory.CreateContext(); - context.Customers.Add(new Customer { Id = "1", PartitionKey = "1" }); - await context.SaveChangesAsync(); + if (defaultContainer) + { + context.Customers.Add(new Customer { Id = "1", PartitionKey = "1" }); + } + else + { + context.OtherContainerCustomers.Add(new OtherContainerCustomer { Id = "1", PartitionKey = "1" }); + } + + if (async) + { + await context.SaveChangesAsync(); + } + else + { + context.SaveChanges(); + } var sessionTokens = context.Database.GetSessionTokens(); - var sessionToken = sessionTokens.GetSessionToken()!; - + + string sessionToken; + if (defaultContainer) + { + sessionToken = sessionTokens.GetSessionToken()!; + } + else + { + sessionToken = sessionTokens.GetSessionToken(OtherContainerName)!; + } + // Only way we can test this is by setting a session token that will fail the request if used.. // This will take a couple of seconds to fail - sessionTokens.SetSessionToken(sessionToken.Substring(0, sessionToken.IndexOf('#') + 1) + int.MaxValue); + var newToken = sessionToken.Substring(0, sessionToken.IndexOf('#') + 1) + int.MaxValue; + + if (defaultContainer) + { + sessionTokens.SetSessionToken(newToken); + } + else + { + sessionTokens.SetSessionToken(OtherContainerName, newToken); + } + + CosmosException ex; + if (async) + { + if (defaultContainer) + { + ex = await Assert.ThrowsAsync(() => context.Customers.Select(x => new { x.Id, x.PartitionKey }).ToListAsync()); + } + else + { + ex = await Assert.ThrowsAsync(() => context.OtherContainerCustomers.Select(x => new { x.Id, x.PartitionKey }).ToListAsync()); + } + } + else + { + if (defaultContainer) + { + ex = Assert.Throws(() => context.Customers.Select(x => new { x.Id, x.PartitionKey }).ToList()); + } + else + { + ex = Assert.Throws(() => context.OtherContainerCustomers.Select(x => new { x.Id, x.PartitionKey }).ToList()); + } + } - var ex = await Assert.ThrowsAsync(() => context.Customers.Select(x => new { x.Id, x.PartitionKey }).ToListAsync()); Assert.Contains("The read session is not available for the input session token.", ex.ResponseBody); } - [ConditionalFact] - // @TODO: and sync.. - public virtual async Task Read_item_uses_session_token() + [ConditionalTheory] + [InlineData(true, true)] + [InlineData(true, false)] + [InlineData(false, true)] + [InlineData(false, false)] + public virtual async Task Read_item_uses_session_token(bool async, bool defaultContainer) { var contextFactory = await InitializeAsync(); - using var context = contextFactory.CreateContext(); - context.Customers.Add(new Customer { Id = "1", PartitionKey = "1" }); - await context.SaveChangesAsync(); + if (defaultContainer) + { + context.Customers.Add(new Customer { Id = "1", PartitionKey = "1" }); + } + else + { + context.OtherContainerCustomers.Add(new OtherContainerCustomer { Id = "1", PartitionKey = "1" }); + } + + if (async) + { + await context.SaveChangesAsync(); + } + else + { + context.SaveChanges(); + } var sessionTokens = context.Database.GetSessionTokens(); - var sessionToken = sessionTokens.GetSessionToken()!; + + string sessionToken; + if (defaultContainer) + { + sessionToken = sessionTokens.GetSessionToken()!; + } + else + { + sessionToken = sessionTokens.GetSessionToken(OtherContainerName)!; + } // Only way we can test this is by setting a session token that will fail the request if used.. // This will take a couple of seconds to fail - sessionTokens.SetSessionToken(sessionToken.Substring(0, sessionToken.IndexOf('#') + 1) + int.MaxValue); + var newToken = sessionToken.Substring(0, sessionToken.IndexOf('#') + 1) + int.MaxValue; + + if (defaultContainer) + { + sessionTokens.SetSessionToken(newToken); + } + else + { + sessionTokens.SetSessionToken(OtherContainerName, newToken); + } + + CosmosException ex; + if (async) + { + if (defaultContainer) + { + ex = await Assert.ThrowsAsync(() => context.Customers.FirstOrDefaultAsync(x => x.Id == "1" && x.PartitionKey == "1")); + } + else + { + ex = await Assert.ThrowsAsync(() => context.OtherContainerCustomers.FirstOrDefaultAsync(x => x.Id == "1" && x.PartitionKey == "1")); + } + } + else + { + if (defaultContainer) + { + ex = Assert.Throws(() => context.Customers.FirstOrDefault(x => x.Id == "1" && x.PartitionKey == "1")); + } + else + { + ex = Assert.Throws(() => context.OtherContainerCustomers.FirstOrDefault(x => x.Id == "1" && x.PartitionKey == "1")); + } + } - var ex = await Assert.ThrowsAsync(() => context.Customers.FirstOrDefaultAsync(x => x.Id == "1" && x.PartitionKey == "1")); Assert.Contains("The read session is not available for the input session token.", ex.ResponseBody); } - [ConditionalFact] - // @TODO: and sync.. - public virtual async Task Query_sets_session_token() + [ConditionalTheory] + [InlineData(true, true)] + [InlineData(true, false)] + [InlineData(false, true)] + [InlineData(false, false)] + public virtual async Task Query_sets_session_token(bool async, bool defaultContainer) { var contextFactory = await InitializeAsync(); - using var context = contextFactory.CreateContext(); - await context.Customers.ToListAsync(); - var sessionToken = context.Database.GetSessionTokens().GetSessionToken(); - Assert.True(!string.IsNullOrWhiteSpace(sessionToken)); + if (async) + { + if (defaultContainer) + { + await context.Customers.ToListAsync(); + } + else + { + await context.OtherContainerCustomers.ToListAsync(); + } + } + else + { + if (defaultContainer) + { + _ = context.Customers.ToList(); + } + else + { + _ = context.OtherContainerCustomers.ToList(); + } + } + + var sessionTokens = context.Database.GetSessionTokens(); + string? sessionToken; + if (defaultContainer) + { + sessionToken = sessionTokens.GetSessionToken(); + } + else + { + sessionToken = sessionTokens.GetSessionToken(OtherContainerName); + } + + Assert.False(string.IsNullOrWhiteSpace(sessionToken)); } - [ConditionalFact] - // @TODO: and sync.. - public virtual async Task PagingQuery_sets_session_token() + [ConditionalTheory] + [InlineData(true, true)] + [InlineData(true, false)] + [InlineData(false, true)] + [InlineData(false, false)] + public virtual async Task PagingQuery_sets_session_token(bool async, bool defaultContainer) { var contextFactory = await InitializeAsync(); - using var context = contextFactory.CreateContext(); - await context.Customers.ToPageAsync(1, null); - var sessionToken = context.Database.GetSessionTokens().GetSessionToken(); - Assert.True(!string.IsNullOrWhiteSpace(sessionToken)); + if (defaultContainer) + { + await context.Customers.ToPageAsync(1, null); + } + else + { + await context.OtherContainerCustomers.ToPageAsync(1, null); + } + + var sessionTokens = context.Database.GetSessionTokens(); + string? sessionToken; + if (defaultContainer) + { + sessionToken = sessionTokens.GetSessionToken(); + } + else + { + sessionToken = sessionTokens.GetSessionToken(OtherContainerName); + } + + Assert.False(string.IsNullOrWhiteSpace(sessionToken)); } - [ConditionalFact] - // @TODO: and sync.. - public virtual async Task Shaped_query_sets_session_token() + [ConditionalTheory] + [InlineData(true, true)] + [InlineData(true, false)] + [InlineData(false, true)] + [InlineData(false, false)] + public virtual async Task Shaped_query_sets_session_token(bool async, bool defaultContainer) { var contextFactory = await InitializeAsync(); - using var context = contextFactory.CreateContext(); - await context.Customers.Select(x => new { x.Id, x.PartitionKey }).ToListAsync(); - var sessionToken = context.Database.GetSessionTokens().GetSessionToken(); - Assert.True(!string.IsNullOrWhiteSpace(sessionToken)); + if (async) + { + if (defaultContainer) + { + await context.Customers.Select(x => new { x.Id, x.PartitionKey }).ToListAsync(); + } + else + { + await context.OtherContainerCustomers.Select(x => new { x.Id, x.PartitionKey }).ToListAsync(); + } + } + else + { + if (defaultContainer) + { + _ = context.Customers.Select(x => new { x.Id, x.PartitionKey }).ToList(); + } + else + { + _ = context.OtherContainerCustomers.Select(x => new { x.Id, x.PartitionKey }).ToList(); + } + } + + var sessionTokens = context.Database.GetSessionTokens(); + string? sessionToken; + if (defaultContainer) + { + sessionToken = sessionTokens.GetSessionToken(); + } + else + { + sessionToken = sessionTokens.GetSessionToken(OtherContainerName); + } + + Assert.False(string.IsNullOrWhiteSpace(sessionToken)); } - [ConditionalFact] - // @TODO: and sync.. - public virtual async Task Read_item_sets_session_token() + [ConditionalTheory] + [InlineData(true, true)] + [InlineData(true, false)] + [InlineData(false, true)] + [InlineData(false, false)] + public virtual async Task Read_item_sets_session_token(bool async, bool defaultContainer) { var contextFactory = await InitializeAsync(); - using var context = contextFactory.CreateContext(); - await context.Customers.FirstOrDefaultAsync(x => x.Id == "1" && x.PartitionKey == "1"); - var sessionToken = context.Database.GetSessionTokens().GetSessionToken(); - Assert.True(!string.IsNullOrWhiteSpace(sessionToken)); + if (async) + { + if (defaultContainer) + { + await context.Customers.FirstOrDefaultAsync(x => x.Id == "1" && x.PartitionKey == "1"); + } + else + { + await context.OtherContainerCustomers.FirstOrDefaultAsync(x => x.Id == "1" && x.PartitionKey == "1"); + } + } + else + { + if (defaultContainer) + { + _ = context.Customers.FirstOrDefault(x => x.Id == "1" && x.PartitionKey == "1"); + } + else + { + _ = context.OtherContainerCustomers.FirstOrDefault(x => x.Id == "1" && x.PartitionKey == "1"); + } + } + + var sessionTokens = context.Database.GetSessionTokens(); + string? sessionToken; + if (defaultContainer) + { + sessionToken = sessionTokens.GetSessionToken(); + } + else + { + sessionToken = sessionTokens.GetSessionToken(OtherContainerName); + } + + Assert.False(string.IsNullOrWhiteSpace(sessionToken)); } - [ConditionalFact] - // @TODO: and sync.. - public virtual async Task Read_item_enumerable_sets_session_token() + [ConditionalTheory] + [InlineData(true, true)] + [InlineData(true, false)] + [InlineData(false, true)] + [InlineData(false, false)] + public virtual async Task Read_item_enumerable_sets_session_token(bool async, bool defaultContainer) { var contextFactory = await InitializeAsync(); - using var context = contextFactory.CreateContext(); - await context.Customers.Where(x => x.Id == "1" && x.PartitionKey == "1").ToListAsync(); - var sessionToken = context.Database.GetSessionTokens().GetSessionToken(); - Assert.True(!string.IsNullOrWhiteSpace(sessionToken)); + if (async) + { + if (defaultContainer) + { + await context.Customers.Where(x => x.Id == "1" && x.PartitionKey == "1").ToListAsync(); + } + else + { + await context.OtherContainerCustomers.Where(x => x.Id == "1" && x.PartitionKey == "1").ToListAsync(); + } + } + else + { + if (defaultContainer) + { + _ = context.Customers.Where(x => x.Id == "1" && x.PartitionKey == "1").ToList(); + } + else + { + _ = context.OtherContainerCustomers.Where(x => x.Id == "1" && x.PartitionKey == "1").ToList(); + } + } + + var sessionTokens = context.Database.GetSessionTokens(); + string? sessionToken; + if (defaultContainer) + { + sessionToken = sessionTokens.GetSessionToken(); + } + else + { + sessionToken = sessionTokens.GetSessionToken(OtherContainerName); + } + + Assert.False(string.IsNullOrWhiteSpace(sessionToken)); } - [ConditionalTheory, InlineData(AutoTransactionBehavior.WhenNeeded), InlineData(AutoTransactionBehavior.Never), InlineData(AutoTransactionBehavior.Always)] - // @TODO: and sync.. - public virtual async Task Add_sets_session_token(AutoTransactionBehavior autoTransactionBehavior) + [ConditionalTheory] + [InlineData(AutoTransactionBehavior.WhenNeeded, true, true)] + [InlineData(AutoTransactionBehavior.WhenNeeded, true, false)] + [InlineData(AutoTransactionBehavior.WhenNeeded, false, true)] + [InlineData(AutoTransactionBehavior.WhenNeeded, false, false)] + [InlineData(AutoTransactionBehavior.Never, true, true)] + [InlineData(AutoTransactionBehavior.Never, true, false)] + [InlineData(AutoTransactionBehavior.Never, false, true)] + [InlineData(AutoTransactionBehavior.Never, false, false)] + [InlineData(AutoTransactionBehavior.Always, true, true)] + [InlineData(AutoTransactionBehavior.Always, true, false)] + [InlineData(AutoTransactionBehavior.Always, false, true)] + [InlineData(AutoTransactionBehavior.Always, false, false)] + public virtual async Task Add_sets_session_token(AutoTransactionBehavior autoTransactionBehavior, bool async, bool defaultContainer) { var contextFactory = await InitializeAsync(); using var context = contextFactory.CreateContext(); context.Database.AutoTransactionBehavior = autoTransactionBehavior; - context.Customers.Add(new Customer { Id = "1", PartitionKey = "1" }); - await context.SaveChangesAsync(); + if (defaultContainer) + { + context.Customers.Add(new Customer { Id = "1", PartitionKey = "1" }); + } + else + { + context.OtherContainerCustomers.Add(new OtherContainerCustomer { Id = "1", PartitionKey = "1" }); + } + + if (async) + { + await context.SaveChangesAsync(); + } + else + { + context.SaveChanges(); + } var sessionTokens = context.Database.GetSessionTokens(); - var sessionToken = sessionTokens.GetSessionToken(); + string? sessionToken; + if (defaultContainer) + { + sessionToken = sessionTokens.GetSessionToken(); + } + else + { + sessionToken = sessionTokens.GetSessionToken(OtherContainerName); + } + Assert.False(string.IsNullOrWhiteSpace(sessionToken)); } - [ConditionalTheory, InlineData(AutoTransactionBehavior.WhenNeeded), InlineData(AutoTransactionBehavior.Never), InlineData(AutoTransactionBehavior.Always)] - // @TODO: and sync.. - public virtual async Task Delete_merges_session_token(AutoTransactionBehavior autoTransactionBehavior) + [ConditionalTheory] + [InlineData(AutoTransactionBehavior.WhenNeeded, true, true)] + [InlineData(AutoTransactionBehavior.WhenNeeded, true, false)] + [InlineData(AutoTransactionBehavior.WhenNeeded, false, true)] + [InlineData(AutoTransactionBehavior.WhenNeeded, false, false)] + [InlineData(AutoTransactionBehavior.Never, true, true)] + [InlineData(AutoTransactionBehavior.Never, true, false)] + [InlineData(AutoTransactionBehavior.Never, false, true)] + [InlineData(AutoTransactionBehavior.Never, false, false)] + [InlineData(AutoTransactionBehavior.Always, true, true)] + [InlineData(AutoTransactionBehavior.Always, true, false)] + [InlineData(AutoTransactionBehavior.Always, false, true)] + [InlineData(AutoTransactionBehavior.Always, false, false)] + public virtual async Task Add_merges_session_token(AutoTransactionBehavior autoTransactionBehavior, bool async, bool defaultContainer) { var contextFactory = await InitializeAsync(); using var context = contextFactory.CreateContext(); + var sessionTokens = context.Database.GetSessionTokens(); context.Database.AutoTransactionBehavior = autoTransactionBehavior; - var customer = new Customer { Id = "1", PartitionKey = "1" }; - context.Customers.Add(customer); + if (defaultContainer) + { + context.Customers.Add(new Customer { Id = "1", PartitionKey = "1" }); + } + else + { + context.OtherContainerCustomers.Add(new OtherContainerCustomer { Id = "1", PartitionKey = "1" }); + } - await context.SaveChangesAsync(); + if (async) + { + await context.SaveChangesAsync(); + } + else + { + context.SaveChanges(); + } - var initialToken = context.Database.GetSessionTokens().GetSessionToken()!; + var initialToken = defaultContainer ? sessionTokens.GetSessionToken()! : sessionTokens.GetSessionToken(OtherContainerName)!; - context.Remove(customer); - await context.SaveChangesAsync(); + if (defaultContainer) + { + context.Customers.Add(new Customer { Id = "2", PartitionKey = "1" }); + } + else + { + context.OtherContainerCustomers.Add(new OtherContainerCustomer { Id = "2", PartitionKey = "1" }); + } + + if (async) + { + await context.SaveChangesAsync(); + } + else + { + context.SaveChanges(); + } + + string? sessionToken; + if (defaultContainer) + { + sessionToken = sessionTokens.GetSessionToken(); + } + else + { + sessionToken = sessionTokens.GetSessionToken(OtherContainerName); + } - var sessionToken = context.Database.GetSessionTokens().GetSessionToken(); Assert.False(string.IsNullOrWhiteSpace(sessionToken)); Assert.NotEqual(sessionToken, initialToken); Assert.StartsWith(initialToken.Substring(0, initialToken.IndexOf('#') + 1), sessionToken); } - [ConditionalTheory, InlineData(AutoTransactionBehavior.WhenNeeded), InlineData(AutoTransactionBehavior.Never), InlineData(AutoTransactionBehavior.Always)] - // @TODO: and sync.. - public virtual async Task Update_merges_session_token(AutoTransactionBehavior autoTransactionBehavior) + [ConditionalTheory] + [InlineData(AutoTransactionBehavior.WhenNeeded, true, true)] + [InlineData(AutoTransactionBehavior.WhenNeeded, true, false)] + [InlineData(AutoTransactionBehavior.WhenNeeded, false, true)] + [InlineData(AutoTransactionBehavior.WhenNeeded, false, false)] + [InlineData(AutoTransactionBehavior.Never, true, true)] + [InlineData(AutoTransactionBehavior.Never, true, false)] + [InlineData(AutoTransactionBehavior.Never, false, true)] + [InlineData(AutoTransactionBehavior.Never, false, false)] + [InlineData(AutoTransactionBehavior.Always, true, true)] + [InlineData(AutoTransactionBehavior.Always, true, false)] + [InlineData(AutoTransactionBehavior.Always, false, true)] + [InlineData(AutoTransactionBehavior.Always, false, false)] + public virtual async Task Delete_merges_session_token(AutoTransactionBehavior autoTransactionBehavior, bool async, bool defaultContainer) { var contextFactory = await InitializeAsync(); using var context = contextFactory.CreateContext(); context.Database.AutoTransactionBehavior = autoTransactionBehavior; - var customer = new Customer { Id = "1", PartitionKey = "1" }; - context.Customers.Add(customer); + string initialToken; + if (defaultContainer) + { + var customer = new Customer { Id = "1", PartitionKey = "1" }; + context.Customers.Add(customer); + + if (async) + { + await context.SaveChangesAsync(); + } + else + { + context.SaveChanges(); + } + + initialToken = context.Database.GetSessionTokens().GetSessionToken()!; + + context.Remove(customer); + + if (async) + { + await context.SaveChangesAsync(); + } + else + { + context.SaveChanges(); + } + } + else + { + var customer = new OtherContainerCustomer { Id = "1", PartitionKey = "1" }; + context.Add(customer); + + if (async) + { + await context.SaveChangesAsync(); + } + else + { + context.SaveChanges(); + } + + initialToken = context.Database.GetSessionTokens().GetSessionToken(OtherContainerName)!; + + context.Remove(customer); + + if (async) + { + await context.SaveChangesAsync(); + } + else + { + context.SaveChanges(); + } + } - await context.SaveChangesAsync(); + var sessionToken = defaultContainer ? context.Database.GetSessionTokens().GetSessionToken() : context.Database.GetSessionTokens().GetSessionToken(OtherContainerName); + Assert.False(string.IsNullOrWhiteSpace(sessionToken)); + Assert.NotEqual(sessionToken, initialToken); + Assert.StartsWith(initialToken.Substring(0, initialToken.IndexOf('#') + 1), sessionToken); + } - var initialToken = context.Database.GetSessionTokens().GetSessionToken()!; + [ConditionalTheory] + [InlineData(AutoTransactionBehavior.WhenNeeded, true, true)] + [InlineData(AutoTransactionBehavior.WhenNeeded, true, false)] + [InlineData(AutoTransactionBehavior.WhenNeeded, false, true)] + [InlineData(AutoTransactionBehavior.WhenNeeded, false, false)] + [InlineData(AutoTransactionBehavior.Never, true, true)] + [InlineData(AutoTransactionBehavior.Never, true, false)] + [InlineData(AutoTransactionBehavior.Never, false, true)] + [InlineData(AutoTransactionBehavior.Never, false, false)] + [InlineData(AutoTransactionBehavior.Always, true, true)] + [InlineData(AutoTransactionBehavior.Always, true, false)] + [InlineData(AutoTransactionBehavior.Always, false, true)] + [InlineData(AutoTransactionBehavior.Always, false, false)] + public virtual async Task Update_merges_session_token(AutoTransactionBehavior autoTransactionBehavior, bool async, bool defaultContainer) + { + var contextFactory = await InitializeAsync(); - customer.Name = "updated"; - await context.SaveChangesAsync(); + using var context = contextFactory.CreateContext(); + context.Database.AutoTransactionBehavior = autoTransactionBehavior; + + string initialToken; + if (defaultContainer) + { + var customer = new Customer { Id = "1", PartitionKey = "1" }; + context.Customers.Add(customer); + + if (async) + { + await context.SaveChangesAsync(); + } + else + { + context.SaveChanges(); + } - var sessionToken = context.Database.GetSessionTokens().GetSessionToken(); + initialToken = context.Database.GetSessionTokens().GetSessionToken()!; + + customer.Name = "updated"; + } + else + { + var customer = new OtherContainerCustomer { Id = "1", PartitionKey = "1" }; + context.Add(customer); + + if (async) + { + await context.SaveChangesAsync(); + } + else + { + context.SaveChanges(); + } + + initialToken = context.Database.GetSessionTokens().GetSessionToken(OtherContainerName)!; + + customer.Name = "updated"; + } + + if (async) + { + await context.SaveChangesAsync(); + } + else + { + context.SaveChanges(); + } + + var sessionToken = defaultContainer ? context.Database.GetSessionTokens().GetSessionToken() : context.Database.GetSessionTokens().GetSessionToken(OtherContainerName); Assert.False(string.IsNullOrWhiteSpace(sessionToken)); Assert.NotEqual(initialToken, sessionToken); Assert.StartsWith(initialToken.Substring(0, initialToken.IndexOf('#') + 1), sessionToken); } - [ConditionalTheory, InlineData(AutoTransactionBehavior.WhenNeeded), InlineData(AutoTransactionBehavior.Never), InlineData(AutoTransactionBehavior.Always)] - // @TODO: And sync.. - public virtual async Task Add_uses_session_token(AutoTransactionBehavior autoTransactionBehavior) + [ConditionalTheory] + [InlineData(AutoTransactionBehavior.WhenNeeded, true, true)] + [InlineData(AutoTransactionBehavior.WhenNeeded, true, false)] + [InlineData(AutoTransactionBehavior.WhenNeeded, false, true)] + [InlineData(AutoTransactionBehavior.WhenNeeded, false, false)] + [InlineData(AutoTransactionBehavior.Never, true, true)] + [InlineData(AutoTransactionBehavior.Never, true, false)] + [InlineData(AutoTransactionBehavior.Never, false, true)] + [InlineData(AutoTransactionBehavior.Never, false, false)] + [InlineData(AutoTransactionBehavior.Always, true, true)] + [InlineData(AutoTransactionBehavior.Always, true, false)] + [InlineData(AutoTransactionBehavior.Always, false, true)] + [InlineData(AutoTransactionBehavior.Always, false, false)] + public virtual async Task Add_uses_session_token(AutoTransactionBehavior autoTransactionBehavior, bool async, bool defaultContainer) { var contextFactory = await InitializeAsync(); @@ -427,19 +1096,46 @@ public virtual async Task Add_uses_session_token(AutoTransactionBehavior autoTra // Only way we can test this is by setting a session token that will fail the request if used.. // Only way to do this for a write is to set an invalid session token.. var internalDictionary = sessionTokens.GetType().GetField("_containerSessionTokens", BindingFlags.NonPublic | BindingFlags.Instance)!.GetValue(sessionTokens)!; - var internalComposite = internalDictionary.GetType().GetProperty("Item", BindingFlags.Public | BindingFlags.Instance)!.GetValue(internalDictionary, new object[] { nameof(CosmosSessionTokenContext) })!; + var internalComposite = internalDictionary.GetType().GetProperty("Item", BindingFlags.Public | BindingFlags.Instance)!.GetValue(internalDictionary, new object[] { defaultContainer ? nameof(CosmosSessionTokenContext) : OtherContainerName })!; internalComposite.GetType().GetField("_string", BindingFlags.NonPublic | BindingFlags.Instance)!.SetValue(internalComposite, "invalidtoken"); internalComposite.GetType().GetField("_isChanged", BindingFlags.NonPublic | BindingFlags.Instance)!.SetValue(internalComposite, false); - context.Customers.Add(new Customer { Id = "1", PartitionKey = "1" }); + if (defaultContainer) + { + context.Customers.Add(new Customer { Id = "1", PartitionKey = "1" }); + } + else + { + context.OtherContainerCustomers.Add(new OtherContainerCustomer { Id = "1", PartitionKey = "1" }); + } + + DbUpdateException ex; + if (async) + { + ex = await Assert.ThrowsAsync(() => context.SaveChangesAsync()); + } + else + { + ex = Assert.Throws(() => context.SaveChanges()); + } - var ex = await Assert.ThrowsAsync(() => context.SaveChangesAsync()); Assert.Contains("The session token provided 'invalidtoken' is invalid.", ((CosmosException)ex.InnerException!).ResponseBody); } - [ConditionalTheory, InlineData(AutoTransactionBehavior.WhenNeeded), InlineData(AutoTransactionBehavior.Never), InlineData(AutoTransactionBehavior.Always)] - // @TODO: And sync.. - public virtual async Task Update_uses_session_token(AutoTransactionBehavior autoTransactionBehavior) + [ConditionalTheory] + [InlineData(AutoTransactionBehavior.WhenNeeded, true, true)] + [InlineData(AutoTransactionBehavior.WhenNeeded, true, false)] + [InlineData(AutoTransactionBehavior.WhenNeeded, false, true)] + [InlineData(AutoTransactionBehavior.WhenNeeded, false, false)] + [InlineData(AutoTransactionBehavior.Never, true, true)] + [InlineData(AutoTransactionBehavior.Never, true, false)] + [InlineData(AutoTransactionBehavior.Never, false, true)] + [InlineData(AutoTransactionBehavior.Never, false, false)] + [InlineData(AutoTransactionBehavior.Always, true, true)] + [InlineData(AutoTransactionBehavior.Always, true, false)] + [InlineData(AutoTransactionBehavior.Always, false, true)] + [InlineData(AutoTransactionBehavior.Always, false, false)] + public virtual async Task Update_uses_session_token(AutoTransactionBehavior autoTransactionBehavior, bool async, bool defaultContainer) { var contextFactory = await InitializeAsync(); @@ -451,19 +1147,46 @@ public virtual async Task Update_uses_session_token(AutoTransactionBehavior auto // Only way we can test this is by setting a session token that will fail the request if used.. // Only way to do this for a write is to set an invalid session token.. var internalDictionary = sessionTokens.GetType().GetField("_containerSessionTokens", BindingFlags.NonPublic | BindingFlags.Instance)!.GetValue(sessionTokens)!; - var internalComposite = internalDictionary.GetType().GetProperty("Item", BindingFlags.Public | BindingFlags.Instance)!.GetValue(internalDictionary, new object[] { nameof(CosmosSessionTokenContext) })!; + var internalComposite = internalDictionary.GetType().GetProperty("Item", BindingFlags.Public | BindingFlags.Instance)!.GetValue(internalDictionary, new object[] { defaultContainer ? nameof(CosmosSessionTokenContext) : OtherContainerName })!; internalComposite.GetType().GetField("_string", BindingFlags.NonPublic | BindingFlags.Instance)!.SetValue(internalComposite, "invalidtoken"); internalComposite.GetType().GetField("_isChanged", BindingFlags.NonPublic | BindingFlags.Instance)!.SetValue(internalComposite, false); - context.Customers.Update(new Customer { Id = "1", PartitionKey = "1" }); + if (defaultContainer) + { + context.Customers.Update(new Customer { Id = "1", PartitionKey = "1" }); + } + else + { + context.OtherContainerCustomers.Update(new OtherContainerCustomer { Id = "1", PartitionKey = "1" }); + } + + DbUpdateException ex; + if (async) + { + ex = await Assert.ThrowsAsync(() => context.SaveChangesAsync()); + } + else + { + ex = Assert.Throws(() => context.SaveChanges()); + } - var ex = await Assert.ThrowsAsync(() => context.SaveChangesAsync()); Assert.Contains("The session token provided 'invalidtoken' is invalid.", ((CosmosException)ex.InnerException!).ResponseBody); } - [ConditionalTheory, InlineData(AutoTransactionBehavior.WhenNeeded), InlineData(AutoTransactionBehavior.Never), InlineData(AutoTransactionBehavior.Always)] - // @TODO: And sync.. - public virtual async Task Delete_uses_session_token(AutoTransactionBehavior autoTransactionBehavior) + [ConditionalTheory] + [InlineData(AutoTransactionBehavior.WhenNeeded, true, true)] + [InlineData(AutoTransactionBehavior.WhenNeeded, true, false)] + [InlineData(AutoTransactionBehavior.WhenNeeded, false, true)] + [InlineData(AutoTransactionBehavior.WhenNeeded, false, false)] + [InlineData(AutoTransactionBehavior.Never, true, true)] + [InlineData(AutoTransactionBehavior.Never, true, false)] + [InlineData(AutoTransactionBehavior.Never, false, true)] + [InlineData(AutoTransactionBehavior.Never, false, false)] + [InlineData(AutoTransactionBehavior.Always, true, true)] + [InlineData(AutoTransactionBehavior.Always, true, false)] + [InlineData(AutoTransactionBehavior.Always, false, true)] + [InlineData(AutoTransactionBehavior.Always, false, false)] + public virtual async Task Delete_uses_session_token(AutoTransactionBehavior autoTransactionBehavior, bool async, bool defaultContainer) { var contextFactory = await InitializeAsync(); @@ -475,13 +1198,29 @@ public virtual async Task Delete_uses_session_token(AutoTransactionBehavior auto // Only way we can test this is by setting a session token that will fail the request if used.. // Only way to do this for a write is to set an invalid session token.. var internalDictionary = sessionTokens.GetType().GetField("_containerSessionTokens", BindingFlags.NonPublic | BindingFlags.Instance)!.GetValue(sessionTokens)!; - var internalComposite = internalDictionary.GetType().GetProperty("Item", BindingFlags.Public | BindingFlags.Instance)!.GetValue(internalDictionary, new object[] { nameof(CosmosSessionTokenContext) })!; + var internalComposite = internalDictionary.GetType().GetProperty("Item", BindingFlags.Public | BindingFlags.Instance)!.GetValue(internalDictionary, new object[] { defaultContainer ? nameof(CosmosSessionTokenContext) : OtherContainerName })!; internalComposite.GetType().GetField("_string", BindingFlags.NonPublic | BindingFlags.Instance)!.SetValue(internalComposite, "invalidtoken"); internalComposite.GetType().GetField("_isChanged", BindingFlags.NonPublic | BindingFlags.Instance)!.SetValue(internalComposite, false); - context.Customers.Remove(new Customer { Id = "1", PartitionKey = "1" }); + if (defaultContainer) + { + context.Customers.Remove(new Customer { Id = "1", PartitionKey = "1" }); + } + else + { + context.OtherContainerCustomers.Remove(new OtherContainerCustomer { Id = "1", PartitionKey = "1" }); + } + + DbUpdateException ex; + if (async) + { + ex = await Assert.ThrowsAsync(() => context.SaveChangesAsync()); + } + else + { + ex = Assert.Throws(() => context.SaveChanges()); + } - var ex = await Assert.ThrowsAsync(() => context.SaveChangesAsync()); Assert.Contains("The session token provided 'invalidtoken' is invalid.", ((CosmosException)ex.InnerException!).ResponseBody); } @@ -533,6 +1272,8 @@ public class OtherContainerCustomer { public string? Id { get; set; } + public string? Name { get; set; } + public string? PartitionKey { get; set; } } } From dadc58bd442cdbb272e28e058ad71ac97cd72142 Mon Sep 17 00:00:00 2001 From: JoasE <32096708+JoasE@users.noreply.github.com> Date: Fri, 24 Oct 2025 16:55:12 +0200 Subject: [PATCH 11/37] Cleanup --- .../Storage/Internal/CosmosClientWrapper.cs | 2 +- .../Storage/Internal/CosmosDatabaseWrapper.cs | 2 +- .../Storage/Internal/CosmosWriteResult.cs | 56 ------------------- 3 files changed, 2 insertions(+), 58 deletions(-) delete mode 100644 src/EFCore.Cosmos/Storage/Internal/CosmosWriteResult.cs diff --git a/src/EFCore.Cosmos/Storage/Internal/CosmosClientWrapper.cs b/src/EFCore.Cosmos/Storage/Internal/CosmosClientWrapper.cs index 7c512b8b916..1fba988ab4c 100644 --- a/src/EFCore.Cosmos/Storage/Internal/CosmosClientWrapper.cs +++ b/src/EFCore.Cosmos/Storage/Internal/CosmosClientWrapper.cs @@ -1360,7 +1360,7 @@ public async ValueTask DisposeAsync() #endregion ResponseMessageEnumerable - private class CosmosFeedIteratorWrapper : FeedIterator + private sealed class CosmosFeedIteratorWrapper : FeedIterator { private readonly FeedIterator _inner; private readonly string _containerName; diff --git a/src/EFCore.Cosmos/Storage/Internal/CosmosDatabaseWrapper.cs b/src/EFCore.Cosmos/Storage/Internal/CosmosDatabaseWrapper.cs index 0955b427289..2ef74edb3eb 100644 --- a/src/EFCore.Cosmos/Storage/Internal/CosmosDatabaseWrapper.cs +++ b/src/EFCore.Cosmos/Storage/Internal/CosmosDatabaseWrapper.cs @@ -62,7 +62,7 @@ public CosmosDatabaseWrapper( /// 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. /// - public ISessionTokenStorage SessionTokenStorage { get; } + public virtual ISessionTokenStorage SessionTokenStorage { get; } /// /// This is an internal API that supports the Entity Framework Core infrastructure and not subject to diff --git a/src/EFCore.Cosmos/Storage/Internal/CosmosWriteResult.cs b/src/EFCore.Cosmos/Storage/Internal/CosmosWriteResult.cs deleted file mode 100644 index 5f8e122ae21..00000000000 --- a/src/EFCore.Cosmos/Storage/Internal/CosmosWriteResult.cs +++ /dev/null @@ -1,56 +0,0 @@ -// 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.Storage.Internal; - -/// -/// 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. -/// -public class CosmosWriteResult -{ - /// - /// 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. - /// - public static CosmosWriteResult Failure { get; } = new CosmosWriteResult(); - - /// - /// 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. - /// - public static CosmosWriteResult Success(string? sessionToken) - => new CosmosWriteResult(sessionToken); - - private CosmosWriteResult() - { - } - - private CosmosWriteResult(string? sessionToken) - { - IsSuccess = true; - SessionToken = sessionToken; - } - - /// - /// 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. - /// - public bool IsSuccess { get; } = true; - - /// - /// 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. - /// - public string? SessionToken { get; } -} From 57f1f372751563e84054d93f47616dbbc63c835a Mon Sep 17 00:00:00 2001 From: JoasE <32096708+JoasE@users.noreply.github.com> Date: Fri, 24 Oct 2025 17:36:48 +0200 Subject: [PATCH 12/37] Add todo --- .../Metadata/Conventions/CosmosRuntimeModelConvention.cs | 1 + 1 file changed, 1 insertion(+) diff --git a/src/EFCore.Cosmos/Metadata/Conventions/CosmosRuntimeModelConvention.cs b/src/EFCore.Cosmos/Metadata/Conventions/CosmosRuntimeModelConvention.cs index f86192daaae..b5d53c3ce2e 100644 --- a/src/EFCore.Cosmos/Metadata/Conventions/CosmosRuntimeModelConvention.cs +++ b/src/EFCore.Cosmos/Metadata/Conventions/CosmosRuntimeModelConvention.cs @@ -45,6 +45,7 @@ protected override void ProcessModelAnnotations( annotations.Remove(CosmosAnnotationNames.Throughput); } + // @TODO: Is this the right place for this? annotations.Add(CosmosAnnotationNames.ContainerNames, GetContainerNames(model)); } From fc8a5801953875995abc4a7e69c27ec85b4f21a5 Mon Sep 17 00:00:00 2001 From: JoasE <32096708+JoasE@users.noreply.github.com> Date: Fri, 24 Oct 2025 19:28:42 +0200 Subject: [PATCH 13/37] Fix empty tokens --- src/EFCore.Cosmos/Storage/Internal/SessionTokenStorage.cs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/EFCore.Cosmos/Storage/Internal/SessionTokenStorage.cs b/src/EFCore.Cosmos/Storage/Internal/SessionTokenStorage.cs index fe42f30345b..e35805083db 100644 --- a/src/EFCore.Cosmos/Storage/Internal/SessionTokenStorage.cs +++ b/src/EFCore.Cosmos/Storage/Internal/SessionTokenStorage.cs @@ -180,7 +180,7 @@ public void Merge(string pkRangeId, VectorSessionToken token) if (_isChanged) { _isChanged = false; - _string = string.Join(",", Tokens.Select(kvp => $"{kvp.Key}:{kvp.Value.ConvertToString()}")); + _string = Tokens.Count == 0 ? null : string.Join(",", Tokens.Select(kvp => $"{kvp.Key}:{kvp.Value.ConvertToString()}")); } return _string; From 364d0f49e5b650e95a529fb69fc2f28d951917b0 Mon Sep 17 00:00:00 2001 From: JoasE <32096708+JoasE@users.noreply.github.com> Date: Fri, 24 Oct 2025 19:29:55 +0200 Subject: [PATCH 14/37] Rename var --- src/EFCore.Cosmos/Storage/Internal/SessionTokenStorage.cs | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/src/EFCore.Cosmos/Storage/Internal/SessionTokenStorage.cs b/src/EFCore.Cosmos/Storage/Internal/SessionTokenStorage.cs index e35805083db..cc553dbdd1b 100644 --- a/src/EFCore.Cosmos/Storage/Internal/SessionTokenStorage.cs +++ b/src/EFCore.Cosmos/Storage/Internal/SessionTokenStorage.cs @@ -115,21 +115,21 @@ public virtual void SetSessionToken(string containerName, string? sessionToken) throw new ArgumentException("sessionToken cannot be whitespace.", sessionToken); } - ref var value = ref CollectionsMarshal.GetValueRefOrNullRef(_containerSessionTokens, containerName); + ref var compositeSessionToken = ref CollectionsMarshal.GetValueRefOrNullRef(_containerSessionTokens, containerName); - if (Unsafe.IsNullRef(ref value)) + if (Unsafe.IsNullRef(ref compositeSessionToken)) { throw new ArgumentException(CosmosStrings.ContainerNameDoesNotExist(containerName), nameof(containerName)); } - value = new CompositeSessionToken(); + compositeSessionToken = new CompositeSessionToken(); if (sessionToken is null) { return; } - ParseAndMerge(value, sessionToken); + ParseAndMerge(compositeSessionToken, sessionToken); } private void ParseAndMerge(CompositeSessionToken compositeSessionToken, string sessionToken) From c9890e127a1fde9cfc32125a81d981682d41165e Mon Sep 17 00:00:00 2001 From: JoasE <32096708+JoasE@users.noreply.github.com> Date: Fri, 24 Oct 2025 19:57:48 +0200 Subject: [PATCH 15/37] Remove some todos --- src/EFCore.Cosmos/Storage/ISessionTokenStorage.cs | 3 +-- .../Storage/Internal/CosmosClientWrapper.cs | 11 +---------- .../Storage/Internal/ICosmosClientWrapper.cs | 1 - 3 files changed, 2 insertions(+), 13 deletions(-) diff --git a/src/EFCore.Cosmos/Storage/ISessionTokenStorage.cs b/src/EFCore.Cosmos/Storage/ISessionTokenStorage.cs index 6fe753b76a1..b18ed566753 100644 --- a/src/EFCore.Cosmos/Storage/ISessionTokenStorage.cs +++ b/src/EFCore.Cosmos/Storage/ISessionTokenStorage.cs @@ -3,9 +3,8 @@ namespace Microsoft.EntityFrameworkCore.Cosmos.Storage; -// @TODO: CosmosSession(Token)Context? /// -/// Defines methods for storing, retrieving, and managing session tokens associated with containers in a . +/// Defines methods for managing session tokens used in a . /// public interface ISessionTokenStorage { diff --git a/src/EFCore.Cosmos/Storage/Internal/CosmosClientWrapper.cs b/src/EFCore.Cosmos/Storage/Internal/CosmosClientWrapper.cs index 1fba988ab4c..e209c351e60 100644 --- a/src/EFCore.Cosmos/Storage/Internal/CosmosClientWrapper.cs +++ b/src/EFCore.Cosmos/Storage/Internal/CosmosClientWrapper.cs @@ -975,14 +975,7 @@ private static async Task CreateSingleItemQueryAsync( using var reader = new StreamReader(responseStream); using var jsonReader = new JsonTextReader(reader); - var jobject = Serializer.Deserialize(jsonReader); - - if (!string.IsNullOrWhiteSpace(responseMessage.Headers.Session)) - { - // @TODO: Set session token.. - } - - return jobject; + return Serializer.Deserialize(jsonReader); } /// @@ -1085,7 +1078,6 @@ public bool MoveNext() queryRequestOptions.PartitionKey = _partitionKeyValue; } - // @TODO: Or should this be inside CreateQuery... queryRequestOptions.SessionToken = _sessionTokenStorage.GetSessionToken(_containerId); _query = _cosmosClientWrapper.CreateQuery( @@ -1195,7 +1187,6 @@ public async ValueTask MoveNextAsync() queryRequestOptions.PartitionKey = _partitionKeyValue; } - // @TODO: Or should this be inside CreateQuery... queryRequestOptions.SessionToken = _sessionTokenStorage.GetSessionToken(_containerId); _query = _cosmosClientWrapper.CreateQuery( diff --git a/src/EFCore.Cosmos/Storage/Internal/ICosmosClientWrapper.cs b/src/EFCore.Cosmos/Storage/Internal/ICosmosClientWrapper.cs index 0337d55c197..620fd3b054b 100644 --- a/src/EFCore.Cosmos/Storage/Internal/ICosmosClientWrapper.cs +++ b/src/EFCore.Cosmos/Storage/Internal/ICosmosClientWrapper.cs @@ -127,7 +127,6 @@ Task ReplaceItemAsync( /// 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. /// - // @TODO: We also need to send session token on writes to keep the same chain or not? Task DeleteItemAsync( string containerId, string documentId, From 255bbe21beaf512dec7e01acfce9dafebdfaa2ad Mon Sep 17 00:00:00 2001 From: JoasE <32096708+JoasE@users.noreply.github.com> Date: Fri, 24 Oct 2025 20:05:14 +0200 Subject: [PATCH 16/37] Move to non-parsing session token management --- .../Properties/CosmosStrings.Designer.cs | 6 + .../Properties/CosmosStrings.resx | 3 + .../Storage/Internal/CosmosClientWrapper.cs | 1 - .../Storage/Internal/SessionTokenStorage.cs | 461 +----------------- .../CosmosSessionTokensTest.cs | 47 +- 5 files changed, 39 insertions(+), 479 deletions(-) diff --git a/src/EFCore.Cosmos/Properties/CosmosStrings.Designer.cs b/src/EFCore.Cosmos/Properties/CosmosStrings.Designer.cs index 6a8bf326061..53dfa712a5f 100644 --- a/src/EFCore.Cosmos/Properties/CosmosStrings.Designer.cs +++ b/src/EFCore.Cosmos/Properties/CosmosStrings.Designer.cs @@ -499,6 +499,12 @@ public static string SaveChangesAutoTransactionBehaviorAlwaysAtomicity public static string SaveChangesAutoTransactionBehaviorAlwaysTriggerAtomicity => GetString("SaveChangesAutoTransactionBehaviorAlwaysTriggerAtomicity"); + /// + /// Session token can not be white space. + /// + public static string SessionTokenCanNotBeWhiteSpace + => GetString("SessionTokenCanNotBeWhiteSpace"); + /// /// SingleOrDefault and FirstOrDefault cannot be used Cosmos SQL does not allow Offset without Limit. Consider specifying a 'Take' operation on the query. /// diff --git a/src/EFCore.Cosmos/Properties/CosmosStrings.resx b/src/EFCore.Cosmos/Properties/CosmosStrings.resx index 9c0e0259e41..e914bee6375 100644 --- a/src/EFCore.Cosmos/Properties/CosmosStrings.resx +++ b/src/EFCore.Cosmos/Properties/CosmosStrings.resx @@ -354,6 +354,9 @@ 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. + + Session token can not be white space. + SingleOrDefault and FirstOrDefault cannot be used Cosmos SQL does not allow Offset without Limit. Consider specifying a 'Take' operation on the query. diff --git a/src/EFCore.Cosmos/Storage/Internal/CosmosClientWrapper.cs b/src/EFCore.Cosmos/Storage/Internal/CosmosClientWrapper.cs index e209c351e60..0c04a201e71 100644 --- a/src/EFCore.Cosmos/Storage/Internal/CosmosClientWrapper.cs +++ b/src/EFCore.Cosmos/Storage/Internal/CosmosClientWrapper.cs @@ -967,7 +967,6 @@ private static async Task CreateSingleItemQueryAsync( } catch (CosmosException ex) when (ex.StatusCode == HttpStatusCode.NotFound && ex.SubStatusCode == resourceNotFoundSubStatusCode) { - // @TODO: Is there a test for this return null? return null; } diff --git a/src/EFCore.Cosmos/Storage/Internal/SessionTokenStorage.cs b/src/EFCore.Cosmos/Storage/Internal/SessionTokenStorage.cs index cc553dbdd1b..2662cea5083 100644 --- a/src/EFCore.Cosmos/Storage/Internal/SessionTokenStorage.cs +++ b/src/EFCore.Cosmos/Storage/Internal/SessionTokenStorage.cs @@ -1,8 +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.Diagnostics.CodeAnalysis; -using System.Globalization; using System.Runtime.CompilerServices; using System.Runtime.InteropServices; using Microsoft.EntityFrameworkCore.Cosmos.Internal; @@ -97,7 +95,7 @@ public virtual void AppendSessionToken(string containerName, string sessionToken throw new ArgumentException(CosmosStrings.ContainerNameDoesNotExist(containerName), nameof(containerName)); } - ParseAndMerge(compositeSessionToken, sessionToken); + compositeSessionToken.Add(sessionToken); } /// @@ -111,8 +109,7 @@ public virtual void SetSessionToken(string containerName, string? sessionToken) ArgumentNullException.ThrowIfNullOrWhiteSpace(containerName, nameof(containerName)); if (sessionToken is not null && string.IsNullOrWhiteSpace(sessionToken)) { - // @TODO: Exception messages in this file. - throw new ArgumentException("sessionToken cannot be whitespace.", sessionToken); + throw new ArgumentException(CosmosStrings.SessionTokenCanNotBeWhiteSpace, nameof(sessionToken)); } ref var compositeSessionToken = ref CollectionsMarshal.GetValueRefOrNullRef(_containerSessionTokens, containerName); @@ -129,50 +126,24 @@ public virtual void SetSessionToken(string containerName, string? sessionToken) return; } - ParseAndMerge(compositeSessionToken, sessionToken); - } - - private void ParseAndMerge(CompositeSessionToken compositeSessionToken, string sessionToken) - { - var parts = sessionToken.Split(','); - foreach (var part in parts) - { - var index = part.IndexOf(':'); - if (index == -1) - { - throw new ArgumentException("CosmosStrings.InvalidSessionToken(sessionToken)", nameof(sessionToken)); - } - - var pkRangeId = sessionToken.Substring(0, index); - var vector = sessionToken.Substring(index + 1); - if (!VectorSessionToken.TryCreate(vector, out var vectorSessionToken)) - { - throw new ArgumentException("CosmosStrings.InvalidSessionToken(sessionToken)", nameof(sessionToken)); - } - - compositeSessionToken.Merge(pkRangeId, vectorSessionToken); - } + compositeSessionToken.Add(sessionToken); } private sealed class CompositeSessionToken { private string? _string; private bool _isChanged; - public Dictionary Tokens { get; } = new(); + private readonly HashSet _tokens = new(); - public void Merge(string pkRangeId, VectorSessionToken token) + public void Add(string token) { - ref var existing = ref CollectionsMarshal.GetValueRefOrAddDefault(Tokens, pkRangeId, out var exists); - if (exists) - { - existing = existing!.Merge(token); - } - else + foreach (var tokenPart in token.Split(',')) { - existing = token; + if (_tokens.Add(tokenPart)) + { + _isChanged = true; + } } - - _isChanged = true; } public string? ConvertToString() @@ -180,420 +151,10 @@ public void Merge(string pkRangeId, VectorSessionToken token) if (_isChanged) { _isChanged = false; - _string = Tokens.Count == 0 ? null : string.Join(",", Tokens.Select(kvp => $"{kvp.Key}:{kvp.Value.ConvertToString()}")); + _string = string.Join(",", _tokens); } return _string; } } - - - /// - private sealed class VectorSessionToken : IEquatable - { - private static readonly IReadOnlyDictionary DefaultLocalLsnByRegion = new Dictionary(0); - - private readonly string sessionToken; - - private readonly long version; - - private readonly long globalLsn; - - private readonly IReadOnlyDictionary localLsnByRegion; - - private static readonly bool isFalseProgressMergeDisabled = string.Equals(Environment.GetEnvironmentVariable("AZURE_COSMOS_SESSION_TOKEN_FALSE_PROGRESS_MERGE_DISABLED"), "true", StringComparison.OrdinalIgnoreCase); - - /// - /// 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. - /// - public long LSN => globalLsn; - - private VectorSessionToken(long version, long globalLsn, IReadOnlyDictionary localLsnByRegion, string? sessionToken = null) - { - this.version = version; - this.globalLsn = globalLsn; - this.localLsnByRegion = localLsnByRegion; - if (sessionToken != null) - { - this.sessionToken = sessionToken; - return; - } - - string? text = null; - if (localLsnByRegion.Any()) - { - text = string.Join("#", localLsnByRegion.Select((KeyValuePair kvp) => string.Format(CultureInfo.InvariantCulture, "{0}{1}{2}", kvp.Key, '=', kvp.Value))); - } - - if (string.IsNullOrEmpty(text)) - { - this.sessionToken = string.Format(CultureInfo.InvariantCulture, "{0}{1}{2}", this.version, "#", this.globalLsn); - return; - } - - this.sessionToken = string.Format(CultureInfo.InvariantCulture, "{0}{1}{2}{3}{4}", this.version, "#", this.globalLsn, "#", text); - } - - /// - /// 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. - /// - public VectorSessionToken(VectorSessionToken other, long globalLSN) - : this(other.version, globalLSN, other.localLsnByRegion) - { - } - - /// - /// 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. - /// - public static bool TryCreate(string sessionToken, [NotNullWhen(true)] out VectorSessionToken? parsedSessionToken) - { - parsedSessionToken = null; - if (TryParseSessionToken(sessionToken, out var num, out var num2, out var readOnlyDictionary)) - { - parsedSessionToken = new VectorSessionToken(num, num2, readOnlyDictionary, sessionToken); - return true; - } - - return false; - } - - /// - /// 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. - /// - public bool Equals(VectorSessionToken? obj) - { - if (!(obj is VectorSessionToken vectorSessionToken)) - { - return false; - } - - if (version == vectorSessionToken.version && globalLsn == vectorSessionToken.globalLsn) - { - return AreRegionProgressEqual(vectorSessionToken.localLsnByRegion); - } - - return false; - } - - /// - /// 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. - /// - public bool IsValid(VectorSessionToken otherSessionToken) - { - if (!(otherSessionToken is VectorSessionToken vectorSessionToken)) - { - throw new ArgumentNullException("otherSessionToken"); - } - - if (isFalseProgressMergeDisabled) - { - if (vectorSessionToken.version < version || vectorSessionToken.globalLsn < globalLsn) - { - return false; - } - } - else if (vectorSessionToken.version < version || (vectorSessionToken.version == version && vectorSessionToken.globalLsn < globalLsn)) - { - return false; - } - - if (vectorSessionToken.version == version && vectorSessionToken.localLsnByRegion.Count != localLsnByRegion.Count) - { - throw new InvalidOperationException("string.Format(CultureInfo.InvariantCulture, RMResources.InvalidRegionsInSessionToken, sessionToken, vectorSessionToken.sessionToken)"); - } - - foreach (KeyValuePair item in vectorSessionToken.localLsnByRegion) - { - uint key = item.Key; - long value = item.Value; - long value2 = -1L; - if (!localLsnByRegion.TryGetValue(key, out value2)) - { - if (version == vectorSessionToken.version) - { - throw new InvalidOperationException("string.Format(CultureInfo.InvariantCulture, RMResources.InvalidRegionsInSessionToken, sessionToken, vectorSessionToken.sessionToken)"); - } - } - else if (value < value2) - { - return false; - } - } - - return true; - } - - /// - /// 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. - /// - public VectorSessionToken Merge(VectorSessionToken vectorSessionToken) - { - if (version == vectorSessionToken.version && localLsnByRegion.Count != vectorSessionToken.localLsnByRegion.Count) - { - throw new InvalidOperationException("string.Format(CultureInfo.InvariantCulture, RMResources.InvalidRegionsInSessionToken, sessionToken, vectorSessionToken.sessionToken)"); - } - - if (version >= vectorSessionToken.version && globalLsn > vectorSessionToken.globalLsn) - { - if (AreAllLocalLsnByRegionsGreaterThanOrEqual(this, vectorSessionToken)) - { - return this; - } - } - else if (vectorSessionToken.version >= version && vectorSessionToken.globalLsn >= globalLsn && AreAllLocalLsnByRegionsGreaterThanOrEqual(vectorSessionToken, this)) - { - return vectorSessionToken; - } - - VectorSessionToken vectorSessionToken2; - VectorSessionToken vectorSessionToken3; - if (version < vectorSessionToken.version) - { - vectorSessionToken2 = this; - vectorSessionToken3 = vectorSessionToken; - } - else - { - vectorSessionToken2 = vectorSessionToken; - vectorSessionToken3 = this; - } - - Dictionary dictionary = new Dictionary(vectorSessionToken3.localLsnByRegion.Count); - foreach (KeyValuePair item in vectorSessionToken3.localLsnByRegion) - { - uint key = item.Key; - long value = item.Value; - long value2 = -1L; - if (vectorSessionToken2.localLsnByRegion.TryGetValue(key, out value2)) - { - dictionary[key] = Math.Max(value, value2); - continue; - } - - if (version == vectorSessionToken.version) - { - throw new InvalidOperationException("string.Format(CultureInfo.InvariantCulture, RMResources.InvalidRegionsInSessionToken, sessionToken, vectorSessionToken.sessionToken)"); - } - - dictionary[key] = value; - } - - long num = Math.Max(version, vectorSessionToken.version); - long num2 = ((version == vectorSessionToken.version || isFalseProgressMergeDisabled) ? Math.Max(globalLsn, vectorSessionToken.globalLsn) : vectorSessionToken3.globalLsn); - return new VectorSessionToken(num, num2, dictionary); - } - - /// - /// 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. - /// - public string ConvertToString() - { - return sessionToken; - } - - private bool AreRegionProgressEqual(IReadOnlyDictionary other) - { - if (localLsnByRegion.Count != other.Count) - { - return false; - } - - foreach (KeyValuePair item in localLsnByRegion) - { - uint key = item.Key; - long value = item.Value; - if (other.TryGetValue(key, out var value2) && value != value2) - { - return false; - } - } - - return true; - } - - private static bool AreAllLocalLsnByRegionsGreaterThanOrEqual(VectorSessionToken higherToken, VectorSessionToken lowerToken) - { - if (higherToken.localLsnByRegion.Count != lowerToken.localLsnByRegion.Count) - { - return false; - } - - if (!higherToken.localLsnByRegion.Any()) - { - return true; - } - - foreach (KeyValuePair item in higherToken.localLsnByRegion) - { - uint key = item.Key; - long value = item.Value; - if (lowerToken.localLsnByRegion.TryGetValue(key, out var value2)) - { - if (value2 > value) - { - return false; - } - - continue; - } - - return false; - } - - return true; - } - - private static bool TryParseSessionToken(string sessionToken, out long version, out long globalLsn, [NotNullWhen(true)] out IReadOnlyDictionary? localLsnByRegion) - { - version = 0L; - localLsnByRegion = null; - globalLsn = -1L; - if (string.IsNullOrEmpty(sessionToken)) - { - return false; - } - - int index = 0; - if (!TryParseLongSegment(sessionToken, ref index, out version)) - { - return false; - } - - if (index >= sessionToken.Length) - { - return false; - } - - if (!TryParseLongSegment(sessionToken, ref index, out globalLsn)) - { - return false; - } - - if (index >= sessionToken.Length) - { - localLsnByRegion = DefaultLocalLsnByRegion; - return true; - } - - Dictionary dictionary = new Dictionary(); - while (index < sessionToken.Length) - { - if (!TryParseUintTillRegionProgressSeparator(sessionToken, ref index, out var value)) - { - return false; - } - - if (!TryParseLongSegment(sessionToken, ref index, out var value2)) - { - return false; - } - - dictionary[value] = value2; - } - - localLsnByRegion = dictionary; - return true; - } - - private static bool TryParseUintTillRegionProgressSeparator(string input, ref int index, out uint value) - { - value = 0u; - if (index >= input.Length) - { - return false; - } - - long num = 0L; - while (index < input.Length) - { - char c = input[index]; - if (c >= '0' && c <= '9') - { - num = num * 10 + (c - 48); - index++; - continue; - } - - if (c == '=') - { - index++; - break; - } - - return false; - } - - if (num > uint.MaxValue || num < 0) - { - return false; - } - - value = (uint)num; - return true; - } - - private static bool TryParseLongSegment(string input, ref int index, out long value) - { - value = 0L; - if (index >= input.Length) - { - return false; - } - - bool flag = false; - if (input[index] == '-') - { - index++; - flag = true; - } - - while (index < input.Length) - { - char c = input[index]; - if (c >= '0' && c <= '9') - { - value = value * 10 + (c - 48); - index++; - continue; - } - - if (c == '#') - { - index++; - break; - } - - return false; - } - - if (flag) - { - value *= -1L; - } - - return true; - } - } - } diff --git a/test/EFCore.Cosmos.FunctionalTests/CosmosSessionTokensTest.cs b/test/EFCore.Cosmos.FunctionalTests/CosmosSessionTokensTest.cs index 62914bb1a7d..5522c61f7b4 100644 --- a/test/EFCore.Cosmos.FunctionalTests/CosmosSessionTokensTest.cs +++ b/test/EFCore.Cosmos.FunctionalTests/CosmosSessionTokensTest.cs @@ -87,7 +87,7 @@ public virtual async Task AppendSessionToken_no_tokens_sets_token(bool defaultCo } [ConditionalTheory, InlineData(true), InlineData(false)] - public virtual async Task AppendSessionToken_append_higher_lsn_same_pkrange_takes_higher_lsn(bool defaultContainer) + public virtual async Task AppendSessionToken_append_token_not_present_adds_token(bool defaultContainer) { var contextFactory = await InitializeAsync(); @@ -106,7 +106,7 @@ public virtual async Task AppendSessionToken_append_higher_lsn_same_pkrange_take var initialToken = defaultContainer ? sessionTokens.GetSessionToken() : sessionTokens.GetSessionToken(OtherContainerName); Assert.False(string.IsNullOrWhiteSpace(initialToken)); - var newToken = initialToken.Substring(0, initialToken.IndexOf('#') + 1) + "999999"; + var newToken = "0:-1#231"; if (defaultContainer) { sessionTokens.AppendSessionToken(newToken); @@ -118,11 +118,11 @@ public virtual async Task AppendSessionToken_append_higher_lsn_same_pkrange_take var updatedToken = defaultContainer ? sessionTokens.GetSessionToken() : sessionTokens.GetSessionToken(OtherContainerName); - Assert.Equal(newToken, updatedToken); + Assert.Equal(initialToken + "," + newToken, updatedToken); } [ConditionalTheory, InlineData(true), InlineData(false)] - public virtual async Task AppendSessionToken_append_lower_lsn_same_pkrange_takes_higher_lsn(bool defaultContainer) + public virtual async Task AppendSessionToken_append_token_already_present_does_not_add_token(bool defaultContainer) { var contextFactory = await InitializeAsync(); @@ -142,7 +142,7 @@ public virtual async Task AppendSessionToken_append_lower_lsn_same_pkrange_takes var initialToken = defaultContainer ? sessionTokens.GetSessionToken() : sessionTokens.GetSessionToken(OtherContainerName); Assert.False(string.IsNullOrWhiteSpace(initialToken)); - var newToken = initialToken.Substring(0, initialToken.IndexOf('#') + 1) + "1"; + var newToken = initialToken; if (defaultContainer) { sessionTokens.AppendSessionToken(newToken); @@ -158,43 +158,34 @@ public virtual async Task AppendSessionToken_append_lower_lsn_same_pkrange_takes } [ConditionalTheory, InlineData(true), InlineData(false)] - public virtual async Task AppendSessionToken_different_pkrange_composites_tokens(bool defaultContainer) + public virtual async Task AppendSessionToken_multiple_tokens_splits_tokens(bool defaultContainer) { var contextFactory = await InitializeAsync(); using var context = contextFactory.CreateContext(); - if (defaultContainer) - { - context.Add(new Customer { Id = "1", PartitionKey = "1" }); - } - else - { - context.Add(new OtherContainerCustomer { Id = "1", PartitionKey = "1" }); - } - - await context.SaveChangesAsync(); var sessionTokens = context.Database.GetSessionTokens(); - var initialToken = defaultContainer ? sessionTokens.GetSessionToken() : sessionTokens.GetSessionToken(OtherContainerName); - Assert.False(string.IsNullOrWhiteSpace(initialToken)); - - var newToken = "99:-1#999999"; + var newToken = "0:-1#123,1:-1#456"; + var appendix = "0:-1#123"; if (defaultContainer) { sessionTokens.AppendSessionToken(newToken); -} + sessionTokens.AppendSessionToken(appendix); + + } else { sessionTokens.AppendSessionToken(OtherContainerName, newToken); + sessionTokens.AppendSessionToken(OtherContainerName, appendix); } var updatedToken = defaultContainer ? sessionTokens.GetSessionToken() : sessionTokens.GetSessionToken(OtherContainerName); - Assert.Equal(initialToken + ",99:-1#999999", updatedToken); + Assert.Equal(newToken, updatedToken); } [ConditionalTheory, InlineData(true), InlineData(false)] - public virtual async Task SetSessionToken_does_not_merge_session_token(bool defaultContainer) + public virtual async Task SetSessionToken_does_not_append_session_token(bool defaultContainer) { var contextFactory = await InitializeAsync(); @@ -214,7 +205,7 @@ public virtual async Task SetSessionToken_does_not_merge_session_token(bool defa var initialToken = defaultContainer ? sessionTokens.GetSessionToken() : sessionTokens.GetSessionToken(OtherContainerName); Assert.False(string.IsNullOrWhiteSpace(initialToken)); - var newToken = "0:-1#1"; + var newToken = "0:-1#1,1:-1#1"; if (defaultContainer) { sessionTokens.SetSessionToken(newToken); @@ -226,7 +217,7 @@ public virtual async Task SetSessionToken_does_not_merge_session_token(bool defa var updatedToken = defaultContainer ? sessionTokens.GetSessionToken() : sessionTokens.GetSessionToken(OtherContainerName); - Assert.Equal("0:-1#1", updatedToken); + Assert.Equal(newToken, updatedToken); } [ConditionalTheory, InlineData(true), InlineData(false)] @@ -914,7 +905,7 @@ public virtual async Task Add_merges_session_token(AutoTransactionBehavior autoT Assert.False(string.IsNullOrWhiteSpace(sessionToken)); Assert.NotEqual(sessionToken, initialToken); - Assert.StartsWith(initialToken.Substring(0, initialToken.IndexOf('#') + 1), sessionToken); + Assert.StartsWith(initialToken + ",", sessionToken); } [ConditionalTheory] @@ -996,7 +987,7 @@ public virtual async Task Delete_merges_session_token(AutoTransactionBehavior au var sessionToken = defaultContainer ? context.Database.GetSessionTokens().GetSessionToken() : context.Database.GetSessionTokens().GetSessionToken(OtherContainerName); Assert.False(string.IsNullOrWhiteSpace(sessionToken)); Assert.NotEqual(sessionToken, initialToken); - Assert.StartsWith(initialToken.Substring(0, initialToken.IndexOf('#') + 1), sessionToken); + Assert.StartsWith(initialToken + ",", sessionToken); } [ConditionalTheory] @@ -1069,7 +1060,7 @@ public virtual async Task Update_merges_session_token(AutoTransactionBehavior au var sessionToken = defaultContainer ? context.Database.GetSessionTokens().GetSessionToken() : context.Database.GetSessionTokens().GetSessionToken(OtherContainerName); Assert.False(string.IsNullOrWhiteSpace(sessionToken)); Assert.NotEqual(initialToken, sessionToken); - Assert.StartsWith(initialToken.Substring(0, initialToken.IndexOf('#') + 1), sessionToken); + Assert.StartsWith(initialToken + ",", sessionToken); } [ConditionalTheory] From 87e0dfd23434ca1b1b0d1bf3829e991b6b5971ac Mon Sep 17 00:00:00 2001 From: JoasE <32096708+JoasE@users.noreply.github.com> Date: Fri, 24 Oct 2025 20:07:01 +0200 Subject: [PATCH 17/37] Revert "Add option ManualSessionTokenManagementEnabled" --- .../CosmosDatabaseFacadeExtensions.cs | 10 ++----- .../CosmosDbContextOptionsBuilder.cs | 16 ---------- .../Internal/CosmosDbOptionExtension.cs | 30 +------------------ .../Internal/CosmosSingletonOptions.cs | 10 ------- .../Internal/ICosmosSingletonOptions.cs | 8 ----- .../CosmosQueryCompilationContextFactory.cs | 3 +- .../Storage/Internal/CosmosDatabaseWrapper.cs | 6 +--- .../CosmosSessionTokensTest.cs | 14 +-------- 8 files changed, 6 insertions(+), 91 deletions(-) diff --git a/src/EFCore.Cosmos/Extensions/CosmosDatabaseFacadeExtensions.cs b/src/EFCore.Cosmos/Extensions/CosmosDatabaseFacadeExtensions.cs index 7bbabdcdabf..65080c57aa1 100644 --- a/src/EFCore.Cosmos/Extensions/CosmosDatabaseFacadeExtensions.cs +++ b/src/EFCore.Cosmos/Extensions/CosmosDatabaseFacadeExtensions.cs @@ -27,7 +27,7 @@ public static CosmosClient GetCosmosClient(this DatabaseFacade databaseFacade) => GetService(databaseFacade).Client; /// - /// Gets the used to manage the session tokens for this . + /// Gets used to manage the session tokens for this . /// /// The for the context. /// The . @@ -39,13 +39,7 @@ public static ISessionTokenStorage GetSessionTokens(this DatabaseFacade database throw new InvalidOperationException(CosmosStrings.CosmosNotInUse); } - if (dbWrapper.SessionTokenStorage is not SessionTokenStorage) - { - // @TODO: string - throw new InvalidOperationException("CosmosStrings.EnableManualSessionTokenManagement"); - } - - return dbWrapper.SessionTokenStorage; + return dbWrapper.SessionTokenStorage!; } private static TService GetService(IInfrastructure databaseFacade) diff --git a/src/EFCore.Cosmos/Infrastructure/CosmosDbContextOptionsBuilder.cs b/src/EFCore.Cosmos/Infrastructure/CosmosDbContextOptionsBuilder.cs index 2dd3ec58ff3..e88fabf78da 100644 --- a/src/EFCore.Cosmos/Infrastructure/CosmosDbContextOptionsBuilder.cs +++ b/src/EFCore.Cosmos/Infrastructure/CosmosDbContextOptionsBuilder.cs @@ -211,22 +211,6 @@ public virtual CosmosDbContextOptionsBuilder MaxRequestsPerTcpConnection(int req public virtual CosmosDbContextOptionsBuilder ContentResponseOnWriteEnabled(bool enabled = true) => WithOption(e => e.ContentResponseOnWriteEnabled(Check.NotNull(enabled))); - - /// - /// Sets the boolean to track and manage session tokens for requests made to Cosmos DB - /// and being able to access them via the method. - /// 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. - /// See Utilize session tokens for more details. - /// - /// - /// See Using DbContextOptions, and - /// Accessing Azure Cosmos DB with EF Core for more information and examples. - /// - /// to track and manually manage session tokens in EF. - public virtual CosmosDbContextOptionsBuilder ManualSessionTokenManagementEnabled(bool enabled = true) - => WithOption(e => e.ManualSessionTokenManagementEnabled(enabled)); - /// /// 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. diff --git a/src/EFCore.Cosmos/Infrastructure/Internal/CosmosDbOptionExtension.cs b/src/EFCore.Cosmos/Infrastructure/Internal/CosmosDbOptionExtension.cs index dbe9e67e273..f2545174ac3 100644 --- a/src/EFCore.Cosmos/Infrastructure/Internal/CosmosDbOptionExtension.cs +++ b/src/EFCore.Cosmos/Infrastructure/Internal/CosmosDbOptionExtension.cs @@ -36,7 +36,6 @@ public class CosmosOptionsExtension : IDbContextOptionsExtension private bool? _enableContentResponseOnWrite; private DbContextOptionsExtensionInfo? _info; private Func? _httpClientFactory; - private bool _enableManualSessionTokenManagement; /// /// This is an internal API that supports the Entity Framework Core infrastructure and not subject to @@ -74,7 +73,6 @@ protected CosmosOptionsExtension(CosmosOptionsExtension copyFrom) _maxTcpConnectionsPerEndpoint = copyFrom._maxTcpConnectionsPerEndpoint; _maxRequestsPerTcpConnection = copyFrom._maxRequestsPerTcpConnection; _httpClientFactory = copyFrom._httpClientFactory; - _enableManualSessionTokenManagement = copyFrom._enableManualSessionTokenManagement; } /// @@ -566,30 +564,6 @@ public virtual CosmosOptionsExtension WithHttpClientFactory(Func? ht return clone; } - /// - /// 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. - /// - public virtual bool EnableManualSessionTokenManagement - => _enableManualSessionTokenManagement; - - /// - /// 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. - /// - public virtual CosmosOptionsExtension ManualSessionTokenManagementEnabled(bool enabled) - { - var clone = Clone(); - - clone._enableManualSessionTokenManagement = enabled; - - return clone; - } - /// /// 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 @@ -658,7 +632,6 @@ public override int GetServiceProviderHashCode() hashCode.Add(Extension._maxTcpConnectionsPerEndpoint); hashCode.Add(Extension._maxRequestsPerTcpConnection); hashCode.Add(Extension._httpClientFactory); - hashCode.Add(Extension._enableManualSessionTokenManagement); _serviceProviderHash = hashCode.ToHashCode(); } @@ -683,8 +656,7 @@ 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._enableManualSessionTokenManagement == otherInfo.Extension._enableManualSessionTokenManagement; + && Extension._httpClientFactory == otherInfo.Extension._httpClientFactory; public override void PopulateDebugInfo(IDictionary debugInfo) { diff --git a/src/EFCore.Cosmos/Infrastructure/Internal/CosmosSingletonOptions.cs b/src/EFCore.Cosmos/Infrastructure/Internal/CosmosSingletonOptions.cs index d5221c0eef3..af29229cfa5 100644 --- a/src/EFCore.Cosmos/Infrastructure/Internal/CosmosSingletonOptions.cs +++ b/src/EFCore.Cosmos/Infrastructure/Internal/CosmosSingletonOptions.cs @@ -151,14 +151,6 @@ public class CosmosSingletonOptions : ICosmosSingletonOptions /// public virtual Func? HttpClientFactory { get; private set; } - /// - /// 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. - /// - public virtual bool EnableManualSessionTokenManagement { get; private set; } - /// /// 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 @@ -186,7 +178,6 @@ public virtual void Initialize(IDbContextOptions options) MaxTcpConnectionsPerEndpoint = cosmosOptions.MaxTcpConnectionsPerEndpoint; MaxRequestsPerTcpConnection = cosmosOptions.MaxRequestsPerTcpConnection; HttpClientFactory = cosmosOptions.HttpClientFactory; - EnableManualSessionTokenManagement = cosmosOptions.EnableManualSessionTokenManagement; } } @@ -217,7 +208,6 @@ public virtual void Validate(IDbContextOptions options) || MaxTcpConnectionsPerEndpoint != cosmosOptions.MaxTcpConnectionsPerEndpoint || MaxRequestsPerTcpConnection != cosmosOptions.MaxRequestsPerTcpConnection || HttpClientFactory != cosmosOptions.HttpClientFactory - || EnableManualSessionTokenManagement != cosmosOptions.EnableManualSessionTokenManagement )) { throw new InvalidOperationException( diff --git a/src/EFCore.Cosmos/Infrastructure/Internal/ICosmosSingletonOptions.cs b/src/EFCore.Cosmos/Infrastructure/Internal/ICosmosSingletonOptions.cs index 8113da2b25f..a26b79a82b3 100644 --- a/src/EFCore.Cosmos/Infrastructure/Internal/ICosmosSingletonOptions.cs +++ b/src/EFCore.Cosmos/Infrastructure/Internal/ICosmosSingletonOptions.cs @@ -155,12 +155,4 @@ public interface ICosmosSingletonOptions : ISingletonOptions /// doing so can result in application failures when updating to a new Entity Framework Core release. /// Func? HttpClientFactory { get; } - - /// - /// 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. - /// - bool EnableManualSessionTokenManagement { get; } } diff --git a/src/EFCore.Cosmos/Query/Internal/CosmosQueryCompilationContextFactory.cs b/src/EFCore.Cosmos/Query/Internal/CosmosQueryCompilationContextFactory.cs index d409c850c97..99a937d82d6 100644 --- a/src/EFCore.Cosmos/Query/Internal/CosmosQueryCompilationContextFactory.cs +++ b/src/EFCore.Cosmos/Query/Internal/CosmosQueryCompilationContextFactory.cs @@ -2,7 +2,6 @@ // 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; @@ -41,7 +40,7 @@ public CosmosQueryCompilationContextFactory(QueryCompilationContextDependencies /// 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. /// - protected virtual ISessionTokenStorage SessionTokenStorage => _sessionTokenStorage ??= ((CosmosDatabaseWrapper)Dependencies.Context.Database.GetService()).SessionTokenStorage; + protected virtual ISessionTokenStorage SessionTokenStorage => _sessionTokenStorage ??= Dependencies.Context.Database.GetSessionTokens(); /// /// This is an internal API that supports the Entity Framework Core infrastructure and not subject to diff --git a/src/EFCore.Cosmos/Storage/Internal/CosmosDatabaseWrapper.cs b/src/EFCore.Cosmos/Storage/Internal/CosmosDatabaseWrapper.cs index 2ef74edb3eb..86f0b5c2763 100644 --- a/src/EFCore.Cosmos/Storage/Internal/CosmosDatabaseWrapper.cs +++ b/src/EFCore.Cosmos/Storage/Internal/CosmosDatabaseWrapper.cs @@ -5,7 +5,6 @@ using System.Runtime.InteropServices; using Microsoft.EntityFrameworkCore.ChangeTracking.Internal; using Microsoft.EntityFrameworkCore.Cosmos.Diagnostics.Internal; -using Microsoft.EntityFrameworkCore.Cosmos.Infrastructure.Internal; using Microsoft.EntityFrameworkCore.Cosmos.Internal; using Microsoft.EntityFrameworkCore.Cosmos.Metadata.Internal; using Microsoft.EntityFrameworkCore.Cosmos.Update.Internal; @@ -39,16 +38,13 @@ public CosmosDatabaseWrapper( DatabaseDependencies dependencies, ICurrentDbContext currentDbContext, ICosmosClientWrapper cosmosClient, - ICosmosSingletonOptions cosmosSingletonOptions, ILoggingOptions loggingOptions) : base(dependencies) { _currentDbContext = currentDbContext; _cosmosClient = cosmosClient; - SessionTokenStorage = cosmosSingletonOptions.EnableManualSessionTokenManagement ? - new SessionTokenStorage(_currentDbContext.Context) - : new NullSessionTokenStorage(); + SessionTokenStorage = new SessionTokenStorage(_currentDbContext.Context); if (loggingOptions.IsSensitiveDataLoggingEnabled) { diff --git a/test/EFCore.Cosmos.FunctionalTests/CosmosSessionTokensTest.cs b/test/EFCore.Cosmos.FunctionalTests/CosmosSessionTokensTest.cs index 5522c61f7b4..5b439b3b31f 100644 --- a/test/EFCore.Cosmos.FunctionalTests/CosmosSessionTokensTest.cs +++ b/test/EFCore.Cosmos.FunctionalTests/CosmosSessionTokensTest.cs @@ -17,19 +17,6 @@ protected override ITestStoreFactory TestStoreFactory protected override DbContextOptionsBuilder AddOptions(DbContextOptionsBuilder builder) => base.AddOptions(builder).ConfigureWarnings(x => x.Ignore(CosmosEventId.SyncNotSupported)); - protected override TestStore CreateTestStore() => CosmosTestStore.Create(StoreName, (c) => c.ManualSessionTokenManagementEnabled()); - - [ConditionalFact] - public virtual async Task GetSessionTokens_throws_if_not_enabled() - { - var contextFactory = await InitializeAsync(createTestStore: () => CosmosTestStore.Create(StoreName)); - - using var context = contextFactory.CreateContext(); - - var exception = Assert.Throws(() => context.Database.GetSessionTokens()); - Assert.Equal("CosmosStrings.EnableManualSessionTokenManagement", exception.Message); - } - [ConditionalFact] public virtual async Task SetSessionToken_ThrowsForNonExistentContainer() { @@ -1084,6 +1071,7 @@ public virtual async Task Add_uses_session_token(AutoTransactionBehavior autoTra context.Database.AutoTransactionBehavior = autoTransactionBehavior; var sessionTokens = context.Database.GetSessionTokens(); + var sessionToken = sessionTokens.GetSessionToken()!; // Only way we can test this is by setting a session token that will fail the request if used.. // Only way to do this for a write is to set an invalid session token.. var internalDictionary = sessionTokens.GetType().GetField("_containerSessionTokens", BindingFlags.NonPublic | BindingFlags.Instance)!.GetValue(sessionTokens)!; From df5809efa17cf1817e12ca87c61df6e99afe7d99 Mon Sep 17 00:00:00 2001 From: JoasE <32096708+JoasE@users.noreply.github.com> Date: Fri, 24 Oct 2025 20:08:04 +0200 Subject: [PATCH 18/37] Cleanup --- test/EFCore.Cosmos.FunctionalTests/CosmosSessionTokensTest.cs | 1 - 1 file changed, 1 deletion(-) diff --git a/test/EFCore.Cosmos.FunctionalTests/CosmosSessionTokensTest.cs b/test/EFCore.Cosmos.FunctionalTests/CosmosSessionTokensTest.cs index 5b439b3b31f..e46c02f9a9d 100644 --- a/test/EFCore.Cosmos.FunctionalTests/CosmosSessionTokensTest.cs +++ b/test/EFCore.Cosmos.FunctionalTests/CosmosSessionTokensTest.cs @@ -1071,7 +1071,6 @@ public virtual async Task Add_uses_session_token(AutoTransactionBehavior autoTra context.Database.AutoTransactionBehavior = autoTransactionBehavior; var sessionTokens = context.Database.GetSessionTokens(); - var sessionToken = sessionTokens.GetSessionToken()!; // Only way we can test this is by setting a session token that will fail the request if used.. // Only way to do this for a write is to set an invalid session token.. var internalDictionary = sessionTokens.GetType().GetField("_containerSessionTokens", BindingFlags.NonPublic | BindingFlags.Instance)!.GetValue(sessionTokens)!; From 2c29285bd23664544638da1b9d105795d0bad379 Mon Sep 17 00:00:00 2001 From: JoasE <32096708+JoasE@users.noreply.github.com> Date: Fri, 24 Oct 2025 20:24:46 +0200 Subject: [PATCH 19/37] Cleanup --- .../Extensions/CosmosDatabaseFacadeExtensions.cs | 2 +- .../TestUtilities}/NullSessionTokenStorage.cs | 4 +++- 2 files changed, 4 insertions(+), 2 deletions(-) rename {src/EFCore.Cosmos/Storage/Internal => test/EFCore.Cosmos.FunctionalTests/TestUtilities}/NullSessionTokenStorage.cs (97%) diff --git a/src/EFCore.Cosmos/Extensions/CosmosDatabaseFacadeExtensions.cs b/src/EFCore.Cosmos/Extensions/CosmosDatabaseFacadeExtensions.cs index 65080c57aa1..5df9cb0c80b 100644 --- a/src/EFCore.Cosmos/Extensions/CosmosDatabaseFacadeExtensions.cs +++ b/src/EFCore.Cosmos/Extensions/CosmosDatabaseFacadeExtensions.cs @@ -39,7 +39,7 @@ public static ISessionTokenStorage GetSessionTokens(this DatabaseFacade database throw new InvalidOperationException(CosmosStrings.CosmosNotInUse); } - return dbWrapper.SessionTokenStorage!; + return dbWrapper.SessionTokenStorage; } private static TService GetService(IInfrastructure databaseFacade) diff --git a/src/EFCore.Cosmos/Storage/Internal/NullSessionTokenStorage.cs b/test/EFCore.Cosmos.FunctionalTests/TestUtilities/NullSessionTokenStorage.cs similarity index 97% rename from src/EFCore.Cosmos/Storage/Internal/NullSessionTokenStorage.cs rename to test/EFCore.Cosmos.FunctionalTests/TestUtilities/NullSessionTokenStorage.cs index 01891e78545..70c2d21d243 100644 --- a/src/EFCore.Cosmos/Storage/Internal/NullSessionTokenStorage.cs +++ b/test/EFCore.Cosmos.FunctionalTests/TestUtilities/NullSessionTokenStorage.cs @@ -1,7 +1,9 @@ // 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.Storage.Internal; +using Microsoft.EntityFrameworkCore.Cosmos.Storage; + +namespace Microsoft.EntityFrameworkCore.TestUtilities; /// /// This is an internal API that supports the Entity Framework Core infrastructure and not subject to From 77a81a3646f06f2732682b5444c6f6a1c175f698 Mon Sep 17 00:00:00 2001 From: JoasE <32096708+JoasE@users.noreply.github.com> Date: Sat, 25 Oct 2025 12:29:36 +0200 Subject: [PATCH 20/37] Add tests for pr comment changes (pooling and query compilation vs execution) --- .../CosmosSessionTokensTest.cs | 206 ++++++++++++++++++ 1 file changed, 206 insertions(+) diff --git a/test/EFCore.Cosmos.FunctionalTests/CosmosSessionTokensTest.cs b/test/EFCore.Cosmos.FunctionalTests/CosmosSessionTokensTest.cs index e46c02f9a9d..064c59784cf 100644 --- a/test/EFCore.Cosmos.FunctionalTests/CosmosSessionTokensTest.cs +++ b/test/EFCore.Cosmos.FunctionalTests/CosmosSessionTokensTest.cs @@ -332,6 +332,212 @@ public virtual async Task Query_uses_session_token(bool async, bool defaultConta Assert.Contains("The read session is not available for the input session token.", ex.ResponseBody); } + [ConditionalTheory] + [InlineData(true, true)] + [InlineData(true, false)] + [InlineData(false, true)] + [InlineData(false, false)] + public virtual async Task Query_on_new_context_does_not_use_same_session_token(bool async, bool defaultContainer) + { + var contextFactory = await InitializeAsync(); + using var context = contextFactory.CreateContext(); + + if (defaultContainer) + { + context.Customers.Add(new Customer { Id = "1", PartitionKey = "1" }); + } + else + { + context.OtherContainerCustomers.Add(new OtherContainerCustomer { Id = "1", PartitionKey = "1" }); + } + + if (async) + { + await context.SaveChangesAsync(); + } + else + { + context.SaveChanges(); + } + + var sessionTokens = context.Database.GetSessionTokens(); + + string sessionToken; + if (defaultContainer) + { + sessionToken = sessionTokens.GetSessionToken()!; + } + else + { + sessionToken = sessionTokens.GetSessionToken(OtherContainerName)!; + } + + // Only way we can test this is by setting a session token that will fail the request if used.. + // This will take a couple of seconds to fail + var newToken = sessionToken.Substring(0, sessionToken.IndexOf('#') + 1) + int.MaxValue; + + if (defaultContainer) + { + sessionTokens.SetSessionToken(newToken); + } + else + { + sessionTokens.SetSessionToken(OtherContainerName, newToken); + } + + if (async) + { + if (defaultContainer) + { + await Assert.ThrowsAsync(() => context.Customers.ToListAsync()); + } + else + { + await Assert.ThrowsAsync(() => context.OtherContainerCustomers.ToListAsync()); + } + } + else + { + if (defaultContainer) + { + Assert.Throws(() => context.Customers.ToList()); + } + else + { + Assert.Throws(() => context.OtherContainerCustomers.ToList()); + } + } + + using var newContext = contextFactory.CreateContext(); + if (async) + { + if (defaultContainer) + { + await newContext.Customers.ToListAsync(); + } + else + { + await newContext.OtherContainerCustomers.ToListAsync(); + } + } + else + { + if (defaultContainer) + { + newContext.Customers.ToList(); + } + else + { + newContext.OtherContainerCustomers.ToList(); + } + } + } + + [ConditionalTheory] + [InlineData(true, true)] + [InlineData(true, false)] + [InlineData(false, true)] + [InlineData(false, false)] + public virtual async Task Query_on_same_newly_pooled_context_does_not_use_same_session_token(bool async, bool defaultContainer) + { + var contextFactory = await InitializeAsync(); + DbContext contextCopy; + using (var context = contextFactory.CreateContext()) + { + contextCopy = context; + if (defaultContainer) + { + context.Customers.Add(new Customer { Id = "1", PartitionKey = "1" }); + } + else + { + context.OtherContainerCustomers.Add(new OtherContainerCustomer { Id = "1", PartitionKey = "1" }); + } + + if (async) + { + await context.SaveChangesAsync(); + } + else + { + context.SaveChanges(); + } + + var sessionTokens = context.Database.GetSessionTokens(); + + string sessionToken; + if (defaultContainer) + { + sessionToken = sessionTokens.GetSessionToken()!; + } + else + { + sessionToken = sessionTokens.GetSessionToken(OtherContainerName)!; + } + + // Only way we can test this is by setting a session token that will fail the request if used.. + // This will take a couple of seconds to fail + var newToken = sessionToken.Substring(0, sessionToken.IndexOf('#') + 1) + int.MaxValue; + + if (defaultContainer) + { + sessionTokens.SetSessionToken(newToken); + } + else + { + sessionTokens.SetSessionToken(OtherContainerName, newToken); + } + + if (async) + { + if (defaultContainer) + { + await Assert.ThrowsAsync(() => context.Customers.ToListAsync()); + } + else + { + await Assert.ThrowsAsync(() => context.OtherContainerCustomers.ToListAsync()); + } + } + else + { + if (defaultContainer) + { + Assert.Throws(() => context.Customers.ToList()); + } + else + { + Assert.Throws(() => context.OtherContainerCustomers.ToList()); + } + } + } + + using var newContext = contextFactory.CreateContext(); + Assert.Same(newContext, contextCopy); + if (async) + { + if (defaultContainer) + { + await newContext.Customers.ToListAsync(); + } + else + { + await newContext.OtherContainerCustomers.ToListAsync(); + } + } + else + { + if (defaultContainer) + { + newContext.Customers.ToList(); + } + else + { + newContext.OtherContainerCustomers.ToList(); + } + } + } + [ConditionalTheory] [InlineData(true), InlineData(false)] public virtual async Task PagingQuery_uses_session_token(bool defaultContainer) From 81e137c0a4dc4efabadc0f568fe341ecc34cbaaa Mon Sep 17 00:00:00 2001 From: JoasE <32096708+JoasE@users.noreply.github.com> Date: Sat, 25 Oct 2025 12:30:03 +0200 Subject: [PATCH 21/37] Clear session token storage on reset Register CosmosDatabaseWrapper as IResettableService --- .../CosmosServiceCollectionExtensions.cs | 1 + src/EFCore.Cosmos/Storage/ISessionTokenStorage.cs | 5 +++++ .../Storage/Internal/CosmosDatabaseWrapper.cs | 13 ++++++++++++- .../Storage/Internal/SessionTokenStorage.cs | 14 ++++++++++++++ .../TestUtilities/NullSessionTokenStorage.cs | 8 ++++++++ 5 files changed, 40 insertions(+), 1 deletion(-) diff --git a/src/EFCore.Cosmos/Extensions/CosmosServiceCollectionExtensions.cs b/src/EFCore.Cosmos/Extensions/CosmosServiceCollectionExtensions.cs index cfff8a16aa9..9a79b99e343 100644 --- a/src/EFCore.Cosmos/Extensions/CosmosServiceCollectionExtensions.cs +++ b/src/EFCore.Cosmos/Extensions/CosmosServiceCollectionExtensions.cs @@ -97,6 +97,7 @@ public static IServiceCollection AddEntityFrameworkCosmos(this IServiceCollectio .TryAdd() .TryAdd>() .TryAdd() + .TryAdd(sp => (CosmosDatabaseWrapper)sp.GetRequiredService()) .TryAdd() .TryAdd() .TryAdd() diff --git a/src/EFCore.Cosmos/Storage/ISessionTokenStorage.cs b/src/EFCore.Cosmos/Storage/ISessionTokenStorage.cs index b18ed566753..8536614b108 100644 --- a/src/EFCore.Cosmos/Storage/ISessionTokenStorage.cs +++ b/src/EFCore.Cosmos/Storage/ISessionTokenStorage.cs @@ -46,4 +46,9 @@ public interface ISessionTokenStorage /// /// The session token to set, or to clear any stored token. void SetSessionToken(string? sessionToken); + + /// + /// Clears all stored session tokens. + /// + void Clear(); } diff --git a/src/EFCore.Cosmos/Storage/Internal/CosmosDatabaseWrapper.cs b/src/EFCore.Cosmos/Storage/Internal/CosmosDatabaseWrapper.cs index 86f0b5c2763..4635f1b5f7b 100644 --- a/src/EFCore.Cosmos/Storage/Internal/CosmosDatabaseWrapper.cs +++ b/src/EFCore.Cosmos/Storage/Internal/CosmosDatabaseWrapper.cs @@ -19,7 +19,7 @@ namespace Microsoft.EntityFrameworkCore.Cosmos.Storage.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. /// -public class CosmosDatabaseWrapper : Database +public class CosmosDatabaseWrapper : Database, IResettableService { private readonly Dictionary _documentCollections = new(); @@ -669,6 +669,17 @@ private DbUpdateException WrapUpdateException(Exception exception, IReadOnlyList }; } + void IResettableService.ResetState() + { + SessionTokenStorage.Clear(); + } + + Task IResettableService.ResetStateAsync(CancellationToken cancellationToken) + { + ((IResettableService)this).ResetState(); + return Task.CompletedTask; + } + private sealed class SaveGroups { public required IEnumerable SingleUpdateEntries { get; init; } diff --git a/src/EFCore.Cosmos/Storage/Internal/SessionTokenStorage.cs b/src/EFCore.Cosmos/Storage/Internal/SessionTokenStorage.cs index 2662cea5083..6434d2b2e72 100644 --- a/src/EFCore.Cosmos/Storage/Internal/SessionTokenStorage.cs +++ b/src/EFCore.Cosmos/Storage/Internal/SessionTokenStorage.cs @@ -129,6 +129,20 @@ public virtual void SetSessionToken(string containerName, string? sessionToken) compositeSessionToken.Add(sessionToken); } + /// + /// 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. + /// + public virtual void Clear() + { + foreach (var key in _containerSessionTokens.Keys) + { + _containerSessionTokens[key] = new CompositeSessionToken(); + } + } + private sealed class CompositeSessionToken { private string? _string; diff --git a/test/EFCore.Cosmos.FunctionalTests/TestUtilities/NullSessionTokenStorage.cs b/test/EFCore.Cosmos.FunctionalTests/TestUtilities/NullSessionTokenStorage.cs index 70c2d21d243..1a9ff6fae52 100644 --- a/test/EFCore.Cosmos.FunctionalTests/TestUtilities/NullSessionTokenStorage.cs +++ b/test/EFCore.Cosmos.FunctionalTests/TestUtilities/NullSessionTokenStorage.cs @@ -29,6 +29,14 @@ public void AppendSessionToken(string sessionToken) { } /// public void AppendSessionToken(string containerName, string sessionToken) { } + /// + /// 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. + /// + public void Clear() { } + /// /// 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 From 0c62ceefe41e33f9d646d742f0ed4c5ef5d97a93 Mon Sep 17 00:00:00 2001 From: JoasE <32096708+JoasE@users.noreply.github.com> Date: Sat, 25 Oct 2025 12:40:26 +0200 Subject: [PATCH 22/37] Store session token storage in query context instead of compilation context --- .../Internal/CosmosQueryCompilationContext.cs | 11 +---------- .../CosmosQueryCompilationContextFactory.cs | 14 +------------- .../Query/Internal/CosmosQueryContext.cs | 11 ++++++++++- .../Query/Internal/CosmosQueryContextFactory.cs | 11 ++++++++++- ...ngExpressionVisitor.PagingQueryingEnumerable.cs | 11 +++-------- ...ompilingExpressionVisitor.QueryingEnumerable.cs | 13 +++---------- ...ExpressionVisitor.ReadItemQueryingEnumerable.cs | 11 +++-------- .../CosmosShapedQueryCompilingExpressionVisitor.cs | 10 +++------- .../CosmosSessionTokensTest.cs | 1 + 9 files changed, 35 insertions(+), 58 deletions(-) diff --git a/src/EFCore.Cosmos/Query/Internal/CosmosQueryCompilationContext.cs b/src/EFCore.Cosmos/Query/Internal/CosmosQueryCompilationContext.cs index 1401fe31cc2..edebd1549e4 100644 --- a/src/EFCore.Cosmos/Query/Internal/CosmosQueryCompilationContext.cs +++ b/src/EFCore.Cosmos/Query/Internal/CosmosQueryCompilationContext.cs @@ -19,9 +19,8 @@ public class CosmosQueryCompilationContext : QueryCompilationContext /// 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. /// - public CosmosQueryCompilationContext(QueryCompilationContextDependencies dependencies, ISessionTokenStorage sessionTokenStorage, bool async) : base(dependencies, async) + public CosmosQueryCompilationContext(QueryCompilationContextDependencies dependencies, bool async) : base(dependencies, async) { - SessionTokenStorage = sessionTokenStorage; } /// @@ -53,12 +52,4 @@ public CosmosQueryCompilationContext(QueryCompilationContextDependencies depende /// doing so can result in application failures when updating to a new Entity Framework Core release. /// public virtual CosmosAliasManager AliasManager { get; } = new(); - - /// - /// 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. - /// - public virtual ISessionTokenStorage SessionTokenStorage { get; } } diff --git a/src/EFCore.Cosmos/Query/Internal/CosmosQueryCompilationContextFactory.cs b/src/EFCore.Cosmos/Query/Internal/CosmosQueryCompilationContextFactory.cs index 99a937d82d6..bd303713db3 100644 --- a/src/EFCore.Cosmos/Query/Internal/CosmosQueryCompilationContextFactory.cs +++ b/src/EFCore.Cosmos/Query/Internal/CosmosQueryCompilationContextFactory.cs @@ -1,8 +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 Microsoft.EntityFrameworkCore.Cosmos.Storage; - namespace Microsoft.EntityFrameworkCore.Cosmos.Query.Internal; /// @@ -13,8 +11,6 @@ namespace Microsoft.EntityFrameworkCore.Cosmos.Query.Internal; /// public class CosmosQueryCompilationContextFactory : IQueryCompilationContextFactory { - private ISessionTokenStorage? _sessionTokenStorage; - /// /// 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 @@ -34,14 +30,6 @@ public CosmosQueryCompilationContextFactory(QueryCompilationContextDependencies /// protected virtual QueryCompilationContextDependencies Dependencies { get; } - /// - /// 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. - /// - protected virtual ISessionTokenStorage SessionTokenStorage => _sessionTokenStorage ??= Dependencies.Context.Database.GetSessionTokens(); - /// /// 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 @@ -49,5 +37,5 @@ public CosmosQueryCompilationContextFactory(QueryCompilationContextDependencies /// doing so can result in application failures when updating to a new Entity Framework Core release. /// public virtual QueryCompilationContext Create(bool async) - => new CosmosQueryCompilationContext(Dependencies, SessionTokenStorage, async); + => new CosmosQueryCompilationContext(Dependencies, async); } diff --git a/src/EFCore.Cosmos/Query/Internal/CosmosQueryContext.cs b/src/EFCore.Cosmos/Query/Internal/CosmosQueryContext.cs index 842f2044047..be7e3e41a38 100644 --- a/src/EFCore.Cosmos/Query/Internal/CosmosQueryContext.cs +++ b/src/EFCore.Cosmos/Query/Internal/CosmosQueryContext.cs @@ -14,7 +14,8 @@ namespace Microsoft.EntityFrameworkCore.Cosmos.Query.Internal; /// public class CosmosQueryContext( QueryContextDependencies dependencies, - ICosmosClientWrapper cosmosClient) + ICosmosClientWrapper cosmosClient, + ISessionTokenStorage sessionTokenStorage) : QueryContext(dependencies) { /// @@ -24,4 +25,12 @@ public class CosmosQueryContext( /// doing so can result in application failures when updating to a new Entity Framework Core release. /// public virtual ICosmosClientWrapper CosmosClient { get; } = cosmosClient; + + /// + /// 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. + /// + public virtual ISessionTokenStorage SessionTokenStorage { get; } = sessionTokenStorage; } diff --git a/src/EFCore.Cosmos/Query/Internal/CosmosQueryContextFactory.cs b/src/EFCore.Cosmos/Query/Internal/CosmosQueryContextFactory.cs index 12c83848a98..2aab8b890f1 100644 --- a/src/EFCore.Cosmos/Query/Internal/CosmosQueryContextFactory.cs +++ b/src/EFCore.Cosmos/Query/Internal/CosmosQueryContextFactory.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; @@ -20,6 +21,14 @@ public class CosmosQueryContextFactory( /// protected virtual QueryContextDependencies Dependencies { get; } = dependencies; + /// + /// 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. + /// + protected virtual ISessionTokenStorage SessionTokenStorage { get; } = ((CosmosDatabaseWrapper)dependencies.CurrentContext.Context.GetService()).SessionTokenStorage; + /// /// 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 @@ -27,5 +36,5 @@ public class CosmosQueryContextFactory( /// doing so can result in application failures when updating to a new Entity Framework Core release. /// public virtual QueryContext Create() - => new CosmosQueryContext(Dependencies, cosmosClient); + => new CosmosQueryContext(Dependencies, cosmosClient, SessionTokenStorage); } diff --git a/src/EFCore.Cosmos/Query/Internal/CosmosShapedQueryCompilingExpressionVisitor.PagingQueryingEnumerable.cs b/src/EFCore.Cosmos/Query/Internal/CosmosShapedQueryCompilingExpressionVisitor.PagingQueryingEnumerable.cs index ba5e8c1a13a..536a973d106 100644 --- a/src/EFCore.Cosmos/Query/Internal/CosmosShapedQueryCompilingExpressionVisitor.PagingQueryingEnumerable.cs +++ b/src/EFCore.Cosmos/Query/Internal/CosmosShapedQueryCompilingExpressionVisitor.PagingQueryingEnumerable.cs @@ -35,7 +35,6 @@ private sealed class PagingQueryingEnumerable : IAsyncEnumerable> private readonly IDiagnosticsLogger _queryLogger; private readonly IDiagnosticsLogger _commandLogger; private readonly bool _standAloneStateManager; - private readonly ISessionTokenStorage _sessionTokenStorage; private readonly CancellationToken _cancellationToken; private readonly IConcurrencyDetector _concurrencyDetector; private readonly IExceptionDetector _exceptionDetector; @@ -114,7 +110,6 @@ public AsyncEnumerator(PagingQueryingEnumerable queryingEnumerable, Cancellat _commandLogger = queryingEnumerable._commandLogger; _standAloneStateManager = queryingEnumerable._standAloneStateManager; _exceptionDetector = _cosmosQueryContext.ExceptionDetector; - _sessionTokenStorage = queryingEnumerable._sessionTokenStorage; _cancellationToken = cancellationToken; _concurrencyDetector = queryingEnumerable._threadSafetyChecksEnabled @@ -159,7 +154,7 @@ public async ValueTask MoveNextAsync() queryRequestOptions.PartitionKey = _cosmosPartitionKey; } - queryRequestOptions.SessionToken = _sessionTokenStorage.GetSessionToken(_cosmosContainer); + queryRequestOptions.SessionToken = _cosmosQueryContext.SessionTokenStorage.GetSessionToken(_cosmosContainer); var cosmosClient = _cosmosQueryContext.CosmosClient; _commandLogger.ExecutingSqlQuery(_cosmosContainer, _cosmosPartitionKey, sqlQuery); @@ -171,7 +166,7 @@ public async ValueTask MoveNextAsync() { queryRequestOptions.MaxItemCount = maxItemCount; using var feedIterator = cosmosClient.CreateQuery( - _cosmosContainer, sqlQuery, _sessionTokenStorage, continuationToken, queryRequestOptions); + _cosmosContainer, sqlQuery, _cosmosQueryContext.SessionTokenStorage, continuationToken, queryRequestOptions); using var responseMessage = await feedIterator.ReadNextAsync(_cancellationToken).ConfigureAwait(false); diff --git a/src/EFCore.Cosmos/Query/Internal/CosmosShapedQueryCompilingExpressionVisitor.QueryingEnumerable.cs b/src/EFCore.Cosmos/Query/Internal/CosmosShapedQueryCompilingExpressionVisitor.QueryingEnumerable.cs index 7cafa5b3805..0246153d9b9 100644 --- a/src/EFCore.Cosmos/Query/Internal/CosmosShapedQueryCompilingExpressionVisitor.QueryingEnumerable.cs +++ b/src/EFCore.Cosmos/Query/Internal/CosmosShapedQueryCompilingExpressionVisitor.QueryingEnumerable.cs @@ -32,7 +32,6 @@ private sealed class QueryingEnumerable : IEnumerable, IAsyncEnumerable private readonly IDiagnosticsLogger _queryLogger; private readonly bool _standAloneStateManager; private readonly bool _threadSafetyChecksEnabled; - private readonly ISessionTokenStorage _sessionTokenStorage; public QueryingEnumerable( CosmosQueryContext cosmosQueryContext, @@ -44,8 +43,7 @@ public QueryingEnumerable( IEntityType rootEntityType, List partitionKeyPropertyValues, bool standAloneStateManager, - bool threadSafetyChecksEnabled, - ISessionTokenStorage sessionTokenStorage) + bool threadSafetyChecksEnabled) { _cosmosQueryContext = cosmosQueryContext; _sqlExpressionFactory = sqlExpressionFactory; @@ -56,7 +54,6 @@ public QueryingEnumerable( _queryLogger = cosmosQueryContext.QueryLogger; _standAloneStateManager = standAloneStateManager; _threadSafetyChecksEnabled = threadSafetyChecksEnabled; - _sessionTokenStorage = sessionTokenStorage; _cosmosContainer = rootEntityType.GetContainer() ?? throw new UnreachableException("Root entity type without a Cosmos container."); @@ -115,7 +112,6 @@ private sealed class Enumerator : IEnumerator private readonly bool _standAloneStateManager; private readonly IConcurrencyDetector _concurrencyDetector; private readonly IExceptionDetector _exceptionDetector; - private readonly ISessionTokenStorage _sessionTokenStorage; private IEnumerator _enumerator; @@ -130,7 +126,6 @@ public Enumerator(QueryingEnumerable queryingEnumerable) _queryLogger = queryingEnumerable._queryLogger; _standAloneStateManager = queryingEnumerable._standAloneStateManager; _exceptionDetector = _cosmosQueryContext.ExceptionDetector; - _sessionTokenStorage = queryingEnumerable._sessionTokenStorage; _concurrencyDetector = queryingEnumerable._threadSafetyChecksEnabled ? _cosmosQueryContext.ConcurrencyDetector @@ -155,7 +150,7 @@ public bool MoveNext() EntityFrameworkMetricsData.ReportQueryExecuting(); _enumerator = _cosmosQueryContext.CosmosClient - .ExecuteSqlQuery(_cosmosContainer, _cosmosPartitionKey, sqlQuery, _sessionTokenStorage) + .ExecuteSqlQuery(_cosmosContainer, _cosmosPartitionKey, sqlQuery, _cosmosQueryContext.SessionTokenStorage) .GetEnumerator(); _cosmosQueryContext.InitializeStateManager(_standAloneStateManager); } @@ -204,7 +199,6 @@ private sealed class AsyncEnumerator : IAsyncEnumerator private readonly PartitionKey _cosmosPartitionKey; private readonly IDiagnosticsLogger _queryLogger; private readonly bool _standAloneStateManager; - private readonly ISessionTokenStorage _sessionTokenStorage; private readonly CancellationToken _cancellationToken; private readonly IConcurrencyDetector _concurrencyDetector; private readonly IExceptionDetector _exceptionDetector; @@ -223,7 +217,6 @@ public AsyncEnumerator(QueryingEnumerable queryingEnumerable, CancellationTok _queryLogger = queryingEnumerable._queryLogger; _standAloneStateManager = queryingEnumerable._standAloneStateManager; _exceptionDetector = _cosmosQueryContext.ExceptionDetector; - _sessionTokenStorage = queryingEnumerable._sessionTokenStorage; _cancellationToken = cancellationToken; @@ -247,7 +240,7 @@ public async ValueTask MoveNextAsync() EntityFrameworkMetricsData.ReportQueryExecuting(); _enumerator = _cosmosQueryContext.CosmosClient - .ExecuteSqlQueryAsync(_cosmosContainer, _cosmosPartitionKey, sqlQuery, _sessionTokenStorage) + .ExecuteSqlQueryAsync(_cosmosContainer, _cosmosPartitionKey, sqlQuery, _cosmosQueryContext.SessionTokenStorage) .GetAsyncEnumerator(_cancellationToken); _cosmosQueryContext.InitializeStateManager(_standAloneStateManager); } diff --git a/src/EFCore.Cosmos/Query/Internal/CosmosShapedQueryCompilingExpressionVisitor.ReadItemQueryingEnumerable.cs b/src/EFCore.Cosmos/Query/Internal/CosmosShapedQueryCompilingExpressionVisitor.ReadItemQueryingEnumerable.cs index 714cf3e3ab0..06e6350ba3f 100644 --- a/src/EFCore.Cosmos/Query/Internal/CosmosShapedQueryCompilingExpressionVisitor.ReadItemQueryingEnumerable.cs +++ b/src/EFCore.Cosmos/Query/Internal/CosmosShapedQueryCompilingExpressionVisitor.ReadItemQueryingEnumerable.cs @@ -31,7 +31,6 @@ private sealed class ReadItemQueryingEnumerable : IEnumerable, IAsyncEnume private readonly IDiagnosticsLogger _queryLogger; private readonly bool _standAloneStateManager; private readonly bool _threadSafetyChecksEnabled; - private readonly ISessionTokenStorage _sessionTokenStorage; public ReadItemQueryingEnumerable( CosmosQueryContext cosmosQueryContext, @@ -41,8 +40,7 @@ public ReadItemQueryingEnumerable( Func shaper, Type contextType, bool standAloneStateManager, - bool threadSafetyChecksEnabled, - ISessionTokenStorage sessionTokenStorage) + bool threadSafetyChecksEnabled) { _cosmosQueryContext = cosmosQueryContext; _rootEntityType = rootEntityType; @@ -52,7 +50,6 @@ public ReadItemQueryingEnumerable( _queryLogger = _cosmosQueryContext.QueryLogger; _standAloneStateManager = standAloneStateManager; _threadSafetyChecksEnabled = threadSafetyChecksEnabled; - _sessionTokenStorage = sessionTokenStorage; _cosmosContainer = rootEntityType.GetContainer() ?? throw new UnreachableException("Root entity type without a Cosmos container."); _cosmosPartitionKey = GeneratePartitionKey( @@ -115,7 +112,6 @@ private sealed class Enumerator : IEnumerator, IAsyncEnumerator private readonly IConcurrencyDetector _concurrencyDetector; private readonly IExceptionDetector _exceptionDetector; private readonly ReadItemQueryingEnumerable _readItemEnumerable; - private readonly ISessionTokenStorage _sessionTokenStorage; private readonly CancellationToken _cancellationToken; private JObject _item; @@ -132,7 +128,6 @@ public Enumerator(ReadItemQueryingEnumerable readItemEnumerable, Cancellation _standAloneStateManager = readItemEnumerable._standAloneStateManager; _exceptionDetector = _cosmosQueryContext.ExceptionDetector; _readItemEnumerable = readItemEnumerable; - _sessionTokenStorage = readItemEnumerable._sessionTokenStorage; _cancellationToken = cancellationToken; _concurrencyDetector = readItemEnumerable._threadSafetyChecksEnabled @@ -167,7 +162,7 @@ public bool MoveNext() _cosmosContainer, _cosmosPartitionKey, resourceId, - _sessionTokenStorage); + _cosmosQueryContext.SessionTokenStorage); return ShapeResult(); } @@ -208,7 +203,7 @@ public async ValueTask MoveNextAsync() _cosmosContainer, _cosmosPartitionKey, resourceId, - _sessionTokenStorage, + _cosmosQueryContext.SessionTokenStorage, _cancellationToken) .ConfigureAwait(false); diff --git a/src/EFCore.Cosmos/Query/Internal/CosmosShapedQueryCompilingExpressionVisitor.cs b/src/EFCore.Cosmos/Query/Internal/CosmosShapedQueryCompilingExpressionVisitor.cs index a8e8b0cb5d4..f79539eb0b7 100644 --- a/src/EFCore.Cosmos/Query/Internal/CosmosShapedQueryCompilingExpressionVisitor.cs +++ b/src/EFCore.Cosmos/Query/Internal/CosmosShapedQueryCompilingExpressionVisitor.cs @@ -83,7 +83,6 @@ protected override Expression VisitShapedQuery(ShapedQueryExpression shapedQuery var threadSafetyConstant = Constant(_threadSafetyChecksEnabled); var standAloneStateManagerConstant = Constant( QueryCompilationContext.QueryTrackingBehavior == QueryTrackingBehavior.NoTrackingWithIdentityResolution); - var sessionTokenStorageConstant = Constant(cosmosQueryCompilationContext.SessionTokenStorage); Check.DebugAssert(!paging || selectExpression.ReadItemInfo is null, "ReadItem is being with paging, impossible."); @@ -98,8 +97,7 @@ protected override Expression VisitShapedQuery(ShapedQueryExpression shapedQuery shaperConstant, contextTypeConstant, standAloneStateManagerConstant, - threadSafetyConstant, - sessionTokenStorageConstant), + threadSafetyConstant), _ when paging => New( typeof(PagingQueryingEnumerable<>).MakeGenericType(shaperLambda.ReturnType).GetConstructors()[0], @@ -115,8 +113,7 @@ protected override Expression VisitShapedQuery(ShapedQueryExpression shapedQuery threadSafetyConstant, Constant(maxItemCount.Name), Constant(continuationToken.Name), - Constant(responseContinuationTokenLimitInKb.Name), - sessionTokenStorageConstant), + Constant(responseContinuationTokenLimitInKb.Name)), _ => New( typeof(QueryingEnumerable<>).MakeGenericType(shaperLambda.ReturnType).GetConstructors()[0], cosmosQueryContextConstant, @@ -128,8 +125,7 @@ protected override Expression VisitShapedQuery(ShapedQueryExpression shapedQuery rootEntityTypeConstant, Constant(cosmosQueryCompilationContext.PartitionKeyPropertyValues), standAloneStateManagerConstant, - threadSafetyConstant, - sessionTokenStorageConstant) + threadSafetyConstant) }; } diff --git a/test/EFCore.Cosmos.FunctionalTests/CosmosSessionTokensTest.cs b/test/EFCore.Cosmos.FunctionalTests/CosmosSessionTokensTest.cs index 064c59784cf..14638cc8f7a 100644 --- a/test/EFCore.Cosmos.FunctionalTests/CosmosSessionTokensTest.cs +++ b/test/EFCore.Cosmos.FunctionalTests/CosmosSessionTokensTest.cs @@ -409,6 +409,7 @@ public virtual async Task Query_on_new_context_does_not_use_same_session_token(b } using var newContext = contextFactory.CreateContext(); + Assert.NotSame(context, newContext); if (async) { if (defaultContainer) From 5c8063712dc5aeb22b7ed4be48550b1edef7cbd2 Mon Sep 17 00:00:00 2001 From: JoasE <32096708+JoasE@users.noreply.github.com> Date: Sat, 25 Oct 2025 12:57:00 +0200 Subject: [PATCH 23/37] Add todo --- test/EFCore.Cosmos.FunctionalTests/CosmosSessionTokensTest.cs | 3 +++ 1 file changed, 3 insertions(+) diff --git a/test/EFCore.Cosmos.FunctionalTests/CosmosSessionTokensTest.cs b/test/EFCore.Cosmos.FunctionalTests/CosmosSessionTokensTest.cs index 14638cc8f7a..8e566d54005 100644 --- a/test/EFCore.Cosmos.FunctionalTests/CosmosSessionTokensTest.cs +++ b/test/EFCore.Cosmos.FunctionalTests/CosmosSessionTokensTest.cs @@ -802,6 +802,9 @@ public virtual async Task Query_sets_session_token(bool async, bool defaultConta Assert.False(string.IsNullOrWhiteSpace(sessionToken)); } + // @TODO set then query appends + // @TODO append then query appends only if not present? + [ConditionalTheory] [InlineData(true, true)] [InlineData(true, false)] From 1f8d6a999d85c2409faa181f220bf72a10c8e9ff Mon Sep 17 00:00:00 2001 From: JoasE <32096708+JoasE@users.noreply.github.com> Date: Sun, 26 Oct 2025 14:57:26 +0100 Subject: [PATCH 24/37] Move to non-parsing but only composing session tokens --- .../CosmosDatabaseFacadeExtensions.cs | 58 +- .../Query/Internal/CosmosQueryContext.cs | 4 +- .../Internal/CosmosQueryContextFactory.cs | 2 +- .../Storage/ISessionTokenStorage.cs | 54 - .../Storage/Internal/CosmosClientWrapper.cs | 68 +- .../Storage/Internal/CosmosDatabaseWrapper.cs | 2 +- .../Storage/Internal/ICosmosClientWrapper.cs | 26 +- .../Storage/Internal/SessionTokenStorage.cs | 53 +- .../CosmosSessionTokensTest.cs | 1253 +++++++---------- .../TestUtilities/CosmosTestStore.cs | 2 +- .../TestUtilities/NullSessionTokenStorage.cs | 71 +- 11 files changed, 682 insertions(+), 911 deletions(-) delete mode 100644 src/EFCore.Cosmos/Storage/ISessionTokenStorage.cs diff --git a/src/EFCore.Cosmos/Extensions/CosmosDatabaseFacadeExtensions.cs b/src/EFCore.Cosmos/Extensions/CosmosDatabaseFacadeExtensions.cs index 5df9cb0c80b..b342720700f 100644 --- a/src/EFCore.Cosmos/Extensions/CosmosDatabaseFacadeExtensions.cs +++ b/src/EFCore.Cosmos/Extensions/CosmosDatabaseFacadeExtensions.cs @@ -27,11 +27,29 @@ public static CosmosClient GetCosmosClient(this DatabaseFacade databaseFacade) => GetService(databaseFacade).Client; /// - /// Gets used to manage the session tokens for this . + /// Gets the composite session token for the default container for this . /// + /// Use this when using only 1 container in the same /// The for the context. - /// The . - public static ISessionTokenStorage GetSessionTokens(this DatabaseFacade databaseFacade) + /// The session token dictionary. + public static string? GetSessionToken(this DatabaseFacade databaseFacade) + { + var db = GetService(databaseFacade); + if (db is not CosmosDatabaseWrapper dbWrapper) + { + throw new InvalidOperationException(CosmosStrings.CosmosNotInUse); + } + + return dbWrapper.SessionTokenStorage.GetSessionToken(); + } + + /// + /// Gets a dictionary that contains the composite session token per container for this . + /// + /// Use this when using multiple containers in the same + /// The for the context. + /// The session token dictionary. + public static IReadOnlyDictionary GetSessionTokens(this DatabaseFacade databaseFacade) { var db = GetService(databaseFacade); if (db is not CosmosDatabaseWrapper dbWrapper) @@ -42,6 +60,40 @@ public static ISessionTokenStorage GetSessionTokens(this DatabaseFacade database return dbWrapper.SessionTokenStorage; } + /// + /// Appends the composite session token for the default container for this . + /// + /// Use this when using only 1 container in the same + /// The for the context. + /// The session token to append. + public static void AppendSessionToken(this DatabaseFacade databaseFacade, string sessionToken) + { + var db = GetService(databaseFacade); + if (db is not CosmosDatabaseWrapper dbWrapper) + { + throw new InvalidOperationException(CosmosStrings.CosmosNotInUse); + } + + dbWrapper.SessionTokenStorage.AppendSessionToken(sessionToken); + } + + /// + /// Appends the composite session token per container for this . + /// + /// Use this when using multiple containers in the same + /// The for the context. + /// The session tokens to append per container. + public static void AppendSessionTokens(this DatabaseFacade databaseFacade, IReadOnlyDictionary sessionTokens) + { + var db = GetService(databaseFacade); + if (db is not CosmosDatabaseWrapper dbWrapper) + { + throw new InvalidOperationException(CosmosStrings.CosmosNotInUse); + } + + dbWrapper.SessionTokenStorage.AppendSessionTokens(sessionTokens); + } + private static TService GetService(IInfrastructure databaseFacade) where TService : class { diff --git a/src/EFCore.Cosmos/Query/Internal/CosmosQueryContext.cs b/src/EFCore.Cosmos/Query/Internal/CosmosQueryContext.cs index be7e3e41a38..91ee46112a8 100644 --- a/src/EFCore.Cosmos/Query/Internal/CosmosQueryContext.cs +++ b/src/EFCore.Cosmos/Query/Internal/CosmosQueryContext.cs @@ -15,7 +15,7 @@ namespace Microsoft.EntityFrameworkCore.Cosmos.Query.Internal; public class CosmosQueryContext( QueryContextDependencies dependencies, ICosmosClientWrapper cosmosClient, - ISessionTokenStorage sessionTokenStorage) + SessionTokenStorage sessionTokenStorage) : QueryContext(dependencies) { /// @@ -32,5 +32,5 @@ public class CosmosQueryContext( /// 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. /// - public virtual ISessionTokenStorage SessionTokenStorage { get; } = sessionTokenStorage; + public virtual SessionTokenStorage SessionTokenStorage { get; } = sessionTokenStorage; } diff --git a/src/EFCore.Cosmos/Query/Internal/CosmosQueryContextFactory.cs b/src/EFCore.Cosmos/Query/Internal/CosmosQueryContextFactory.cs index 2aab8b890f1..ec9c6bf1558 100644 --- a/src/EFCore.Cosmos/Query/Internal/CosmosQueryContextFactory.cs +++ b/src/EFCore.Cosmos/Query/Internal/CosmosQueryContextFactory.cs @@ -27,7 +27,7 @@ public class CosmosQueryContextFactory( /// 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. /// - protected virtual ISessionTokenStorage SessionTokenStorage { get; } = ((CosmosDatabaseWrapper)dependencies.CurrentContext.Context.GetService()).SessionTokenStorage; + protected virtual SessionTokenStorage SessionTokenStorage { get; } = ((CosmosDatabaseWrapper)dependencies.CurrentContext.Context.GetService()).SessionTokenStorage; /// /// This is an internal API that supports the Entity Framework Core infrastructure and not subject to diff --git a/src/EFCore.Cosmos/Storage/ISessionTokenStorage.cs b/src/EFCore.Cosmos/Storage/ISessionTokenStorage.cs deleted file mode 100644 index 8536614b108..00000000000 --- a/src/EFCore.Cosmos/Storage/ISessionTokenStorage.cs +++ /dev/null @@ -1,54 +0,0 @@ -// 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.Storage; - -/// -/// Defines methods for managing session tokens used in a . -/// -public interface ISessionTokenStorage -{ - /// - /// Appends or merges the session token specified to any composite already stored in the storage for the default container. - /// - /// The session token to append or merge. - void AppendSessionToken(string sessionToken); - - /// - /// Appends or merges the session token specified to any composite already stored in the storage for the specified container. - /// - /// The name of the container to append the session token for. - /// The session token to append or merge. - void AppendSessionToken(string containerName, string sessionToken); - - /// - /// Gets the composite session token for the default container. - /// - /// The composite session token, or if none is stored. - string? GetSessionToken(); - - /// - /// Gets the composite session token for the specified container. - /// - /// The name of the container to get the session token for. - /// The composite session token, or if none is stored. - string? GetSessionToken(string containerName); - - /// - /// Overwrites the session token for the container. - /// - /// The name of the container to set the session token for. - /// The session token to set, or to clear any stored token. - void SetSessionToken(string containerName, string? sessionToken); - - /// - /// Overwrites the session token for the default container. - /// - /// The session token to set, or to clear any stored token. - void SetSessionToken(string? sessionToken); - - /// - /// Clears all stored session tokens. - /// - void Clear(); -} diff --git a/src/EFCore.Cosmos/Storage/Internal/CosmosClientWrapper.cs b/src/EFCore.Cosmos/Storage/Internal/CosmosClientWrapper.cs index 0c04a201e71..7641d176f12 100644 --- a/src/EFCore.Cosmos/Storage/Internal/CosmosClientWrapper.cs +++ b/src/EFCore.Cosmos/Storage/Internal/CosmosClientWrapper.cs @@ -389,7 +389,7 @@ public virtual bool CreateItem( string containerId, JToken document, IUpdateEntry entry, - ISessionTokenStorage sessionTokenStorage) + SessionTokenStorage sessionTokenStorage) { _databaseLogger.SyncNotSupported(); @@ -398,7 +398,7 @@ public virtual bool CreateItem( private static bool CreateItemOnce( DbContext context, - (string ContainerId, JToken Document, IUpdateEntry Entry, ISessionTokenStorage SessionTokenStorage, CosmosClientWrapper Wrapper) parameters) + (string ContainerId, JToken Document, IUpdateEntry Entry, SessionTokenStorage SessionTokenStorage, CosmosClientWrapper Wrapper) parameters) => CreateItemOnceAsync(context, parameters).GetAwaiter().GetResult(); /// @@ -411,13 +411,13 @@ public virtual Task CreateItemAsync( string containerId, JToken document, IUpdateEntry updateEntry, - ISessionTokenStorage sessionTokenStorage, + SessionTokenStorage sessionTokenStorage, CancellationToken cancellationToken = default) => _executionStrategy.ExecuteAsync((containerId, document, updateEntry, sessionTokenStorage, this), CreateItemOnceAsync, null, cancellationToken); private static async Task CreateItemOnceAsync( DbContext _, - (string ContainerId, JToken Document, IUpdateEntry Entry, ISessionTokenStorage SessionTokenStorage, CosmosClientWrapper Wrapper) parameters, + (string ContainerId, JToken Document, IUpdateEntry Entry, SessionTokenStorage SessionTokenStorage, CosmosClientWrapper Wrapper) parameters, CancellationToken cancellationToken = default) { using var stream = Serialize(parameters.Document); @@ -475,7 +475,7 @@ public virtual bool ReplaceItem( string documentId, JObject document, IUpdateEntry entry, - ISessionTokenStorage sessionTokenStorage) + SessionTokenStorage sessionTokenStorage) { _databaseLogger.SyncNotSupported(); @@ -484,7 +484,7 @@ public virtual bool ReplaceItem( private static bool ReplaceItemOnce( DbContext context, - (string ContainerId, string ItemId, JObject Document, IUpdateEntry Entry, ISessionTokenStorage SessionTokenStorage, CosmosClientWrapper Wrapper) parameters) + (string ContainerId, string ItemId, JObject Document, IUpdateEntry Entry, SessionTokenStorage SessionTokenStorage, CosmosClientWrapper Wrapper) parameters) => ReplaceItemOnceAsync(context, parameters).GetAwaiter().GetResult(); /// @@ -498,14 +498,14 @@ public virtual Task ReplaceItemAsync( string documentId, JObject document, IUpdateEntry updateEntry, - ISessionTokenStorage sessionTokenStorage, + SessionTokenStorage sessionTokenStorage, CancellationToken cancellationToken = default) => _executionStrategy.ExecuteAsync( (collectionId, documentId, document, updateEntry, sessionTokenStorage, this), ReplaceItemOnceAsync, null, cancellationToken); private static async Task ReplaceItemOnceAsync( DbContext _, - (string ContainerId, string ResourceId, JObject Document, IUpdateEntry Entry, ISessionTokenStorage SessionTokenStorage, CosmosClientWrapper Wrapper) parameters, + (string ContainerId, string ResourceId, JObject Document, IUpdateEntry Entry, SessionTokenStorage SessionTokenStorage, CosmosClientWrapper Wrapper) parameters, CancellationToken cancellationToken = default) { using var stream = Serialize(parameters.Document); @@ -563,7 +563,7 @@ public virtual bool DeleteItem( string containerId, string documentId, IUpdateEntry entry, - ISessionTokenStorage sessionTokenStorage) + SessionTokenStorage sessionTokenStorage) { _databaseLogger.SyncNotSupported(); @@ -572,7 +572,7 @@ public virtual bool DeleteItem( private static bool DeleteItemOnce( DbContext context, - (string ContainerId, string DocumentId, IUpdateEntry Entry, ISessionTokenStorage SessionTokenStorage, CosmosClientWrapper Wrapper) parameters) + (string ContainerId, string DocumentId, IUpdateEntry Entry, SessionTokenStorage SessionTokenStorage, CosmosClientWrapper Wrapper) parameters) => DeleteItemOnceAsync(context, parameters).GetAwaiter().GetResult(); /// @@ -585,13 +585,13 @@ public virtual Task DeleteItemAsync( string containerId, string documentId, IUpdateEntry entry, - ISessionTokenStorage sessionTokenStorage, + SessionTokenStorage sessionTokenStorage, CancellationToken cancellationToken = default) => _executionStrategy.ExecuteAsync((containerId, documentId, entry, sessionTokenStorage, this), DeleteItemOnceAsync, null, cancellationToken); private static async Task DeleteItemOnceAsync( DbContext? _, - (string ContainerId, string ResourceId, IUpdateEntry Entry, ISessionTokenStorage SessionTokenStorage, CosmosClientWrapper Wrapper) parameters, + (string ContainerId, string ResourceId, IUpdateEntry Entry, SessionTokenStorage SessionTokenStorage, CosmosClientWrapper Wrapper) parameters, CancellationToken cancellationToken = default) { var containerId = parameters.ContainerId; @@ -667,7 +667,7 @@ public virtual ICosmosTransactionalBatchWrapper CreateTransactionalBatch(string /// 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. /// - public virtual CosmosTransactionalBatchResult ExecuteTransactionalBatch(ICosmosTransactionalBatchWrapper batch, ISessionTokenStorage sessionTokenStorage) + public virtual CosmosTransactionalBatchResult ExecuteTransactionalBatch(ICosmosTransactionalBatchWrapper batch, SessionTokenStorage sessionTokenStorage) { _databaseLogger.SyncNotSupported(); @@ -675,7 +675,7 @@ public virtual CosmosTransactionalBatchResult ExecuteTransactionalBatch(ICosmosT } private static CosmosTransactionalBatchResult ExecuteBatchOnce(DbContext _, - (ICosmosTransactionalBatchWrapper Batch, ISessionTokenStorage SessionTokenStorage, CosmosClientWrapper Wrapper) parameters) + (ICosmosTransactionalBatchWrapper Batch, SessionTokenStorage SessionTokenStorage, CosmosClientWrapper Wrapper) parameters) => ExecuteBatchOnceAsync(_, parameters).GetAwaiter().GetResult(); /// @@ -684,11 +684,11 @@ private static CosmosTransactionalBatchResult ExecuteBatchOnce(DbContext _, /// 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. /// - public virtual Task ExecuteTransactionalBatchAsync(ICosmosTransactionalBatchWrapper batch, ISessionTokenStorage sessionTokenStorage, CancellationToken cancellationToken = default) + public virtual Task ExecuteTransactionalBatchAsync(ICosmosTransactionalBatchWrapper batch, SessionTokenStorage sessionTokenStorage, CancellationToken cancellationToken = default) => _executionStrategy.ExecuteAsync((batch, sessionTokenStorage, this), ExecuteBatchOnceAsync, null, cancellationToken); private static async Task ExecuteBatchOnceAsync(DbContext _, - (ICosmosTransactionalBatchWrapper Batch, ISessionTokenStorage SessionTokenStorage, CosmosClientWrapper Wrapper) parameters, + (ICosmosTransactionalBatchWrapper Batch, SessionTokenStorage SessionTokenStorage, CosmosClientWrapper Wrapper) parameters, CancellationToken cancellationToken = default) { var batch = parameters.Batch; @@ -780,7 +780,7 @@ private static PartitionKey ExtractPartitionKeyValue(IUpdateEntry entry) return builder.Build(); } - private static void ProcessResponse(string containerId, ResponseMessage response, IUpdateEntry entry, ISessionTokenStorage sessionTokenStorage) + private static void ProcessResponse(string containerId, ResponseMessage response, IUpdateEntry entry, SessionTokenStorage sessionTokenStorage) { response.EnsureSuccessStatusCode(); @@ -792,7 +792,7 @@ private static void ProcessResponse(string containerId, ResponseMessage response ProcessResponse(entry, response.Headers.ETag, response.Content); } - private static void ProcessResponse(string containerId, TransactionalBatchResponse batchResponse, IReadOnlyList entries, ISessionTokenStorage sessionTokenStorage) + private static void ProcessResponse(string containerId, TransactionalBatchResponse batchResponse, IReadOnlyList entries, SessionTokenStorage sessionTokenStorage) { if (!string.IsNullOrWhiteSpace(batchResponse.Headers.Session)) { @@ -840,7 +840,7 @@ public virtual IEnumerable ExecuteSqlQuery( string containerId, PartitionKey partitionKeyValue, CosmosSqlQuery query, - ISessionTokenStorage sessionTokenStorage) + SessionTokenStorage sessionTokenStorage) { _databaseLogger.SyncNotSupported(); @@ -859,7 +859,7 @@ public virtual IAsyncEnumerable ExecuteSqlQueryAsync( string containerId, PartitionKey partitionKeyValue, CosmosSqlQuery query, - ISessionTokenStorage sessionTokenStorage) + SessionTokenStorage sessionTokenStorage) { _commandLogger.ExecutingSqlQuery(containerId, partitionKeyValue, query); @@ -876,7 +876,7 @@ public virtual IAsyncEnumerable ExecuteSqlQueryAsync( string containerId, PartitionKey partitionKeyValue, string resourceId, - ISessionTokenStorage sessionTokenStorage) + SessionTokenStorage sessionTokenStorage) { _databaseLogger.SyncNotSupported(); @@ -905,7 +905,7 @@ public virtual IAsyncEnumerable ExecuteSqlQueryAsync( string containerId, PartitionKey partitionKeyValue, string resourceId, - ISessionTokenStorage sessionTokenStorage, + SessionTokenStorage sessionTokenStorage, CancellationToken cancellationToken = default) { _commandLogger.ExecutingReadItem(containerId, partitionKeyValue, resourceId); @@ -930,12 +930,12 @@ public virtual IAsyncEnumerable ExecuteSqlQueryAsync( private static ResponseMessage CreateSingleItemQuery( DbContext? context, - (string ContainerId, PartitionKey PartitionKeyValue, string ResourceId, ISessionTokenStorage SessionTokenStorage, CosmosClientWrapper Wrapper) parameters) + (string ContainerId, PartitionKey PartitionKeyValue, string ResourceId, SessionTokenStorage SessionTokenStorage, CosmosClientWrapper Wrapper) parameters) => CreateSingleItemQueryAsync(context, parameters).GetAwaiter().GetResult(); private static async Task CreateSingleItemQueryAsync( DbContext? _, - (string ContainerId, PartitionKey PartitionKeyValue, string ResourceId, ISessionTokenStorage SessionTokenStorage, CosmosClientWrapper Wrapper) parameters, + (string ContainerId, PartitionKey PartitionKeyValue, string ResourceId, SessionTokenStorage SessionTokenStorage, CosmosClientWrapper Wrapper) parameters, CancellationToken cancellationToken = default) { var (containerId, partitionKeyValue, resourceId, sessionTokenStorage, wrapper) = parameters; @@ -986,7 +986,7 @@ private static async Task CreateSingleItemQueryAsync( public virtual FeedIterator CreateQuery( string containerId, CosmosSqlQuery query, - ISessionTokenStorage sessionTokenStorage, + SessionTokenStorage sessionTokenStorage, string? continuationToken = null, QueryRequestOptions? queryRequestOptions = null) { @@ -1029,14 +1029,14 @@ private sealed class DocumentEnumerable( string containerId, PartitionKey partitionKeyValue, CosmosSqlQuery cosmosSqlQuery, - ISessionTokenStorage sessionTokenStorage) + SessionTokenStorage sessionTokenStorage) : IEnumerable { private readonly CosmosClientWrapper _cosmosClient = cosmosClient; private readonly string _containerId = containerId; private readonly PartitionKey _partitionKeyValue = partitionKeyValue; private readonly CosmosSqlQuery _cosmosSqlQuery = cosmosSqlQuery; - private readonly ISessionTokenStorage _sessionTokenStorage = sessionTokenStorage; + private readonly SessionTokenStorage _sessionTokenStorage = sessionTokenStorage; public IEnumerator GetEnumerator() => new Enumerator(this); @@ -1050,7 +1050,7 @@ private sealed class Enumerator(DocumentEnumerable documentEnumerable) : IEnumer private readonly string _containerId = documentEnumerable._containerId; private readonly PartitionKey _partitionKeyValue = documentEnumerable._partitionKeyValue; private readonly CosmosSqlQuery _cosmosSqlQuery = documentEnumerable._cosmosSqlQuery; - private readonly ISessionTokenStorage _sessionTokenStorage = documentEnumerable._sessionTokenStorage; + private readonly SessionTokenStorage _sessionTokenStorage = documentEnumerable._sessionTokenStorage; private JToken? _current; private ResponseMessage? _responseMessage; @@ -1141,14 +1141,14 @@ private sealed class DocumentAsyncEnumerable( string containerId, PartitionKey partitionKeyValue, CosmosSqlQuery cosmosSqlQuery, - ISessionTokenStorage sessionTokenStorage) + SessionTokenStorage sessionTokenStorage) : IAsyncEnumerable { private readonly CosmosClientWrapper _cosmosClient = cosmosClient; private readonly string _containerId = containerId; private readonly PartitionKey _partitionKeyValue = partitionKeyValue; private readonly CosmosSqlQuery _cosmosSqlQuery = cosmosSqlQuery; - private readonly ISessionTokenStorage _sessionTokenStorage = sessionTokenStorage; + private readonly SessionTokenStorage _sessionTokenStorage = sessionTokenStorage; public IAsyncEnumerator GetAsyncEnumerator(CancellationToken cancellationToken = default) => new AsyncEnumerator(this, cancellationToken); @@ -1160,7 +1160,7 @@ private sealed class AsyncEnumerator(DocumentAsyncEnumerable documentEnumerable, private readonly string _containerId = documentEnumerable._containerId; private readonly PartitionKey _partitionKeyValue = documentEnumerable._partitionKeyValue; private readonly CosmosSqlQuery _cosmosSqlQuery = documentEnumerable._cosmosSqlQuery; - private readonly ISessionTokenStorage _sessionTokenStorage = documentEnumerable._sessionTokenStorage; + private readonly SessionTokenStorage _sessionTokenStorage = documentEnumerable._sessionTokenStorage; private JToken? _current; private ResponseMessage? _responseMessage; @@ -1354,9 +1354,9 @@ private sealed class CosmosFeedIteratorWrapper : FeedIterator { private readonly FeedIterator _inner; private readonly string _containerName; - private readonly ISessionTokenStorage _sessionTokenStorage; + private readonly SessionTokenStorage _sessionTokenStorage; - public CosmosFeedIteratorWrapper(FeedIterator inner, string containerName, ISessionTokenStorage sessionTokenStorage) + public CosmosFeedIteratorWrapper(FeedIterator inner, string containerName, SessionTokenStorage sessionTokenStorage) { _inner = inner; _containerName = containerName; @@ -1370,7 +1370,7 @@ public override async Task ReadNextAsync(CancellationToken canc var response = await _inner.ReadNextAsync(cancellationToken).ConfigureAwait(false); if (!string.IsNullOrWhiteSpace(response.Headers.Session)) { - _sessionTokenStorage.SetSessionToken(_containerName, response.Headers.Session); + _sessionTokenStorage.AppendSessionToken(_containerName, response.Headers.Session); } return response; } diff --git a/src/EFCore.Cosmos/Storage/Internal/CosmosDatabaseWrapper.cs b/src/EFCore.Cosmos/Storage/Internal/CosmosDatabaseWrapper.cs index 4635f1b5f7b..f7f1fbddcf6 100644 --- a/src/EFCore.Cosmos/Storage/Internal/CosmosDatabaseWrapper.cs +++ b/src/EFCore.Cosmos/Storage/Internal/CosmosDatabaseWrapper.cs @@ -58,7 +58,7 @@ public CosmosDatabaseWrapper( /// 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. /// - public virtual ISessionTokenStorage SessionTokenStorage { get; } + public virtual SessionTokenStorage SessionTokenStorage { get; } /// /// This is an internal API that supports the Entity Framework Core infrastructure and not subject to diff --git a/src/EFCore.Cosmos/Storage/Internal/ICosmosClientWrapper.cs b/src/EFCore.Cosmos/Storage/Internal/ICosmosClientWrapper.cs index 620fd3b054b..54e2cea7a63 100644 --- a/src/EFCore.Cosmos/Storage/Internal/ICosmosClientWrapper.cs +++ b/src/EFCore.Cosmos/Storage/Internal/ICosmosClientWrapper.cs @@ -67,7 +67,7 @@ public interface ICosmosClientWrapper /// 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. /// - bool CreateItem(string containerId, JToken document, IUpdateEntry entry, ISessionTokenStorage sessionTokenStorage); + bool CreateItem(string containerId, JToken document, IUpdateEntry entry, SessionTokenStorage sessionTokenStorage); /// /// This is an internal API that supports the Entity Framework Core infrastructure and not subject to @@ -80,7 +80,7 @@ bool ReplaceItem( string documentId, JObject document, IUpdateEntry entry, - ISessionTokenStorage sessionTokenStorage); + SessionTokenStorage sessionTokenStorage); /// /// This is an internal API that supports the Entity Framework Core infrastructure and not subject to @@ -92,7 +92,7 @@ bool DeleteItem( string containerId, string documentId, IUpdateEntry entry, - ISessionTokenStorage sessionTokenStorage); + SessionTokenStorage sessionTokenStorage); /// /// This is an internal API that supports the Entity Framework Core infrastructure and not subject to @@ -104,7 +104,7 @@ Task CreateItemAsync( string containerId, JToken document, IUpdateEntry updateEntry, - ISessionTokenStorage sessionTokenStorage, + SessionTokenStorage sessionTokenStorage, CancellationToken cancellationToken = default); /// @@ -118,7 +118,7 @@ Task ReplaceItemAsync( string documentId, JObject document, IUpdateEntry updateEntry, - ISessionTokenStorage sessionTokenStorage, + SessionTokenStorage sessionTokenStorage, CancellationToken cancellationToken = default); /// @@ -131,7 +131,7 @@ Task DeleteItemAsync( string containerId, string documentId, IUpdateEntry entry, - ISessionTokenStorage sessionTokenStorage, + SessionTokenStorage sessionTokenStorage, CancellationToken cancellationToken = default); /// @@ -143,7 +143,7 @@ Task DeleteItemAsync( FeedIterator CreateQuery( string containerId, CosmosSqlQuery query, - ISessionTokenStorage sessionTokenStorage, + SessionTokenStorage sessionTokenStorage, string? continuationToken = null, QueryRequestOptions? queryRequestOptions = null); @@ -157,7 +157,7 @@ FeedIterator CreateQuery( string containerId, PartitionKey partitionKeyValue, string resourceId, - ISessionTokenStorage sessionTokenStorage); + SessionTokenStorage sessionTokenStorage); /// /// This is an internal API that supports the Entity Framework Core infrastructure and not subject to @@ -169,7 +169,7 @@ FeedIterator CreateQuery( string containerId, PartitionKey partitionKeyValue, string resourceId, - ISessionTokenStorage sessionTokenStorage, + SessionTokenStorage sessionTokenStorage, CancellationToken cancellationToken = default); /// @@ -182,7 +182,7 @@ IEnumerable ExecuteSqlQuery( string containerId, PartitionKey partitionKeyValue, CosmosSqlQuery query, - ISessionTokenStorage sessionTokenStorage); + SessionTokenStorage sessionTokenStorage); /// /// This is an internal API that supports the Entity Framework Core infrastructure and not subject to @@ -194,7 +194,7 @@ IAsyncEnumerable ExecuteSqlQueryAsync( string containerId, PartitionKey partitionKeyValue, CosmosSqlQuery query, - ISessionTokenStorage sessionTokenStorage); + SessionTokenStorage sessionTokenStorage); /// /// This is an internal API that supports the Entity Framework Core infrastructure and not subject to @@ -226,7 +226,7 @@ IAsyncEnumerable ExecuteSqlQueryAsync( /// 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. /// - Task ExecuteTransactionalBatchAsync(ICosmosTransactionalBatchWrapper batch, ISessionTokenStorage sessionTokenStorage, CancellationToken cancellationToken = default); + Task ExecuteTransactionalBatchAsync(ICosmosTransactionalBatchWrapper batch, SessionTokenStorage sessionTokenStorage, CancellationToken cancellationToken = default); /// /// This is an internal API that supports the Entity Framework Core infrastructure and not subject to @@ -234,5 +234,5 @@ IAsyncEnumerable ExecuteSqlQueryAsync( /// 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. /// - CosmosTransactionalBatchResult ExecuteTransactionalBatch(ICosmosTransactionalBatchWrapper batch, ISessionTokenStorage sessionTokenStorage); + CosmosTransactionalBatchResult ExecuteTransactionalBatch(ICosmosTransactionalBatchWrapper batch, SessionTokenStorage sessionTokenStorage); } diff --git a/src/EFCore.Cosmos/Storage/Internal/SessionTokenStorage.cs b/src/EFCore.Cosmos/Storage/Internal/SessionTokenStorage.cs index 6434d2b2e72..fccbb1d920f 100644 --- a/src/EFCore.Cosmos/Storage/Internal/SessionTokenStorage.cs +++ b/src/EFCore.Cosmos/Storage/Internal/SessionTokenStorage.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 System.Collections; using System.Runtime.CompilerServices; using System.Runtime.InteropServices; using Microsoft.EntityFrameworkCore.Cosmos.Internal; @@ -14,7 +15,7 @@ namespace Microsoft.EntityFrameworkCore.Cosmos.Storage.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. /// -public class SessionTokenStorage : ISessionTokenStorage +public class SessionTokenStorage : IReadOnlyDictionary { private readonly Dictionary _containerSessionTokens; private readonly string _defaultContainerName; @@ -40,17 +41,7 @@ public SessionTokenStorage(DbContext dbContext) /// 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. /// - public virtual string? GetSessionToken() - => GetSessionToken(_defaultContainerName); - - /// - /// 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. - /// - public virtual void SetSessionToken(string? sessionToken) - => SetSessionToken(_defaultContainerName, sessionToken); + public virtual string? GetSessionToken() => GetSessionToken(_defaultContainerName); /// /// This is an internal API that supports the Entity Framework Core infrastructure and not subject to @@ -79,6 +70,22 @@ public virtual void AppendSessionToken(string sessionToken) return value.ConvertToString(); } + /// + /// 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. + /// + public virtual void AppendSessionTokens(IReadOnlyDictionary sessionTokens) + { + ArgumentNullException.ThrowIfNull(sessionTokens, nameof(sessionTokens)); + foreach (var (containerName, sessionToken) in sessionTokens) + { + // @TODO: null checks? + AppendSessionToken(containerName, sessionToken); + } + } + /// /// 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 @@ -171,4 +178,26 @@ public void Add(string token) return _string; } } + + + #region IReadOnlyDictionary + IEnumerable IReadOnlyDictionary.Keys => _containerSessionTokens.Keys; + + IEnumerable IReadOnlyDictionary.Values => _containerSessionTokens.Values.Select(v => v.ConvertToString()); + + int IReadOnlyCollection>.Count => _containerSessionTokens.Count; + + string? IReadOnlyDictionary.this[string key] => _containerSessionTokens[key].ConvertToString(); + + bool IReadOnlyDictionary.ContainsKey(string key) => _containerSessionTokens.ContainsKey(key); + + bool IReadOnlyDictionary.TryGetValue(string key, out string? value) + => _containerSessionTokens.TryGetValue(key, out var compositeSessionToken) ? + (value = compositeSessionToken.ConvertToString()) != null + : (value = null) == null; + + IEnumerator> IEnumerable>.GetEnumerator() => _containerSessionTokens.Select(x => new KeyValuePair(x.Key, x.Value.ConvertToString())).GetEnumerator(); + + IEnumerator IEnumerable.GetEnumerator() => ((IEnumerable>)this).GetEnumerator(); + #endregion } diff --git a/test/EFCore.Cosmos.FunctionalTests/CosmosSessionTokensTest.cs b/test/EFCore.Cosmos.FunctionalTests/CosmosSessionTokensTest.cs index 8e566d54005..e61cd871746 100644 --- a/test/EFCore.Cosmos.FunctionalTests/CosmosSessionTokensTest.cs +++ b/test/EFCore.Cosmos.FunctionalTests/CosmosSessionTokensTest.cs @@ -17,135 +17,91 @@ protected override ITestStoreFactory TestStoreFactory protected override DbContextOptionsBuilder AddOptions(DbContextOptionsBuilder builder) => base.AddOptions(builder).ConfigureWarnings(x => x.Ignore(CosmosEventId.SyncNotSupported)); - [ConditionalFact] - public virtual async Task SetSessionToken_ThrowsForNonExistentContainer() - { - var contextFactory = await InitializeAsync(); - - using var context = contextFactory.CreateContext(); - var sessionTokens = context.Database.GetSessionTokens(); - var exception = Assert.Throws(() => sessionTokens.SetSessionToken("Not the container name", "0:-1#231")); - Assert.Equal(CosmosStrings.ContainerNameDoesNotExist("Not the container name") + " (Parameter 'containerName')", exception.Message); - } - [ConditionalFact] public virtual async Task AppendSessionToken_ThrowsForNonExistentContainer() { var contextFactory = await InitializeAsync(); using var context = contextFactory.CreateContext(); - var sessionTokens = context.Database.GetSessionTokens(); - var exception = Assert.Throws(() => sessionTokens.AppendSessionToken("Not the container name", "0:-1#231")); + var exception = Assert.Throws(() => context.Database.AppendSessionTokens(new Dictionary { { OtherContainerName, "0:-1#231"}, { "Not the container name", "0:-1#231" } })); Assert.Equal(CosmosStrings.ContainerNameDoesNotExist("Not the container name") + " (Parameter 'containerName')", exception.Message); } [ConditionalFact] - public virtual async Task GetSessionToken_ThrowsForNonExistentContainer() + public virtual async Task AppendSessionToken_no_token_sets_token() { var contextFactory = await InitializeAsync(); using var context = contextFactory.CreateContext(); - var sessionTokens = context.Database.GetSessionTokens(); - var exception = Assert.Throws(() => sessionTokens.GetSessionToken("Not the container name")); - Assert.Equal(CosmosStrings.ContainerNameDoesNotExist("Not the container name") + " (Parameter 'containerName')", exception.Message); - } + IReadOnlyDictionary sessionTokens; + context.Database.AppendSessionToken("0:-1#231"); + sessionTokens = context.Database.GetSessionTokens(); + Assert.Equal("0:-1#231", sessionTokens[nameof(CosmosSessionTokenContext)]); + Assert.Equal(nameof(CosmosSessionTokenContext), sessionTokens.First().Key); + Assert.Equal(sessionTokens[nameof(CosmosSessionTokenContext)], sessionTokens.First().Value); + } - [ConditionalTheory, InlineData(true), InlineData(false)] - public virtual async Task AppendSessionToken_no_tokens_sets_token(bool defaultContainer) + [ConditionalFact] + public virtual async Task AppendSessionTokens_no_tokens_sets_tokens() { var contextFactory = await InitializeAsync(); using var context = contextFactory.CreateContext(); + IReadOnlyDictionary sessionTokens; - var sessionTokens = context.Database.GetSessionTokens(); - if (defaultContainer) - { - sessionTokens.AppendSessionToken("0:-1#231"); - } - else - { - sessionTokens.AppendSessionToken(OtherContainerName, "0:-1#231"); - } - - var updatedToken = defaultContainer ? sessionTokens.GetSessionToken() : sessionTokens.GetSessionToken(OtherContainerName); + context.Database.AppendSessionTokens(new Dictionary { { OtherContainerName, "0:-1#123" }, { nameof(CosmosSessionTokenContext), "0:-1#231" } }); + sessionTokens = context.Database.GetSessionTokens(); - Assert.Equal("0:-1#231", updatedToken); + Assert.Equal("0:-1#123", sessionTokens[OtherContainerName]); + Assert.Equal("0:-1#231", sessionTokens[nameof(CosmosSessionTokenContext)]); + Assert.Equal(nameof(CosmosSessionTokenContext), sessionTokens.First().Key); + Assert.Equal(sessionTokens[nameof(CosmosSessionTokenContext)], sessionTokens.First().Value); } - [ConditionalTheory, InlineData(true), InlineData(false)] - public virtual async Task AppendSessionToken_append_token_not_present_adds_token(bool defaultContainer) + [ConditionalFact] + public virtual async Task AppendSessionToken_append_token_already_present_does_not_add_token() { var contextFactory = await InitializeAsync(); - + using var context = contextFactory.CreateContext(); - if (defaultContainer) - { - context.Add(new Customer { Id = "1", PartitionKey = "1" }); - } - else - { - context.Add(new OtherContainerCustomer { Id = "1", PartitionKey = "1" }); - } + context.Add(new Customer { Id = "1", PartitionKey = "1" }); + await context.SaveChangesAsync(); - var sessionTokens = context.Database.GetSessionTokens(); - var initialToken = defaultContainer ? sessionTokens.GetSessionToken() : sessionTokens.GetSessionToken(OtherContainerName); + var initialToken = context.Database.GetSessionToken(); Assert.False(string.IsNullOrWhiteSpace(initialToken)); + context.Database.AppendSessionToken(initialToken); - var newToken = "0:-1#231"; - if (defaultContainer) - { - sessionTokens.AppendSessionToken(newToken); - } - else - { - sessionTokens.AppendSessionToken(OtherContainerName, newToken); - } - - var updatedToken = defaultContainer ? sessionTokens.GetSessionToken() : sessionTokens.GetSessionToken(OtherContainerName); + var updatedToken = context.Database.GetSessionToken(); - Assert.Equal(initialToken + "," + newToken, updatedToken); + Assert.Equal(initialToken, updatedToken); } - [ConditionalTheory, InlineData(true), InlineData(false)] - public virtual async Task AppendSessionToken_append_token_already_present_does_not_add_token(bool defaultContainer) + [ConditionalFact] + public virtual async Task AppendSessionTokens_append_token_already_present_does_not_add_token() { var contextFactory = await InitializeAsync(); - + using var context = contextFactory.CreateContext(); - if (defaultContainer) - { - context.Add(new Customer { Id = "1", PartitionKey = "1" }); - } - else - { - context.Add(new OtherContainerCustomer { Id = "1", PartitionKey = "1" }); - } + context.Add(new Customer { Id = "1", PartitionKey = "1" }); + context.Add(new OtherContainerCustomer { Id = "1", PartitionKey = "1" }); await context.SaveChangesAsync(); - var sessionTokens = context.Database.GetSessionTokens(); - var initialToken = defaultContainer ? sessionTokens.GetSessionToken() : sessionTokens.GetSessionToken(OtherContainerName); - Assert.False(string.IsNullOrWhiteSpace(initialToken)); + var initialTokens = context.Database.GetSessionTokens(); + context.Database.AppendSessionTokens(initialTokens!); - var newToken = initialToken; - if (defaultContainer) - { - sessionTokens.AppendSessionToken(newToken); - } - else + var updatedTokens = context.Database.GetSessionTokens(); + + foreach (var pair in updatedTokens) { - sessionTokens.AppendSessionToken(OtherContainerName, newToken); + Assert.Equal(initialTokens[pair.Key], pair.Value); } - - var updatedToken = defaultContainer ? sessionTokens.GetSessionToken() : sessionTokens.GetSessionToken(OtherContainerName); - - Assert.Equal(initialToken, updatedToken); } - [ConditionalTheory, InlineData(true), InlineData(false)] - public virtual async Task AppendSessionToken_multiple_tokens_splits_tokens(bool defaultContainer) + [ConditionalFact] + public virtual async Task AppendSessionToken_multiple_tokens_splits_tokens() { var contextFactory = await InitializeAsync(); @@ -154,122 +110,57 @@ public virtual async Task AppendSessionToken_multiple_tokens_splits_tokens(bool var sessionTokens = context.Database.GetSessionTokens(); var newToken = "0:-1#123,1:-1#456"; var appendix = "0:-1#123"; - if (defaultContainer) - { - sessionTokens.AppendSessionToken(newToken); - sessionTokens.AppendSessionToken(appendix); - - } - else - { - sessionTokens.AppendSessionToken(OtherContainerName, newToken); - sessionTokens.AppendSessionToken(OtherContainerName, appendix); - } + context.Database.AppendSessionToken(newToken); + context.Database.AppendSessionToken(appendix); - var updatedToken = defaultContainer ? sessionTokens.GetSessionToken() : sessionTokens.GetSessionToken(OtherContainerName); + var updatedToken = context.Database.GetSessionToken(); Assert.Equal(newToken, updatedToken); } - [ConditionalTheory, InlineData(true), InlineData(false)] - public virtual async Task SetSessionToken_does_not_append_session_token(bool defaultContainer) + [ConditionalFact] + public virtual async Task AppendSessionTokens_multiple_tokens_splits_tokens() { var contextFactory = await InitializeAsync(); using var context = contextFactory.CreateContext(); - if (defaultContainer) - { - context.Add(new Customer { Id = "1", PartitionKey = "1" }); - } - else - { - context.Add(new OtherContainerCustomer { Id = "1", PartitionKey = "1" }); - } - - await context.SaveChangesAsync(); var sessionTokens = context.Database.GetSessionTokens(); - var initialToken = defaultContainer ? sessionTokens.GetSessionToken() : sessionTokens.GetSessionToken(OtherContainerName); - Assert.False(string.IsNullOrWhiteSpace(initialToken)); - - var newToken = "0:-1#1,1:-1#1"; - if (defaultContainer) - { - sessionTokens.SetSessionToken(newToken); - } - else - { - sessionTokens.SetSessionToken(OtherContainerName, newToken); - } - - var updatedToken = defaultContainer ? sessionTokens.GetSessionToken() : sessionTokens.GetSessionToken(OtherContainerName); - - Assert.Equal(newToken, updatedToken); - } - - [ConditionalTheory, InlineData(true), InlineData(false)] - public virtual async Task SetSessionToken_null_sets_session_token_null(bool defaultContainer) - { - var contextFactory = await InitializeAsync(); - - using var context = contextFactory.CreateContext(); - if (defaultContainer) - { - context.Add(new Customer { Id = "1", PartitionKey = "1" }); - } - else - { - context.Add(new OtherContainerCustomer { Id = "1", PartitionKey = "1" }); - } + var newToken = "0:-1#123,1:-1#456"; + var appendix = "0:-1#123"; - await context.SaveChangesAsync(); + context.Database.AppendSessionTokens(new Dictionary { { OtherContainerName, newToken }, { nameof(CosmosSessionTokenContext), newToken } }); + context.Database.AppendSessionTokens(new Dictionary { { OtherContainerName, appendix }, { nameof(CosmosSessionTokenContext), appendix } }); - var sessionTokens = context.Database.GetSessionTokens(); - var initialToken = defaultContainer ? sessionTokens.GetSessionToken() : sessionTokens.GetSessionToken(OtherContainerName); - Assert.False(string.IsNullOrWhiteSpace(initialToken)); + var updatedTokens = context.Database.GetSessionTokens(); - if (defaultContainer) - { - sessionTokens.SetSessionToken(null); - } - else + foreach (var pair in updatedTokens) { - sessionTokens.SetSessionToken(OtherContainerName, null); + Assert.Equal(newToken, pair.Value); } - - var updatedToken = defaultContainer ? sessionTokens.GetSessionToken() : sessionTokens.GetSessionToken(OtherContainerName); - - Assert.Null(updatedToken); } - [ConditionalTheory, InlineData(true), InlineData(false)] - public virtual async Task GetSessionToken_no_token_returns_null(bool defaultContainer) + [ConditionalFact] + public virtual async Task GetSessionToken_no_token_returns_null() { var contextFactory = await InitializeAsync(); using var context = contextFactory.CreateContext(); var sessionTokens = context.Database.GetSessionTokens(); - var initialToken = defaultContainer ? sessionTokens.GetSessionToken() : sessionTokens.GetSessionToken(OtherContainerName); - Assert.Null(initialToken); + Assert.Equal(2, sessionTokens.Count); + Assert.Equal(nameof(CosmosSessionTokenContext), sessionTokens.First().Key); + Assert.Null(sessionTokens.First().Value); + Assert.Equal(OtherContainerName, sessionTokens.Last().Key); + Assert.Null(sessionTokens.Last().Value); } - [ConditionalTheory] - [InlineData(true, true)] - [InlineData(true, false)] - [InlineData(false, true)] - [InlineData(false, false)] - public virtual async Task Query_uses_session_token(bool async, bool defaultContainer) + [ConditionalTheory, InlineData(true), InlineData(false)] + public virtual async Task Query_uses_session_token(bool async) { var contextFactory = await InitializeAsync(); using var context = contextFactory.CreateContext(); - - if (defaultContainer) - { - context.Customers.Add(new Customer { Id = "1", PartitionKey = "1" }); - } - else - { - context.OtherContainerCustomers.Add(new OtherContainerCustomer { Id = "1", PartitionKey = "1" }); - } + + context.Customers.Add(new Customer { Id = "1", PartitionKey = "1" }); + context.OtherContainerCustomers.Add(new OtherContainerCustomer { Id = "1", PartitionKey = "1" }); if (async) { @@ -282,74 +173,38 @@ public virtual async Task Query_uses_session_token(bool async, bool defaultConta var sessionTokens = context.Database.GetSessionTokens(); - string sessionToken; - if (defaultContainer) - { - sessionToken = sessionTokens.GetSessionToken()!; - } - else - { - sessionToken = sessionTokens.GetSessionToken(OtherContainerName)!; - } - // Only way we can test this is by setting a session token that will fail the request if used.. // This will take a couple of seconds to fail - var newToken = sessionToken.Substring(0, sessionToken.IndexOf('#') + 1) + int.MaxValue; + var newTokens = sessionTokens.ToDictionary(x => x.Key, x => x.Value!.Substring(0, x.Value.IndexOf('#') + 1) + int.MaxValue); - if (defaultContainer) - { - sessionTokens.SetSessionToken(newToken); - } - else - { - sessionTokens.SetSessionToken(OtherContainerName, newToken); - } + context.Database.AppendSessionTokens(newTokens); + + CosmosException ex1; + CosmosException ex2; - CosmosException ex; if (async) { - if (defaultContainer) - { - ex = await Assert.ThrowsAsync(() => context.Customers.ToListAsync()); - } - else - { - ex = await Assert.ThrowsAsync(() => context.OtherContainerCustomers.ToListAsync()); - } + ex1 = await Assert.ThrowsAsync(() => context.Customers.ToListAsync()); + ex2 = await Assert.ThrowsAsync(() => context.OtherContainerCustomers.ToListAsync()); } else { - if (defaultContainer) - { - ex = Assert.Throws(() => context.Customers.ToList()); - } - else - { - ex = Assert.Throws(() => context.OtherContainerCustomers.ToList()); - } + ex1 = Assert.Throws(() => context.Customers.ToList()); + ex2 = Assert.Throws(() => context.OtherContainerCustomers.ToList()); } - Assert.Contains("The read session is not available for the input session token.", ex.ResponseBody); + Assert.Contains("The read session is not available for the input session token.", ex1.ResponseBody); + Assert.Contains("The read session is not available for the input session token.", ex2.ResponseBody); } - [ConditionalTheory] - [InlineData(true, true)] - [InlineData(true, false)] - [InlineData(false, true)] - [InlineData(false, false)] - public virtual async Task Query_on_new_context_does_not_use_same_session_token(bool async, bool defaultContainer) + [ConditionalTheory, InlineData(true), InlineData(false)] + public virtual async Task Query_on_new_context_does_not_use_same_session_token(bool async) { var contextFactory = await InitializeAsync(); using var context = contextFactory.CreateContext(); - - if (defaultContainer) - { - context.Customers.Add(new Customer { Id = "1", PartitionKey = "1" }); - } - else - { - context.OtherContainerCustomers.Add(new OtherContainerCustomer { Id = "1", PartitionKey = "1" }); - } + + context.Customers.Add(new Customer { Id = "1", PartitionKey = "1" }); + context.OtherContainerCustomers.Add(new OtherContainerCustomer { Id = "1", PartitionKey = "1" }); if (async) { @@ -362,98 +217,47 @@ public virtual async Task Query_on_new_context_does_not_use_same_session_token(b var sessionTokens = context.Database.GetSessionTokens(); - string sessionToken; - if (defaultContainer) - { - sessionToken = sessionTokens.GetSessionToken()!; - } - else - { - sessionToken = sessionTokens.GetSessionToken(OtherContainerName)!; - } - // Only way we can test this is by setting a session token that will fail the request if used.. // This will take a couple of seconds to fail - var newToken = sessionToken.Substring(0, sessionToken.IndexOf('#') + 1) + int.MaxValue; + var newTokens = sessionTokens.ToDictionary(x => x.Key, x => x.Value!.Substring(0, x.Value.IndexOf('#') + 1) + int.MaxValue); - if (defaultContainer) - { - sessionTokens.SetSessionToken(newToken); - } - else - { - sessionTokens.SetSessionToken(OtherContainerName, newToken); - } + context.Database.AppendSessionTokens(newTokens); if (async) { - if (defaultContainer) - { - await Assert.ThrowsAsync(() => context.Customers.ToListAsync()); - } - else - { - await Assert.ThrowsAsync(() => context.OtherContainerCustomers.ToListAsync()); - } + await Assert.ThrowsAsync(() => context.Customers.ToListAsync()); + await Assert.ThrowsAsync(() => context.OtherContainerCustomers.ToListAsync()); } else { - if (defaultContainer) - { - Assert.Throws(() => context.Customers.ToList()); - } - else - { - Assert.Throws(() => context.OtherContainerCustomers.ToList()); - } + Assert.Throws(() => context.Customers.ToList()); + Assert.Throws(() => context.OtherContainerCustomers.ToList()); } using var newContext = contextFactory.CreateContext(); Assert.NotSame(context, newContext); if (async) { - if (defaultContainer) - { - await newContext.Customers.ToListAsync(); - } - else - { - await newContext.OtherContainerCustomers.ToListAsync(); - } + await newContext.Customers.ToListAsync(); + await newContext.OtherContainerCustomers.ToListAsync(); } else { - if (defaultContainer) - { - newContext.Customers.ToList(); - } - else - { - newContext.OtherContainerCustomers.ToList(); - } + newContext.Customers.ToList(); + newContext.OtherContainerCustomers.ToList(); } } - [ConditionalTheory] - [InlineData(true, true)] - [InlineData(true, false)] - [InlineData(false, true)] - [InlineData(false, false)] - public virtual async Task Query_on_same_newly_pooled_context_does_not_use_same_session_token(bool async, bool defaultContainer) + [ConditionalTheory, InlineData(true), InlineData(false)] + public virtual async Task Query_on_same_newly_pooled_context_does_not_use_same_session_token(bool async) { var contextFactory = await InitializeAsync(); DbContext contextCopy; using (var context = contextFactory.CreateContext()) { contextCopy = context; - if (defaultContainer) - { - context.Customers.Add(new Customer { Id = "1", PartitionKey = "1" }); - } - else - { - context.OtherContainerCustomers.Add(new OtherContainerCustomer { Id = "1", PartitionKey = "1" }); - } + context.Customers.Add(new Customer { Id = "1", PartitionKey = "1" }); + context.OtherContainerCustomers.Add(new OtherContainerCustomer { Id = "1", PartitionKey = "1" }); if (async) { @@ -466,376 +270,440 @@ public virtual async Task Query_on_same_newly_pooled_context_does_not_use_same_s var sessionTokens = context.Database.GetSessionTokens(); - string sessionToken; - if (defaultContainer) - { - sessionToken = sessionTokens.GetSessionToken()!; - } - else - { - sessionToken = sessionTokens.GetSessionToken(OtherContainerName)!; - } - // Only way we can test this is by setting a session token that will fail the request if used.. // This will take a couple of seconds to fail - var newToken = sessionToken.Substring(0, sessionToken.IndexOf('#') + 1) + int.MaxValue; + var newTokens = sessionTokens.ToDictionary(x => x.Key, x => x.Value!.Substring(0, x.Value.IndexOf('#') + 1) + int.MaxValue); - if (defaultContainer) - { - sessionTokens.SetSessionToken(newToken); - } - else - { - sessionTokens.SetSessionToken(OtherContainerName, newToken); - } + context.Database.AppendSessionTokens(newTokens); if (async) { - if (defaultContainer) - { - await Assert.ThrowsAsync(() => context.Customers.ToListAsync()); - } - else - { - await Assert.ThrowsAsync(() => context.OtherContainerCustomers.ToListAsync()); - } + await Assert.ThrowsAsync(() => context.Customers.ToListAsync()); + await Assert.ThrowsAsync(() => context.OtherContainerCustomers.ToListAsync()); } else { - if (defaultContainer) - { - Assert.Throws(() => context.Customers.ToList()); - } - else - { - Assert.Throws(() => context.OtherContainerCustomers.ToList()); - } + Assert.Throws(() => context.Customers.ToList()); + Assert.Throws(() => context.OtherContainerCustomers.ToList()); } } using var newContext = contextFactory.CreateContext(); Assert.Same(newContext, contextCopy); - if (async) + if (async) // @TODO: Maybe even no if here? { - if (defaultContainer) - { - await newContext.Customers.ToListAsync(); - } - else - { - await newContext.OtherContainerCustomers.ToListAsync(); - } + await newContext.Customers.ToListAsync(); + await newContext.OtherContainerCustomers.ToListAsync(); } else { - if (defaultContainer) - { - newContext.Customers.ToList(); - } - else - { - newContext.OtherContainerCustomers.ToList(); - } + newContext.Customers.ToList(); + newContext.OtherContainerCustomers.ToList(); } } - [ConditionalTheory] - [InlineData(true), InlineData(false)] - public virtual async Task PagingQuery_uses_session_token(bool defaultContainer) + [ConditionalTheory, InlineData(true), InlineData(false)] + public virtual async Task PagingQuery_uses_session_token(bool async) { var contextFactory = await InitializeAsync(); using var context = contextFactory.CreateContext(); - if (defaultContainer) + context.Customers.Add(new Customer { Id = "1", PartitionKey = "1" }); + context.OtherContainerCustomers.Add(new OtherContainerCustomer { Id = "1", PartitionKey = "1" }); + + if (async) { - context.Customers.Add(new Customer { Id = "1", PartitionKey = "1" }); + await context.SaveChangesAsync(); } else { - context.OtherContainerCustomers.Add(new OtherContainerCustomer { Id = "1", PartitionKey = "1" }); + context.SaveChanges(); } - await context.SaveChangesAsync(); - var sessionTokens = context.Database.GetSessionTokens(); - string sessionToken; - if (defaultContainer) + // Only way we can test this is by setting a session token that will fail the request if used.. + // This will take a couple of seconds to fail + var newTokens = sessionTokens.ToDictionary(x => x.Key, x => x.Value!.Substring(0, x.Value.IndexOf('#') + 1) + int.MaxValue); + + context.Database.AppendSessionTokens(newTokens); + + var ex1 = await Assert.ThrowsAsync(() => context.Customers.ToPageAsync(1, null)); + var ex2 = await Assert.ThrowsAsync(() => context.OtherContainerCustomers.ToPageAsync(1, null)); + + Assert.Contains("The read session is not available for the input session token.", ex1.ResponseBody); + Assert.Contains("The read session is not available for the input session token.", ex2.ResponseBody); + + } + + [ConditionalTheory, InlineData(true), InlineData(false)] + public virtual async Task Shaped_query_uses_session_token(bool async) + { + var contextFactory = await InitializeAsync(); + using var context = contextFactory.CreateContext(); + + context.Customers.Add(new Customer { Id = "1", PartitionKey = "1" }); + context.OtherContainerCustomers.Add(new OtherContainerCustomer { Id = "1", PartitionKey = "1" }); + + if (async) { - sessionToken = sessionTokens.GetSessionToken()!; + await context.SaveChangesAsync(); } else { - sessionToken = sessionTokens.GetSessionToken(OtherContainerName)!; + context.SaveChanges(); } + var sessionTokens = context.Database.GetSessionTokens(); + // Only way we can test this is by setting a session token that will fail the request if used.. // This will take a couple of seconds to fail - var newToken = sessionToken.Substring(0, sessionToken.IndexOf('#') + 1) + int.MaxValue; + var newTokens = sessionTokens.ToDictionary(x => x.Key, x => x.Value!.Substring(0, x.Value.IndexOf('#') + 1) + int.MaxValue); - if (defaultContainer) + context.Database.AppendSessionTokens(newTokens); + + CosmosException ex1; + CosmosException ex2; + if (async) { - sessionTokens.SetSessionToken(newToken); + ex1 = await Assert.ThrowsAsync(() => context.Customers.Select(x => new { x.Id, x.PartitionKey }).ToListAsync()); + ex2 = await Assert.ThrowsAsync(() => context.OtherContainerCustomers.Select(x => new { x.Id, x.PartitionKey }).ToListAsync()); } else { - sessionTokens.SetSessionToken(OtherContainerName, newToken); + ex1 = Assert.Throws(() => context.Customers.Select(x => new { x.Id, x.PartitionKey }).ToList()); + ex2 = Assert.Throws(() => context.OtherContainerCustomers.Select(x => new { x.Id, x.PartitionKey }).ToList()); } - CosmosException ex; - if (defaultContainer) + Assert.Contains("The read session is not available for the input session token.", ex1.ResponseBody); + Assert.Contains("The read session is not available for the input session token.", ex2.ResponseBody); + + } + + [ConditionalTheory, InlineData(true), InlineData(false)] + public virtual async Task Read_item_uses_session_token(bool async) + { + var contextFactory = await InitializeAsync(); + using var context = contextFactory.CreateContext(); + + context.Customers.Add(new Customer { Id = "1", PartitionKey = "1" }); + context.OtherContainerCustomers.Add(new OtherContainerCustomer { Id = "1", PartitionKey = "1" }); + + if (async) { - ex = await Assert.ThrowsAsync(() => context.Customers.ToPageAsync(1, null)); + await context.SaveChangesAsync(); } else { - ex = await Assert.ThrowsAsync(() => context.OtherContainerCustomers.ToPageAsync(1, null)); + context.SaveChanges(); } - Assert.Contains("The read session is not available for the input session token.", ex.ResponseBody); - } + var sessionTokens = context.Database.GetSessionTokens(); - [ConditionalTheory] - [InlineData(true, true)] - [InlineData(true, false)] - [InlineData(false, true)] - [InlineData(false, false)] - public virtual async Task Shaped_query_uses_session_token(bool async, bool defaultContainer) - { - var contextFactory = await InitializeAsync(); - using var context = contextFactory.CreateContext(); + // Only way we can test this is by setting a session token that will fail the request if used.. + // This will take a couple of seconds to fail + var newTokens = sessionTokens.ToDictionary(x => x.Key, x => x.Value!.Substring(0, x.Value.IndexOf('#') + 1) + int.MaxValue); - if (defaultContainer) + context.Database.AppendSessionTokens(newTokens); + + CosmosException ex1; + CosmosException ex2; + if (async) { - context.Customers.Add(new Customer { Id = "1", PartitionKey = "1" }); + ex1 = await Assert.ThrowsAsync(() => context.Customers.FirstOrDefaultAsync(x => x.Id == "1" && x.PartitionKey == "1")); + ex2 = await Assert.ThrowsAsync(() => context.OtherContainerCustomers.FirstOrDefaultAsync(x => x.Id == "1" && x.PartitionKey == "1")); } else { - context.OtherContainerCustomers.Add(new OtherContainerCustomer { Id = "1", PartitionKey = "1" }); + ex1 = Assert.Throws(() => context.Customers.FirstOrDefault(x => x.Id == "1" && x.PartitionKey == "1")); + ex2 = Assert.Throws(() => context.OtherContainerCustomers.FirstOrDefault(x => x.Id == "1" && x.PartitionKey == "1")); } + Assert.Contains("The read session is not available for the input session token.", ex1.ResponseBody); + Assert.Contains("The read session is not available for the input session token.", ex2.ResponseBody); + + } + + [ConditionalTheory, InlineData(true), InlineData(false)] + public virtual async Task Query_sets_session_token(bool async) + { + var contextFactory = await InitializeAsync(); + using var context = contextFactory.CreateContext(); + if (async) { - await context.SaveChangesAsync(); + await context.Customers.ToListAsync(); + await context.OtherContainerCustomers.ToListAsync(); } else { - context.SaveChanges(); + _ = context.Customers.ToList(); + _ = context.OtherContainerCustomers.ToList(); } var sessionTokens = context.Database.GetSessionTokens(); + Assert.False(string.IsNullOrWhiteSpace(sessionTokens.First().Value)); + Assert.False(string.IsNullOrWhiteSpace(sessionTokens.Last().Value)); + } - string sessionToken; - if (defaultContainer) + [ConditionalTheory, InlineData(true), InlineData(false)] + public virtual async Task Query_appends_session_token(bool async) + { + var contextFactory = await InitializeAsync(); + using var context = contextFactory.CreateContext(); + + if (async) { - sessionToken = sessionTokens.GetSessionToken()!; + await context.Customers.ToListAsync(); + await context.OtherContainerCustomers.ToListAsync(); } else { - sessionToken = sessionTokens.GetSessionToken(OtherContainerName)!; + _ = context.Customers.ToList(); + _ = context.OtherContainerCustomers.ToList(); } - // Only way we can test this is by setting a session token that will fail the request if used.. - // This will take a couple of seconds to fail - var newToken = sessionToken.Substring(0, sessionToken.IndexOf('#') + 1) + int.MaxValue; + var initialTokens = context.Database.GetSessionTokens().ToDictionary(); // Make a copy - if (defaultContainer) + using var otherContext = contextFactory.CreateContext(); + otherContext.Customers.Add(new Customer { Id = "1", PartitionKey = "1" }); + otherContext.OtherContainerCustomers.Add(new OtherContainerCustomer { Id = "1", PartitionKey = "1" }); + + if (async) { - sessionTokens.SetSessionToken(newToken); + await otherContext.SaveChangesAsync(); } else { - sessionTokens.SetSessionToken(OtherContainerName, newToken); + otherContext.SaveChanges(); } - CosmosException ex; + var otherTokens = otherContext.Database.GetSessionTokens(); + if (async) { - if (defaultContainer) - { - ex = await Assert.ThrowsAsync(() => context.Customers.Select(x => new { x.Id, x.PartitionKey }).ToListAsync()); - } - else - { - ex = await Assert.ThrowsAsync(() => context.OtherContainerCustomers.Select(x => new { x.Id, x.PartitionKey }).ToListAsync()); - } + await context.Customers.ToListAsync(); + await context.OtherContainerCustomers.ToListAsync(); } else { - if (defaultContainer) - { - ex = Assert.Throws(() => context.Customers.Select(x => new { x.Id, x.PartitionKey }).ToList()); - } - else - { - ex = Assert.Throws(() => context.OtherContainerCustomers.Select(x => new { x.Id, x.PartitionKey }).ToList()); - } + _ = context.Customers.ToList(); + _ = context.OtherContainerCustomers.ToList(); } - Assert.Contains("The read session is not available for the input session token.", ex.ResponseBody); + var sessionTokens = context.Database.GetSessionTokens(); + foreach (var token in sessionTokens) + { + var initialToken = initialTokens[token.Key]; + var otherToken = otherTokens[token.Key]; + Assert.Equal(initialToken + "," + otherToken, token.Value); + } } - [ConditionalTheory] - [InlineData(true, true)] - [InlineData(true, false)] - [InlineData(false, true)] - [InlineData(false, false)] - public virtual async Task Read_item_uses_session_token(bool async, bool defaultContainer) + [ConditionalTheory, InlineData(true), InlineData(false)] + public virtual async Task Query_same_session_does_not_append_session_token(bool async) { var contextFactory = await InitializeAsync(); using var context = contextFactory.CreateContext(); - if (defaultContainer) + if (async) { - context.Customers.Add(new Customer { Id = "1", PartitionKey = "1" }); + await context.Customers.ToListAsync(); + await context.OtherContainerCustomers.ToListAsync(); } else { - context.OtherContainerCustomers.Add(new OtherContainerCustomer { Id = "1", PartitionKey = "1" }); + _ = context.Customers.ToList(); + _ = context.OtherContainerCustomers.ToList(); } + var initialTokens = context.Database.GetSessionTokens().ToDictionary(); // Make a copy + if (async) { - await context.SaveChangesAsync(); + await context.Customers.ToListAsync(); + await context.OtherContainerCustomers.ToListAsync(); } else { - context.SaveChanges(); + _ = context.Customers.ToList(); + _ = context.OtherContainerCustomers.ToList(); } var sessionTokens = context.Database.GetSessionTokens(); + foreach (var token in sessionTokens) + { + var initialToken = initialTokens[token.Key]; + Assert.Equal(initialToken, token.Value); + } + } - string sessionToken; - if (defaultContainer) + [ConditionalFact] + public virtual async Task PagingQuery_appends_session_token() + { + var contextFactory = await InitializeAsync(); + using var context = contextFactory.CreateContext(); + + await context.Customers.ToPageAsync(1, null); + await context.OtherContainerCustomers.ToPageAsync(1, null); + + var initialTokens = context.Database.GetSessionTokens().ToDictionary(); // Make a copy + + using var otherContext = contextFactory.CreateContext(); + otherContext.Customers.Add(new Customer { Id = "1", PartitionKey = "1" }); + otherContext.OtherContainerCustomers.Add(new OtherContainerCustomer { Id = "1", PartitionKey = "1" }); + + await otherContext.SaveChangesAsync(); + otherContext.SaveChanges(); + + var otherTokens = otherContext.Database.GetSessionTokens(); + + await context.Customers.ToPageAsync(1, null); + await context.OtherContainerCustomers.ToPageAsync(1, null); + + var sessionTokens = context.Database.GetSessionTokens(); + foreach (var token in sessionTokens) + { + var initialToken = initialTokens[token.Key]; + var otherToken = otherTokens[token.Key]; + Assert.Equal(initialToken + "," + otherToken, token.Value); + } + } + + [ConditionalTheory, InlineData(true), InlineData(false)] + public virtual async Task Read_item_appends_session_token(bool async) + { + var contextFactory = await InitializeAsync(); + using var context = contextFactory.CreateContext(); + + if (async) { - sessionToken = sessionTokens.GetSessionToken()!; + await context.Customers.FirstOrDefaultAsync(x => x.Id == "1" && x.PartitionKey == "1"); + await context.OtherContainerCustomers.FirstOrDefaultAsync(x => x.Id == "1" && x.PartitionKey == "1"); } else { - sessionToken = sessionTokens.GetSessionToken(OtherContainerName)!; + _ = context.Customers.ToList(); + _ = context.OtherContainerCustomers.ToList(); } - // Only way we can test this is by setting a session token that will fail the request if used.. - // This will take a couple of seconds to fail - var newToken = sessionToken.Substring(0, sessionToken.IndexOf('#') + 1) + int.MaxValue; + var initialTokens = context.Database.GetSessionTokens().ToDictionary(); // Make a copy - if (defaultContainer) + using var otherContext = contextFactory.CreateContext(); + otherContext.Customers.Add(new Customer { Id = "1", PartitionKey = "1" }); + otherContext.OtherContainerCustomers.Add(new OtherContainerCustomer { Id = "1", PartitionKey = "1" }); + + if (async) { - sessionTokens.SetSessionToken(newToken); + await otherContext.SaveChangesAsync(); } else { - sessionTokens.SetSessionToken(OtherContainerName, newToken); + otherContext.SaveChanges(); } - CosmosException ex; + var otherTokens = otherContext.Database.GetSessionTokens(); + if (async) { - if (defaultContainer) - { - ex = await Assert.ThrowsAsync(() => context.Customers.FirstOrDefaultAsync(x => x.Id == "1" && x.PartitionKey == "1")); - } - else - { - ex = await Assert.ThrowsAsync(() => context.OtherContainerCustomers.FirstOrDefaultAsync(x => x.Id == "1" && x.PartitionKey == "1")); - } + await context.Customers.FirstOrDefaultAsync(x => x.Id == "1" && x.PartitionKey == "1"); + await context.OtherContainerCustomers.FirstOrDefaultAsync(x => x.Id == "1" && x.PartitionKey == "1"); } else { - if (defaultContainer) - { - ex = Assert.Throws(() => context.Customers.FirstOrDefault(x => x.Id == "1" && x.PartitionKey == "1")); - } - else - { - ex = Assert.Throws(() => context.OtherContainerCustomers.FirstOrDefault(x => x.Id == "1" && x.PartitionKey == "1")); - } + _ = context.Customers.ToList(); + _ = context.OtherContainerCustomers.ToList(); } - Assert.Contains("The read session is not available for the input session token.", ex.ResponseBody); + var sessionTokens = context.Database.GetSessionTokens(); + foreach (var token in sessionTokens) + { + var initialToken = initialTokens[token.Key]; + var otherToken = otherTokens[token.Key]; + Assert.Equal(initialToken + "," + otherToken, token.Value); + } } - [ConditionalTheory] - [InlineData(true, true)] - [InlineData(true, false)] - [InlineData(false, true)] - [InlineData(false, false)] - public virtual async Task Query_sets_session_token(bool async, bool defaultContainer) + [ConditionalTheory, InlineData(true), InlineData(false)] + public virtual async Task Read_item_enumerable_sets_session_token(bool async) { var contextFactory = await InitializeAsync(); using var context = contextFactory.CreateContext(); if (async) { - if (defaultContainer) - { - await context.Customers.ToListAsync(); - } - else - { - await context.OtherContainerCustomers.ToListAsync(); - } + await context.Customers.Where(x => x.Id == "1" && x.PartitionKey == "1").ToListAsync(); + await context.OtherContainerCustomers.Where(x => x.Id == "1" && x.PartitionKey == "1").ToListAsync(); } else { - if (defaultContainer) - { - _ = context.Customers.ToList(); - } - else - { - _ = context.OtherContainerCustomers.ToList(); - } + _ = context.Customers.ToList(); + _ = context.OtherContainerCustomers.ToList(); } - var sessionTokens = context.Database.GetSessionTokens(); - string? sessionToken; - if (defaultContainer) + var initialTokens = context.Database.GetSessionTokens().ToDictionary(); // Make a copy + + using var otherContext = contextFactory.CreateContext(); + otherContext.Customers.Add(new Customer { Id = "1", PartitionKey = "1" }); + otherContext.OtherContainerCustomers.Add(new OtherContainerCustomer { Id = "1", PartitionKey = "1" }); + + if (async) { - sessionToken = sessionTokens.GetSessionToken(); + await otherContext.SaveChangesAsync(); } else { - sessionToken = sessionTokens.GetSessionToken(OtherContainerName); + otherContext.SaveChanges(); } - Assert.False(string.IsNullOrWhiteSpace(sessionToken)); - } + var otherTokens = otherContext.Database.GetSessionTokens(); - // @TODO set then query appends - // @TODO append then query appends only if not present? + if (async) + { + await context.Customers.Where(x => x.Id == "1" && x.PartitionKey == "1").ToListAsync(); + await context.OtherContainerCustomers.Where(x => x.Id == "1" && x.PartitionKey == "1").ToListAsync(); + } + else + { + _ = context.Customers.ToList(); + _ = context.OtherContainerCustomers.ToList(); + } + + var sessionTokens = context.Database.GetSessionTokens(); + foreach (var token in sessionTokens) + { + var initialToken = initialTokens[token.Key]; + var otherToken = otherTokens[token.Key]; + Assert.Equal(initialToken + "," + otherToken, token.Value); + } + } [ConditionalTheory] - [InlineData(true, true)] - [InlineData(true, false)] - [InlineData(false, true)] - [InlineData(false, false)] - public virtual async Task PagingQuery_sets_session_token(bool async, bool defaultContainer) + [InlineData(true)] + [InlineData(false)] + public virtual async Task Add_AutoTransactionBehavior_never_sets_session_token(bool async) { var contextFactory = await InitializeAsync(); + using var context = contextFactory.CreateContext(); + context.Database.AutoTransactionBehavior = AutoTransactionBehavior.Never; + context.Customers.Add(new Customer { Id = "1", PartitionKey = "1" }); + context.OtherContainerCustomers.Add(new OtherContainerCustomer { Id = "1", PartitionKey = "1" }); - if (defaultContainer) + if (async) { - await context.Customers.ToPageAsync(1, null); + await context.SaveChangesAsync(); } else { - await context.OtherContainerCustomers.ToPageAsync(1, null); + context.SaveChanges(); } var sessionTokens = context.Database.GetSessionTokens(); - string? sessionToken; - if (defaultContainer) + foreach (var sessionToken in sessionTokens.Values) { - sessionToken = sessionTokens.GetSessionToken(); + Assert.False(string.IsNullOrWhiteSpace(sessionToken)); } - else - { - sessionToken = sessionTokens.GetSessionToken(OtherContainerName); - } - - Assert.False(string.IsNullOrWhiteSpace(sessionToken)); } [ConditionalTheory] @@ -843,93 +711,79 @@ public virtual async Task PagingQuery_sets_session_token(bool async, bool defaul [InlineData(true, false)] [InlineData(false, true)] [InlineData(false, false)] - public virtual async Task Shaped_query_sets_session_token(bool async, bool defaultContainer) + public virtual async Task Add_AutoTransactionBehavior_always_sets_session_token(bool async, bool defaultContainer) { var contextFactory = await InitializeAsync(); - using var context = contextFactory.CreateContext(); - if (async) + using var context = contextFactory.CreateContext(); + context.Database.AutoTransactionBehavior = AutoTransactionBehavior.Always; + if (defaultContainer) { - if (defaultContainer) - { - await context.Customers.Select(x => new { x.Id, x.PartitionKey }).ToListAsync(); - } - else - { - await context.OtherContainerCustomers.Select(x => new { x.Id, x.PartitionKey }).ToListAsync(); - } + context.Customers.Add(new Customer { Id = "1", PartitionKey = "1" }); } else { - if (defaultContainer) - { - _ = context.Customers.Select(x => new { x.Id, x.PartitionKey }).ToList(); - } - else - { - _ = context.OtherContainerCustomers.Select(x => new { x.Id, x.PartitionKey }).ToList(); - } + context.OtherContainerCustomers.Add(new OtherContainerCustomer { Id = "1", PartitionKey = "1" }); } - var sessionTokens = context.Database.GetSessionTokens(); - string? sessionToken; - if (defaultContainer) + if (async) { - sessionToken = sessionTokens.GetSessionToken(); + await context.SaveChangesAsync(); } else { - sessionToken = sessionTokens.GetSessionToken(OtherContainerName); + context.SaveChanges(); } + var sessionTokens = context.Database.GetSessionTokens(); + var sessionToken = defaultContainer ? sessionTokens[nameof(CosmosSessionTokenContext)]! : sessionTokens[OtherContainerName]!; Assert.False(string.IsNullOrWhiteSpace(sessionToken)); } [ConditionalTheory] - [InlineData(true, true)] - [InlineData(true, false)] - [InlineData(false, true)] - [InlineData(false, false)] - public virtual async Task Read_item_sets_session_token(bool async, bool defaultContainer) + [InlineData(true)] + [InlineData(false)] + public virtual async Task Add_never_merges_session_token(bool async) { var contextFactory = await InitializeAsync(); + using var context = contextFactory.CreateContext(); + context.Database.AutoTransactionBehavior = AutoTransactionBehavior.Never; + + context.Customers.Add(new Customer { Id = "1", PartitionKey = "1" }); + context.OtherContainerCustomers.Add(new OtherContainerCustomer { Id = "1", PartitionKey = "1" }); if (async) { - if (defaultContainer) - { - await context.Customers.FirstOrDefaultAsync(x => x.Id == "1" && x.PartitionKey == "1"); - } - else - { - await context.OtherContainerCustomers.FirstOrDefaultAsync(x => x.Id == "1" && x.PartitionKey == "1"); - } + await context.SaveChangesAsync(); } else { - if (defaultContainer) - { - _ = context.Customers.FirstOrDefault(x => x.Id == "1" && x.PartitionKey == "1"); - } - else - { - _ = context.OtherContainerCustomers.FirstOrDefault(x => x.Id == "1" && x.PartitionKey == "1"); - } + context.SaveChanges(); } - var sessionTokens = context.Database.GetSessionTokens(); - string? sessionToken; - if (defaultContainer) + var initialTokens = context.Database.GetSessionTokens().ToDictionary(); + + context.Customers.Add(new Customer { Id = "2", PartitionKey = "1" }); + context.OtherContainerCustomers.Add(new OtherContainerCustomer { Id = "2", PartitionKey = "1" }); + + if (async) { - sessionToken = sessionTokens.GetSessionToken(); + await context.SaveChangesAsync(); } else { - sessionToken = sessionTokens.GetSessionToken(OtherContainerName); + context.SaveChanges(); } - Assert.False(string.IsNullOrWhiteSpace(sessionToken)); + var sessionTokens = context.Database.GetSessionTokens(); + foreach (var sessionToken in sessionTokens) + { + var initialToken = initialTokens[sessionToken.Key]; + Assert.NotEqual(sessionToken.Value, initialToken); + Assert.StartsWith(initialToken + ",", sessionToken.Value); + Assert.False(string.IsNullOrWhiteSpace(sessionToken.Value!.Substring(sessionToken.Value.IndexOf(",") + 1))); + } } [ConditionalTheory] @@ -937,77 +791,85 @@ public virtual async Task Read_item_sets_session_token(bool async, bool defaultC [InlineData(true, false)] [InlineData(false, true)] [InlineData(false, false)] - public virtual async Task Read_item_enumerable_sets_session_token(bool async, bool defaultContainer) + public virtual async Task Add_always_merges_session_token(bool async, bool defaultContainer) { var contextFactory = await InitializeAsync(); + using var context = contextFactory.CreateContext(); + context.Database.AutoTransactionBehavior = AutoTransactionBehavior.Always; + + if (defaultContainer) + { + context.Customers.Add(new Customer { Id = "1", PartitionKey = "1" }); + } + else + { + context.OtherContainerCustomers.Add(new OtherContainerCustomer { Id = "1", PartitionKey = "1" }); + } if (async) { - if (defaultContainer) - { - await context.Customers.Where(x => x.Id == "1" && x.PartitionKey == "1").ToListAsync(); - } - else - { - await context.OtherContainerCustomers.Where(x => x.Id == "1" && x.PartitionKey == "1").ToListAsync(); - } + await context.SaveChangesAsync(); } else { - if (defaultContainer) - { - _ = context.Customers.Where(x => x.Id == "1" && x.PartitionKey == "1").ToList(); - } - else - { - _ = context.OtherContainerCustomers.Where(x => x.Id == "1" && x.PartitionKey == "1").ToList(); - } + context.SaveChanges(); } - var sessionTokens = context.Database.GetSessionTokens(); - string? sessionToken; + var initialTokens = context.Database.GetSessionTokens().ToDictionary(); + if (defaultContainer) { - sessionToken = sessionTokens.GetSessionToken(); + context.Customers.Add(new Customer { Id = "2", PartitionKey = "1" }); + } + else + { + context.OtherContainerCustomers.Add(new OtherContainerCustomer { Id = "2", PartitionKey = "1" }); + } + + if (async) + { + await context.SaveChangesAsync(); } else { - sessionToken = sessionTokens.GetSessionToken(OtherContainerName); + context.SaveChanges(); } + var sessionTokens = context.Database.GetSessionTokens(); + var sessionToken = defaultContainer ? sessionTokens[nameof(CosmosSessionTokenContext)]! : sessionTokens[OtherContainerName]!; Assert.False(string.IsNullOrWhiteSpace(sessionToken)); } [ConditionalTheory] - [InlineData(AutoTransactionBehavior.WhenNeeded, true, true)] - [InlineData(AutoTransactionBehavior.WhenNeeded, true, false)] - [InlineData(AutoTransactionBehavior.WhenNeeded, false, true)] - [InlineData(AutoTransactionBehavior.WhenNeeded, false, false)] - [InlineData(AutoTransactionBehavior.Never, true, true)] - [InlineData(AutoTransactionBehavior.Never, true, false)] - [InlineData(AutoTransactionBehavior.Never, false, true)] - [InlineData(AutoTransactionBehavior.Never, false, false)] - [InlineData(AutoTransactionBehavior.Always, true, true)] - [InlineData(AutoTransactionBehavior.Always, true, false)] - [InlineData(AutoTransactionBehavior.Always, false, true)] - [InlineData(AutoTransactionBehavior.Always, false, false)] - public virtual async Task Add_sets_session_token(AutoTransactionBehavior autoTransactionBehavior, bool async, bool defaultContainer) + [InlineData(true)] + [InlineData(false)] + public virtual async Task Delete_never_merges_session_token(bool async) { var contextFactory = await InitializeAsync(); using var context = contextFactory.CreateContext(); - context.Database.AutoTransactionBehavior = autoTransactionBehavior; + context.Database.AutoTransactionBehavior = AutoTransactionBehavior.Never; - if (defaultContainer) + var customer = new Customer { Id = "1", PartitionKey = "1" }; + var otherContainerCustomer = new OtherContainerCustomer { Id = "1", PartitionKey = "1" }; + context.Customers.Add(customer); + context.OtherContainerCustomers.Add(otherContainerCustomer); + + if (async) { - context.Customers.Add(new Customer { Id = "1", PartitionKey = "1" }); + await context.SaveChangesAsync(); } else { - context.OtherContainerCustomers.Add(new OtherContainerCustomer { Id = "1", PartitionKey = "1" }); + context.SaveChanges(); } + var initialTokens = context.Database.GetSessionTokens().ToDictionary(); + + context.Customers.Remove(customer); + context.OtherContainerCustomers.Remove(otherContainerCustomer); + if (async) { await context.SaveChangesAsync(); @@ -1018,39 +880,26 @@ public virtual async Task Add_sets_session_token(AutoTransactionBehavior autoTra } var sessionTokens = context.Database.GetSessionTokens(); - string? sessionToken; - if (defaultContainer) + foreach (var sessionToken in sessionTokens) { - sessionToken = sessionTokens.GetSessionToken(); + var initialToken = initialTokens[sessionToken.Key]; + Assert.NotEqual(sessionToken.Value, initialToken); + Assert.StartsWith(initialToken + ",", sessionToken.Value); + Assert.False(string.IsNullOrWhiteSpace(sessionToken.Value!.Substring(sessionToken.Value.IndexOf(",") + 1))); } - else - { - sessionToken = sessionTokens.GetSessionToken(OtherContainerName); - } - - Assert.False(string.IsNullOrWhiteSpace(sessionToken)); } [ConditionalTheory] - [InlineData(AutoTransactionBehavior.WhenNeeded, true, true)] - [InlineData(AutoTransactionBehavior.WhenNeeded, true, false)] - [InlineData(AutoTransactionBehavior.WhenNeeded, false, true)] - [InlineData(AutoTransactionBehavior.WhenNeeded, false, false)] - [InlineData(AutoTransactionBehavior.Never, true, true)] - [InlineData(AutoTransactionBehavior.Never, true, false)] - [InlineData(AutoTransactionBehavior.Never, false, true)] - [InlineData(AutoTransactionBehavior.Never, false, false)] - [InlineData(AutoTransactionBehavior.Always, true, true)] - [InlineData(AutoTransactionBehavior.Always, true, false)] - [InlineData(AutoTransactionBehavior.Always, false, true)] - [InlineData(AutoTransactionBehavior.Always, false, false)] - public virtual async Task Add_merges_session_token(AutoTransactionBehavior autoTransactionBehavior, bool async, bool defaultContainer) + [InlineData(true, true)] + [InlineData(true, false)] + [InlineData(false, true)] + [InlineData(false, false)] + public virtual async Task Delete_always_merges_session_token(bool async, bool defaultContainer) { var contextFactory = await InitializeAsync(); using var context = contextFactory.CreateContext(); - var sessionTokens = context.Database.GetSessionTokens(); - context.Database.AutoTransactionBehavior = autoTransactionBehavior; + context.Database.AutoTransactionBehavior = AutoTransactionBehavior.Always; if (defaultContainer) { @@ -1070,15 +919,16 @@ public virtual async Task Add_merges_session_token(AutoTransactionBehavior autoT context.SaveChanges(); } - var initialToken = defaultContainer ? sessionTokens.GetSessionToken()! : sessionTokens.GetSessionToken(OtherContainerName)!; + context.ChangeTracker.Clear(); + var initialTokens = context.Database.GetSessionTokens().ToDictionary(); if (defaultContainer) { - context.Customers.Add(new Customer { Id = "2", PartitionKey = "1" }); + context.Customers.Remove(new Customer { Id = "1", PartitionKey = "1" }); } else { - context.OtherContainerCustomers.Add(new OtherContainerCustomer { Id = "2", PartitionKey = "1" }); + context.OtherContainerCustomers.Remove(new OtherContainerCustomer { Id = "1", PartitionKey = "1" }); } if (async) @@ -1090,159 +940,99 @@ public virtual async Task Add_merges_session_token(AutoTransactionBehavior autoT context.SaveChanges(); } - string? sessionToken; - if (defaultContainer) - { - sessionToken = sessionTokens.GetSessionToken(); - } - else - { - sessionToken = sessionTokens.GetSessionToken(OtherContainerName); - } - + var sessionTokens = context.Database.GetSessionTokens(); + var sessionToken = defaultContainer ? sessionTokens[nameof(CosmosSessionTokenContext)]! : sessionTokens[OtherContainerName]!; Assert.False(string.IsNullOrWhiteSpace(sessionToken)); - Assert.NotEqual(sessionToken, initialToken); - Assert.StartsWith(initialToken + ",", sessionToken); } [ConditionalTheory] - [InlineData(AutoTransactionBehavior.WhenNeeded, true, true)] - [InlineData(AutoTransactionBehavior.WhenNeeded, true, false)] - [InlineData(AutoTransactionBehavior.WhenNeeded, false, true)] - [InlineData(AutoTransactionBehavior.WhenNeeded, false, false)] - [InlineData(AutoTransactionBehavior.Never, true, true)] - [InlineData(AutoTransactionBehavior.Never, true, false)] - [InlineData(AutoTransactionBehavior.Never, false, true)] - [InlineData(AutoTransactionBehavior.Never, false, false)] - [InlineData(AutoTransactionBehavior.Always, true, true)] - [InlineData(AutoTransactionBehavior.Always, true, false)] - [InlineData(AutoTransactionBehavior.Always, false, true)] - [InlineData(AutoTransactionBehavior.Always, false, false)] - public virtual async Task Delete_merges_session_token(AutoTransactionBehavior autoTransactionBehavior, bool async, bool defaultContainer) + [InlineData(true)] + [InlineData(false)] + public virtual async Task Update_never_merges_session_token(bool async) { var contextFactory = await InitializeAsync(); using var context = contextFactory.CreateContext(); - context.Database.AutoTransactionBehavior = autoTransactionBehavior; - - string initialToken; - if (defaultContainer) - { - var customer = new Customer { Id = "1", PartitionKey = "1" }; - context.Customers.Add(customer); - - if (async) - { - await context.SaveChangesAsync(); - } - else - { - context.SaveChanges(); - } + context.Database.AutoTransactionBehavior = AutoTransactionBehavior.Never; - initialToken = context.Database.GetSessionTokens().GetSessionToken()!; + var customer = new Customer { Id = "1", PartitionKey = "1" }; + var otherContainerCustomer = new OtherContainerCustomer { Id = "1", PartitionKey = "1" }; + context.Customers.Add(customer); + context.OtherContainerCustomers.Add(otherContainerCustomer); - context.Remove(customer); - - if (async) - { - await context.SaveChangesAsync(); - } - else - { - context.SaveChanges(); - } + if (async) + { + await context.SaveChangesAsync(); } else { - var customer = new OtherContainerCustomer { Id = "1", PartitionKey = "1" }; - context.Add(customer); - - if (async) - { - await context.SaveChangesAsync(); - } - else - { - context.SaveChanges(); - } + context.SaveChanges(); + } - initialToken = context.Database.GetSessionTokens().GetSessionToken(OtherContainerName)!; + var initialTokens = context.Database.GetSessionTokens().ToDictionary(); - context.Remove(customer); + customer.Name = "updated"; + otherContainerCustomer.Name = "updated"; - if (async) - { - await context.SaveChangesAsync(); - } - else - { - context.SaveChanges(); - } + if (async) + { + await context.SaveChangesAsync(); + } + else + { + context.SaveChanges(); } - var sessionToken = defaultContainer ? context.Database.GetSessionTokens().GetSessionToken() : context.Database.GetSessionTokens().GetSessionToken(OtherContainerName); - Assert.False(string.IsNullOrWhiteSpace(sessionToken)); - Assert.NotEqual(sessionToken, initialToken); - Assert.StartsWith(initialToken + ",", sessionToken); + var sessionTokens = context.Database.GetSessionTokens(); + foreach (var sessionToken in sessionTokens) + { + var initialToken = initialTokens[sessionToken.Key]; + Assert.NotEqual(sessionToken.Value, initialToken); + Assert.StartsWith(initialToken + ",", sessionToken.Value); + Assert.False(string.IsNullOrWhiteSpace(sessionToken.Value!.Substring(sessionToken.Value.IndexOf(",") + 1))); + } } [ConditionalTheory] - [InlineData(AutoTransactionBehavior.WhenNeeded, true, true)] - [InlineData(AutoTransactionBehavior.WhenNeeded, true, false)] - [InlineData(AutoTransactionBehavior.WhenNeeded, false, true)] - [InlineData(AutoTransactionBehavior.WhenNeeded, false, false)] - [InlineData(AutoTransactionBehavior.Never, true, true)] - [InlineData(AutoTransactionBehavior.Never, true, false)] - [InlineData(AutoTransactionBehavior.Never, false, true)] - [InlineData(AutoTransactionBehavior.Never, false, false)] - [InlineData(AutoTransactionBehavior.Always, true, true)] - [InlineData(AutoTransactionBehavior.Always, true, false)] - [InlineData(AutoTransactionBehavior.Always, false, true)] - [InlineData(AutoTransactionBehavior.Always, false, false)] - public virtual async Task Update_merges_session_token(AutoTransactionBehavior autoTransactionBehavior, bool async, bool defaultContainer) + [InlineData(true, true)] + [InlineData(true, false)] + [InlineData(false, true)] + [InlineData(false, false)] + public virtual async Task Update_always_merges_session_token(bool async, bool defaultContainer) { var contextFactory = await InitializeAsync(); using var context = contextFactory.CreateContext(); - context.Database.AutoTransactionBehavior = autoTransactionBehavior; + context.Database.AutoTransactionBehavior = AutoTransactionBehavior.Always; - string initialToken; if (defaultContainer) { - var customer = new Customer { Id = "1", PartitionKey = "1" }; - context.Customers.Add(customer); - - if (async) - { - await context.SaveChangesAsync(); - } - else - { - context.SaveChanges(); - } - - initialToken = context.Database.GetSessionTokens().GetSessionToken()!; - - customer.Name = "updated"; + context.Customers.Add(new Customer { Id = "1", PartitionKey = "1" }); } else { - var customer = new OtherContainerCustomer { Id = "1", PartitionKey = "1" }; - context.Add(customer); + context.OtherContainerCustomers.Add(new OtherContainerCustomer { Id = "1", PartitionKey = "1" }); + } - if (async) - { - await context.SaveChangesAsync(); - } - else - { - context.SaveChanges(); - } + if (async) + { + await context.SaveChangesAsync(); + } + else + { + context.SaveChanges(); + } - initialToken = context.Database.GetSessionTokens().GetSessionToken(OtherContainerName)!; + context.ChangeTracker.Clear(); + var initialTokens = context.Database.GetSessionTokens().ToDictionary(); - customer.Name = "updated"; + if (defaultContainer) + { + context.Customers.Update(new Customer { Id = "1", Name = "updated", PartitionKey = "1" }); + } + else + { + context.OtherContainerCustomers.Update(new OtherContainerCustomer { Id = "1", Name = "updated", PartitionKey = "1" }); } if (async) @@ -1254,10 +1044,9 @@ public virtual async Task Update_merges_session_token(AutoTransactionBehavior au context.SaveChanges(); } - var sessionToken = defaultContainer ? context.Database.GetSessionTokens().GetSessionToken() : context.Database.GetSessionTokens().GetSessionToken(OtherContainerName); + var sessionTokens = context.Database.GetSessionTokens(); + var sessionToken = defaultContainer ? sessionTokens[nameof(CosmosSessionTokenContext)]! : sessionTokens[OtherContainerName]!; Assert.False(string.IsNullOrWhiteSpace(sessionToken)); - Assert.NotEqual(initialToken, sessionToken); - Assert.StartsWith(initialToken + ",", sessionToken); } [ConditionalTheory] @@ -1331,7 +1120,6 @@ public virtual async Task Update_uses_session_token(AutoTransactionBehavior auto context.Database.AutoTransactionBehavior = autoTransactionBehavior; var sessionTokens = context.Database.GetSessionTokens(); - var sessionToken = sessionTokens.GetSessionToken()!; // Only way we can test this is by setting a session token that will fail the request if used.. // Only way to do this for a write is to set an invalid session token.. var internalDictionary = sessionTokens.GetType().GetField("_containerSessionTokens", BindingFlags.NonPublic | BindingFlags.Instance)!.GetValue(sessionTokens)!; @@ -1382,7 +1170,6 @@ public virtual async Task Delete_uses_session_token(AutoTransactionBehavior auto context.Database.AutoTransactionBehavior = autoTransactionBehavior; var sessionTokens = context.Database.GetSessionTokens(); - var sessionToken = sessionTokens.GetSessionToken()!; // Only way we can test this is by setting a session token that will fail the request if used.. // Only way to do this for a write is to set an invalid session token.. var internalDictionary = sessionTokens.GetType().GetField("_containerSessionTokens", BindingFlags.NonPublic | BindingFlags.Instance)!.GetValue(sessionTokens)!; diff --git a/test/EFCore.Cosmos.FunctionalTests/TestUtilities/CosmosTestStore.cs b/test/EFCore.Cosmos.FunctionalTests/TestUtilities/CosmosTestStore.cs index 7c4ff895950..e24d9f7b42c 100644 --- a/test/EFCore.Cosmos.FunctionalTests/TestUtilities/CosmosTestStore.cs +++ b/test/EFCore.Cosmos.FunctionalTests/TestUtilities/CosmosTestStore.cs @@ -241,7 +241,7 @@ private async Task CreateFromFile(DbContext context) document["$type"] = entityName; await cosmosClient.CreateItemAsync( - containerName!, document, new FakeUpdateEntry(), new NullSessionTokenStorage()).ConfigureAwait(false); + containerName!, document, new FakeUpdateEntry(), new NullSessionTokenStorage(context)).ConfigureAwait(false); } else if (reader.TokenType == JsonToken.EndObject) { diff --git a/test/EFCore.Cosmos.FunctionalTests/TestUtilities/NullSessionTokenStorage.cs b/test/EFCore.Cosmos.FunctionalTests/TestUtilities/NullSessionTokenStorage.cs index 1a9ff6fae52..96d99ceabd3 100644 --- a/test/EFCore.Cosmos.FunctionalTests/TestUtilities/NullSessionTokenStorage.cs +++ b/test/EFCore.Cosmos.FunctionalTests/TestUtilities/NullSessionTokenStorage.cs @@ -1,70 +1,27 @@ // 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.TestUtilities; -/// -/// 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. -/// -public sealed class NullSessionTokenStorage : ISessionTokenStorage +public sealed class NullSessionTokenStorage : SessionTokenStorage { - /// - /// 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. - /// - public void AppendSessionToken(string sessionToken) { } + public NullSessionTokenStorage(DbContext dbContext) : base(dbContext) + { + } - /// - /// 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. - /// - public void AppendSessionToken(string containerName, string sessionToken) { } + public override void AppendSessionToken(string sessionToken) { } - /// - /// 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. - /// - public void Clear() { } + public override void AppendSessionToken(string containerName, string sessionToken) { } - /// - /// 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. - /// - public string? GetSessionToken() => null; - /// - /// 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. - /// - public string? GetSessionToken(string containerName) => null; + public override void Clear() { } - /// - /// 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. - /// - public void SetSessionToken(string containerName, string? sessionToken) { } + public override string? GetSessionToken() => null; - /// - /// 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. - /// - public void SetSessionToken(string? sessionToken) { } + public override string? GetSessionToken(string containerName) => null; + + public override void SetSessionToken(string containerName, string? sessionToken) { } + + public override void AppendSessionTokens(IReadOnlyDictionary sessionTokens) { } } From e729af3a2c1ac435b95b640468eb850b759ce75b Mon Sep 17 00:00:00 2001 From: JoasE <32096708+JoasE@users.noreply.github.com> Date: Tue, 28 Oct 2025 19:09:13 +0100 Subject: [PATCH 25/37] Readd option ManualSessionTokenManagementEnabled --- .../CosmosDatabaseFacadeExtensions.cs | 40 ++++------- .../CosmosDbContextOptionsBuilder.cs | 17 +++++ .../Internal/CosmosDbOptionExtension.cs | 30 +++++++- .../Internal/CosmosSingletonOptions.cs | 10 +++ .../Internal/ICosmosSingletonOptions.cs | 8 +++ .../Properties/CosmosStrings.Designer.cs | 6 ++ .../Properties/CosmosStrings.resx | 3 + .../Query/Internal/CosmosQueryContext.cs | 5 +- .../Internal/CosmosQueryContextFactory.cs | 3 +- .../Storage/Internal/CosmosClientWrapper.cs | 66 ++++++++--------- .../Storage/Internal/CosmosDatabaseWrapper.cs | 8 ++- .../Storage/Internal/ICosmosClientWrapper.cs | 26 +++---- .../Storage/Internal/ISessionTokenStorage.cs | 70 +++++++++++++++++++ .../Internal/NullSessionTokenStorage.cs | 69 ++++++++++++++++++ .../Storage/Internal/SessionTokenStorage.cs | 2 +- .../CosmosSessionTokensTest.cs | 13 ++++ .../TestUtilities/CosmosTestStore.cs | 2 +- .../TestUtilities/NullSessionTokenStorage.cs | 27 ------- .../CosmosDbContextOptionsExtensionsTests.cs | 1 + 19 files changed, 295 insertions(+), 111 deletions(-) create mode 100644 src/EFCore.Cosmos/Storage/Internal/ISessionTokenStorage.cs create mode 100644 src/EFCore.Cosmos/Storage/Internal/NullSessionTokenStorage.cs delete mode 100644 test/EFCore.Cosmos.FunctionalTests/TestUtilities/NullSessionTokenStorage.cs diff --git a/src/EFCore.Cosmos/Extensions/CosmosDatabaseFacadeExtensions.cs b/src/EFCore.Cosmos/Extensions/CosmosDatabaseFacadeExtensions.cs index b342720700f..83f889fdb4e 100644 --- a/src/EFCore.Cosmos/Extensions/CosmosDatabaseFacadeExtensions.cs +++ b/src/EFCore.Cosmos/Extensions/CosmosDatabaseFacadeExtensions.cs @@ -33,15 +33,7 @@ public static CosmosClient GetCosmosClient(this DatabaseFacade databaseFacade) /// The for the context. /// The session token dictionary. public static string? GetSessionToken(this DatabaseFacade databaseFacade) - { - var db = GetService(databaseFacade); - if (db is not CosmosDatabaseWrapper dbWrapper) - { - throw new InvalidOperationException(CosmosStrings.CosmosNotInUse); - } - - return dbWrapper.SessionTokenStorage.GetSessionToken(); - } + => GetSessionTokenStorage(databaseFacade).GetSessionToken(); /// /// Gets a dictionary that contains the composite session token per container for this . @@ -50,15 +42,7 @@ public static CosmosClient GetCosmosClient(this DatabaseFacade databaseFacade) /// The for the context. /// The session token dictionary. public static IReadOnlyDictionary GetSessionTokens(this DatabaseFacade databaseFacade) - { - var db = GetService(databaseFacade); - if (db is not CosmosDatabaseWrapper dbWrapper) - { - throw new InvalidOperationException(CosmosStrings.CosmosNotInUse); - } - - return dbWrapper.SessionTokenStorage; - } + => GetSessionTokenStorage(databaseFacade); /// /// Appends the composite session token for the default container for this . @@ -67,15 +51,7 @@ public static CosmosClient GetCosmosClient(this DatabaseFacade databaseFacade) /// The for the context. /// The session token to append. public static void AppendSessionToken(this DatabaseFacade databaseFacade, string sessionToken) - { - var db = GetService(databaseFacade); - if (db is not CosmosDatabaseWrapper dbWrapper) - { - throw new InvalidOperationException(CosmosStrings.CosmosNotInUse); - } - - dbWrapper.SessionTokenStorage.AppendSessionToken(sessionToken); - } + => GetSessionTokenStorage(databaseFacade).AppendSessionToken(sessionToken); /// /// Appends the composite session token per container for this . @@ -84,6 +60,9 @@ public static void AppendSessionToken(this DatabaseFacade databaseFacade, string /// The for the context. /// The session tokens to append per container. public static void AppendSessionTokens(this DatabaseFacade databaseFacade, IReadOnlyDictionary sessionTokens) + => GetSessionTokenStorage(databaseFacade).AppendSessionTokens(sessionTokens); + + private static SessionTokenStorage GetSessionTokenStorage(DatabaseFacade databaseFacade) { var db = GetService(databaseFacade); if (db is not CosmosDatabaseWrapper dbWrapper) @@ -91,7 +70,12 @@ public static void AppendSessionTokens(this DatabaseFacade databaseFacade, IRead throw new InvalidOperationException(CosmosStrings.CosmosNotInUse); } - dbWrapper.SessionTokenStorage.AppendSessionTokens(sessionTokens); + if (dbWrapper.SessionTokenStorage is not SessionTokenStorage sts) + { + throw new InvalidOperationException(CosmosStrings.EnableManualSessionTokenManagement); + } + + return sts; } private static TService GetService(IInfrastructure databaseFacade) diff --git a/src/EFCore.Cosmos/Infrastructure/CosmosDbContextOptionsBuilder.cs b/src/EFCore.Cosmos/Infrastructure/CosmosDbContextOptionsBuilder.cs index e88fabf78da..c7494edcfe1 100644 --- a/src/EFCore.Cosmos/Infrastructure/CosmosDbContextOptionsBuilder.cs +++ b/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))); + + /// + /// Sets the boolean to track and manage session tokens for requests made to Cosmos DB + /// and being able to access them via the and 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 Utilize session tokens for more details. + /// + /// + /// See Using DbContextOptions, and + /// Accessing Azure Cosmos DB with EF Core for more information and examples. + /// + /// to track and manually manage session tokens in EF. + public virtual CosmosDbContextOptionsBuilder ManualSessionTokenManagementEnabled(bool enabled = true) + => WithOption(e => e.ManualSessionTokenManagementEnabled(enabled)); + /// /// 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. diff --git a/src/EFCore.Cosmos/Infrastructure/Internal/CosmosDbOptionExtension.cs b/src/EFCore.Cosmos/Infrastructure/Internal/CosmosDbOptionExtension.cs index f2545174ac3..dbe9e67e273 100644 --- a/src/EFCore.Cosmos/Infrastructure/Internal/CosmosDbOptionExtension.cs +++ b/src/EFCore.Cosmos/Infrastructure/Internal/CosmosDbOptionExtension.cs @@ -36,6 +36,7 @@ public class CosmosOptionsExtension : IDbContextOptionsExtension private bool? _enableContentResponseOnWrite; private DbContextOptionsExtensionInfo? _info; private Func? _httpClientFactory; + private bool _enableManualSessionTokenManagement; /// /// 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; } /// @@ -564,6 +566,30 @@ public virtual CosmosOptionsExtension WithHttpClientFactory(Func? ht return clone; } + /// + /// 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. + /// + public virtual bool EnableManualSessionTokenManagement + => _enableManualSessionTokenManagement; + + /// + /// 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. + /// + public virtual CosmosOptionsExtension ManualSessionTokenManagementEnabled(bool enabled) + { + var clone = Clone(); + + clone._enableManualSessionTokenManagement = enabled; + + return clone; + } + /// /// 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 debugInfo) { diff --git a/src/EFCore.Cosmos/Infrastructure/Internal/CosmosSingletonOptions.cs b/src/EFCore.Cosmos/Infrastructure/Internal/CosmosSingletonOptions.cs index af29229cfa5..d5221c0eef3 100644 --- a/src/EFCore.Cosmos/Infrastructure/Internal/CosmosSingletonOptions.cs +++ b/src/EFCore.Cosmos/Infrastructure/Internal/CosmosSingletonOptions.cs @@ -151,6 +151,14 @@ public class CosmosSingletonOptions : ICosmosSingletonOptions /// public virtual Func? HttpClientFactory { get; private set; } + /// + /// 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. + /// + public virtual bool EnableManualSessionTokenManagement { get; private set; } + /// /// 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( diff --git a/src/EFCore.Cosmos/Infrastructure/Internal/ICosmosSingletonOptions.cs b/src/EFCore.Cosmos/Infrastructure/Internal/ICosmosSingletonOptions.cs index a26b79a82b3..8113da2b25f 100644 --- a/src/EFCore.Cosmos/Infrastructure/Internal/ICosmosSingletonOptions.cs +++ b/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. /// Func? HttpClientFactory { get; } + + /// + /// 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. + /// + bool EnableManualSessionTokenManagement { get; } } diff --git a/src/EFCore.Cosmos/Properties/CosmosStrings.Designer.cs b/src/EFCore.Cosmos/Properties/CosmosStrings.Designer.cs index 53dfa712a5f..62e0ae5afc5 100644 --- a/src/EFCore.Cosmos/Properties/CosmosStrings.Designer.cs +++ b/src/EFCore.Cosmos/Properties/CosmosStrings.Designer.cs @@ -165,6 +165,12 @@ public static string ElementWithValueConverter(object? propertyType, object? str GetString("ElementWithValueConverter", nameof(propertyType), nameof(structuralType), nameof(property), nameof(elementType)), propertyType, structuralType, property, elementType); + /// + /// Enable manual session token management using CosmosDbContextOptionsBuilder.ManualSessionTokenManagementEnabled to use this method. + /// + public static string EnableManualSessionTokenManagement + => GetString("EnableManualSessionTokenManagement"); + /// /// The type of the etag property '{property}' on '{entityType}' is '{propertyType}'. All etag properties must be strings or have a string value converter. /// diff --git a/src/EFCore.Cosmos/Properties/CosmosStrings.resx b/src/EFCore.Cosmos/Properties/CosmosStrings.resx index e914bee6375..9a5f4df5d38 100644 --- a/src/EFCore.Cosmos/Properties/CosmosStrings.resx +++ b/src/EFCore.Cosmos/Properties/CosmosStrings.resx @@ -175,6 +175,9 @@ 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. + + Enable manual session token management using CosmosDbContextOptionsBuilder.ManualSessionTokenManagementEnabled to use this method. + The type of the etag property '{property}' on '{entityType}' is '{propertyType}'. All etag properties must be strings or have a string value converter. diff --git a/src/EFCore.Cosmos/Query/Internal/CosmosQueryContext.cs b/src/EFCore.Cosmos/Query/Internal/CosmosQueryContext.cs index 91ee46112a8..0249b3ebdc2 100644 --- a/src/EFCore.Cosmos/Query/Internal/CosmosQueryContext.cs +++ b/src/EFCore.Cosmos/Query/Internal/CosmosQueryContext.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 Microsoft.EntityFrameworkCore.Cosmos.Storage; using Microsoft.EntityFrameworkCore.Cosmos.Storage.Internal; namespace Microsoft.EntityFrameworkCore.Cosmos.Query.Internal; @@ -15,7 +14,7 @@ namespace Microsoft.EntityFrameworkCore.Cosmos.Query.Internal; public class CosmosQueryContext( QueryContextDependencies dependencies, ICosmosClientWrapper cosmosClient, - SessionTokenStorage sessionTokenStorage) + ISessionTokenStorage sessionTokenStorage) : QueryContext(dependencies) { /// @@ -32,5 +31,5 @@ public class CosmosQueryContext( /// 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. /// - public virtual SessionTokenStorage SessionTokenStorage { get; } = sessionTokenStorage; + public virtual ISessionTokenStorage SessionTokenStorage { get; } = sessionTokenStorage; } diff --git a/src/EFCore.Cosmos/Query/Internal/CosmosQueryContextFactory.cs b/src/EFCore.Cosmos/Query/Internal/CosmosQueryContextFactory.cs index ec9c6bf1558..245a89a1700 100644 --- a/src/EFCore.Cosmos/Query/Internal/CosmosQueryContextFactory.cs +++ b/src/EFCore.Cosmos/Query/Internal/CosmosQueryContextFactory.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 Microsoft.EntityFrameworkCore.Cosmos.Storage; using Microsoft.EntityFrameworkCore.Cosmos.Storage.Internal; namespace Microsoft.EntityFrameworkCore.Cosmos.Query.Internal; @@ -27,7 +26,7 @@ public class CosmosQueryContextFactory( /// 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. /// - protected virtual SessionTokenStorage SessionTokenStorage { get; } = ((CosmosDatabaseWrapper)dependencies.CurrentContext.Context.GetService()).SessionTokenStorage; + protected virtual ISessionTokenStorage SessionTokenStorage { get; } = ((CosmosDatabaseWrapper)dependencies.CurrentContext.Context.GetService()).SessionTokenStorage; /// /// This is an internal API that supports the Entity Framework Core infrastructure and not subject to diff --git a/src/EFCore.Cosmos/Storage/Internal/CosmosClientWrapper.cs b/src/EFCore.Cosmos/Storage/Internal/CosmosClientWrapper.cs index 7641d176f12..7f06009f230 100644 --- a/src/EFCore.Cosmos/Storage/Internal/CosmosClientWrapper.cs +++ b/src/EFCore.Cosmos/Storage/Internal/CosmosClientWrapper.cs @@ -389,7 +389,7 @@ public virtual bool CreateItem( string containerId, JToken document, IUpdateEntry entry, - SessionTokenStorage sessionTokenStorage) + ISessionTokenStorage sessionTokenStorage) { _databaseLogger.SyncNotSupported(); @@ -398,7 +398,7 @@ public virtual bool CreateItem( private static bool CreateItemOnce( DbContext context, - (string ContainerId, JToken Document, IUpdateEntry Entry, SessionTokenStorage SessionTokenStorage, CosmosClientWrapper Wrapper) parameters) + (string ContainerId, JToken Document, IUpdateEntry Entry, ISessionTokenStorage SessionTokenStorage, CosmosClientWrapper Wrapper) parameters) => CreateItemOnceAsync(context, parameters).GetAwaiter().GetResult(); /// @@ -411,13 +411,13 @@ public virtual Task CreateItemAsync( string containerId, JToken document, IUpdateEntry updateEntry, - SessionTokenStorage sessionTokenStorage, + ISessionTokenStorage sessionTokenStorage, CancellationToken cancellationToken = default) => _executionStrategy.ExecuteAsync((containerId, document, updateEntry, sessionTokenStorage, this), CreateItemOnceAsync, null, cancellationToken); private static async Task CreateItemOnceAsync( DbContext _, - (string ContainerId, JToken Document, IUpdateEntry Entry, SessionTokenStorage SessionTokenStorage, CosmosClientWrapper Wrapper) parameters, + (string ContainerId, JToken Document, IUpdateEntry Entry, ISessionTokenStorage SessionTokenStorage, CosmosClientWrapper Wrapper) parameters, CancellationToken cancellationToken = default) { using var stream = Serialize(parameters.Document); @@ -475,7 +475,7 @@ public virtual bool ReplaceItem( string documentId, JObject document, IUpdateEntry entry, - SessionTokenStorage sessionTokenStorage) + ISessionTokenStorage sessionTokenStorage) { _databaseLogger.SyncNotSupported(); @@ -484,7 +484,7 @@ public virtual bool ReplaceItem( private static bool ReplaceItemOnce( DbContext context, - (string ContainerId, string ItemId, JObject Document, IUpdateEntry Entry, SessionTokenStorage SessionTokenStorage, CosmosClientWrapper Wrapper) parameters) + (string ContainerId, string ItemId, JObject Document, IUpdateEntry Entry, ISessionTokenStorage SessionTokenStorage, CosmosClientWrapper Wrapper) parameters) => ReplaceItemOnceAsync(context, parameters).GetAwaiter().GetResult(); /// @@ -498,14 +498,14 @@ public virtual Task ReplaceItemAsync( string documentId, JObject document, IUpdateEntry updateEntry, - SessionTokenStorage sessionTokenStorage, + ISessionTokenStorage sessionTokenStorage, CancellationToken cancellationToken = default) => _executionStrategy.ExecuteAsync( (collectionId, documentId, document, updateEntry, sessionTokenStorage, this), ReplaceItemOnceAsync, null, cancellationToken); private static async Task ReplaceItemOnceAsync( DbContext _, - (string ContainerId, string ResourceId, JObject Document, IUpdateEntry Entry, SessionTokenStorage SessionTokenStorage, CosmosClientWrapper Wrapper) parameters, + (string ContainerId, string ResourceId, JObject Document, IUpdateEntry Entry, ISessionTokenStorage SessionTokenStorage, CosmosClientWrapper Wrapper) parameters, CancellationToken cancellationToken = default) { using var stream = Serialize(parameters.Document); @@ -563,7 +563,7 @@ public virtual bool DeleteItem( string containerId, string documentId, IUpdateEntry entry, - SessionTokenStorage sessionTokenStorage) + ISessionTokenStorage sessionTokenStorage) { _databaseLogger.SyncNotSupported(); @@ -572,7 +572,7 @@ public virtual bool DeleteItem( private static bool DeleteItemOnce( DbContext context, - (string ContainerId, string DocumentId, IUpdateEntry Entry, SessionTokenStorage SessionTokenStorage, CosmosClientWrapper Wrapper) parameters) + (string ContainerId, string DocumentId, IUpdateEntry Entry, ISessionTokenStorage SessionTokenStorage, CosmosClientWrapper Wrapper) parameters) => DeleteItemOnceAsync(context, parameters).GetAwaiter().GetResult(); /// @@ -585,13 +585,13 @@ public virtual Task DeleteItemAsync( string containerId, string documentId, IUpdateEntry entry, - SessionTokenStorage sessionTokenStorage, + ISessionTokenStorage sessionTokenStorage, CancellationToken cancellationToken = default) => _executionStrategy.ExecuteAsync((containerId, documentId, entry, sessionTokenStorage, this), DeleteItemOnceAsync, null, cancellationToken); private static async Task DeleteItemOnceAsync( DbContext? _, - (string ContainerId, string ResourceId, IUpdateEntry Entry, SessionTokenStorage SessionTokenStorage, CosmosClientWrapper Wrapper) parameters, + (string ContainerId, string ResourceId, IUpdateEntry Entry, ISessionTokenStorage SessionTokenStorage, CosmosClientWrapper Wrapper) parameters, CancellationToken cancellationToken = default) { var containerId = parameters.ContainerId; @@ -667,7 +667,7 @@ public virtual ICosmosTransactionalBatchWrapper CreateTransactionalBatch(string /// 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. /// - public virtual CosmosTransactionalBatchResult ExecuteTransactionalBatch(ICosmosTransactionalBatchWrapper batch, SessionTokenStorage sessionTokenStorage) + public virtual CosmosTransactionalBatchResult ExecuteTransactionalBatch(ICosmosTransactionalBatchWrapper batch, ISessionTokenStorage sessionTokenStorage) { _databaseLogger.SyncNotSupported(); @@ -675,7 +675,7 @@ public virtual CosmosTransactionalBatchResult ExecuteTransactionalBatch(ICosmosT } private static CosmosTransactionalBatchResult ExecuteBatchOnce(DbContext _, - (ICosmosTransactionalBatchWrapper Batch, SessionTokenStorage SessionTokenStorage, CosmosClientWrapper Wrapper) parameters) + (ICosmosTransactionalBatchWrapper Batch, ISessionTokenStorage SessionTokenStorage, CosmosClientWrapper Wrapper) parameters) => ExecuteBatchOnceAsync(_, parameters).GetAwaiter().GetResult(); /// @@ -684,11 +684,11 @@ private static CosmosTransactionalBatchResult ExecuteBatchOnce(DbContext _, /// 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. /// - public virtual Task ExecuteTransactionalBatchAsync(ICosmosTransactionalBatchWrapper batch, SessionTokenStorage sessionTokenStorage, CancellationToken cancellationToken = default) + public virtual Task ExecuteTransactionalBatchAsync(ICosmosTransactionalBatchWrapper batch, ISessionTokenStorage sessionTokenStorage, CancellationToken cancellationToken = default) => _executionStrategy.ExecuteAsync((batch, sessionTokenStorage, this), ExecuteBatchOnceAsync, null, cancellationToken); private static async Task ExecuteBatchOnceAsync(DbContext _, - (ICosmosTransactionalBatchWrapper Batch, SessionTokenStorage SessionTokenStorage, CosmosClientWrapper Wrapper) parameters, + (ICosmosTransactionalBatchWrapper Batch, ISessionTokenStorage SessionTokenStorage, CosmosClientWrapper Wrapper) parameters, CancellationToken cancellationToken = default) { var batch = parameters.Batch; @@ -780,7 +780,7 @@ private static PartitionKey ExtractPartitionKeyValue(IUpdateEntry entry) return builder.Build(); } - private static void ProcessResponse(string containerId, ResponseMessage response, IUpdateEntry entry, SessionTokenStorage sessionTokenStorage) + private static void ProcessResponse(string containerId, ResponseMessage response, IUpdateEntry entry, ISessionTokenStorage sessionTokenStorage) { response.EnsureSuccessStatusCode(); @@ -792,7 +792,7 @@ private static void ProcessResponse(string containerId, ResponseMessage response ProcessResponse(entry, response.Headers.ETag, response.Content); } - private static void ProcessResponse(string containerId, TransactionalBatchResponse batchResponse, IReadOnlyList entries, SessionTokenStorage sessionTokenStorage) + private static void ProcessResponse(string containerId, TransactionalBatchResponse batchResponse, IReadOnlyList entries, ISessionTokenStorage sessionTokenStorage) { if (!string.IsNullOrWhiteSpace(batchResponse.Headers.Session)) { @@ -840,7 +840,7 @@ public virtual IEnumerable ExecuteSqlQuery( string containerId, PartitionKey partitionKeyValue, CosmosSqlQuery query, - SessionTokenStorage sessionTokenStorage) + ISessionTokenStorage sessionTokenStorage) { _databaseLogger.SyncNotSupported(); @@ -859,7 +859,7 @@ public virtual IAsyncEnumerable ExecuteSqlQueryAsync( string containerId, PartitionKey partitionKeyValue, CosmosSqlQuery query, - SessionTokenStorage sessionTokenStorage) + ISessionTokenStorage sessionTokenStorage) { _commandLogger.ExecutingSqlQuery(containerId, partitionKeyValue, query); @@ -876,7 +876,7 @@ public virtual IAsyncEnumerable ExecuteSqlQueryAsync( string containerId, PartitionKey partitionKeyValue, string resourceId, - SessionTokenStorage sessionTokenStorage) + ISessionTokenStorage sessionTokenStorage) { _databaseLogger.SyncNotSupported(); @@ -905,7 +905,7 @@ public virtual IAsyncEnumerable ExecuteSqlQueryAsync( string containerId, PartitionKey partitionKeyValue, string resourceId, - SessionTokenStorage sessionTokenStorage, + ISessionTokenStorage sessionTokenStorage, CancellationToken cancellationToken = default) { _commandLogger.ExecutingReadItem(containerId, partitionKeyValue, resourceId); @@ -930,12 +930,12 @@ public virtual IAsyncEnumerable ExecuteSqlQueryAsync( private static ResponseMessage CreateSingleItemQuery( DbContext? context, - (string ContainerId, PartitionKey PartitionKeyValue, string ResourceId, SessionTokenStorage SessionTokenStorage, CosmosClientWrapper Wrapper) parameters) + (string ContainerId, PartitionKey PartitionKeyValue, string ResourceId, ISessionTokenStorage SessionTokenStorage, CosmosClientWrapper Wrapper) parameters) => CreateSingleItemQueryAsync(context, parameters).GetAwaiter().GetResult(); private static async Task CreateSingleItemQueryAsync( DbContext? _, - (string ContainerId, PartitionKey PartitionKeyValue, string ResourceId, SessionTokenStorage SessionTokenStorage, CosmosClientWrapper Wrapper) parameters, + (string ContainerId, PartitionKey PartitionKeyValue, string ResourceId, ISessionTokenStorage SessionTokenStorage, CosmosClientWrapper Wrapper) parameters, CancellationToken cancellationToken = default) { var (containerId, partitionKeyValue, resourceId, sessionTokenStorage, wrapper) = parameters; @@ -986,7 +986,7 @@ private static async Task CreateSingleItemQueryAsync( public virtual FeedIterator CreateQuery( string containerId, CosmosSqlQuery query, - SessionTokenStorage sessionTokenStorage, + ISessionTokenStorage sessionTokenStorage, string? continuationToken = null, QueryRequestOptions? queryRequestOptions = null) { @@ -1029,14 +1029,14 @@ private sealed class DocumentEnumerable( string containerId, PartitionKey partitionKeyValue, CosmosSqlQuery cosmosSqlQuery, - SessionTokenStorage sessionTokenStorage) + ISessionTokenStorage sessionTokenStorage) : IEnumerable { private readonly CosmosClientWrapper _cosmosClient = cosmosClient; private readonly string _containerId = containerId; private readonly PartitionKey _partitionKeyValue = partitionKeyValue; private readonly CosmosSqlQuery _cosmosSqlQuery = cosmosSqlQuery; - private readonly SessionTokenStorage _sessionTokenStorage = sessionTokenStorage; + private readonly ISessionTokenStorage _sessionTokenStorage = sessionTokenStorage; public IEnumerator GetEnumerator() => new Enumerator(this); @@ -1050,7 +1050,7 @@ private sealed class Enumerator(DocumentEnumerable documentEnumerable) : IEnumer private readonly string _containerId = documentEnumerable._containerId; private readonly PartitionKey _partitionKeyValue = documentEnumerable._partitionKeyValue; private readonly CosmosSqlQuery _cosmosSqlQuery = documentEnumerable._cosmosSqlQuery; - private readonly SessionTokenStorage _sessionTokenStorage = documentEnumerable._sessionTokenStorage; + private readonly ISessionTokenStorage _sessionTokenStorage = documentEnumerable._sessionTokenStorage; private JToken? _current; private ResponseMessage? _responseMessage; @@ -1141,14 +1141,14 @@ private sealed class DocumentAsyncEnumerable( string containerId, PartitionKey partitionKeyValue, CosmosSqlQuery cosmosSqlQuery, - SessionTokenStorage sessionTokenStorage) + ISessionTokenStorage sessionTokenStorage) : IAsyncEnumerable { private readonly CosmosClientWrapper _cosmosClient = cosmosClient; private readonly string _containerId = containerId; private readonly PartitionKey _partitionKeyValue = partitionKeyValue; private readonly CosmosSqlQuery _cosmosSqlQuery = cosmosSqlQuery; - private readonly SessionTokenStorage _sessionTokenStorage = sessionTokenStorage; + private readonly ISessionTokenStorage _sessionTokenStorage = sessionTokenStorage; public IAsyncEnumerator GetAsyncEnumerator(CancellationToken cancellationToken = default) => new AsyncEnumerator(this, cancellationToken); @@ -1160,7 +1160,7 @@ private sealed class AsyncEnumerator(DocumentAsyncEnumerable documentEnumerable, private readonly string _containerId = documentEnumerable._containerId; private readonly PartitionKey _partitionKeyValue = documentEnumerable._partitionKeyValue; private readonly CosmosSqlQuery _cosmosSqlQuery = documentEnumerable._cosmosSqlQuery; - private readonly SessionTokenStorage _sessionTokenStorage = documentEnumerable._sessionTokenStorage; + private readonly ISessionTokenStorage _sessionTokenStorage = documentEnumerable._sessionTokenStorage; private JToken? _current; private ResponseMessage? _responseMessage; @@ -1354,9 +1354,9 @@ private sealed class CosmosFeedIteratorWrapper : FeedIterator { private readonly FeedIterator _inner; private readonly string _containerName; - private readonly SessionTokenStorage _sessionTokenStorage; + private readonly ISessionTokenStorage _sessionTokenStorage; - public CosmosFeedIteratorWrapper(FeedIterator inner, string containerName, SessionTokenStorage sessionTokenStorage) + public CosmosFeedIteratorWrapper(FeedIterator inner, string containerName, ISessionTokenStorage sessionTokenStorage) { _inner = inner; _containerName = containerName; diff --git a/src/EFCore.Cosmos/Storage/Internal/CosmosDatabaseWrapper.cs b/src/EFCore.Cosmos/Storage/Internal/CosmosDatabaseWrapper.cs index f7f1fbddcf6..7cc6997ee4a 100644 --- a/src/EFCore.Cosmos/Storage/Internal/CosmosDatabaseWrapper.cs +++ b/src/EFCore.Cosmos/Storage/Internal/CosmosDatabaseWrapper.cs @@ -5,6 +5,7 @@ using System.Runtime.InteropServices; using Microsoft.EntityFrameworkCore.ChangeTracking.Internal; using Microsoft.EntityFrameworkCore.Cosmos.Diagnostics.Internal; +using Microsoft.EntityFrameworkCore.Cosmos.Infrastructure.Internal; using Microsoft.EntityFrameworkCore.Cosmos.Internal; using Microsoft.EntityFrameworkCore.Cosmos.Metadata.Internal; using Microsoft.EntityFrameworkCore.Cosmos.Update.Internal; @@ -38,13 +39,16 @@ public CosmosDatabaseWrapper( DatabaseDependencies dependencies, ICurrentDbContext currentDbContext, ICosmosClientWrapper cosmosClient, + ICosmosSingletonOptions cosmosSingletonOptions, ILoggingOptions loggingOptions) : base(dependencies) { _currentDbContext = currentDbContext; _cosmosClient = cosmosClient; - SessionTokenStorage = new SessionTokenStorage(_currentDbContext.Context); + SessionTokenStorage = cosmosSingletonOptions.EnableManualSessionTokenManagement ? + new SessionTokenStorage(_currentDbContext.Context) + : new NullSessionTokenStorage(); if (loggingOptions.IsSensitiveDataLoggingEnabled) { @@ -58,7 +62,7 @@ public CosmosDatabaseWrapper( /// 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. /// - public virtual SessionTokenStorage SessionTokenStorage { get; } + public virtual ISessionTokenStorage SessionTokenStorage { get; } /// /// This is an internal API that supports the Entity Framework Core infrastructure and not subject to diff --git a/src/EFCore.Cosmos/Storage/Internal/ICosmosClientWrapper.cs b/src/EFCore.Cosmos/Storage/Internal/ICosmosClientWrapper.cs index 54e2cea7a63..620fd3b054b 100644 --- a/src/EFCore.Cosmos/Storage/Internal/ICosmosClientWrapper.cs +++ b/src/EFCore.Cosmos/Storage/Internal/ICosmosClientWrapper.cs @@ -67,7 +67,7 @@ public interface ICosmosClientWrapper /// 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. /// - bool CreateItem(string containerId, JToken document, IUpdateEntry entry, SessionTokenStorage sessionTokenStorage); + bool CreateItem(string containerId, JToken document, IUpdateEntry entry, ISessionTokenStorage sessionTokenStorage); /// /// This is an internal API that supports the Entity Framework Core infrastructure and not subject to @@ -80,7 +80,7 @@ bool ReplaceItem( string documentId, JObject document, IUpdateEntry entry, - SessionTokenStorage sessionTokenStorage); + ISessionTokenStorage sessionTokenStorage); /// /// This is an internal API that supports the Entity Framework Core infrastructure and not subject to @@ -92,7 +92,7 @@ bool DeleteItem( string containerId, string documentId, IUpdateEntry entry, - SessionTokenStorage sessionTokenStorage); + ISessionTokenStorage sessionTokenStorage); /// /// This is an internal API that supports the Entity Framework Core infrastructure and not subject to @@ -104,7 +104,7 @@ Task CreateItemAsync( string containerId, JToken document, IUpdateEntry updateEntry, - SessionTokenStorage sessionTokenStorage, + ISessionTokenStorage sessionTokenStorage, CancellationToken cancellationToken = default); /// @@ -118,7 +118,7 @@ Task ReplaceItemAsync( string documentId, JObject document, IUpdateEntry updateEntry, - SessionTokenStorage sessionTokenStorage, + ISessionTokenStorage sessionTokenStorage, CancellationToken cancellationToken = default); /// @@ -131,7 +131,7 @@ Task DeleteItemAsync( string containerId, string documentId, IUpdateEntry entry, - SessionTokenStorage sessionTokenStorage, + ISessionTokenStorage sessionTokenStorage, CancellationToken cancellationToken = default); /// @@ -143,7 +143,7 @@ Task DeleteItemAsync( FeedIterator CreateQuery( string containerId, CosmosSqlQuery query, - SessionTokenStorage sessionTokenStorage, + ISessionTokenStorage sessionTokenStorage, string? continuationToken = null, QueryRequestOptions? queryRequestOptions = null); @@ -157,7 +157,7 @@ FeedIterator CreateQuery( string containerId, PartitionKey partitionKeyValue, string resourceId, - SessionTokenStorage sessionTokenStorage); + ISessionTokenStorage sessionTokenStorage); /// /// This is an internal API that supports the Entity Framework Core infrastructure and not subject to @@ -169,7 +169,7 @@ FeedIterator CreateQuery( string containerId, PartitionKey partitionKeyValue, string resourceId, - SessionTokenStorage sessionTokenStorage, + ISessionTokenStorage sessionTokenStorage, CancellationToken cancellationToken = default); /// @@ -182,7 +182,7 @@ IEnumerable ExecuteSqlQuery( string containerId, PartitionKey partitionKeyValue, CosmosSqlQuery query, - SessionTokenStorage sessionTokenStorage); + ISessionTokenStorage sessionTokenStorage); /// /// This is an internal API that supports the Entity Framework Core infrastructure and not subject to @@ -194,7 +194,7 @@ IAsyncEnumerable ExecuteSqlQueryAsync( string containerId, PartitionKey partitionKeyValue, CosmosSqlQuery query, - SessionTokenStorage sessionTokenStorage); + ISessionTokenStorage sessionTokenStorage); /// /// This is an internal API that supports the Entity Framework Core infrastructure and not subject to @@ -226,7 +226,7 @@ IAsyncEnumerable ExecuteSqlQueryAsync( /// 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. /// - Task ExecuteTransactionalBatchAsync(ICosmosTransactionalBatchWrapper batch, SessionTokenStorage sessionTokenStorage, CancellationToken cancellationToken = default); + Task ExecuteTransactionalBatchAsync(ICosmosTransactionalBatchWrapper batch, ISessionTokenStorage sessionTokenStorage, CancellationToken cancellationToken = default); /// /// This is an internal API that supports the Entity Framework Core infrastructure and not subject to @@ -234,5 +234,5 @@ IAsyncEnumerable ExecuteSqlQueryAsync( /// 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. /// - CosmosTransactionalBatchResult ExecuteTransactionalBatch(ICosmosTransactionalBatchWrapper batch, SessionTokenStorage sessionTokenStorage); + CosmosTransactionalBatchResult ExecuteTransactionalBatch(ICosmosTransactionalBatchWrapper batch, ISessionTokenStorage sessionTokenStorage); } diff --git a/src/EFCore.Cosmos/Storage/Internal/ISessionTokenStorage.cs b/src/EFCore.Cosmos/Storage/Internal/ISessionTokenStorage.cs new file mode 100644 index 00000000000..fbbeea80a46 --- /dev/null +++ b/src/EFCore.Cosmos/Storage/Internal/ISessionTokenStorage.cs @@ -0,0 +1,70 @@ +// 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.Storage.Internal; + +/// +/// 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. +/// +public interface ISessionTokenStorage +{ + /// + /// 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. + /// + public void AppendSessionToken(string sessionToken); + + /// + /// 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. + /// + public void AppendSessionToken(string containerName, string sessionToken); + + /// + /// 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. + /// + public void AppendSessionTokens(IReadOnlyDictionary sessionTokens); + + /// + /// 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. + /// + public void Clear(); + + /// + /// 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. + /// + public string? GetSessionToken(); + + /// + /// 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. + /// + public string? GetSessionToken(string containerName); + + /// + /// 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. + /// + public void SetSessionToken(string containerName, string? sessionToken); +} diff --git a/src/EFCore.Cosmos/Storage/Internal/NullSessionTokenStorage.cs b/src/EFCore.Cosmos/Storage/Internal/NullSessionTokenStorage.cs new file mode 100644 index 00000000000..9524a69392c --- /dev/null +++ b/src/EFCore.Cosmos/Storage/Internal/NullSessionTokenStorage.cs @@ -0,0 +1,69 @@ +// 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.Storage.Internal; + +/// +/// 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. +/// +public class NullSessionTokenStorage : ISessionTokenStorage +{ + /// + /// 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. + /// + public void AppendSessionToken(string sessionToken) { } + + /// + /// 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. + /// + public void AppendSessionToken(string containerName, string sessionToken) { } + + /// + /// 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. + /// + public void Clear() { } + + /// + /// 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. + /// + public string? GetSessionToken() => null; + + /// + /// 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. + /// + public string? GetSessionToken(string containerName) => null; + + /// + /// 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. + /// + public void SetSessionToken(string containerName, string? sessionToken) { } + + /// + /// 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. + /// + public void AppendSessionTokens(IReadOnlyDictionary sessionTokens) { } +} diff --git a/src/EFCore.Cosmos/Storage/Internal/SessionTokenStorage.cs b/src/EFCore.Cosmos/Storage/Internal/SessionTokenStorage.cs index fccbb1d920f..43264369018 100644 --- a/src/EFCore.Cosmos/Storage/Internal/SessionTokenStorage.cs +++ b/src/EFCore.Cosmos/Storage/Internal/SessionTokenStorage.cs @@ -15,7 +15,7 @@ namespace Microsoft.EntityFrameworkCore.Cosmos.Storage.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. /// -public class SessionTokenStorage : IReadOnlyDictionary +public class SessionTokenStorage : ISessionTokenStorage, IReadOnlyDictionary { private readonly Dictionary _containerSessionTokens; private readonly string _defaultContainerName; diff --git a/test/EFCore.Cosmos.FunctionalTests/CosmosSessionTokensTest.cs b/test/EFCore.Cosmos.FunctionalTests/CosmosSessionTokensTest.cs index e61cd871746..eb552039f4b 100644 --- a/test/EFCore.Cosmos.FunctionalTests/CosmosSessionTokensTest.cs +++ b/test/EFCore.Cosmos.FunctionalTests/CosmosSessionTokensTest.cs @@ -17,6 +17,19 @@ protected override ITestStoreFactory TestStoreFactory protected override DbContextOptionsBuilder AddOptions(DbContextOptionsBuilder builder) => base.AddOptions(builder).ConfigureWarnings(x => x.Ignore(CosmosEventId.SyncNotSupported)); + protected override TestStore CreateTestStore() => CosmosTestStore.Create(StoreName, (c) => c.ManualSessionTokenManagementEnabled()); + + [ConditionalFact] + public virtual async Task GetSessionTokens_throws_if_not_enabled() + { + var contextFactory = await InitializeAsync(createTestStore: () => CosmosTestStore.Create(StoreName)); + + using var context = contextFactory.CreateContext(); + + var exception = Assert.Throws(() => context.Database.GetSessionTokens()); + Assert.Equal(CosmosStrings.EnableManualSessionTokenManagement, exception.Message); + } + [ConditionalFact] public virtual async Task AppendSessionToken_ThrowsForNonExistentContainer() { diff --git a/test/EFCore.Cosmos.FunctionalTests/TestUtilities/CosmosTestStore.cs b/test/EFCore.Cosmos.FunctionalTests/TestUtilities/CosmosTestStore.cs index e24d9f7b42c..7c4ff895950 100644 --- a/test/EFCore.Cosmos.FunctionalTests/TestUtilities/CosmosTestStore.cs +++ b/test/EFCore.Cosmos.FunctionalTests/TestUtilities/CosmosTestStore.cs @@ -241,7 +241,7 @@ private async Task CreateFromFile(DbContext context) document["$type"] = entityName; await cosmosClient.CreateItemAsync( - containerName!, document, new FakeUpdateEntry(), new NullSessionTokenStorage(context)).ConfigureAwait(false); + containerName!, document, new FakeUpdateEntry(), new NullSessionTokenStorage()).ConfigureAwait(false); } else if (reader.TokenType == JsonToken.EndObject) { diff --git a/test/EFCore.Cosmos.FunctionalTests/TestUtilities/NullSessionTokenStorage.cs b/test/EFCore.Cosmos.FunctionalTests/TestUtilities/NullSessionTokenStorage.cs deleted file mode 100644 index 96d99ceabd3..00000000000 --- a/test/EFCore.Cosmos.FunctionalTests/TestUtilities/NullSessionTokenStorage.cs +++ /dev/null @@ -1,27 +0,0 @@ -// 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.Internal; - -namespace Microsoft.EntityFrameworkCore.TestUtilities; - -public sealed class NullSessionTokenStorage : SessionTokenStorage -{ - public NullSessionTokenStorage(DbContext dbContext) : base(dbContext) - { - } - - public override void AppendSessionToken(string sessionToken) { } - - public override void AppendSessionToken(string containerName, string sessionToken) { } - - public override void Clear() { } - - public override string? GetSessionToken() => null; - - public override string? GetSessionToken(string containerName) => null; - - public override void SetSessionToken(string containerName, string? sessionToken) { } - - public override void AppendSessionTokens(IReadOnlyDictionary sessionTokens) { } -} diff --git a/test/EFCore.Cosmos.Tests/Extensions/CosmosDbContextOptionsExtensionsTests.cs b/test/EFCore.Cosmos.Tests/Extensions/CosmosDbContextOptionsExtensionsTests.cs index 895f45889bf..50b04a60a93 100644 --- a/test/EFCore.Cosmos.Tests/Extensions/CosmosDbContextOptionsExtensionsTests.cs +++ b/test/EFCore.Cosmos.Tests/Extensions/CosmosDbContextOptionsExtensionsTests.cs @@ -65,6 +65,7 @@ public void Can_create_options_with_valid_values() Test(o => o.MaxTcpConnectionsPerEndpoint(3), o => Assert.Equal(3, o.MaxTcpConnectionsPerEndpoint)); Test(o => o.LimitToEndpoint(), o => Assert.True(o.LimitToEndpoint)); Test(o => o.ContentResponseOnWriteEnabled(), o => Assert.True(o.EnableContentResponseOnWrite)); + Test(o => o.ManualSessionTokenManagementEnabled(), o => Assert.True(o.EnableManualSessionTokenManagement)); var webProxy = new WebProxy(); Test(o => o.WebProxy(webProxy), o => Assert.Same(webProxy, o.WebProxy)); From 9974e8fbfd1cac24b12af899215917840793a4df Mon Sep 17 00:00:00 2001 From: JoasE <32096708+JoasE@users.noreply.github.com> Date: Sat, 1 Nov 2025 12:03:44 +0100 Subject: [PATCH 26/37] Move container checks to extension method and other small improvements and cleanup --- .../CosmosDatabaseFacadeExtensions.cs | 26 ++++- .../Extensions/CosmosModelExtensions.cs | 8 -- .../CosmosRuntimeModelConvention.cs | 11 --- .../Internal/CosmosAnnotationNames.cs | 8 -- .../Properties/CosmosStrings.Designer.cs | 2 +- .../Properties/CosmosStrings.resx | 2 +- .../Storage/Internal/ISessionTokenStorage.cs | 21 ++-- .../Internal/NullSessionTokenStorage.cs | 8 ++ .../Storage/Internal/SessionTokenStorage.cs | 76 +++------------ .../CosmosSessionTokensTest.cs | 96 ++++++++++--------- 10 files changed, 104 insertions(+), 154 deletions(-) diff --git a/src/EFCore.Cosmos/Extensions/CosmosDatabaseFacadeExtensions.cs b/src/EFCore.Cosmos/Extensions/CosmosDatabaseFacadeExtensions.cs index 83f889fdb4e..2a7d3e6a78a 100644 --- a/src/EFCore.Cosmos/Extensions/CosmosDatabaseFacadeExtensions.cs +++ b/src/EFCore.Cosmos/Extensions/CosmosDatabaseFacadeExtensions.cs @@ -3,7 +3,6 @@ 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 @@ -42,7 +41,7 @@ public static CosmosClient GetCosmosClient(this DatabaseFacade databaseFacade) /// The for the context. /// The session token dictionary. public static IReadOnlyDictionary GetSessionTokens(this DatabaseFacade databaseFacade) - => GetSessionTokenStorage(databaseFacade); + => GetSessionTokenStorage(databaseFacade).ToDictionary(); /// /// Appends the composite session token for the default container for this . @@ -60,7 +59,28 @@ public static void AppendSessionToken(this DatabaseFacade databaseFacade, string /// The for the context. /// The session tokens to append per container. public static void AppendSessionTokens(this DatabaseFacade databaseFacade, IReadOnlyDictionary sessionTokens) - => GetSessionTokenStorage(databaseFacade).AppendSessionTokens(sessionTokens); + { + var sessionTokenStorage = GetSessionTokenStorage(databaseFacade); + + var containerNames = GetContainerNames(databaseFacade.GetService()); + foreach (var sessionToken in sessionTokens) + { + if (!containerNames.Contains(sessionToken.Key)) + { + throw new InvalidOperationException(CosmosStrings.ContainerNameDoesNotExist(sessionToken.Key)); + } + } + + sessionTokenStorage.AppendSessionTokens(sessionTokens); + } + + private static HashSet GetContainerNames(IModel model) + => model.GetEntityTypes() + .Where(et => et.FindPrimaryKey() != null) + .Select(et => et.GetContainer()) + .Where(container => container != null) + .Distinct()! + .ToHashSet()!; private static SessionTokenStorage GetSessionTokenStorage(DatabaseFacade databaseFacade) { diff --git a/src/EFCore.Cosmos/Extensions/CosmosModelExtensions.cs b/src/EFCore.Cosmos/Extensions/CosmosModelExtensions.cs index 7d52fc47071..be6f8f79366 100644 --- a/src/EFCore.Cosmos/Extensions/CosmosModelExtensions.cs +++ b/src/EFCore.Cosmos/Extensions/CosmosModelExtensions.cs @@ -23,14 +23,6 @@ public static class CosmosModelExtensions public static string? GetDefaultContainer(this IReadOnlyModel model) => (string?)model[CosmosAnnotationNames.ContainerName]; - /// - /// Returns the all container names used in the model. - /// - /// The model. - /// A set of the names of the containers used in the model. - public static HashSet GetContainerNames(this IReadOnlyModel model) - => (HashSet)model[CosmosAnnotationNames.ContainerNames]!; - /// /// Sets the default container name. /// diff --git a/src/EFCore.Cosmos/Metadata/Conventions/CosmosRuntimeModelConvention.cs b/src/EFCore.Cosmos/Metadata/Conventions/CosmosRuntimeModelConvention.cs index b5d53c3ce2e..4ae848ac437 100644 --- a/src/EFCore.Cosmos/Metadata/Conventions/CosmosRuntimeModelConvention.cs +++ b/src/EFCore.Cosmos/Metadata/Conventions/CosmosRuntimeModelConvention.cs @@ -44,19 +44,8 @@ protected override void ProcessModelAnnotations( { annotations.Remove(CosmosAnnotationNames.Throughput); } - - // @TODO: Is this the right place for this? - annotations.Add(CosmosAnnotationNames.ContainerNames, GetContainerNames(model)); } - private HashSet GetContainerNames(IModel model) - => model.GetEntityTypes() - .Where(et => et.FindPrimaryKey() != null) - .Select(et => et.GetContainer()) - .Where(container => container != null) - .Distinct()! - .ToHashSet()!; - /// /// Updates the entity type annotations that will be set on the read-only object. /// diff --git a/src/EFCore.Cosmos/Metadata/Internal/CosmosAnnotationNames.cs b/src/EFCore.Cosmos/Metadata/Internal/CosmosAnnotationNames.cs index 85f259dfa12..5aac3bede02 100644 --- a/src/EFCore.Cosmos/Metadata/Internal/CosmosAnnotationNames.cs +++ b/src/EFCore.Cosmos/Metadata/Internal/CosmosAnnotationNames.cs @@ -27,14 +27,6 @@ public static class CosmosAnnotationNames /// public const string ContainerName = Prefix + "ContainerName"; - /// - /// 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. - /// - public const string ContainerNames = Prefix + "ContainerNames"; // @TODO: is this the right way? - /// /// 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 diff --git a/src/EFCore.Cosmos/Properties/CosmosStrings.Designer.cs b/src/EFCore.Cosmos/Properties/CosmosStrings.Designer.cs index 62e0ae5afc5..22fcb2e18cd 100644 --- a/src/EFCore.Cosmos/Properties/CosmosStrings.Designer.cs +++ b/src/EFCore.Cosmos/Properties/CosmosStrings.Designer.cs @@ -166,7 +166,7 @@ public static string ElementWithValueConverter(object? propertyType, object? str propertyType, structuralType, property, elementType); /// - /// Enable manual session token management using CosmosDbContextOptionsBuilder.ManualSessionTokenManagementEnabled to use this method. + /// Enable manual session token management using 'options.ManualSessionTokenManagementEnabled' to use this method. /// public static string EnableManualSessionTokenManagement => GetString("EnableManualSessionTokenManagement"); diff --git a/src/EFCore.Cosmos/Properties/CosmosStrings.resx b/src/EFCore.Cosmos/Properties/CosmosStrings.resx index 9a5f4df5d38..71bd33f52cf 100644 --- a/src/EFCore.Cosmos/Properties/CosmosStrings.resx +++ b/src/EFCore.Cosmos/Properties/CosmosStrings.resx @@ -176,7 +176,7 @@ 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. - Enable manual session token management using CosmosDbContextOptionsBuilder.ManualSessionTokenManagementEnabled to use this method. + Enable manual session token management using 'options.ManualSessionTokenManagementEnabled' to use this method. The type of the etag property '{property}' on '{entityType}' is '{propertyType}'. All etag properties must be strings or have a string value converter. diff --git a/src/EFCore.Cosmos/Storage/Internal/ISessionTokenStorage.cs b/src/EFCore.Cosmos/Storage/Internal/ISessionTokenStorage.cs index fbbeea80a46..97a1a4927a8 100644 --- a/src/EFCore.Cosmos/Storage/Internal/ISessionTokenStorage.cs +++ b/src/EFCore.Cosmos/Storage/Internal/ISessionTokenStorage.cs @@ -2,6 +2,7 @@ // The .NET Foundation licenses this file to you under the MIT license. + namespace Microsoft.EntityFrameworkCore.Cosmos.Storage.Internal; /// @@ -18,15 +19,7 @@ public interface ISessionTokenStorage /// 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. /// - public void AppendSessionToken(string sessionToken); - - /// - /// 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. - /// - public void AppendSessionToken(string containerName, string sessionToken); + string? GetSessionToken(); /// /// This is an internal API that supports the Entity Framework Core infrastructure and not subject to @@ -34,7 +27,7 @@ public interface ISessionTokenStorage /// 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. /// - public void AppendSessionTokens(IReadOnlyDictionary sessionTokens); + void AppendSessionToken(string sessionToken); /// /// This is an internal API that supports the Entity Framework Core infrastructure and not subject to @@ -42,7 +35,7 @@ public interface ISessionTokenStorage /// 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. /// - public void Clear(); + public void AppendSessionToken(string containerName, string sessionToken); /// /// This is an internal API that supports the Entity Framework Core infrastructure and not subject to @@ -50,7 +43,7 @@ public interface ISessionTokenStorage /// 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. /// - public string? GetSessionToken(); + public string? GetSessionToken(string containerName); /// /// This is an internal API that supports the Entity Framework Core infrastructure and not subject to @@ -58,7 +51,7 @@ public interface ISessionTokenStorage /// 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. /// - public string? GetSessionToken(string containerName); + public void Clear(); /// /// This is an internal API that supports the Entity Framework Core infrastructure and not subject to @@ -66,5 +59,5 @@ public interface ISessionTokenStorage /// 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. /// - public void SetSessionToken(string containerName, string? sessionToken); + public IReadOnlyDictionary ToDictionary(); } diff --git a/src/EFCore.Cosmos/Storage/Internal/NullSessionTokenStorage.cs b/src/EFCore.Cosmos/Storage/Internal/NullSessionTokenStorage.cs index 9524a69392c..dbcffa10e9f 100644 --- a/src/EFCore.Cosmos/Storage/Internal/NullSessionTokenStorage.cs +++ b/src/EFCore.Cosmos/Storage/Internal/NullSessionTokenStorage.cs @@ -66,4 +66,12 @@ public void SetSessionToken(string containerName, string? sessionToken) { } /// doing so can result in application failures when updating to a new Entity Framework Core release. /// public void AppendSessionTokens(IReadOnlyDictionary sessionTokens) { } + + /// + /// 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. + /// + public IReadOnlyDictionary ToDictionary() => new Dictionary(); } diff --git a/src/EFCore.Cosmos/Storage/Internal/SessionTokenStorage.cs b/src/EFCore.Cosmos/Storage/Internal/SessionTokenStorage.cs index 43264369018..f7750c7c066 100644 --- a/src/EFCore.Cosmos/Storage/Internal/SessionTokenStorage.cs +++ b/src/EFCore.Cosmos/Storage/Internal/SessionTokenStorage.cs @@ -1,10 +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 System.Collections; -using System.Runtime.CompilerServices; using System.Runtime.InteropServices; -using Microsoft.EntityFrameworkCore.Cosmos.Internal; using Microsoft.EntityFrameworkCore.Cosmos.Metadata.Internal; namespace Microsoft.EntityFrameworkCore.Cosmos.Storage.Internal; @@ -15,9 +12,9 @@ namespace Microsoft.EntityFrameworkCore.Cosmos.Storage.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. /// -public class SessionTokenStorage : ISessionTokenStorage, IReadOnlyDictionary +public class SessionTokenStorage : ISessionTokenStorage { - private readonly Dictionary _containerSessionTokens; + private readonly Dictionary _containerSessionTokens = new(); private readonly string _defaultContainerName; /// @@ -28,11 +25,7 @@ public class SessionTokenStorage : ISessionTokenStorage, IReadOnlyDictionary public SessionTokenStorage(DbContext dbContext) { - var defaultContainerName = (string)dbContext.Model.GetAnnotation(CosmosAnnotationNames.ContainerName).Value!; - var containerNames = (HashSet)dbContext.Model.GetAnnotation(CosmosAnnotationNames.ContainerNames).Value!; - - _defaultContainerName = defaultContainerName; - _containerSessionTokens = containerNames.ToDictionary(containerName => containerName, _ => new CompositeSessionToken()); + _defaultContainerName = (string)dbContext.Model.GetAnnotation(CosmosAnnotationNames.ContainerName).Value!; } /// @@ -64,7 +57,7 @@ public virtual void AppendSessionToken(string sessionToken) if (!_containerSessionTokens.TryGetValue(containerName, out var value)) { - throw new ArgumentException(CosmosStrings.ContainerNameDoesNotExist(containerName), nameof(containerName)); + return null; } return value.ConvertToString(); @@ -81,7 +74,6 @@ public virtual void AppendSessionTokens(IReadOnlyDictionary sess ArgumentNullException.ThrowIfNull(sessionTokens, nameof(sessionTokens)); foreach (var (containerName, sessionToken) in sessionTokens) { - // @TODO: null checks? AppendSessionToken(containerName, sessionToken); } } @@ -97,12 +89,13 @@ public virtual void AppendSessionToken(string containerName, string sessionToken ArgumentNullException.ThrowIfNullOrWhiteSpace(containerName, nameof(containerName)); ArgumentNullException.ThrowIfNullOrWhiteSpace(sessionToken, nameof(sessionToken)); - if (!_containerSessionTokens.TryGetValue(containerName, out var compositeSessionToken)) + ref var compositeSessionToken = ref CollectionsMarshal.GetValueRefOrAddDefault(_containerSessionTokens, containerName, out var exists); + if (!exists) { - throw new ArgumentException(CosmosStrings.ContainerNameDoesNotExist(containerName), nameof(containerName)); + compositeSessionToken = new(); } - compositeSessionToken.Add(sessionToken); + compositeSessionToken!.Add(sessionToken); } /// @@ -111,29 +104,12 @@ public virtual void AppendSessionToken(string containerName, string sessionToken /// 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. /// - public virtual void SetSessionToken(string containerName, string? sessionToken) + public virtual void Clear() { - ArgumentNullException.ThrowIfNullOrWhiteSpace(containerName, nameof(containerName)); - if (sessionToken is not null && string.IsNullOrWhiteSpace(sessionToken)) - { - throw new ArgumentException(CosmosStrings.SessionTokenCanNotBeWhiteSpace, nameof(sessionToken)); - } - - ref var compositeSessionToken = ref CollectionsMarshal.GetValueRefOrNullRef(_containerSessionTokens, containerName); - - if (Unsafe.IsNullRef(ref compositeSessionToken)) - { - throw new ArgumentException(CosmosStrings.ContainerNameDoesNotExist(containerName), nameof(containerName)); - } - - compositeSessionToken = new CompositeSessionToken(); - - if (sessionToken is null) + foreach (var key in _containerSessionTokens.Keys) { - return; + _containerSessionTokens[key] = new CompositeSessionToken(); } - - compositeSessionToken.Add(sessionToken); } /// @@ -142,13 +118,7 @@ public virtual void SetSessionToken(string containerName, string? sessionToken) /// 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. /// - public virtual void Clear() - { - foreach (var key in _containerSessionTokens.Keys) - { - _containerSessionTokens[key] = new CompositeSessionToken(); - } - } + public virtual IReadOnlyDictionary ToDictionary() => _containerSessionTokens.ToDictionary(x => x.Key, x => x.Value.ConvertToString()); private sealed class CompositeSessionToken { @@ -178,26 +148,4 @@ public void Add(string token) return _string; } } - - - #region IReadOnlyDictionary - IEnumerable IReadOnlyDictionary.Keys => _containerSessionTokens.Keys; - - IEnumerable IReadOnlyDictionary.Values => _containerSessionTokens.Values.Select(v => v.ConvertToString()); - - int IReadOnlyCollection>.Count => _containerSessionTokens.Count; - - string? IReadOnlyDictionary.this[string key] => _containerSessionTokens[key].ConvertToString(); - - bool IReadOnlyDictionary.ContainsKey(string key) => _containerSessionTokens.ContainsKey(key); - - bool IReadOnlyDictionary.TryGetValue(string key, out string? value) - => _containerSessionTokens.TryGetValue(key, out var compositeSessionToken) ? - (value = compositeSessionToken.ConvertToString()) != null - : (value = null) == null; - - IEnumerator> IEnumerable>.GetEnumerator() => _containerSessionTokens.Select(x => new KeyValuePair(x.Key, x.Value.ConvertToString())).GetEnumerator(); - - IEnumerator IEnumerable.GetEnumerator() => ((IEnumerable>)this).GetEnumerator(); - #endregion } diff --git a/test/EFCore.Cosmos.FunctionalTests/CosmosSessionTokensTest.cs b/test/EFCore.Cosmos.FunctionalTests/CosmosSessionTokensTest.cs index eb552039f4b..8e60578f73e 100644 --- a/test/EFCore.Cosmos.FunctionalTests/CosmosSessionTokensTest.cs +++ b/test/EFCore.Cosmos.FunctionalTests/CosmosSessionTokensTest.cs @@ -31,13 +31,13 @@ public virtual async Task GetSessionTokens_throws_if_not_enabled() } [ConditionalFact] - public virtual async Task AppendSessionToken_ThrowsForNonExistentContainer() + public virtual async Task AppendSessionTokens_throws_for_non_existent_container() { var contextFactory = await InitializeAsync(); using var context = contextFactory.CreateContext(); - var exception = Assert.Throws(() => context.Database.AppendSessionTokens(new Dictionary { { OtherContainerName, "0:-1#231"}, { "Not the container name", "0:-1#231" } })); - Assert.Equal(CosmosStrings.ContainerNameDoesNotExist("Not the container name") + " (Parameter 'containerName')", exception.Message); + var exception = Assert.Throws(() => context.Database.AppendSessionTokens(new Dictionary { { OtherContainerName, "0:-1#231"}, { "Not the container name", "0:-1#231" } })); + Assert.Equal(CosmosStrings.ContainerNameDoesNotExist("Not the container name"), exception.Message); } [ConditionalFact] @@ -68,8 +68,6 @@ public virtual async Task AppendSessionTokens_no_tokens_sets_tokens() Assert.Equal("0:-1#123", sessionTokens[OtherContainerName]); Assert.Equal("0:-1#231", sessionTokens[nameof(CosmosSessionTokenContext)]); - Assert.Equal(nameof(CosmosSessionTokenContext), sessionTokens.First().Key); - Assert.Equal(sessionTokens[nameof(CosmosSessionTokenContext)], sessionTokens.First().Value); } [ConditionalFact] @@ -154,16 +152,21 @@ public virtual async Task AppendSessionTokens_multiple_tokens_splits_tokens() } [ConditionalFact] - public virtual async Task GetSessionToken_no_token_returns_null() + public virtual async Task GetSessionTokens_no_token_returns_empty() { var contextFactory = await InitializeAsync(); using var context = contextFactory.CreateContext(); var sessionTokens = context.Database.GetSessionTokens(); - Assert.Equal(2, sessionTokens.Count); - Assert.Equal(nameof(CosmosSessionTokenContext), sessionTokens.First().Key); - Assert.Null(sessionTokens.First().Value); - Assert.Equal(OtherContainerName, sessionTokens.Last().Key); - Assert.Null(sessionTokens.Last().Value); + Assert.Equal(0, sessionTokens.Count); + } + + [ConditionalFact] + public virtual async Task GetSessionToken_no_token_returns_null() + { + var contextFactory = await InitializeAsync(); + using var context = contextFactory.CreateContext(); + var sessionToken = context.Database.GetSessionToken(); + Assert.Null(sessionToken); } [ConditionalTheory, InlineData(true), InlineData(false)] @@ -303,16 +306,8 @@ public virtual async Task Query_on_same_newly_pooled_context_does_not_use_same_s using var newContext = contextFactory.CreateContext(); Assert.Same(newContext, contextCopy); - if (async) // @TODO: Maybe even no if here? - { - await newContext.Customers.ToListAsync(); - await newContext.OtherContainerCustomers.ToListAsync(); - } - else - { - newContext.Customers.ToList(); - newContext.OtherContainerCustomers.ToList(); - } + await newContext.Customers.ToListAsync(); + await newContext.OtherContainerCustomers.ToListAsync(); } [ConditionalTheory, InlineData(true), InlineData(false)] @@ -476,7 +471,7 @@ public virtual async Task Query_appends_session_token(bool async) _ = context.OtherContainerCustomers.ToList(); } - var initialTokens = context.Database.GetSessionTokens().ToDictionary(); // Make a copy + var initialTokens = context.Database.GetSessionTokens(); using var otherContext = contextFactory.CreateContext(); otherContext.Customers.Add(new Customer { Id = "1", PartitionKey = "1" }); @@ -530,7 +525,7 @@ public virtual async Task Query_same_session_does_not_append_session_token(bool _ = context.OtherContainerCustomers.ToList(); } - var initialTokens = context.Database.GetSessionTokens().ToDictionary(); // Make a copy + var initialTokens = context.Database.GetSessionTokens(); if (async) { @@ -560,7 +555,7 @@ public virtual async Task PagingQuery_appends_session_token() await context.Customers.ToPageAsync(1, null); await context.OtherContainerCustomers.ToPageAsync(1, null); - var initialTokens = context.Database.GetSessionTokens().ToDictionary(); // Make a copy + var initialTokens = context.Database.GetSessionTokens(); using var otherContext = contextFactory.CreateContext(); otherContext.Customers.Add(new Customer { Id = "1", PartitionKey = "1" }); @@ -600,7 +595,7 @@ public virtual async Task Read_item_appends_session_token(bool async) _ = context.OtherContainerCustomers.ToList(); } - var initialTokens = context.Database.GetSessionTokens().ToDictionary(); // Make a copy + var initialTokens = context.Database.GetSessionTokens(); using var otherContext = contextFactory.CreateContext(); otherContext.Customers.Add(new Customer { Id = "1", PartitionKey = "1" }); @@ -654,7 +649,7 @@ public virtual async Task Read_item_enumerable_sets_session_token(bool async) _ = context.OtherContainerCustomers.ToList(); } - var initialTokens = context.Database.GetSessionTokens().ToDictionary(); // Make a copy + var initialTokens = context.Database.GetSessionTokens(); using var otherContext = contextFactory.CreateContext(); otherContext.Customers.Add(new Customer { Id = "1", PartitionKey = "1" }); @@ -775,7 +770,7 @@ public virtual async Task Add_never_merges_session_token(bool async) context.SaveChanges(); } - var initialTokens = context.Database.GetSessionTokens().ToDictionary(); + var initialTokens = context.Database.GetSessionTokens(); context.Customers.Add(new Customer { Id = "2", PartitionKey = "1" }); context.OtherContainerCustomers.Add(new OtherContainerCustomer { Id = "2", PartitionKey = "1" }); @@ -829,7 +824,7 @@ public virtual async Task Add_always_merges_session_token(bool async, bool defau context.SaveChanges(); } - var initialTokens = context.Database.GetSessionTokens().ToDictionary(); + var initialTokens = context.Database.GetSessionTokens(); if (defaultContainer) { @@ -878,7 +873,7 @@ public virtual async Task Delete_never_merges_session_token(bool async) context.SaveChanges(); } - var initialTokens = context.Database.GetSessionTokens().ToDictionary(); + var initialTokens = context.Database.GetSessionTokens(); context.Customers.Remove(customer); context.OtherContainerCustomers.Remove(otherContainerCustomer); @@ -933,7 +928,7 @@ public virtual async Task Delete_always_merges_session_token(bool async, bool de } context.ChangeTracker.Clear(); - var initialTokens = context.Database.GetSessionTokens().ToDictionary(); + var initialTokens = context.Database.GetSessionTokens(); if (defaultContainer) { @@ -982,7 +977,7 @@ public virtual async Task Update_never_merges_session_token(bool async) context.SaveChanges(); } - var initialTokens = context.Database.GetSessionTokens().ToDictionary(); + var initialTokens = context.Database.GetSessionTokens(); customer.Name = "updated"; otherContainerCustomer.Name = "updated"; @@ -1037,7 +1032,7 @@ public virtual async Task Update_always_merges_session_token(bool async, bool de } context.ChangeTracker.Clear(); - var initialTokens = context.Database.GetSessionTokens().ToDictionary(); + var initialTokens = context.Database.GetSessionTokens(); if (defaultContainer) { @@ -1085,10 +1080,15 @@ public virtual async Task Add_uses_session_token(AutoTransactionBehavior autoTra var sessionTokens = context.Database.GetSessionTokens(); // Only way we can test this is by setting a session token that will fail the request if used.. // Only way to do this for a write is to set an invalid session token.. - var internalDictionary = sessionTokens.GetType().GetField("_containerSessionTokens", BindingFlags.NonPublic | BindingFlags.Instance)!.GetValue(sessionTokens)!; - var internalComposite = internalDictionary.GetType().GetProperty("Item", BindingFlags.Public | BindingFlags.Instance)!.GetValue(internalDictionary, new object[] { defaultContainer ? nameof(CosmosSessionTokenContext) : OtherContainerName })!; - internalComposite.GetType().GetField("_string", BindingFlags.NonPublic | BindingFlags.Instance)!.SetValue(internalComposite, "invalidtoken"); - internalComposite.GetType().GetField("_isChanged", BindingFlags.NonPublic | BindingFlags.Instance)!.SetValue(internalComposite, false); + + if (defaultContainer) + { + context.Database.AppendSessionToken("invalidtoken"); + } + else + { + context.Database.AppendSessionTokens(new Dictionary { { OtherContainerName, "invalidtoken" } }); + } if (defaultContainer) { @@ -1135,10 +1135,14 @@ public virtual async Task Update_uses_session_token(AutoTransactionBehavior auto var sessionTokens = context.Database.GetSessionTokens(); // Only way we can test this is by setting a session token that will fail the request if used.. // Only way to do this for a write is to set an invalid session token.. - var internalDictionary = sessionTokens.GetType().GetField("_containerSessionTokens", BindingFlags.NonPublic | BindingFlags.Instance)!.GetValue(sessionTokens)!; - var internalComposite = internalDictionary.GetType().GetProperty("Item", BindingFlags.Public | BindingFlags.Instance)!.GetValue(internalDictionary, new object[] { defaultContainer ? nameof(CosmosSessionTokenContext) : OtherContainerName })!; - internalComposite.GetType().GetField("_string", BindingFlags.NonPublic | BindingFlags.Instance)!.SetValue(internalComposite, "invalidtoken"); - internalComposite.GetType().GetField("_isChanged", BindingFlags.NonPublic | BindingFlags.Instance)!.SetValue(internalComposite, false); + if (defaultContainer) + { + context.Database.AppendSessionToken("invalidtoken"); + } + else + { + context.Database.AppendSessionTokens(new Dictionary { { OtherContainerName, "invalidtoken" } }); + } if (defaultContainer) { @@ -1185,10 +1189,14 @@ public virtual async Task Delete_uses_session_token(AutoTransactionBehavior auto var sessionTokens = context.Database.GetSessionTokens(); // Only way we can test this is by setting a session token that will fail the request if used.. // Only way to do this for a write is to set an invalid session token.. - var internalDictionary = sessionTokens.GetType().GetField("_containerSessionTokens", BindingFlags.NonPublic | BindingFlags.Instance)!.GetValue(sessionTokens)!; - var internalComposite = internalDictionary.GetType().GetProperty("Item", BindingFlags.Public | BindingFlags.Instance)!.GetValue(internalDictionary, new object[] { defaultContainer ? nameof(CosmosSessionTokenContext) : OtherContainerName })!; - internalComposite.GetType().GetField("_string", BindingFlags.NonPublic | BindingFlags.Instance)!.SetValue(internalComposite, "invalidtoken"); - internalComposite.GetType().GetField("_isChanged", BindingFlags.NonPublic | BindingFlags.Instance)!.SetValue(internalComposite, false); + if (defaultContainer) + { + context.Database.AppendSessionToken("invalidtoken"); + } + else + { + context.Database.AppendSessionTokens(new Dictionary { { OtherContainerName, "invalidtoken" } }); + } if (defaultContainer) { From 017ffc73719b4f1304dea9296ef6ac516cb820a1 Mon Sep 17 00:00:00 2001 From: JoasE <32096708+JoasE@users.noreply.github.com> Date: Sat, 1 Nov 2025 12:19:08 +0100 Subject: [PATCH 27/37] Cleanup --- .../Properties/CosmosStrings.Designer.cs | 6 ------ src/EFCore.Cosmos/Properties/CosmosStrings.resx | 3 --- .../Internal/CosmosQueryCompilationContext.cs | 15 ++------------- .../CosmosQueryCompilationContextFactory.cs | 9 ++------- .../Storage/Internal/ISessionTokenStorage.cs | 2 -- 5 files changed, 4 insertions(+), 31 deletions(-) diff --git a/src/EFCore.Cosmos/Properties/CosmosStrings.Designer.cs b/src/EFCore.Cosmos/Properties/CosmosStrings.Designer.cs index 22fcb2e18cd..3c5f5e3c41f 100644 --- a/src/EFCore.Cosmos/Properties/CosmosStrings.Designer.cs +++ b/src/EFCore.Cosmos/Properties/CosmosStrings.Designer.cs @@ -505,12 +505,6 @@ public static string SaveChangesAutoTransactionBehaviorAlwaysAtomicity public static string SaveChangesAutoTransactionBehaviorAlwaysTriggerAtomicity => GetString("SaveChangesAutoTransactionBehaviorAlwaysTriggerAtomicity"); - /// - /// Session token can not be white space. - /// - public static string SessionTokenCanNotBeWhiteSpace - => GetString("SessionTokenCanNotBeWhiteSpace"); - /// /// SingleOrDefault and FirstOrDefault cannot be used Cosmos SQL does not allow Offset without Limit. Consider specifying a 'Take' operation on the query. /// diff --git a/src/EFCore.Cosmos/Properties/CosmosStrings.resx b/src/EFCore.Cosmos/Properties/CosmosStrings.resx index 71bd33f52cf..9622d0028bc 100644 --- a/src/EFCore.Cosmos/Properties/CosmosStrings.resx +++ b/src/EFCore.Cosmos/Properties/CosmosStrings.resx @@ -357,9 +357,6 @@ 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. - - Session token can not be white space. - SingleOrDefault and FirstOrDefault cannot be used Cosmos SQL does not allow Offset without Limit. Consider specifying a 'Take' operation on the query. diff --git a/src/EFCore.Cosmos/Query/Internal/CosmosQueryCompilationContext.cs b/src/EFCore.Cosmos/Query/Internal/CosmosQueryCompilationContext.cs index edebd1549e4..bc804f1f869 100644 --- a/src/EFCore.Cosmos/Query/Internal/CosmosQueryCompilationContext.cs +++ b/src/EFCore.Cosmos/Query/Internal/CosmosQueryCompilationContext.cs @@ -1,8 +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 Microsoft.EntityFrameworkCore.Cosmos.Storage; - namespace Microsoft.EntityFrameworkCore.Cosmos.Query.Internal; /// @@ -11,18 +9,9 @@ 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. /// -public class CosmosQueryCompilationContext : QueryCompilationContext +public class CosmosQueryCompilationContext(QueryCompilationContextDependencies dependencies, bool async) + : QueryCompilationContext(dependencies, async) { - /// - /// 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. - /// - public CosmosQueryCompilationContext(QueryCompilationContextDependencies dependencies, bool async) : base(dependencies, async) - { - } - /// /// The root entity type being queried. /// diff --git a/src/EFCore.Cosmos/Query/Internal/CosmosQueryCompilationContextFactory.cs b/src/EFCore.Cosmos/Query/Internal/CosmosQueryCompilationContextFactory.cs index bd303713db3..dc5b69cb0f4 100644 --- a/src/EFCore.Cosmos/Query/Internal/CosmosQueryCompilationContextFactory.cs +++ b/src/EFCore.Cosmos/Query/Internal/CosmosQueryCompilationContextFactory.cs @@ -18,15 +18,10 @@ public class CosmosQueryCompilationContextFactory : IQueryCompilationContextFact /// doing so can result in application failures when updating to a new Entity Framework Core release. /// public CosmosQueryCompilationContextFactory(QueryCompilationContextDependencies dependencies) - { - Dependencies = dependencies; - } + => Dependencies = dependencies; /// - /// 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. + /// Dependencies for this service. /// protected virtual QueryCompilationContextDependencies Dependencies { get; } diff --git a/src/EFCore.Cosmos/Storage/Internal/ISessionTokenStorage.cs b/src/EFCore.Cosmos/Storage/Internal/ISessionTokenStorage.cs index 97a1a4927a8..f506854ee46 100644 --- a/src/EFCore.Cosmos/Storage/Internal/ISessionTokenStorage.cs +++ b/src/EFCore.Cosmos/Storage/Internal/ISessionTokenStorage.cs @@ -1,8 +1,6 @@ // 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.Storage.Internal; /// From 0342b9ccc6cb8f8078d193b6e5ad672cf80cccce Mon Sep 17 00:00:00 2001 From: JoasE <32096708+JoasE@users.noreply.github.com> Date: Sat, 1 Nov 2025 12:52:23 +0100 Subject: [PATCH 28/37] Fix make public api virtual --- .../Storage/Internal/NullSessionTokenStorage.cs | 16 ++++++++-------- 1 file changed, 8 insertions(+), 8 deletions(-) diff --git a/src/EFCore.Cosmos/Storage/Internal/NullSessionTokenStorage.cs b/src/EFCore.Cosmos/Storage/Internal/NullSessionTokenStorage.cs index dbcffa10e9f..8bf584cda54 100644 --- a/src/EFCore.Cosmos/Storage/Internal/NullSessionTokenStorage.cs +++ b/src/EFCore.Cosmos/Storage/Internal/NullSessionTokenStorage.cs @@ -17,7 +17,7 @@ public class NullSessionTokenStorage : ISessionTokenStorage /// 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. /// - public void AppendSessionToken(string sessionToken) { } + public virtual void AppendSessionToken(string sessionToken) { } /// /// This is an internal API that supports the Entity Framework Core infrastructure and not subject to @@ -25,7 +25,7 @@ public void AppendSessionToken(string sessionToken) { } /// 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. /// - public void AppendSessionToken(string containerName, string sessionToken) { } + public virtual void AppendSessionToken(string containerName, string sessionToken) { } /// /// This is an internal API that supports the Entity Framework Core infrastructure and not subject to @@ -33,7 +33,7 @@ public void AppendSessionToken(string containerName, string sessionToken) { } /// 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. /// - public void Clear() { } + public virtual void Clear() { } /// /// This is an internal API that supports the Entity Framework Core infrastructure and not subject to @@ -41,7 +41,7 @@ public void Clear() { } /// 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. /// - public string? GetSessionToken() => null; + public virtual string? GetSessionToken() => null; /// /// This is an internal API that supports the Entity Framework Core infrastructure and not subject to @@ -49,7 +49,7 @@ public void Clear() { } /// 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. /// - public string? GetSessionToken(string containerName) => null; + public virtual string? GetSessionToken(string containerName) => null; /// /// This is an internal API that supports the Entity Framework Core infrastructure and not subject to @@ -57,7 +57,7 @@ public void Clear() { } /// 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. /// - public void SetSessionToken(string containerName, string? sessionToken) { } + public virtual void SetSessionToken(string containerName, string? sessionToken) { } /// /// This is an internal API that supports the Entity Framework Core infrastructure and not subject to @@ -65,7 +65,7 @@ public void SetSessionToken(string containerName, string? sessionToken) { } /// 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. /// - public void AppendSessionTokens(IReadOnlyDictionary sessionTokens) { } + public virtual void AppendSessionTokens(IReadOnlyDictionary sessionTokens) { } /// /// This is an internal API that supports the Entity Framework Core infrastructure and not subject to @@ -73,5 +73,5 @@ public void AppendSessionTokens(IReadOnlyDictionary sessionToken /// 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. /// - public IReadOnlyDictionary ToDictionary() => new Dictionary(); + public virtual IReadOnlyDictionary ToDictionary() => new Dictionary(); } From 8bd06b43e1d5e610cae5a50ef1396db7cb244fdb Mon Sep 17 00:00:00 2001 From: JoasE <32096708+JoasE@users.noreply.github.com> Date: Sat, 1 Nov 2025 16:54:04 +0100 Subject: [PATCH 29/37] Cleanup & don't return null values in dictionary --- .../CosmosDatabaseFacadeExtensions.cs | 16 ++++++------- .../Storage/Internal/ISessionTokenStorage.cs | 6 ++--- .../Internal/NullSessionTokenStorage.cs | 2 +- .../Storage/Internal/SessionTokenStorage.cs | 24 +++++++++---------- .../CosmosSessionTokensTest.cs | 4 ++-- 5 files changed, 26 insertions(+), 26 deletions(-) diff --git a/src/EFCore.Cosmos/Extensions/CosmosDatabaseFacadeExtensions.cs b/src/EFCore.Cosmos/Extensions/CosmosDatabaseFacadeExtensions.cs index 2a7d3e6a78a..f3387ed781d 100644 --- a/src/EFCore.Cosmos/Extensions/CosmosDatabaseFacadeExtensions.cs +++ b/src/EFCore.Cosmos/Extensions/CosmosDatabaseFacadeExtensions.cs @@ -28,34 +28,34 @@ public static CosmosClient GetCosmosClient(this DatabaseFacade databaseFacade) /// /// Gets the composite session token for the default container for this . /// - /// Use this when using only 1 container in the same + /// Use this when using only 1 container in the same . /// The for the context. - /// The session token dictionary. + /// The session token for the default container in the context, or if none present. public static string? GetSessionToken(this DatabaseFacade databaseFacade) => GetSessionTokenStorage(databaseFacade).GetSessionToken(); /// /// Gets a dictionary that contains the composite session token per container for this . /// - /// Use this when using multiple containers in the same + /// Use this when using multiple containers in the same . /// The for the context. /// The session token dictionary. - public static IReadOnlyDictionary GetSessionTokens(this DatabaseFacade databaseFacade) + public static IReadOnlyDictionary GetSessionTokens(this DatabaseFacade databaseFacade) => GetSessionTokenStorage(databaseFacade).ToDictionary(); /// - /// Appends the composite session token for the default container for this . + /// Appends the composite session token for the default container for this . /// - /// Use this when using only 1 container in the same + /// Use this when using only 1 container in the same . /// The for the context. /// The session token to append. public static void AppendSessionToken(this DatabaseFacade databaseFacade, string sessionToken) => GetSessionTokenStorage(databaseFacade).AppendSessionToken(sessionToken); /// - /// Appends the composite session token per container for this . + /// Appends the composite sessions token per container for this with the tokens specified in . /// - /// Use this when using multiple containers in the same + /// Use this when using multiple containers in the same . /// The for the context. /// The session tokens to append per container. public static void AppendSessionTokens(this DatabaseFacade databaseFacade, IReadOnlyDictionary sessionTokens) diff --git a/src/EFCore.Cosmos/Storage/Internal/ISessionTokenStorage.cs b/src/EFCore.Cosmos/Storage/Internal/ISessionTokenStorage.cs index f506854ee46..1bd6659fe3f 100644 --- a/src/EFCore.Cosmos/Storage/Internal/ISessionTokenStorage.cs +++ b/src/EFCore.Cosmos/Storage/Internal/ISessionTokenStorage.cs @@ -17,7 +17,7 @@ public interface ISessionTokenStorage /// 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. /// - string? GetSessionToken(); + public string? GetSessionToken(); /// /// This is an internal API that supports the Entity Framework Core infrastructure and not subject to @@ -25,7 +25,7 @@ public interface ISessionTokenStorage /// 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. /// - void AppendSessionToken(string sessionToken); + public void AppendSessionToken(string sessionToken); /// /// This is an internal API that supports the Entity Framework Core infrastructure and not subject to @@ -57,5 +57,5 @@ public interface ISessionTokenStorage /// 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. /// - public IReadOnlyDictionary ToDictionary(); + public IReadOnlyDictionary ToDictionary(); } diff --git a/src/EFCore.Cosmos/Storage/Internal/NullSessionTokenStorage.cs b/src/EFCore.Cosmos/Storage/Internal/NullSessionTokenStorage.cs index 8bf584cda54..bc111042bef 100644 --- a/src/EFCore.Cosmos/Storage/Internal/NullSessionTokenStorage.cs +++ b/src/EFCore.Cosmos/Storage/Internal/NullSessionTokenStorage.cs @@ -73,5 +73,5 @@ public virtual void AppendSessionTokens(IReadOnlyDictionary sess /// 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. /// - public virtual IReadOnlyDictionary ToDictionary() => new Dictionary(); + public virtual IReadOnlyDictionary ToDictionary() => new Dictionary(); } diff --git a/src/EFCore.Cosmos/Storage/Internal/SessionTokenStorage.cs b/src/EFCore.Cosmos/Storage/Internal/SessionTokenStorage.cs index f7750c7c066..a267c982f53 100644 --- a/src/EFCore.Cosmos/Storage/Internal/SessionTokenStorage.cs +++ b/src/EFCore.Cosmos/Storage/Internal/SessionTokenStorage.cs @@ -92,10 +92,12 @@ public virtual void AppendSessionToken(string containerName, string sessionToken ref var compositeSessionToken = ref CollectionsMarshal.GetValueRefOrAddDefault(_containerSessionTokens, containerName, out var exists); if (!exists) { - compositeSessionToken = new(); + compositeSessionToken = new(sessionToken); + } + else + { + compositeSessionToken!.Add(sessionToken); } - - compositeSessionToken!.Add(sessionToken); } /// @@ -105,12 +107,7 @@ public virtual void AppendSessionToken(string containerName, string sessionToken /// doing so can result in application failures when updating to a new Entity Framework Core release. /// public virtual void Clear() - { - foreach (var key in _containerSessionTokens.Keys) - { - _containerSessionTokens[key] = new CompositeSessionToken(); - } - } + => _containerSessionTokens.Clear(); /// /// This is an internal API that supports the Entity Framework Core infrastructure and not subject to @@ -118,7 +115,7 @@ public virtual void Clear() /// 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. /// - public virtual IReadOnlyDictionary ToDictionary() => _containerSessionTokens.ToDictionary(x => x.Key, x => x.Value.ConvertToString()); + public virtual IReadOnlyDictionary ToDictionary() => _containerSessionTokens.ToDictionary(x => x.Key, x => x.Value.ConvertToString()); private sealed class CompositeSessionToken { @@ -126,6 +123,9 @@ private sealed class CompositeSessionToken private bool _isChanged; private readonly HashSet _tokens = new(); + public CompositeSessionToken(string token) + => Add(token); + public void Add(string token) { foreach (var tokenPart in token.Split(',')) @@ -137,7 +137,7 @@ public void Add(string token) } } - public string? ConvertToString() + public string ConvertToString() { if (_isChanged) { @@ -145,7 +145,7 @@ public void Add(string token) _string = string.Join(",", _tokens); } - return _string; + return _string!; } } } diff --git a/test/EFCore.Cosmos.FunctionalTests/CosmosSessionTokensTest.cs b/test/EFCore.Cosmos.FunctionalTests/CosmosSessionTokensTest.cs index 8e60578f73e..3d396536d72 100644 --- a/test/EFCore.Cosmos.FunctionalTests/CosmosSessionTokensTest.cs +++ b/test/EFCore.Cosmos.FunctionalTests/CosmosSessionTokensTest.cs @@ -46,7 +46,7 @@ public virtual async Task AppendSessionToken_no_token_sets_token() var contextFactory = await InitializeAsync(); using var context = contextFactory.CreateContext(); - IReadOnlyDictionary sessionTokens; + IReadOnlyDictionary sessionTokens; context.Database.AppendSessionToken("0:-1#231"); sessionTokens = context.Database.GetSessionTokens(); @@ -61,7 +61,7 @@ public virtual async Task AppendSessionTokens_no_tokens_sets_tokens() var contextFactory = await InitializeAsync(); using var context = contextFactory.CreateContext(); - IReadOnlyDictionary sessionTokens; + IReadOnlyDictionary sessionTokens; context.Database.AppendSessionTokens(new Dictionary { { OtherContainerName, "0:-1#123" }, { nameof(CosmosSessionTokenContext), "0:-1#231" } }); sessionTokens = context.Database.GetSessionTokens(); From 34eeeea3dafbdb8a11305fbfd73228fe7bbb84ba Mon Sep 17 00:00:00 2001 From: JoasE <32096708+JoasE@users.noreply.github.com> Date: Thu, 6 Nov 2025 16:21:36 +0100 Subject: [PATCH 30/37] WIP --- .../CosmosDatabaseFacadeExtensions.cs | 53 +++- .../CosmosDbContextOptionsBuilder.cs | 15 +- .../Internal/CosmosDbOptionExtension.cs | 16 +- .../Internal/CosmosSingletonOptions.cs | 6 +- .../Internal/ICosmosSingletonOptions.cs | 2 +- .../SessionTokenManagementMode.cs | 40 +++ .../Storage/Internal/CosmosClientWrapper.cs | 8 +- .../Storage/Internal/CosmosDatabaseWrapper.cs | 7 +- .../Storage/Internal/ISessionTokenStorage.cs | 35 ++- .../Internal/NullSessionTokenStorage.cs | 76 ++---- .../Storage/Internal/SessionTokenStorage.cs | 104 +++++--- .../CosmosSessionTokensEnforcedManualTest.cs | 19 ++ .../CosmosSessionTokensFullyAutomaticTest.cs | 78 ++++++ .../CosmosSessionTokensManualTest.cs | 11 + .../CosmosSessionTokensSemiAutomaticTest.cs | 11 + ...Test.cs => CosmosSessionTokensTestBase.cs} | 239 +++++++++++------- 16 files changed, 505 insertions(+), 215 deletions(-) create mode 100644 src/EFCore.Cosmos/Infrastructure/SessionTokenManagementMode.cs create mode 100644 test/EFCore.Cosmos.FunctionalTests/CosmosSessionTokensEnforcedManualTest.cs create mode 100644 test/EFCore.Cosmos.FunctionalTests/CosmosSessionTokensFullyAutomaticTest.cs create mode 100644 test/EFCore.Cosmos.FunctionalTests/CosmosSessionTokensManualTest.cs create mode 100644 test/EFCore.Cosmos.FunctionalTests/CosmosSessionTokensSemiAutomaticTest.cs rename test/EFCore.Cosmos.FunctionalTests/{CosmosSessionTokensTest.cs => CosmosSessionTokensTestBase.cs} (91%) diff --git a/src/EFCore.Cosmos/Extensions/CosmosDatabaseFacadeExtensions.cs b/src/EFCore.Cosmos/Extensions/CosmosDatabaseFacadeExtensions.cs index f3387ed781d..78c010b450b 100644 --- a/src/EFCore.Cosmos/Extensions/CosmosDatabaseFacadeExtensions.cs +++ b/src/EFCore.Cosmos/Extensions/CosmosDatabaseFacadeExtensions.cs @@ -32,7 +32,7 @@ public static CosmosClient GetCosmosClient(this DatabaseFacade databaseFacade) /// The for the context. /// The session token for the default container in the context, or if none present. public static string? GetSessionToken(this DatabaseFacade databaseFacade) - => GetSessionTokenStorage(databaseFacade).GetSessionToken(); + => GetSessionTokenStorage(databaseFacade).GetDefaultContainerTrackedToken(); /// /// Gets a dictionary that contains the composite session token per container for this . @@ -41,7 +41,16 @@ public static CosmosClient GetCosmosClient(this DatabaseFacade databaseFacade) /// The for the context. /// The session token dictionary. public static IReadOnlyDictionary GetSessionTokens(this DatabaseFacade databaseFacade) - => GetSessionTokenStorage(databaseFacade).ToDictionary(); + => GetSessionTokenStorage(databaseFacade).GetTrackedTokens(); + + /// + /// Sets the composite session token for the default container for this . + /// + /// Use this when using only 1 container in the same . + /// The for the context. + /// The session token to set. + public static void UseSessionToken(this DatabaseFacade databaseFacade, string sessionToken) + => GetSessionTokenStorage(databaseFacade).SetDefaultContainerSessionToken(sessionToken); /// /// Appends the composite session token for the default container for this . @@ -50,7 +59,20 @@ public static IReadOnlyDictionary GetSessionTokens(this Database /// The for the context. /// The session token to append. public static void AppendSessionToken(this DatabaseFacade databaseFacade, string sessionToken) - => GetSessionTokenStorage(databaseFacade).AppendSessionToken(sessionToken); + => GetSessionTokenStorage(databaseFacade).AppendDefaultContainerSessionToken(sessionToken); + + /// + /// Sets the composite sessions token per container for this with the tokens specified in . + /// + /// Use this when using multiple containers in the same . + /// The for the context. + /// The session tokens to set per container. + public static void UseSessionTokens(this DatabaseFacade databaseFacade, IReadOnlyDictionary sessionTokens) + { + var sessionTokenStorage = GetSessionTokenStorage(databaseFacade, sessionTokens); + + sessionTokenStorage.SetSessionTokens(sessionTokens); + } /// /// Appends the composite sessions token per container for this with the tokens specified in . @@ -60,16 +82,7 @@ public static void AppendSessionToken(this DatabaseFacade databaseFacade, string /// The session tokens to append per container. public static void AppendSessionTokens(this DatabaseFacade databaseFacade, IReadOnlyDictionary sessionTokens) { - var sessionTokenStorage = GetSessionTokenStorage(databaseFacade); - - var containerNames = GetContainerNames(databaseFacade.GetService()); - foreach (var sessionToken in sessionTokens) - { - if (!containerNames.Contains(sessionToken.Key)) - { - throw new InvalidOperationException(CosmosStrings.ContainerNameDoesNotExist(sessionToken.Key)); - } - } + var sessionTokenStorage = GetSessionTokenStorage(databaseFacade, sessionTokens); sessionTokenStorage.AppendSessionTokens(sessionTokens); } @@ -82,7 +95,7 @@ private static HashSet GetContainerNames(IModel model) .Distinct()! .ToHashSet()!; - private static SessionTokenStorage GetSessionTokenStorage(DatabaseFacade databaseFacade) + private static SessionTokenStorage GetSessionTokenStorage(DatabaseFacade databaseFacade, IReadOnlyDictionary? sessionTokens = null) { var db = GetService(databaseFacade); if (db is not CosmosDatabaseWrapper dbWrapper) @@ -95,6 +108,18 @@ private static SessionTokenStorage GetSessionTokenStorage(DatabaseFacade databas throw new InvalidOperationException(CosmosStrings.EnableManualSessionTokenManagement); } + if (sessionTokens != null) + { + var containerNames = GetContainerNames(databaseFacade.GetService()); + foreach (var sessionToken in sessionTokens) + { + if (!containerNames.Contains(sessionToken.Key)) + { + throw new InvalidOperationException(CosmosStrings.ContainerNameDoesNotExist(sessionToken.Key)); + } + } + } + return sts; } diff --git a/src/EFCore.Cosmos/Infrastructure/CosmosDbContextOptionsBuilder.cs b/src/EFCore.Cosmos/Infrastructure/CosmosDbContextOptionsBuilder.cs index c7494edcfe1..605c4551034 100644 --- a/src/EFCore.Cosmos/Infrastructure/CosmosDbContextOptionsBuilder.cs +++ b/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; @@ -213,20 +214,20 @@ public virtual CosmosDbContextOptionsBuilder ContentResponseOnWriteEnabled(bool /// - /// Sets the boolean to track and manage session tokens for requests made to Cosmos DB - /// and being able to access them via the and methods. - /// This is only relevant when your application needs to manage session tokens manually. + /// Sets the to use. + /// By default, 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. - /// Enabling manual session token management can break session consistency when not handled properly. + /// Manual session token management can break session consistency when not handled properly. /// See Utilize session tokens for more details. /// /// /// See Using DbContextOptions, and /// Accessing Azure Cosmos DB with EF Core for more information and examples. /// - /// to track and manually manage session tokens in EF. - public virtual CosmosDbContextOptionsBuilder ManualSessionTokenManagementEnabled(bool enabled = true) - => WithOption(e => e.ManualSessionTokenManagementEnabled(enabled)); + /// The to use. + public virtual CosmosDbContextOptionsBuilder SessionTokenManagementMode(SessionTokenManagementMode mode) + => WithOption(e => e.WithSessionTokenManagementMode(mode)); /// /// Sets an option by cloning the extension used to store the settings. This ensures the builder diff --git a/src/EFCore.Cosmos/Infrastructure/Internal/CosmosDbOptionExtension.cs b/src/EFCore.Cosmos/Infrastructure/Internal/CosmosDbOptionExtension.cs index dbe9e67e273..6758b48ce0a 100644 --- a/src/EFCore.Cosmos/Infrastructure/Internal/CosmosDbOptionExtension.cs +++ b/src/EFCore.Cosmos/Infrastructure/Internal/CosmosDbOptionExtension.cs @@ -36,7 +36,7 @@ public class CosmosOptionsExtension : IDbContextOptionsExtension private bool? _enableContentResponseOnWrite; private DbContextOptionsExtensionInfo? _info; private Func? _httpClientFactory; - private bool _enableManualSessionTokenManagement; + private SessionTokenManagementMode _sessionTokenManagementMode = SessionTokenManagementMode.FullyAutomatic; /// /// This is an internal API that supports the Entity Framework Core infrastructure and not subject to @@ -74,7 +74,7 @@ protected CosmosOptionsExtension(CosmosOptionsExtension copyFrom) _maxTcpConnectionsPerEndpoint = copyFrom._maxTcpConnectionsPerEndpoint; _maxRequestsPerTcpConnection = copyFrom._maxRequestsPerTcpConnection; _httpClientFactory = copyFrom._httpClientFactory; - _enableManualSessionTokenManagement = copyFrom._enableManualSessionTokenManagement; + _sessionTokenManagementMode = copyFrom._sessionTokenManagementMode; } /// @@ -572,8 +572,8 @@ public virtual CosmosOptionsExtension WithHttpClientFactory(Func? ht /// 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. /// - public virtual bool EnableManualSessionTokenManagement - => _enableManualSessionTokenManagement; + public virtual SessionTokenManagementMode SessionTokenManagementMode + => _sessionTokenManagementMode; /// /// This is an internal API that supports the Entity Framework Core infrastructure and not subject to @@ -581,11 +581,11 @@ public virtual bool EnableManualSessionTokenManagement /// 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. /// - public virtual CosmosOptionsExtension ManualSessionTokenManagementEnabled(bool enabled) + public virtual CosmosOptionsExtension WithSessionTokenManagementMode(SessionTokenManagementMode mode) { var clone = Clone(); - clone._enableManualSessionTokenManagement = enabled; + clone._sessionTokenManagementMode = mode; return clone; } @@ -658,7 +658,7 @@ public override int GetServiceProviderHashCode() hashCode.Add(Extension._maxTcpConnectionsPerEndpoint); hashCode.Add(Extension._maxRequestsPerTcpConnection); hashCode.Add(Extension._httpClientFactory); - hashCode.Add(Extension._enableManualSessionTokenManagement); + hashCode.Add(Extension._sessionTokenManagementMode); _serviceProviderHash = hashCode.ToHashCode(); } @@ -684,7 +684,7 @@ public override bool ShouldUseSameServiceProvider(DbContextOptionsExtensionInfo && Extension._maxTcpConnectionsPerEndpoint == otherInfo.Extension._maxTcpConnectionsPerEndpoint && Extension._maxRequestsPerTcpConnection == otherInfo.Extension._maxRequestsPerTcpConnection && Extension._httpClientFactory == otherInfo.Extension._httpClientFactory - && Extension._enableManualSessionTokenManagement == otherInfo.Extension._enableManualSessionTokenManagement; + && Extension._sessionTokenManagementMode == otherInfo.Extension._sessionTokenManagementMode; public override void PopulateDebugInfo(IDictionary debugInfo) { diff --git a/src/EFCore.Cosmos/Infrastructure/Internal/CosmosSingletonOptions.cs b/src/EFCore.Cosmos/Infrastructure/Internal/CosmosSingletonOptions.cs index d5221c0eef3..42e5ea687a9 100644 --- a/src/EFCore.Cosmos/Infrastructure/Internal/CosmosSingletonOptions.cs +++ b/src/EFCore.Cosmos/Infrastructure/Internal/CosmosSingletonOptions.cs @@ -157,7 +157,7 @@ public class CosmosSingletonOptions : ICosmosSingletonOptions /// 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. /// - public virtual bool EnableManualSessionTokenManagement { get; private set; } + public virtual SessionTokenManagementMode SessionTokenManagementMode { get; private set; } /// /// This is an internal API that supports the Entity Framework Core infrastructure and not subject to @@ -186,7 +186,7 @@ public virtual void Initialize(IDbContextOptions options) MaxTcpConnectionsPerEndpoint = cosmosOptions.MaxTcpConnectionsPerEndpoint; MaxRequestsPerTcpConnection = cosmosOptions.MaxRequestsPerTcpConnection; HttpClientFactory = cosmosOptions.HttpClientFactory; - EnableManualSessionTokenManagement = cosmosOptions.EnableManualSessionTokenManagement; + SessionTokenManagementMode = cosmosOptions.SessionTokenManagementMode; } } @@ -217,7 +217,7 @@ public virtual void Validate(IDbContextOptions options) || MaxTcpConnectionsPerEndpoint != cosmosOptions.MaxTcpConnectionsPerEndpoint || MaxRequestsPerTcpConnection != cosmosOptions.MaxRequestsPerTcpConnection || HttpClientFactory != cosmosOptions.HttpClientFactory - || EnableManualSessionTokenManagement != cosmosOptions.EnableManualSessionTokenManagement + || SessionTokenManagementMode != cosmosOptions.SessionTokenManagementMode )) { throw new InvalidOperationException( diff --git a/src/EFCore.Cosmos/Infrastructure/Internal/ICosmosSingletonOptions.cs b/src/EFCore.Cosmos/Infrastructure/Internal/ICosmosSingletonOptions.cs index 8113da2b25f..cbdbf8c1d79 100644 --- a/src/EFCore.Cosmos/Infrastructure/Internal/ICosmosSingletonOptions.cs +++ b/src/EFCore.Cosmos/Infrastructure/Internal/ICosmosSingletonOptions.cs @@ -162,5 +162,5 @@ public interface ICosmosSingletonOptions : ISingletonOptions /// 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. /// - bool EnableManualSessionTokenManagement { get; } + SessionTokenManagementMode SessionTokenManagementMode { get; } } diff --git a/src/EFCore.Cosmos/Infrastructure/SessionTokenManagementMode.cs b/src/EFCore.Cosmos/Infrastructure/SessionTokenManagementMode.cs new file mode 100644 index 00000000000..e0cb402f513 --- /dev/null +++ b/src/EFCore.Cosmos/Infrastructure/SessionTokenManagementMode.cs @@ -0,0 +1,40 @@ +// 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; + +/// +/// Defines the behaviour of EF regarding the management of Cosmos DB session tokens. +/// +/// +/// See Consistency level choices for more info. +/// +public enum SessionTokenManagementMode +{ + /// + /// 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. UseSessionTokens and GetSessionTokens 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. + /// + FullyAutomatic, + + /// + /// Allows the usage of UseSessionTokens to overwrite the default Cosmos DB SDK automatic session token management by use of the UseSessionTokens method on a instance. + /// EF will track and parse session tokens returned from Cosmos DB, which can be retrieved via . + /// + SemiAutomatic, + + /// + /// Fully overwrites the Cosmos DB SDK automatic session token management, and only uses session tokens specified via UseSessionTokens. + /// EF will track and parse session tokens returned from Cosmos DB, which can be retrieved via . + /// + Manual, + + /// + /// Same as , but will throw an exception if UseSessionTokens was not invoked before executong a read. + /// + EnforcedManual +} diff --git a/src/EFCore.Cosmos/Storage/Internal/CosmosClientWrapper.cs b/src/EFCore.Cosmos/Storage/Internal/CosmosClientWrapper.cs index 7f06009f230..3a7cb5709c8 100644 --- a/src/EFCore.Cosmos/Storage/Internal/CosmosClientWrapper.cs +++ b/src/EFCore.Cosmos/Storage/Internal/CosmosClientWrapper.cs @@ -786,7 +786,7 @@ private static void ProcessResponse(string containerId, ResponseMessage response if (!string.IsNullOrWhiteSpace(response.Headers.Session)) { - sessionTokenStorage.AppendSessionToken(containerId, response.Headers.Session); + sessionTokenStorage.TrackSessionToken(containerId, response.Headers.Session); } ProcessResponse(entry, response.Headers.ETag, response.Content); @@ -796,7 +796,7 @@ private static void ProcessResponse(string containerId, TransactionalBatchRespon { if (!string.IsNullOrWhiteSpace(batchResponse.Headers.Session)) { - sessionTokenStorage.AppendSessionToken(containerId, batchResponse.Headers.Session); + sessionTokenStorage.TrackSessionToken(containerId, batchResponse.Headers.Session); } for (var i = 0; i < batchResponse.Count; i++) @@ -951,7 +951,7 @@ private static async Task CreateSingleItemQueryAsync( if (!string.IsNullOrWhiteSpace(response.Headers.Session)) { - sessionTokenStorage.AppendSessionToken(containerId, response.Headers.Session); + sessionTokenStorage.TrackSessionToken(containerId, response.Headers.Session); } return response; @@ -1370,7 +1370,7 @@ public override async Task ReadNextAsync(CancellationToken canc var response = await _inner.ReadNextAsync(cancellationToken).ConfigureAwait(false); if (!string.IsNullOrWhiteSpace(response.Headers.Session)) { - _sessionTokenStorage.AppendSessionToken(_containerName, response.Headers.Session); + _sessionTokenStorage.TrackSessionToken(_containerName, response.Headers.Session); } return response; } diff --git a/src/EFCore.Cosmos/Storage/Internal/CosmosDatabaseWrapper.cs b/src/EFCore.Cosmos/Storage/Internal/CosmosDatabaseWrapper.cs index 7cc6997ee4a..f1d2f430818 100644 --- a/src/EFCore.Cosmos/Storage/Internal/CosmosDatabaseWrapper.cs +++ b/src/EFCore.Cosmos/Storage/Internal/CosmosDatabaseWrapper.cs @@ -5,6 +5,7 @@ using System.Runtime.InteropServices; using Microsoft.EntityFrameworkCore.ChangeTracking.Internal; using Microsoft.EntityFrameworkCore.Cosmos.Diagnostics.Internal; +using Microsoft.EntityFrameworkCore.Cosmos.Infrastructure; using Microsoft.EntityFrameworkCore.Cosmos.Infrastructure.Internal; using Microsoft.EntityFrameworkCore.Cosmos.Internal; using Microsoft.EntityFrameworkCore.Cosmos.Metadata.Internal; @@ -46,9 +47,9 @@ public CosmosDatabaseWrapper( _currentDbContext = currentDbContext; _cosmosClient = cosmosClient; - SessionTokenStorage = cosmosSingletonOptions.EnableManualSessionTokenManagement ? - new SessionTokenStorage(_currentDbContext.Context) - : new NullSessionTokenStorage(); + SessionTokenStorage = cosmosSingletonOptions.SessionTokenManagementMode == SessionTokenManagementMode.FullyAutomatic ? + new NullSessionTokenStorage() : + new SessionTokenStorage((string)_currentDbContext.Context.Model.GetAnnotation(CosmosAnnotationNames.ContainerName).Value!, cosmosSingletonOptions.SessionTokenManagementMode); if (loggingOptions.IsSensitiveDataLoggingEnabled) { diff --git a/src/EFCore.Cosmos/Storage/Internal/ISessionTokenStorage.cs b/src/EFCore.Cosmos/Storage/Internal/ISessionTokenStorage.cs index 1bd6659fe3f..ba6e5c56ed2 100644 --- a/src/EFCore.Cosmos/Storage/Internal/ISessionTokenStorage.cs +++ b/src/EFCore.Cosmos/Storage/Internal/ISessionTokenStorage.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. + namespace Microsoft.EntityFrameworkCore.Cosmos.Storage.Internal; /// @@ -17,7 +18,15 @@ public interface ISessionTokenStorage /// 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. /// - public string? GetSessionToken(); + public void AppendDefaultContainerSessionToken(string sessionToken); + + /// + /// 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. + /// + public void AppendSessionTokens(IReadOnlyDictionary sessionTokens); /// /// This is an internal API that supports the Entity Framework Core infrastructure and not subject to @@ -25,7 +34,7 @@ public interface ISessionTokenStorage /// 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. /// - public void AppendSessionToken(string sessionToken); + public void Clear(); /// /// This is an internal API that supports the Entity Framework Core infrastructure and not subject to @@ -33,7 +42,7 @@ public interface ISessionTokenStorage /// 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. /// - public void AppendSessionToken(string containerName, string sessionToken); + public string? GetDefaultContainerTrackedToken(); /// /// This is an internal API that supports the Entity Framework Core infrastructure and not subject to @@ -49,7 +58,23 @@ public interface ISessionTokenStorage /// 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. /// - public void Clear(); + public IReadOnlyDictionary GetTrackedTokens(); + + /// + /// 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. + /// + public void TrackSessionToken(string containerName, string sessionToken); + + /// + /// 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. + /// + public void SetDefaultContainerSessionToken(string sessionToken); /// /// This is an internal API that supports the Entity Framework Core infrastructure and not subject to @@ -57,5 +82,5 @@ public interface ISessionTokenStorage /// 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. /// - public IReadOnlyDictionary ToDictionary(); + public void SetSessionTokens(IReadOnlyDictionary sessionTokens); } diff --git a/src/EFCore.Cosmos/Storage/Internal/NullSessionTokenStorage.cs b/src/EFCore.Cosmos/Storage/Internal/NullSessionTokenStorage.cs index bc111042bef..5cee1da05e2 100644 --- a/src/EFCore.Cosmos/Storage/Internal/NullSessionTokenStorage.cs +++ b/src/EFCore.Cosmos/Storage/Internal/NullSessionTokenStorage.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. + namespace Microsoft.EntityFrameworkCore.Cosmos.Storage.Internal; /// @@ -11,67 +12,30 @@ namespace Microsoft.EntityFrameworkCore.Cosmos.Storage.Internal; /// public class NullSessionTokenStorage : ISessionTokenStorage { - /// - /// 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. - /// - public virtual void AppendSessionToken(string sessionToken) { } + /// + public void AppendDefaultContainerSessionToken(string sessionToken) { } + + /// + public void AppendSessionTokens(IReadOnlyDictionary sessionTokens) {} - /// - /// 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. - /// - public virtual void AppendSessionToken(string containerName, string sessionToken) { } + /// + public void Clear() {} - /// - /// 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. - /// - public virtual void Clear() { } + /// + public string? GetDefaultContainerTrackedToken() => null; - /// - /// 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. - /// - public virtual string? GetSessionToken() => null; + /// + public string? GetSessionToken(string containerName) => null; - /// - /// 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. - /// - public virtual string? GetSessionToken(string containerName) => null; + /// + public IReadOnlyDictionary GetTrackedTokens() => null!; - /// - /// 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. - /// - public virtual void SetSessionToken(string containerName, string? sessionToken) { } + /// + public void TrackSessionToken(string containerName, string sessionToken) {} - /// - /// 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. - /// - public virtual void AppendSessionTokens(IReadOnlyDictionary sessionTokens) { } + /// + public void SetDefaultContainerSessionToken(string sessionToken) {} - /// - /// 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. - /// - public virtual IReadOnlyDictionary ToDictionary() => new Dictionary(); + /// + public void SetSessionTokens(IReadOnlyDictionary sessionTokens) {} } diff --git a/src/EFCore.Cosmos/Storage/Internal/SessionTokenStorage.cs b/src/EFCore.Cosmos/Storage/Internal/SessionTokenStorage.cs index a267c982f53..7809fe8451a 100644 --- a/src/EFCore.Cosmos/Storage/Internal/SessionTokenStorage.cs +++ b/src/EFCore.Cosmos/Storage/Internal/SessionTokenStorage.cs @@ -2,7 +2,7 @@ // The .NET Foundation licenses this file to you under the MIT license. using System.Runtime.InteropServices; -using Microsoft.EntityFrameworkCore.Cosmos.Metadata.Internal; +using Microsoft.EntityFrameworkCore.Cosmos.Infrastructure; namespace Microsoft.EntityFrameworkCore.Cosmos.Storage.Internal; @@ -14,8 +14,11 @@ namespace Microsoft.EntityFrameworkCore.Cosmos.Storage.Internal; /// public class SessionTokenStorage : ISessionTokenStorage { - private readonly Dictionary _containerSessionTokens = new(); + private bool _useSessionTokens = false; + + private Dictionary _containerSessionTokens = new(); private readonly string _defaultContainerName; + private readonly SessionTokenManagementMode _mode; /// /// This is an internal API that supports the Entity Framework Core infrastructure and not subject to @@ -23,9 +26,12 @@ public class SessionTokenStorage : ISessionTokenStorage /// 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. /// - public SessionTokenStorage(DbContext dbContext) + public SessionTokenStorage(string defaultContainerName, SessionTokenManagementMode mode) { - _defaultContainerName = (string)dbContext.Model.GetAnnotation(CosmosAnnotationNames.ContainerName).Value!; + Debug.Assert(mode != SessionTokenManagementMode.FullyAutomatic, $"Use {nameof(NullSessionTokenStorage)} instead."); + + _defaultContainerName = defaultContainerName; + _mode = mode; } /// @@ -34,7 +40,11 @@ public SessionTokenStorage(DbContext dbContext) /// 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. /// - public virtual string? GetSessionToken() => GetSessionToken(_defaultContainerName); + public void SetSessionTokens(IReadOnlyDictionary sessionTokens) + { + _useSessionTokens = true; + _containerSessionTokens = sessionTokens.ToDictionary(x => x.Key, x => new CompositeSessionToken(x.Value)); + } /// /// This is an internal API that supports the Entity Framework Core infrastructure and not subject to @@ -42,8 +52,14 @@ public SessionTokenStorage(DbContext dbContext) /// 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. /// - public virtual void AppendSessionToken(string sessionToken) - => AppendSessionToken(_defaultContainerName, sessionToken); + public void AppendSessionTokens(IReadOnlyDictionary sessionTokens) + { + _useSessionTokens = true; + foreach (var sessionToken in sessionTokens) + { + TrackSessionToken(sessionToken.Key, sessionToken.Value); + } + } /// /// This is an internal API that supports the Entity Framework Core infrastructure and not subject to @@ -51,16 +67,22 @@ public virtual void AppendSessionToken(string sessionToken) /// 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. /// - public virtual string? GetSessionToken(string containerName) + public virtual void AppendDefaultContainerSessionToken(string sessionToken) { - ArgumentNullException.ThrowIfNullOrWhiteSpace(containerName, nameof(containerName)); - - if (!_containerSessionTokens.TryGetValue(containerName, out var value)) - { - return null; - } + _useSessionTokens = true; + TrackSessionToken(_defaultContainerName, sessionToken); + } - return value.ConvertToString(); + /// + /// 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. + /// + public virtual void SetDefaultContainerSessionToken(string sessionToken) + { + _useSessionTokens = true; + _containerSessionTokens[_defaultContainerName] = new CompositeSessionToken(sessionToken); } /// @@ -69,13 +91,42 @@ public virtual void AppendSessionToken(string sessionToken) /// 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. /// - public virtual void AppendSessionTokens(IReadOnlyDictionary sessionTokens) + public virtual IReadOnlyDictionary GetTrackedTokens() => _containerSessionTokens.ToDictionary(x => x.Key, x => x.Value.ConvertToString()); + + /// + /// 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. + /// + public virtual string? GetDefaultContainerTrackedToken() => _containerSessionTokens.GetValueOrDefault(_defaultContainerName)?.ConvertToString(); + + /// + /// 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. + /// + public virtual string? GetSessionToken(string containerName) { - ArgumentNullException.ThrowIfNull(sessionTokens, nameof(sessionTokens)); - foreach (var (containerName, sessionToken) in sessionTokens) + ArgumentNullException.ThrowIfNullOrWhiteSpace(containerName, nameof(containerName)); + + if (_mode == SessionTokenManagementMode.SemiAutomatic && !_useSessionTokens) + { + return null; + } + + if (!_containerSessionTokens.TryGetValue(containerName, out var value)) { - AppendSessionToken(containerName, sessionToken); + if (_mode == SessionTokenManagementMode.EnforcedManual) + { + throw new InvalidOperationException(); + } + + return null; } + + return value.ConvertToString(); } /// @@ -84,7 +135,7 @@ public virtual void AppendSessionTokens(IReadOnlyDictionary sess /// 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. /// - public virtual void AppendSessionToken(string containerName, string sessionToken) + public virtual void TrackSessionToken(string containerName, string sessionToken) { ArgumentNullException.ThrowIfNullOrWhiteSpace(containerName, nameof(containerName)); ArgumentNullException.ThrowIfNullOrWhiteSpace(sessionToken, nameof(sessionToken)); @@ -107,15 +158,10 @@ public virtual void AppendSessionToken(string containerName, string sessionToken /// doing so can result in application failures when updating to a new Entity Framework Core release. /// public virtual void Clear() - => _containerSessionTokens.Clear(); - - /// - /// 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. - /// - public virtual IReadOnlyDictionary ToDictionary() => _containerSessionTokens.ToDictionary(x => x.Key, x => x.Value.ConvertToString()); + { + _useSessionTokens = false; + _containerSessionTokens.Clear(); + } private sealed class CompositeSessionToken { diff --git a/test/EFCore.Cosmos.FunctionalTests/CosmosSessionTokensEnforcedManualTest.cs b/test/EFCore.Cosmos.FunctionalTests/CosmosSessionTokensEnforcedManualTest.cs new file mode 100644 index 00000000000..81950ce0d3f --- /dev/null +++ b/test/EFCore.Cosmos.FunctionalTests/CosmosSessionTokensEnforcedManualTest.cs @@ -0,0 +1,19 @@ +// 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.Infrastructure; + +namespace Microsoft.EntityFrameworkCore; + +public class CosmosSessionTokensEnforcedManualTest(NonSharedFixture fixture) : CosmosSessionTokensNonFullyAutomaticTestBase(fixture) +{ + protected override SessionTokenManagementMode Mode => SessionTokenManagementMode.EnforcedManual; + + [ConditionalTheory, InlineData(true), InlineData(false)] + public async virtual Task SaveChanges_throws_without_token(bool async) + { + var contextFactory = await InitializeAsync(); + + + } +} diff --git a/test/EFCore.Cosmos.FunctionalTests/CosmosSessionTokensFullyAutomaticTest.cs b/test/EFCore.Cosmos.FunctionalTests/CosmosSessionTokensFullyAutomaticTest.cs new file mode 100644 index 00000000000..b1b17165c4d --- /dev/null +++ b/test/EFCore.Cosmos.FunctionalTests/CosmosSessionTokensFullyAutomaticTest.cs @@ -0,0 +1,78 @@ +// 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.Infrastructure; +using Microsoft.EntityFrameworkCore.Cosmos.Internal; + +namespace Microsoft.EntityFrameworkCore; + +public class CosmosSessionTokensFullyAutomaticTest(NonSharedFixture fixture) : CosmosSessionTokensTestBase(fixture) +{ + protected override SessionTokenManagementMode Mode => SessionTokenManagementMode.FullyAutomatic; + + [ConditionalFact] + public virtual async Task GetSessionToken_throws() + { + var contextFactory = await InitializeAsync(); + + using var context = contextFactory.CreateContext(); + + var exception = Assert.Throws(() => context.Database.GetSessionToken()); + Assert.Equal(CosmosStrings.EnableManualSessionTokenManagement, exception.Message); + } + + [ConditionalFact] + public virtual async Task UseSessionToken_throws() + { + var contextFactory = await InitializeAsync(); + + using var context = contextFactory.CreateContext(); + + var exception = Assert.Throws(() => context.Database.UseSessionToken("0:-1#231")); + Assert.Equal(CosmosStrings.EnableManualSessionTokenManagement, exception.Message); + } + + [ConditionalFact] + public virtual async Task AppendSessionToken_throws() + { + var contextFactory = await InitializeAsync(); + + using var context = contextFactory.CreateContext(); + + var exception = Assert.Throws(() => context.Database.AppendSessionToken("0:-1#231")); + Assert.Equal(CosmosStrings.EnableManualSessionTokenManagement, exception.Message); + } + + [ConditionalFact] + public virtual async Task GetSessionTokens_throws() + { + var contextFactory = await InitializeAsync(); + + using var context = contextFactory.CreateContext(); + + var exception = Assert.Throws(() => context.Database.GetSessionTokens()); + Assert.Equal(CosmosStrings.EnableManualSessionTokenManagement, exception.Message); + } + + [ConditionalFact] + public virtual async Task UseSessionTokens_throws() + { + var contextFactory = await InitializeAsync(); + + using var context = contextFactory.CreateContext(); + + var exception = Assert.Throws(() => context.Database.UseSessionTokens(new Dictionary() { { nameof(CosmosSessionTokenContext), "0:-1#231" } })); + Assert.Equal(CosmosStrings.EnableManualSessionTokenManagement, exception.Message); + } + + [ConditionalFact] + public virtual async Task AppendSessionTokens_throws() + { + var contextFactory = await InitializeAsync(); + + using var context = contextFactory.CreateContext(); + + var exception = Assert.Throws(() => context.Database.AppendSessionTokens(new Dictionary() { { nameof(CosmosSessionTokenContext), "0:-1#231" } })); + Assert.Equal(CosmosStrings.EnableManualSessionTokenManagement, exception.Message); + } +} diff --git a/test/EFCore.Cosmos.FunctionalTests/CosmosSessionTokensManualTest.cs b/test/EFCore.Cosmos.FunctionalTests/CosmosSessionTokensManualTest.cs new file mode 100644 index 00000000000..8ad70e07b25 --- /dev/null +++ b/test/EFCore.Cosmos.FunctionalTests/CosmosSessionTokensManualTest.cs @@ -0,0 +1,11 @@ +// 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.Infrastructure; + +namespace Microsoft.EntityFrameworkCore; + +public class CosmosSessionTokensManualTest(NonSharedFixture fixture) : CosmosSessionTokensNonFullyAutomaticTestBase(fixture) +{ + protected override SessionTokenManagementMode Mode => SessionTokenManagementMode.Manual; +} diff --git a/test/EFCore.Cosmos.FunctionalTests/CosmosSessionTokensSemiAutomaticTest.cs b/test/EFCore.Cosmos.FunctionalTests/CosmosSessionTokensSemiAutomaticTest.cs new file mode 100644 index 00000000000..15900c48f60 --- /dev/null +++ b/test/EFCore.Cosmos.FunctionalTests/CosmosSessionTokensSemiAutomaticTest.cs @@ -0,0 +1,11 @@ +// 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.Infrastructure; + +namespace Microsoft.EntityFrameworkCore; + +public class CosmosSessionTokensSemiAutomaticTest(NonSharedFixture fixture) : CosmosSessionTokensNonFullyAutomaticTestBase(fixture) +{ + protected override SessionTokenManagementMode Mode => SessionTokenManagementMode.SemiAutomatic; +} diff --git a/test/EFCore.Cosmos.FunctionalTests/CosmosSessionTokensTest.cs b/test/EFCore.Cosmos.FunctionalTests/CosmosSessionTokensTestBase.cs similarity index 91% rename from test/EFCore.Cosmos.FunctionalTests/CosmosSessionTokensTest.cs rename to test/EFCore.Cosmos.FunctionalTests/CosmosSessionTokensTestBase.cs index 3d396536d72..7937df854ac 100644 --- a/test/EFCore.Cosmos.FunctionalTests/CosmosSessionTokensTest.cs +++ b/test/EFCore.Cosmos.FunctionalTests/CosmosSessionTokensTestBase.cs @@ -2,34 +2,82 @@ // The .NET Foundation licenses this file to you under the MIT license. using Microsoft.Azure.Cosmos; +using Microsoft.EntityFrameworkCore.Cosmos.Infrastructure; using Microsoft.EntityFrameworkCore.Cosmos.Internal; namespace Microsoft.EntityFrameworkCore; -public class CosmosSessionTokensTest(NonSharedFixture fixture) : NonSharedModelTestBase(fixture), IClassFixture +public abstract class CosmosSessionTokensTestBase(NonSharedFixture fixture) : NonSharedModelTestBase(fixture), IClassFixture { - private const string OtherContainerName = "Other"; - protected override string StoreName - => "CosmosSessionTokensTest"; + protected const string OtherContainerName = "Other"; protected override ITestStoreFactory TestStoreFactory => CosmosTestStoreFactory.Instance; + protected override string StoreName => nameof(CosmosSessionTokensTestBase); + + protected abstract SessionTokenManagementMode Mode { get; } + protected override DbContextOptionsBuilder AddOptions(DbContextOptionsBuilder builder) => base.AddOptions(builder).ConfigureWarnings(x => x.Ignore(CosmosEventId.SyncNotSupported)); - protected override TestStore CreateTestStore() => CosmosTestStore.Create(StoreName, (c) => c.ManualSessionTokenManagementEnabled()); + protected override TestStore CreateTestStore() => CosmosTestStore.Create(StoreName, (c) => c.SessionTokenManagementMode(Mode)); - [ConditionalFact] - public virtual async Task GetSessionTokens_throws_if_not_enabled() + public class CosmosSessionTokenContext(DbContextOptions options) : PoolableDbContext(options) { - var contextFactory = await InitializeAsync(createTestStore: () => CosmosTestStore.Create(StoreName)); + public DbSet Customers { get; set; } = null!; + public DbSet OtherContainerCustomers { get; set; } = null!; - using var context = contextFactory.CreateContext(); + protected override void OnModelCreating(ModelBuilder builder) + { + builder.Entity( + b => + { + b.HasKey(c => c.Id); + b.Property(c => c.ETag).IsETagConcurrency(); + b.OwnsMany(x => x.Children); + b.HasPartitionKey(c => c.PartitionKey); + }); + + builder.Entity( + b => + { + b.HasKey(c => c.Id); + b.HasPartitionKey(c => c.PartitionKey); + b.ToContainer(OtherContainerName); + }); + } + } + + public class Customer + { + public string? Id { get; set; } + + public string? Name { get; set; } + + public string? ETag { get; set; } + + public string? PartitionKey { get; set; } + + public ICollection Children { get; } = new HashSet(); + } + + public class DummyChild + { + public string? Id { get; init; } + } + + public class OtherContainerCustomer + { + public string? Id { get; set; } - var exception = Assert.Throws(() => context.Database.GetSessionTokens()); - Assert.Equal(CosmosStrings.EnableManualSessionTokenManagement, exception.Message); + public string? Name { get; set; } + + public string? PartitionKey { get; set; } } +} +public abstract class CosmosSessionTokensNonFullyAutomaticTestBase(NonSharedFixture fixture) : CosmosSessionTokensTestBase(fixture) +{ [ConditionalFact] public virtual async Task AppendSessionTokens_throws_for_non_existent_container() { @@ -93,7 +141,7 @@ public virtual async Task AppendSessionToken_append_token_already_present_does_n public virtual async Task AppendSessionTokens_append_token_already_present_does_not_add_token() { var contextFactory = await InitializeAsync(); - + using var context = contextFactory.CreateContext(); context.Add(new Customer { Id = "1", PartitionKey = "1" }); context.Add(new OtherContainerCustomer { Id = "1", PartitionKey = "1" }); @@ -151,6 +199,86 @@ public virtual async Task AppendSessionTokens_multiple_tokens_splits_tokens() } } + [ConditionalFact] + public virtual async Task UseSessionToken_sets_tokens() + { + var contextFactory = await InitializeAsync(); + + using var context = contextFactory.CreateContext(); + + var sessionTokens = context.Database.GetSessionTokens(); + var newToken = "0:-1#123,1:-1#456"; + var overwrite = "0:-1#123"; + + context.Database.UseSessionToken(newToken); + context.Database.UseSessionToken(overwrite); + + var updatedToken = context.Database.GetSessionToken(); + Assert.Equal(overwrite, updatedToken); + } + + [ConditionalFact] + public virtual async Task UseSessionToken_multiple_tokens_splits_tokens() + { + var contextFactory = await InitializeAsync(); + + using var context = contextFactory.CreateContext(); + + var sessionTokens = context.Database.GetSessionTokens(); + var newToken = "0:-1#123,1:-1#456"; + var appendix = "0:-1#123"; + + context.Database.UseSessionToken(newToken); + context.Database.AppendSessionToken(appendix); + + var updatedToken = context.Database.GetSessionToken(); + Assert.Equal(newToken, updatedToken); + } + + [ConditionalFact] + public virtual async Task UseSessionTokens_sets_tokens() + { + var contextFactory = await InitializeAsync(); + + using var context = contextFactory.CreateContext(); + + var sessionTokens = context.Database.GetSessionTokens(); + var newToken = "0:-1#123,1:-1#456"; + var overwrite = "0:-1#123"; + + context.Database.UseSessionTokens(new Dictionary { { OtherContainerName, newToken }, { nameof(CosmosSessionTokenContext), newToken } }); + context.Database.UseSessionTokens(new Dictionary { { OtherContainerName, overwrite }, { nameof(CosmosSessionTokenContext), overwrite } }); + + var updatedTokens = context.Database.GetSessionTokens(); + + foreach (var pair in updatedTokens) + { + Assert.Equal(overwrite, pair.Value); + } + } + + [ConditionalFact] + public virtual async Task UseSessionTokens_multiple_tokens_splits_tokens() + { + var contextFactory = await InitializeAsync(); + + using var context = contextFactory.CreateContext(); + + var sessionTokens = context.Database.GetSessionTokens(); + var newToken = "0:-1#123,1:-1#456"; + var appendix = "0:-1#123"; + + context.Database.UseSessionTokens(new Dictionary { { OtherContainerName, newToken }, { nameof(CosmosSessionTokenContext), newToken } }); + context.Database.AppendSessionTokens(new Dictionary { { OtherContainerName, appendix }, { nameof(CosmosSessionTokenContext), appendix } }); + + var updatedTokens = context.Database.GetSessionTokens(); + + foreach (var pair in updatedTokens) + { + Assert.Equal(newToken, pair.Value); + } + } + [ConditionalFact] public virtual async Task GetSessionTokens_no_token_returns_empty() { @@ -190,8 +318,7 @@ public virtual async Task Query_uses_session_token(bool async) var sessionTokens = context.Database.GetSessionTokens(); // Only way we can test this is by setting a session token that will fail the request if used.. - // This will take a couple of seconds to fail - var newTokens = sessionTokens.ToDictionary(x => x.Key, x => x.Value!.Substring(0, x.Value.IndexOf('#') + 1) + int.MaxValue); + var newTokens = sessionTokens.ToDictionary(x => x.Key, x => "invalidtoken"); context.Database.AppendSessionTokens(newTokens); @@ -209,8 +336,8 @@ public virtual async Task Query_uses_session_token(bool async) ex2 = Assert.Throws(() => context.OtherContainerCustomers.ToList()); } - Assert.Contains("The read session is not available for the input session token.", ex1.ResponseBody); - Assert.Contains("The read session is not available for the input session token.", ex2.ResponseBody); + Assert.Contains("The session token provided 'invalidtoken' is invalid", ex1.ResponseBody); + Assert.Contains("The session token provided 'invalidtoken' is invalid", ex2.ResponseBody); } [ConditionalTheory, InlineData(true), InlineData(false)] @@ -234,8 +361,7 @@ public virtual async Task Query_on_new_context_does_not_use_same_session_token(b var sessionTokens = context.Database.GetSessionTokens(); // Only way we can test this is by setting a session token that will fail the request if used.. - // This will take a couple of seconds to fail - var newTokens = sessionTokens.ToDictionary(x => x.Key, x => x.Value!.Substring(0, x.Value.IndexOf('#') + 1) + int.MaxValue); + var newTokens = sessionTokens.ToDictionary(x => x.Key, x => "invalidtoken"); context.Database.AppendSessionTokens(newTokens); @@ -287,8 +413,7 @@ public virtual async Task Query_on_same_newly_pooled_context_does_not_use_same_s var sessionTokens = context.Database.GetSessionTokens(); // Only way we can test this is by setting a session token that will fail the request if used.. - // This will take a couple of seconds to fail - var newTokens = sessionTokens.ToDictionary(x => x.Key, x => x.Value!.Substring(0, x.Value.IndexOf('#') + 1) + int.MaxValue); + var newTokens = sessionTokens.ToDictionary(x => x.Key, x => "invalidtoken"); context.Database.AppendSessionTokens(newTokens); @@ -331,16 +456,15 @@ public virtual async Task PagingQuery_uses_session_token(bool async) var sessionTokens = context.Database.GetSessionTokens(); // Only way we can test this is by setting a session token that will fail the request if used.. - // This will take a couple of seconds to fail - var newTokens = sessionTokens.ToDictionary(x => x.Key, x => x.Value!.Substring(0, x.Value.IndexOf('#') + 1) + int.MaxValue); + var newTokens = sessionTokens.ToDictionary(x => x.Key, x => "invalidtoken"); context.Database.AppendSessionTokens(newTokens); var ex1 = await Assert.ThrowsAsync(() => context.Customers.ToPageAsync(1, null)); var ex2 = await Assert.ThrowsAsync(() => context.OtherContainerCustomers.ToPageAsync(1, null)); - Assert.Contains("The read session is not available for the input session token.", ex1.ResponseBody); - Assert.Contains("The read session is not available for the input session token.", ex2.ResponseBody); + Assert.Contains("The session token provided 'invalidtoken' is invalid", ex1.ResponseBody); + Assert.Contains("The session token provided 'invalidtoken' is invalid", ex2.ResponseBody); } @@ -365,8 +489,7 @@ public virtual async Task Shaped_query_uses_session_token(bool async) var sessionTokens = context.Database.GetSessionTokens(); // Only way we can test this is by setting a session token that will fail the request if used.. - // This will take a couple of seconds to fail - var newTokens = sessionTokens.ToDictionary(x => x.Key, x => x.Value!.Substring(0, x.Value.IndexOf('#') + 1) + int.MaxValue); + var newTokens = sessionTokens.ToDictionary(x => x.Key, x => "invalidtoken"); context.Database.AppendSessionTokens(newTokens); @@ -383,8 +506,8 @@ public virtual async Task Shaped_query_uses_session_token(bool async) ex2 = Assert.Throws(() => context.OtherContainerCustomers.Select(x => new { x.Id, x.PartitionKey }).ToList()); } - Assert.Contains("The read session is not available for the input session token.", ex1.ResponseBody); - Assert.Contains("The read session is not available for the input session token.", ex2.ResponseBody); + Assert.Contains("The session token provided 'invalidtoken' is invalid", ex1.ResponseBody); + Assert.Contains("The session token provided 'invalidtoken' is invalid", ex2.ResponseBody); } @@ -409,8 +532,7 @@ public virtual async Task Read_item_uses_session_token(bool async) var sessionTokens = context.Database.GetSessionTokens(); // Only way we can test this is by setting a session token that will fail the request if used.. - // This will take a couple of seconds to fail - var newTokens = sessionTokens.ToDictionary(x => x.Key, x => x.Value!.Substring(0, x.Value.IndexOf('#') + 1) + int.MaxValue); + var newTokens = sessionTokens.ToDictionary(x => x.Key, x => "invalidtoken"); context.Database.AppendSessionTokens(newTokens); @@ -427,8 +549,8 @@ public virtual async Task Read_item_uses_session_token(bool async) ex2 = Assert.Throws(() => context.OtherContainerCustomers.FirstOrDefault(x => x.Id == "1" && x.PartitionKey == "1")); } - Assert.Contains("The read session is not available for the input session token.", ex1.ResponseBody); - Assert.Contains("The read session is not available for the input session token.", ex2.ResponseBody); + Assert.Contains("The session token provided 'invalidtoken' is invalid", ex1.ResponseBody); + Assert.Contains("The session token provided 'invalidtoken' is invalid", ex2.ResponseBody); } @@ -1219,57 +1341,4 @@ public virtual async Task Delete_uses_session_token(AutoTransactionBehavior auto Assert.Contains("The session token provided 'invalidtoken' is invalid.", ((CosmosException)ex.InnerException!).ResponseBody); } - - public class CosmosSessionTokenContext(DbContextOptions options) : PoolableDbContext(options) - { - public DbSet Customers { get; set; } = null!; - public DbSet OtherContainerCustomers { get; set; } = null!; - - protected override void OnModelCreating(ModelBuilder builder) - { - builder.Entity( - b => - { - b.HasKey(c => c.Id); - b.Property(c => c.ETag).IsETagConcurrency(); - b.OwnsMany(x => x.Children); - b.HasPartitionKey(c => c.PartitionKey); - }); - - builder.Entity( - b => - { - b.HasKey(c => c.Id); - b.HasPartitionKey(c => c.PartitionKey); - b.ToContainer(OtherContainerName); - }); - } - } - - public class Customer - { - public string? Id { get; set; } - - public string? Name { get; set; } - - public string? ETag { get; set; } - - public string? PartitionKey { get; set; } - - public ICollection Children { get; } = new HashSet(); - } - - public class DummyChild - { - public string? Id { get; init; } - } - - public class OtherContainerCustomer - { - public string? Id { get; set; } - - public string? Name { get; set; } - - public string? PartitionKey { get; set; } - } } From e0b9cedb5db2e74f11266df886ad63b5d6a14deb Mon Sep 17 00:00:00 2001 From: JoasE <32096708+JoasE@users.noreply.github.com> Date: Thu, 6 Nov 2025 22:29:58 +0100 Subject: [PATCH 31/37] Improve tests and sessiontokenstorage --- .../CosmosDatabaseFacadeExtensions.cs | 35 +- .../CosmosServiceCollectionExtensions.cs | 3 +- .../Storage/Internal/CosmosDatabaseWrapper.cs | 9 +- .../Storage/Internal/ISessionTokenStorage.cs | 4 +- .../Internal/ISessionTokenStorageFactory.cs | 21 + .../Internal/NullSessionTokenStorage.cs | 4 +- .../Storage/Internal/SessionTokenStorage.cs | 133 +- .../Internal/SessionTokenStorageFactory.cs | 52 + .../CosmosSessionTokensEnforcedManualTest.cs | 19 - .../CosmosSessionTokensFullyAutomaticTest.cs | 78 - .../CosmosSessionTokensManualTest.cs | 11 - .../CosmosSessionTokensSemiAutomaticTest.cs | 11 - .../CosmosSessionTokensTest.cs | 699 +++++++++ .../CosmosSessionTokensTestBase.cs | 1344 ----------------- .../CosmosDbContextOptionsExtensionsTests.cs | 2 +- .../Internal/SessionTokenStorageTest.cs | 785 ++++++++++ 16 files changed, 1660 insertions(+), 1550 deletions(-) create mode 100644 src/EFCore.Cosmos/Storage/Internal/ISessionTokenStorageFactory.cs create mode 100644 src/EFCore.Cosmos/Storage/Internal/SessionTokenStorageFactory.cs delete mode 100644 test/EFCore.Cosmos.FunctionalTests/CosmosSessionTokensEnforcedManualTest.cs delete mode 100644 test/EFCore.Cosmos.FunctionalTests/CosmosSessionTokensFullyAutomaticTest.cs delete mode 100644 test/EFCore.Cosmos.FunctionalTests/CosmosSessionTokensManualTest.cs delete mode 100644 test/EFCore.Cosmos.FunctionalTests/CosmosSessionTokensSemiAutomaticTest.cs create mode 100644 test/EFCore.Cosmos.FunctionalTests/CosmosSessionTokensTest.cs delete mode 100644 test/EFCore.Cosmos.FunctionalTests/CosmosSessionTokensTestBase.cs create mode 100644 test/EFCore.Cosmos.Tests/Storage/Internal/SessionTokenStorageTest.cs diff --git a/src/EFCore.Cosmos/Extensions/CosmosDatabaseFacadeExtensions.cs b/src/EFCore.Cosmos/Extensions/CosmosDatabaseFacadeExtensions.cs index 78c010b450b..fcdbf80fca5 100644 --- a/src/EFCore.Cosmos/Extensions/CosmosDatabaseFacadeExtensions.cs +++ b/src/EFCore.Cosmos/Extensions/CosmosDatabaseFacadeExtensions.cs @@ -40,7 +40,7 @@ public static CosmosClient GetCosmosClient(this DatabaseFacade databaseFacade) /// Use this when using multiple containers in the same . /// The for the context. /// The session token dictionary. - public static IReadOnlyDictionary GetSessionTokens(this DatabaseFacade databaseFacade) + public static IReadOnlyDictionary GetSessionTokens(this DatabaseFacade databaseFacade) => GetSessionTokenStorage(databaseFacade).GetTrackedTokens(); /// @@ -67,7 +67,7 @@ public static void AppendSessionToken(this DatabaseFacade databaseFacade, string /// Use this when using multiple containers in the same . /// The for the context. /// The session tokens to set per container. - public static void UseSessionTokens(this DatabaseFacade databaseFacade, IReadOnlyDictionary sessionTokens) + public static void UseSessionTokens(this DatabaseFacade databaseFacade, IReadOnlyDictionary sessionTokens) { var sessionTokenStorage = GetSessionTokenStorage(databaseFacade, sessionTokens); @@ -82,20 +82,12 @@ public static void UseSessionTokens(this DatabaseFacade databaseFacade, IReadOnl /// The session tokens to append per container. public static void AppendSessionTokens(this DatabaseFacade databaseFacade, IReadOnlyDictionary sessionTokens) { - var sessionTokenStorage = GetSessionTokenStorage(databaseFacade, sessionTokens); + var sessionTokenStorage = GetSessionTokenStorage(databaseFacade, (IReadOnlyDictionary)sessionTokens); sessionTokenStorage.AppendSessionTokens(sessionTokens); } - private static HashSet GetContainerNames(IModel model) - => model.GetEntityTypes() - .Where(et => et.FindPrimaryKey() != null) - .Select(et => et.GetContainer()) - .Where(container => container != null) - .Distinct()! - .ToHashSet()!; - - private static SessionTokenStorage GetSessionTokenStorage(DatabaseFacade databaseFacade, IReadOnlyDictionary? sessionTokens = null) + private static ISessionTokenStorage GetSessionTokenStorage(DatabaseFacade databaseFacade, IReadOnlyDictionary? sessionTokens = null) { var db = GetService(databaseFacade); if (db is not CosmosDatabaseWrapper dbWrapper) @@ -103,24 +95,7 @@ private static SessionTokenStorage GetSessionTokenStorage(DatabaseFacade databas throw new InvalidOperationException(CosmosStrings.CosmosNotInUse); } - if (dbWrapper.SessionTokenStorage is not SessionTokenStorage sts) - { - throw new InvalidOperationException(CosmosStrings.EnableManualSessionTokenManagement); - } - - if (sessionTokens != null) - { - var containerNames = GetContainerNames(databaseFacade.GetService()); - foreach (var sessionToken in sessionTokens) - { - if (!containerNames.Contains(sessionToken.Key)) - { - throw new InvalidOperationException(CosmosStrings.ContainerNameDoesNotExist(sessionToken.Key)); - } - } - } - - return sts; + return dbWrapper.SessionTokenStorage; } private static TService GetService(IInfrastructure databaseFacade) diff --git a/src/EFCore.Cosmos/Extensions/CosmosServiceCollectionExtensions.cs b/src/EFCore.Cosmos/Extensions/CosmosServiceCollectionExtensions.cs index 9a79b99e343..f560d9b4ae6 100644 --- a/src/EFCore.Cosmos/Extensions/CosmosServiceCollectionExtensions.cs +++ b/src/EFCore.Cosmos/Extensions/CosmosServiceCollectionExtensions.cs @@ -122,7 +122,8 @@ public static IServiceCollection AddEntityFrameworkCosmos(this IServiceCollectio .TryAddScoped() .TryAddScoped() .TryAddScoped() - .TryAddScoped()); + .TryAddScoped() + .TryAddScoped()); builder.TryAddCoreServices(); diff --git a/src/EFCore.Cosmos/Storage/Internal/CosmosDatabaseWrapper.cs b/src/EFCore.Cosmos/Storage/Internal/CosmosDatabaseWrapper.cs index f1d2f430818..f9d742d9c78 100644 --- a/src/EFCore.Cosmos/Storage/Internal/CosmosDatabaseWrapper.cs +++ b/src/EFCore.Cosmos/Storage/Internal/CosmosDatabaseWrapper.cs @@ -5,8 +5,6 @@ using System.Runtime.InteropServices; using Microsoft.EntityFrameworkCore.ChangeTracking.Internal; using Microsoft.EntityFrameworkCore.Cosmos.Diagnostics.Internal; -using Microsoft.EntityFrameworkCore.Cosmos.Infrastructure; -using Microsoft.EntityFrameworkCore.Cosmos.Infrastructure.Internal; using Microsoft.EntityFrameworkCore.Cosmos.Internal; using Microsoft.EntityFrameworkCore.Cosmos.Metadata.Internal; using Microsoft.EntityFrameworkCore.Cosmos.Update.Internal; @@ -40,16 +38,13 @@ public CosmosDatabaseWrapper( DatabaseDependencies dependencies, ICurrentDbContext currentDbContext, ICosmosClientWrapper cosmosClient, - ICosmosSingletonOptions cosmosSingletonOptions, + ISessionTokenStorageFactory sessionTokenStorageFactory, ILoggingOptions loggingOptions) : base(dependencies) { _currentDbContext = currentDbContext; _cosmosClient = cosmosClient; - - SessionTokenStorage = cosmosSingletonOptions.SessionTokenManagementMode == SessionTokenManagementMode.FullyAutomatic ? - new NullSessionTokenStorage() : - new SessionTokenStorage((string)_currentDbContext.Context.Model.GetAnnotation(CosmosAnnotationNames.ContainerName).Value!, cosmosSingletonOptions.SessionTokenManagementMode); + SessionTokenStorage = sessionTokenStorageFactory.Create(); if (loggingOptions.IsSensitiveDataLoggingEnabled) { diff --git a/src/EFCore.Cosmos/Storage/Internal/ISessionTokenStorage.cs b/src/EFCore.Cosmos/Storage/Internal/ISessionTokenStorage.cs index ba6e5c56ed2..d0c160b6e61 100644 --- a/src/EFCore.Cosmos/Storage/Internal/ISessionTokenStorage.cs +++ b/src/EFCore.Cosmos/Storage/Internal/ISessionTokenStorage.cs @@ -58,7 +58,7 @@ public interface ISessionTokenStorage /// 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. /// - public IReadOnlyDictionary GetTrackedTokens(); + public IReadOnlyDictionary GetTrackedTokens(); /// /// This is an internal API that supports the Entity Framework Core infrastructure and not subject to @@ -82,5 +82,5 @@ public interface ISessionTokenStorage /// 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. /// - public void SetSessionTokens(IReadOnlyDictionary sessionTokens); + public void SetSessionTokens(IReadOnlyDictionary sessionTokens); } diff --git a/src/EFCore.Cosmos/Storage/Internal/ISessionTokenStorageFactory.cs b/src/EFCore.Cosmos/Storage/Internal/ISessionTokenStorageFactory.cs new file mode 100644 index 00000000000..25311d88d48 --- /dev/null +++ b/src/EFCore.Cosmos/Storage/Internal/ISessionTokenStorageFactory.cs @@ -0,0 +1,21 @@ +// 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.Storage.Internal; + +/// +/// 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. +/// +public interface ISessionTokenStorageFactory +{ + /// + /// 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. + /// + public ISessionTokenStorage Create(); +} diff --git a/src/EFCore.Cosmos/Storage/Internal/NullSessionTokenStorage.cs b/src/EFCore.Cosmos/Storage/Internal/NullSessionTokenStorage.cs index 5cee1da05e2..1fb63959e8c 100644 --- a/src/EFCore.Cosmos/Storage/Internal/NullSessionTokenStorage.cs +++ b/src/EFCore.Cosmos/Storage/Internal/NullSessionTokenStorage.cs @@ -28,7 +28,7 @@ public void Clear() {} public string? GetSessionToken(string containerName) => null; /// - public IReadOnlyDictionary GetTrackedTokens() => null!; + public IReadOnlyDictionary GetTrackedTokens() => null!; /// public void TrackSessionToken(string containerName, string sessionToken) {} @@ -37,5 +37,5 @@ public void TrackSessionToken(string containerName, string sessionToken) {} public void SetDefaultContainerSessionToken(string sessionToken) {} /// - public void SetSessionTokens(IReadOnlyDictionary sessionTokens) {} + public void SetSessionTokens(IReadOnlyDictionary sessionTokens) {} } diff --git a/src/EFCore.Cosmos/Storage/Internal/SessionTokenStorage.cs b/src/EFCore.Cosmos/Storage/Internal/SessionTokenStorage.cs index 7809fe8451a..42aff6f34ac 100644 --- a/src/EFCore.Cosmos/Storage/Internal/SessionTokenStorage.cs +++ b/src/EFCore.Cosmos/Storage/Internal/SessionTokenStorage.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.Runtime.InteropServices; using Microsoft.EntityFrameworkCore.Cosmos.Infrastructure; namespace Microsoft.EntityFrameworkCore.Cosmos.Storage.Internal; @@ -14,10 +13,9 @@ namespace Microsoft.EntityFrameworkCore.Cosmos.Storage.Internal; /// public class SessionTokenStorage : ISessionTokenStorage { - private bool _useSessionTokens = false; - - private Dictionary _containerSessionTokens = new(); + private readonly Dictionary _containerSessionTokens; private readonly string _defaultContainerName; + private readonly HashSet _containerNames; private readonly SessionTokenManagementMode _mode; /// @@ -26,12 +24,14 @@ public class SessionTokenStorage : ISessionTokenStorage /// 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. /// - public SessionTokenStorage(string defaultContainerName, SessionTokenManagementMode mode) + public SessionTokenStorage(string defaultContainerName, HashSet containerNames, SessionTokenManagementMode mode) { - Debug.Assert(mode != SessionTokenManagementMode.FullyAutomatic, $"Use {nameof(NullSessionTokenStorage)} instead."); - + Debug.Assert(containerNames.Contains(defaultContainerName)); _defaultContainerName = defaultContainerName; + _containerNames = containerNames; _mode = mode; + + _containerSessionTokens = containerNames.ToDictionary(x => x, x => new CompositeSessionToken()); } /// @@ -40,10 +40,18 @@ public SessionTokenStorage(string defaultContainerName, SessionTokenManagementMo /// 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. /// - public void SetSessionTokens(IReadOnlyDictionary sessionTokens) + public virtual void SetSessionTokens(IReadOnlyDictionary sessionTokens) { - _useSessionTokens = true; - _containerSessionTokens = sessionTokens.ToDictionary(x => x.Key, x => new CompositeSessionToken(x.Value)); + CheckMode(); + foreach (var sessionToken in sessionTokens) + { + if (!_containerNames.Contains(sessionToken.Key)) + { + throw new InvalidOperationException("invalid container name"); + } + + _containerSessionTokens[sessionToken.Key] = new CompositeSessionToken(sessionToken.Value, true); + } } /// @@ -52,12 +60,17 @@ public void SetSessionTokens(IReadOnlyDictionary sessionTokens) /// 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. /// - public void AppendSessionTokens(IReadOnlyDictionary sessionTokens) + public virtual void AppendSessionTokens(IReadOnlyDictionary sessionTokens) { - _useSessionTokens = true; + CheckMode(); foreach (var sessionToken in sessionTokens) { - TrackSessionToken(sessionToken.Key, sessionToken.Value); + if (!_containerNames.Contains(sessionToken.Key)) + { + throw new InvalidOperationException("invalid container name"); + } + + _containerSessionTokens[sessionToken.Key].Add(sessionToken.Value, true); } } @@ -69,8 +82,9 @@ public void AppendSessionTokens(IReadOnlyDictionary sessionToken /// public virtual void AppendDefaultContainerSessionToken(string sessionToken) { - _useSessionTokens = true; - TrackSessionToken(_defaultContainerName, sessionToken); + ArgumentException.ThrowIfNullOrWhiteSpace(sessionToken, nameof(sessionToken)); + CheckMode(); + _containerSessionTokens[_defaultContainerName].Add(sessionToken, true); } /// @@ -79,10 +93,10 @@ public virtual void AppendDefaultContainerSessionToken(string sessionToken) /// 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. /// - public virtual void SetDefaultContainerSessionToken(string sessionToken) + public virtual void SetDefaultContainerSessionToken(string? sessionToken) { - _useSessionTokens = true; - _containerSessionTokens[_defaultContainerName] = new CompositeSessionToken(sessionToken); + CheckMode(); + _containerSessionTokens[_defaultContainerName] = new CompositeSessionToken(sessionToken, true); } /// @@ -91,7 +105,11 @@ public virtual void SetDefaultContainerSessionToken(string sessionToken) /// 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. /// - public virtual IReadOnlyDictionary GetTrackedTokens() => _containerSessionTokens.ToDictionary(x => x.Key, x => x.Value.ConvertToString()); + public virtual IReadOnlyDictionary GetTrackedTokens() + { + CheckMode(); + return _containerSessionTokens.ToDictionary(x => x.Key, x => x.Value.ConvertToString()); + } /// /// This is an internal API that supports the Entity Framework Core infrastructure and not subject to @@ -99,7 +117,11 @@ public virtual void SetDefaultContainerSessionToken(string sessionToken) /// 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. /// - public virtual string? GetDefaultContainerTrackedToken() => _containerSessionTokens.GetValueOrDefault(_defaultContainerName)?.ConvertToString(); + public virtual string? GetDefaultContainerTrackedToken() + { + CheckMode(); + return _containerSessionTokens[_defaultContainerName].ConvertToString(); + } /// /// This is an internal API that supports the Entity Framework Core infrastructure and not subject to @@ -111,22 +133,19 @@ public virtual void SetDefaultContainerSessionToken(string sessionToken) { ArgumentNullException.ThrowIfNullOrWhiteSpace(containerName, nameof(containerName)); - if (_mode == SessionTokenManagementMode.SemiAutomatic && !_useSessionTokens) + if (_mode == SessionTokenManagementMode.FullyAutomatic) { return null; } - if (!_containerSessionTokens.TryGetValue(containerName, out var value)) - { - if (_mode == SessionTokenManagementMode.EnforcedManual) - { - throw new InvalidOperationException(); - } + var sessionToken = _containerSessionTokens[containerName]; - return null; + if (!sessionToken.IsSet && _mode == SessionTokenManagementMode.EnforcedManual) + { + throw new InvalidOperationException("No session token set for container while EnforcedManual"); } - return value.ConvertToString(); + return sessionToken.ConvertToString(); } /// @@ -140,15 +159,12 @@ public virtual void TrackSessionToken(string containerName, string sessionToken) ArgumentNullException.ThrowIfNullOrWhiteSpace(containerName, nameof(containerName)); ArgumentNullException.ThrowIfNullOrWhiteSpace(sessionToken, nameof(sessionToken)); - ref var compositeSessionToken = ref CollectionsMarshal.GetValueRefOrAddDefault(_containerSessionTokens, containerName, out var exists); - if (!exists) + if (_mode == SessionTokenManagementMode.FullyAutomatic) { - compositeSessionToken = new(sessionToken); - } - else - { - compositeSessionToken!.Add(sessionToken); + return; } + + _containerSessionTokens[containerName].Add(sessionToken); } /// @@ -159,8 +175,18 @@ public virtual void TrackSessionToken(string containerName, string sessionToken) /// public virtual void Clear() { - _useSessionTokens = false; - _containerSessionTokens.Clear(); + foreach (var key in _containerSessionTokens.Keys) + { + _containerSessionTokens[key] = new CompositeSessionToken(); + } + } + + private void CheckMode() + { + if (_mode == SessionTokenManagementMode.FullyAutomatic) + { + throw new InvalidOperationException("Can't use session tokens with FullyAutomatic"); + } } private sealed class CompositeSessionToken @@ -169,11 +195,30 @@ private sealed class CompositeSessionToken private bool _isChanged; private readonly HashSet _tokens = new(); - public CompositeSessionToken(string token) - => Add(token); + public CompositeSessionToken(string? token, bool isSet = false) + { + if (token != null) + { + Add(token); + } + IsSet = isSet; + } + + public CompositeSessionToken() + { + } - public void Add(string token) + public bool IsSet { get; private set; } + + public void Add(string token, bool isSet = false) { + IsSet = IsSet || isSet; + + if (token == null) + { + return; + } + foreach (var tokenPart in token.Split(',')) { if (_tokens.Add(tokenPart)) @@ -183,15 +228,15 @@ public void Add(string token) } } - public string ConvertToString() + public string? ConvertToString() { if (_isChanged) { _isChanged = false; - _string = string.Join(",", _tokens); + _string = IsSet && _tokens.Count == 0 ? null : string.Join(",", _tokens); } - return _string!; + return _string; } } } diff --git a/src/EFCore.Cosmos/Storage/Internal/SessionTokenStorageFactory.cs b/src/EFCore.Cosmos/Storage/Internal/SessionTokenStorageFactory.cs new file mode 100644 index 00000000000..4db1d15de7e --- /dev/null +++ b/src/EFCore.Cosmos/Storage/Internal/SessionTokenStorageFactory.cs @@ -0,0 +1,52 @@ +// 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.Infrastructure; +using Microsoft.EntityFrameworkCore.Cosmos.Infrastructure.Internal; +using Microsoft.EntityFrameworkCore.Cosmos.Metadata.Internal; + +namespace Microsoft.EntityFrameworkCore.Cosmos.Storage.Internal; + +/// +/// 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. +/// +public class SessionTokenStorageFactory : ISessionTokenStorageFactory +{ + private readonly string _defaultContainerName; + private readonly HashSet _containerNames; + private readonly SessionTokenManagementMode _mode; + + /// + /// 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. + /// + public SessionTokenStorageFactory(ICurrentDbContext currentDbContext, ICosmosSingletonOptions options) + { + _defaultContainerName = (string)currentDbContext.Context.Model.GetAnnotation(CosmosAnnotationNames.ContainerName).Value!; + _containerNames = new HashSet([_defaultContainerName, ..GetContainerNames(currentDbContext.Context.Model)]); + _mode = options.SessionTokenManagementMode; + } + + /// + /// 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. + /// + public ISessionTokenStorage Create() + => _mode == SessionTokenManagementMode.FullyAutomatic ? + new NullSessionTokenStorage() : + new SessionTokenStorage(_defaultContainerName, _containerNames, _mode); + + + private static IEnumerable GetContainerNames(IModel model) + => model.GetEntityTypes() + .Where(et => et.FindPrimaryKey() != null) + .Select(et => et.GetContainer()) + .Where(container => container != null)!; +} diff --git a/test/EFCore.Cosmos.FunctionalTests/CosmosSessionTokensEnforcedManualTest.cs b/test/EFCore.Cosmos.FunctionalTests/CosmosSessionTokensEnforcedManualTest.cs deleted file mode 100644 index 81950ce0d3f..00000000000 --- a/test/EFCore.Cosmos.FunctionalTests/CosmosSessionTokensEnforcedManualTest.cs +++ /dev/null @@ -1,19 +0,0 @@ -// 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.Infrastructure; - -namespace Microsoft.EntityFrameworkCore; - -public class CosmosSessionTokensEnforcedManualTest(NonSharedFixture fixture) : CosmosSessionTokensNonFullyAutomaticTestBase(fixture) -{ - protected override SessionTokenManagementMode Mode => SessionTokenManagementMode.EnforcedManual; - - [ConditionalTheory, InlineData(true), InlineData(false)] - public async virtual Task SaveChanges_throws_without_token(bool async) - { - var contextFactory = await InitializeAsync(); - - - } -} diff --git a/test/EFCore.Cosmos.FunctionalTests/CosmosSessionTokensFullyAutomaticTest.cs b/test/EFCore.Cosmos.FunctionalTests/CosmosSessionTokensFullyAutomaticTest.cs deleted file mode 100644 index b1b17165c4d..00000000000 --- a/test/EFCore.Cosmos.FunctionalTests/CosmosSessionTokensFullyAutomaticTest.cs +++ /dev/null @@ -1,78 +0,0 @@ -// 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.Infrastructure; -using Microsoft.EntityFrameworkCore.Cosmos.Internal; - -namespace Microsoft.EntityFrameworkCore; - -public class CosmosSessionTokensFullyAutomaticTest(NonSharedFixture fixture) : CosmosSessionTokensTestBase(fixture) -{ - protected override SessionTokenManagementMode Mode => SessionTokenManagementMode.FullyAutomatic; - - [ConditionalFact] - public virtual async Task GetSessionToken_throws() - { - var contextFactory = await InitializeAsync(); - - using var context = contextFactory.CreateContext(); - - var exception = Assert.Throws(() => context.Database.GetSessionToken()); - Assert.Equal(CosmosStrings.EnableManualSessionTokenManagement, exception.Message); - } - - [ConditionalFact] - public virtual async Task UseSessionToken_throws() - { - var contextFactory = await InitializeAsync(); - - using var context = contextFactory.CreateContext(); - - var exception = Assert.Throws(() => context.Database.UseSessionToken("0:-1#231")); - Assert.Equal(CosmosStrings.EnableManualSessionTokenManagement, exception.Message); - } - - [ConditionalFact] - public virtual async Task AppendSessionToken_throws() - { - var contextFactory = await InitializeAsync(); - - using var context = contextFactory.CreateContext(); - - var exception = Assert.Throws(() => context.Database.AppendSessionToken("0:-1#231")); - Assert.Equal(CosmosStrings.EnableManualSessionTokenManagement, exception.Message); - } - - [ConditionalFact] - public virtual async Task GetSessionTokens_throws() - { - var contextFactory = await InitializeAsync(); - - using var context = contextFactory.CreateContext(); - - var exception = Assert.Throws(() => context.Database.GetSessionTokens()); - Assert.Equal(CosmosStrings.EnableManualSessionTokenManagement, exception.Message); - } - - [ConditionalFact] - public virtual async Task UseSessionTokens_throws() - { - var contextFactory = await InitializeAsync(); - - using var context = contextFactory.CreateContext(); - - var exception = Assert.Throws(() => context.Database.UseSessionTokens(new Dictionary() { { nameof(CosmosSessionTokenContext), "0:-1#231" } })); - Assert.Equal(CosmosStrings.EnableManualSessionTokenManagement, exception.Message); - } - - [ConditionalFact] - public virtual async Task AppendSessionTokens_throws() - { - var contextFactory = await InitializeAsync(); - - using var context = contextFactory.CreateContext(); - - var exception = Assert.Throws(() => context.Database.AppendSessionTokens(new Dictionary() { { nameof(CosmosSessionTokenContext), "0:-1#231" } })); - Assert.Equal(CosmosStrings.EnableManualSessionTokenManagement, exception.Message); - } -} diff --git a/test/EFCore.Cosmos.FunctionalTests/CosmosSessionTokensManualTest.cs b/test/EFCore.Cosmos.FunctionalTests/CosmosSessionTokensManualTest.cs deleted file mode 100644 index 8ad70e07b25..00000000000 --- a/test/EFCore.Cosmos.FunctionalTests/CosmosSessionTokensManualTest.cs +++ /dev/null @@ -1,11 +0,0 @@ -// 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.Infrastructure; - -namespace Microsoft.EntityFrameworkCore; - -public class CosmosSessionTokensManualTest(NonSharedFixture fixture) : CosmosSessionTokensNonFullyAutomaticTestBase(fixture) -{ - protected override SessionTokenManagementMode Mode => SessionTokenManagementMode.Manual; -} diff --git a/test/EFCore.Cosmos.FunctionalTests/CosmosSessionTokensSemiAutomaticTest.cs b/test/EFCore.Cosmos.FunctionalTests/CosmosSessionTokensSemiAutomaticTest.cs deleted file mode 100644 index 15900c48f60..00000000000 --- a/test/EFCore.Cosmos.FunctionalTests/CosmosSessionTokensSemiAutomaticTest.cs +++ /dev/null @@ -1,11 +0,0 @@ -// 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.Infrastructure; - -namespace Microsoft.EntityFrameworkCore; - -public class CosmosSessionTokensSemiAutomaticTest(NonSharedFixture fixture) : CosmosSessionTokensNonFullyAutomaticTestBase(fixture) -{ - protected override SessionTokenManagementMode Mode => SessionTokenManagementMode.SemiAutomatic; -} diff --git a/test/EFCore.Cosmos.FunctionalTests/CosmosSessionTokensTest.cs b/test/EFCore.Cosmos.FunctionalTests/CosmosSessionTokensTest.cs new file mode 100644 index 00000000000..e4a98ececba --- /dev/null +++ b/test/EFCore.Cosmos.FunctionalTests/CosmosSessionTokensTest.cs @@ -0,0 +1,699 @@ +// Licensed to the .NET Foundation under one or more agreements. +// The .NET Foundation licenses this file to you under the MIT license. + +using Microsoft.Azure.Cosmos; +using Microsoft.EntityFrameworkCore.Cosmos.Storage.Internal; +using Microsoft.Extensions.DependencyInjection.Extensions; + +namespace Microsoft.EntityFrameworkCore; + +public class CosmosSessionTokensTest(NonSharedFixture fixture) : NonSharedModelTestBase(fixture), IClassFixture +{ + protected const string OtherContainerName = "Other"; + + protected override ITestStoreFactory TestStoreFactory + => CosmosTestStoreFactory.Instance; + + protected override string StoreName => nameof(CosmosSessionTokensTest); + + protected override IServiceCollection AddServices(IServiceCollection serviceCollection) + => base.AddServices(serviceCollection).Replace(ServiceDescriptor.Singleton()); + + private static TestSessionTokenStorage _sessionTokenStorage = null!; + + private class TestSessionTokenStorageFactory : ISessionTokenStorageFactory + { + public ISessionTokenStorage Create() + => _sessionTokenStorage = new(); + } + + private class TestSessionTokenStorage : ISessionTokenStorage + { + public Dictionary SessionTokens { get; set; } = new() { { nameof(CosmosSessionTokenContext), null }, { OtherContainerName, null } }; + + public List AppendDefaultContainerSessionTokenCalls { get; set; } = new(); + public List> AppendSessionTokensCalls { get; set; } = new(); + public List SetDefaultContainerSessionTokenCalls { get; set; } = new(); + + public List> SetSessionTokensCalls { get; set; } = new(); + public List<(string containerName, string sessionToken)> TrackSessionTokenCalls { get; set; } = new(); + public bool ClearCalled { get; set; } + + + + public void AppendDefaultContainerSessionToken(string sessionToken) => AppendDefaultContainerSessionTokenCalls.Add(sessionToken); + + public void AppendSessionTokens(IReadOnlyDictionary sessionTokens) => AppendSessionTokensCalls.Add(sessionTokens); + public void Clear() => ClearCalled = true; + public string? GetDefaultContainerTrackedToken() => SessionTokens.FirstOrDefault().Value; + public string? GetSessionToken(string containerName) => SessionTokens[containerName]; + public IReadOnlyDictionary GetTrackedTokens() => SessionTokens; + public void SetDefaultContainerSessionToken(string sessionToken) => SetDefaultContainerSessionTokenCalls.Add(sessionToken); + public void SetSessionTokens(IReadOnlyDictionary sessionTokens) => SetSessionTokensCalls.Add(sessionTokens); + public void TrackSessionToken(string containerName, string sessionToken) => TrackSessionTokenCalls.Add((containerName, sessionToken)); + } + + + [ConditionalFact] + public virtual async Task AppendSessionToken_uses_AppendDefaultContainerSessionToken() + { + var contextFactory = await InitializeAsync(); + + using var context = contextFactory.CreateContext(); + var arg = "0:-1#231"; + context.Database.AppendSessionToken(arg); + Assert.Equal(arg, _sessionTokenStorage.AppendDefaultContainerSessionTokenCalls.Single()); + } + + [ConditionalFact] + public virtual async Task AppendSessionTokens_uses_AppendSessionTokens() + { + var contextFactory = await InitializeAsync(); + + using var context = contextFactory.CreateContext(); + + var arg = new Dictionary { { OtherContainerName, "0:-1#123" }, { nameof(CosmosSessionTokenContext), "0:-1#231" } }; + context.Database.AppendSessionTokens(arg); + Assert.Equal(arg, _sessionTokenStorage.AppendSessionTokensCalls.Single()); + } + + [ConditionalFact] + public virtual async Task UseSessionToken_uses_SetDefaultContainerSessionToken() + { + var contextFactory = await InitializeAsync(); + + using var context = contextFactory.CreateContext(); + var arg = "0:-1#231"; + context.Database.UseSessionToken(arg); + Assert.Equal(arg, _sessionTokenStorage.SetDefaultContainerSessionTokenCalls.Single()); + } + + [ConditionalFact] + public virtual async Task UseSessionTokens_uses_SetSessionTokens() + { + var contextFactory = await InitializeAsync(); + + using var context = contextFactory.CreateContext(); + + var arg = new Dictionary { { OtherContainerName, "0:-1#123" }, { nameof(CosmosSessionTokenContext), "0:-1#231" } }; + context.Database.UseSessionTokens(arg); + Assert.Equal(arg, _sessionTokenStorage.SetSessionTokensCalls.Single()); + } + + [ConditionalFact] + public virtual async Task GetSessionTokens_uses_GetTrackedSessionTokens() + { + var contextFactory = await InitializeAsync(); + using var context = contextFactory.CreateContext(); + _sessionTokenStorage.SessionTokens = new Dictionary { { OtherContainerName, "0:-1#123" }, { nameof(CosmosSessionTokenContext), "0:-1#231" } }; + var sessionTokens = context.Database.GetSessionTokens(); + Assert.Equal(_sessionTokenStorage.SessionTokens, sessionTokens); + } + + [ConditionalFact] + public virtual async Task Query_uses_session_token() + { + var contextFactory = await InitializeAsync(); + using var context = contextFactory.CreateContext(); + + _sessionTokenStorage.SessionTokens = new Dictionary { { OtherContainerName, "invalidtoken" }, { nameof(CosmosSessionTokenContext), "invalidtoken" } }; + + var exes = new List + { + await Assert.ThrowsAsync(() => context.Customers.ToListAsync()), + await Assert.ThrowsAsync(() => context.OtherContainerCustomers.ToListAsync()) + }; + + foreach (var ex in exes) + { + Assert.Contains("The session token provided 'invalidtoken' is invalid", ex.ResponseBody); + } + } + + [ConditionalFact] + public virtual async Task New_context_does_not_use_same_SessionTokenStorage() + { + var contextFactory = await InitializeAsync(); + using var context = contextFactory.CreateContext(); + + var oldSessionTokenStorage = _sessionTokenStorage; + + using var newContext = contextFactory.CreateContext(); + Assert.NotSame(context, newContext); + Assert.NotSame(oldSessionTokenStorage, _sessionTokenStorage); + } + + [ConditionalFact] + public virtual async Task Pooled_context_clears_SessionTokenStorage() + { + var contextFactory = await InitializeAsync(); + DbContext contextCopy; + ISessionTokenStorage sessionTokenStorageCopy; + using (var context = contextFactory.CreateContext()) + { + contextCopy = context; + sessionTokenStorageCopy = _sessionTokenStorage; + Assert.False(_sessionTokenStorage.ClearCalled); + } + + using var newContext = contextFactory.CreateContext(); + Assert.Same(newContext, contextCopy); + Assert.Same(_sessionTokenStorage, sessionTokenStorageCopy); + Assert.True(_sessionTokenStorage.ClearCalled); + } + + [ConditionalFact] + public virtual async Task PagingQuery_uses_session_token() + { + var contextFactory = await InitializeAsync(); + using var context = contextFactory.CreateContext(); + + _sessionTokenStorage.SessionTokens = new Dictionary { { OtherContainerName, "invalidtoken" }, { nameof(CosmosSessionTokenContext), "invalidtoken" } }; + + var exes = new List() + { + await Assert.ThrowsAsync(() => context.Customers.ToPageAsync(1, null)), + await Assert.ThrowsAsync(() => context.OtherContainerCustomers.ToPageAsync(1, null)), + }; + + foreach (var ex in exes) + { + Assert.Contains("The session token provided 'invalidtoken' is invalid", ex.ResponseBody); + } + } + + [ConditionalFact] + public virtual async Task Shaped_query_uses_session_token() + { + var contextFactory = await InitializeAsync(); + using var context = contextFactory.CreateContext(); + + _sessionTokenStorage.SessionTokens = new Dictionary { { OtherContainerName, "invalidtoken" }, { nameof(CosmosSessionTokenContext), "invalidtoken" } }; + + var exes = new List() + { + await Assert.ThrowsAsync(() => context.Customers.Select(x => new { x.Id, x.PartitionKey }).ToListAsync()), + await Assert.ThrowsAsync(() => context.OtherContainerCustomers.Select(x => new { x.Id, x.PartitionKey }).ToListAsync()) + }; + + foreach (var ex in exes) + { + Assert.Contains("The session token provided 'invalidtoken' is invalid", ex.ResponseBody); + } + } + + [ConditionalFact] + public virtual async Task Read_item_uses_session_token() + { + var contextFactory = await InitializeAsync(); + using var context = contextFactory.CreateContext(); + + _sessionTokenStorage.SessionTokens = new Dictionary { { OtherContainerName, "invalidtoken" }, { nameof(CosmosSessionTokenContext), "invalidtoken" } }; + + var exes = new List() + { + await Assert.ThrowsAsync(() => context.Customers.FirstOrDefaultAsync(x => x.Id == "1" && x.PartitionKey == "1")), + await Assert.ThrowsAsync(() => context.OtherContainerCustomers.FirstOrDefaultAsync(x => x.Id == "1" && x.PartitionKey == "1")) + }; + + foreach (var ex in exes) + { + Assert.Contains("The session token provided 'invalidtoken' is invalid", ex.ResponseBody); + } + } + + [ConditionalFact] + public virtual async Task Query_uses_TrackSessionToken() + { + var contextFactory = await InitializeAsync(); + using var context = contextFactory.CreateContext(); + + await context.Customers.ToListAsync(); + await context.OtherContainerCustomers.ToListAsync(); + + Assert.Equal(2, _sessionTokenStorage.TrackSessionTokenCalls.Count); + var defaultContainerCall = _sessionTokenStorage.TrackSessionTokenCalls.First(); + var otherContainerCall = _sessionTokenStorage.TrackSessionTokenCalls.Last(); + + Assert.Equal(nameof(CosmosSessionTokenContext), defaultContainerCall.containerName); + Assert.NotEmpty(defaultContainerCall.sessionToken); + + Assert.Equal(OtherContainerName, otherContainerCall.containerName); + Assert.NotEmpty(otherContainerCall.sessionToken); + } + + [ConditionalFact] + public virtual async Task PagingQuery_uses_TrackSessionToken() + { + var contextFactory = await InitializeAsync(); + using var context = contextFactory.CreateContext(); + + await context.Customers.ToPageAsync(1, null); + await context.OtherContainerCustomers.ToPageAsync(1, null); + + Assert.Equal(2, _sessionTokenStorage.TrackSessionTokenCalls.Count); + var defaultContainerCall = _sessionTokenStorage.TrackSessionTokenCalls.First(); + var otherContainerCall = _sessionTokenStorage.TrackSessionTokenCalls.Last(); + + Assert.Equal(nameof(CosmosSessionTokenContext), defaultContainerCall.containerName); + Assert.NotEmpty(defaultContainerCall.sessionToken); + + Assert.Equal(OtherContainerName, otherContainerCall.containerName); + Assert.NotEmpty(otherContainerCall.sessionToken); + } + + [ConditionalFact] + public virtual async Task Read_item_uses_TrackSessionToken() + { + var contextFactory = await InitializeAsync(); + using var context = contextFactory.CreateContext(); + + await context.Customers.FirstOrDefaultAsync(x => x.Id == "1" && x.PartitionKey == "1"); + await context.OtherContainerCustomers.FirstOrDefaultAsync(x => x.Id == "1" && x.PartitionKey == "1"); + + Assert.Equal(2, _sessionTokenStorage.TrackSessionTokenCalls.Count); + var defaultContainerCall = _sessionTokenStorage.TrackSessionTokenCalls.First(); + var otherContainerCall = _sessionTokenStorage.TrackSessionTokenCalls.Last(); + + Assert.Equal(nameof(CosmosSessionTokenContext), defaultContainerCall.containerName); + Assert.NotEmpty(defaultContainerCall.sessionToken); + + Assert.Equal(OtherContainerName, otherContainerCall.containerName); + Assert.NotEmpty(otherContainerCall.sessionToken); + } + + [ConditionalFact] + public virtual async Task Read_item_enumerable_uses_TrackSessionToken() + { + var contextFactory = await InitializeAsync(); + using var context = contextFactory.CreateContext(); + + await context.Customers.Where(x => x.Id == "1" && x.PartitionKey == "1").ToListAsync(); + await context.OtherContainerCustomers.Where(x => x.Id == "1" && x.PartitionKey == "1").ToListAsync(); + + Assert.Equal(2, _sessionTokenStorage.TrackSessionTokenCalls.Count); + var defaultContainerCall = _sessionTokenStorage.TrackSessionTokenCalls.First(); + var otherContainerCall = _sessionTokenStorage.TrackSessionTokenCalls.Last(); + + Assert.Equal(nameof(CosmosSessionTokenContext), defaultContainerCall.containerName); + Assert.NotEmpty(defaultContainerCall.sessionToken); + + Assert.Equal(OtherContainerName, otherContainerCall.containerName); + Assert.NotEmpty(otherContainerCall.sessionToken); + } + + [ConditionalFact] + public virtual async Task Add_AutoTransactionBehavior_Never_uses_TrackSessionToken() + { + var contextFactory = await InitializeAsync(); + + using var context = contextFactory.CreateContext(); + context.Database.AutoTransactionBehavior = AutoTransactionBehavior.Never; + context.Customers.Add(new Customer { Id = "1", PartitionKey = "1" }); + context.OtherContainerCustomers.Add(new OtherContainerCustomer { Id = "1", PartitionKey = "1" }); + + await context.SaveChangesAsync(); + + Assert.Equal(2, _sessionTokenStorage.TrackSessionTokenCalls.Count); + var defaultContainerCall = _sessionTokenStorage.TrackSessionTokenCalls.First(); + var otherContainerCall = _sessionTokenStorage.TrackSessionTokenCalls.Last(); + + Assert.Equal(nameof(CosmosSessionTokenContext), defaultContainerCall.containerName); + Assert.NotEmpty(defaultContainerCall.sessionToken); + + Assert.Equal(OtherContainerName, otherContainerCall.containerName); + Assert.NotEmpty(otherContainerCall.sessionToken); + } + + [ConditionalTheory] + [InlineData(true)] + [InlineData(false)] + public virtual async Task Add_AutoTransactionBehavior_Always_uses_TrackSessionToken(bool defaultContainer) + { + var contextFactory = await InitializeAsync(); + + using var context = contextFactory.CreateContext(); + context.Database.AutoTransactionBehavior = AutoTransactionBehavior.Always; + if (defaultContainer) + { + context.Customers.Add(new Customer { Id = "1", PartitionKey = "1" }); + } + else + { + context.OtherContainerCustomers.Add(new OtherContainerCustomer { Id = "1", PartitionKey = "1" }); + } + + await context.SaveChangesAsync(); + + Assert.Equal(1, _sessionTokenStorage.TrackSessionTokenCalls.Count); + var call = _sessionTokenStorage.TrackSessionTokenCalls.First(); + + if (defaultContainer) + { + Assert.Equal(nameof(CosmosSessionTokenContext), call.containerName); + } + else + { + Assert.Equal(OtherContainerName, call.containerName); + } + + Assert.NotEmpty(call.sessionToken); + } + + [ConditionalFact] + public virtual async Task Delete_never_uses_TrackSessionToken() + { + var contextFactory = await InitializeAsync(); + + using var context = contextFactory.CreateContext(); + context.Database.AutoTransactionBehavior = AutoTransactionBehavior.Never; + + var customer = new Customer { Id = "1", PartitionKey = "1" }; + var otherContainerCustomer = new OtherContainerCustomer { Id = "1", PartitionKey = "1" }; + context.Customers.Add(customer); + context.OtherContainerCustomers.Add(otherContainerCustomer); + + await context.SaveChangesAsync(); + + var initialDefaultContainerCall = _sessionTokenStorage.TrackSessionTokenCalls[0]; + var initialOtherContainerCall = _sessionTokenStorage.TrackSessionTokenCalls[1]; + + context.Customers.Remove(customer); + context.OtherContainerCustomers.Remove(otherContainerCustomer); + + await context.SaveChangesAsync(); + + Assert.Equal(4, _sessionTokenStorage.TrackSessionTokenCalls.Count); + var defaultContainerCall = _sessionTokenStorage.TrackSessionTokenCalls[2]; + var otherContainerCall = _sessionTokenStorage.TrackSessionTokenCalls[3]; + + Assert.Equal(nameof(CosmosSessionTokenContext), defaultContainerCall.containerName); + Assert.NotEmpty(defaultContainerCall.sessionToken); + + Assert.Equal(OtherContainerName, otherContainerCall.containerName); + Assert.NotEmpty(otherContainerCall.sessionToken); + + Assert.Equal(initialDefaultContainerCall.containerName, defaultContainerCall.containerName); + Assert.Equal(initialOtherContainerCall.containerName, otherContainerCall.containerName); + + Assert.NotEqual(initialDefaultContainerCall.sessionToken, defaultContainerCall.sessionToken); + Assert.NotEqual(initialOtherContainerCall.sessionToken, otherContainerCall.sessionToken); + } + + [ConditionalTheory] + [InlineData(true)] + [InlineData(false)] + public virtual async Task Delete_always_uses_TrackSessionToken(bool defaultContainer) + { + var contextFactory = await InitializeAsync(); + + using var context = contextFactory.CreateContext(); + context.Database.AutoTransactionBehavior = AutoTransactionBehavior.Always; + + if (defaultContainer) + { + context.Customers.Add(new Customer { Id = "1", PartitionKey = "1" }); + } + else + { + context.OtherContainerCustomers.Add(new OtherContainerCustomer { Id = "1", PartitionKey = "1" }); + } + + await context.SaveChangesAsync(); + + context.ChangeTracker.Clear(); + var initialCall = _sessionTokenStorage.TrackSessionTokenCalls[0]; + + if (defaultContainer) + { + context.Customers.Remove(new Customer { Id = "1", PartitionKey = "1" }); + } + else + { + context.OtherContainerCustomers.Remove(new OtherContainerCustomer { Id = "1", PartitionKey = "1" }); + } + + await context.SaveChangesAsync(); + + Assert.Equal(2, _sessionTokenStorage.TrackSessionTokenCalls.Count); + var call = _sessionTokenStorage.TrackSessionTokenCalls[1]; + + if (defaultContainer) + { + Assert.Equal(nameof(CosmosSessionTokenContext), call.containerName); + } + else + { + Assert.Equal(OtherContainerName, call.containerName); + + } + Assert.NotEmpty(call.sessionToken); + + Assert.Equal(initialCall.containerName, call.containerName); + Assert.NotEqual(initialCall.sessionToken, call.sessionToken); + } + + [ConditionalFact] + public virtual async Task Update_never_uses_TrackSessionToken() + { + var contextFactory = await InitializeAsync(); + + using var context = contextFactory.CreateContext(); + context.Database.AutoTransactionBehavior = AutoTransactionBehavior.Never; + + var customer = new Customer { Id = "1", PartitionKey = "1" }; + var otherContainerCustomer = new OtherContainerCustomer { Id = "1", PartitionKey = "1" }; + context.Customers.Add(customer); + context.OtherContainerCustomers.Add(otherContainerCustomer); + + await context.SaveChangesAsync(); + + var initialDefaultContainerCall = _sessionTokenStorage.TrackSessionTokenCalls[0]; + var initialOtherContainerCall = _sessionTokenStorage.TrackSessionTokenCalls[1]; + + customer.Name = "updated"; + otherContainerCustomer.Name = "updated"; + + await context.SaveChangesAsync(); + + Assert.Equal(4, _sessionTokenStorage.TrackSessionTokenCalls.Count); + var defaultContainerCall = _sessionTokenStorage.TrackSessionTokenCalls[2]; + var otherContainerCall = _sessionTokenStorage.TrackSessionTokenCalls[3]; + + Assert.Equal(nameof(CosmosSessionTokenContext), defaultContainerCall.containerName); + Assert.NotEmpty(defaultContainerCall.sessionToken); + + Assert.Equal(OtherContainerName, otherContainerCall.containerName); + Assert.NotEmpty(otherContainerCall.sessionToken); + + Assert.Equal(initialDefaultContainerCall.containerName, defaultContainerCall.containerName); + Assert.Equal(initialOtherContainerCall.containerName, otherContainerCall.containerName); + + Assert.NotEqual(initialDefaultContainerCall.sessionToken, defaultContainerCall.sessionToken); + Assert.NotEqual(initialOtherContainerCall.sessionToken, otherContainerCall.sessionToken); + } + + [ConditionalTheory] + [InlineData(true)] + [InlineData(false)] + public virtual async Task Update_always_uses_TrackSessionToken(bool defaultContainer) + { + var contextFactory = await InitializeAsync(); + + using var context = contextFactory.CreateContext(); + context.Database.AutoTransactionBehavior = AutoTransactionBehavior.Always; + + if (defaultContainer) + { + context.Customers.Add(new Customer { Id = "1", PartitionKey = "1" }); + } + else + { + context.OtherContainerCustomers.Add(new OtherContainerCustomer { Id = "1", PartitionKey = "1" }); + } + + await context.SaveChangesAsync(); + + context.ChangeTracker.Clear(); + var initialCall = _sessionTokenStorage.TrackSessionTokenCalls[0]; + + if (defaultContainer) + { + context.Customers.Update(new Customer { Id = "1", Name = "updated", PartitionKey = "1" }); + } + else + { + context.OtherContainerCustomers.Update(new OtherContainerCustomer { Id = "1", Name = "updated", PartitionKey = "1" }); + } + + await context.SaveChangesAsync(); + + Assert.Equal(2, _sessionTokenStorage.TrackSessionTokenCalls.Count); + var call = _sessionTokenStorage.TrackSessionTokenCalls[1]; + + if (defaultContainer) + { + Assert.Equal(nameof(CosmosSessionTokenContext), call.containerName); + } + else + { + Assert.Equal(OtherContainerName, call.containerName); + + } + Assert.NotEmpty(call.sessionToken); + + Assert.Equal(initialCall.containerName, call.containerName); + Assert.NotEqual(initialCall.sessionToken, call.sessionToken); + } + + [ConditionalTheory] + [InlineData(AutoTransactionBehavior.WhenNeeded, true)] + [InlineData(AutoTransactionBehavior.WhenNeeded, false)] + [InlineData(AutoTransactionBehavior.Never, false)] + [InlineData(AutoTransactionBehavior.Never, true)] + [InlineData(AutoTransactionBehavior.Always, false)] + [InlineData(AutoTransactionBehavior.Always, true)] + public virtual async Task Add_uses_GetSessionToken(AutoTransactionBehavior autoTransactionBehavior, bool defaultContainer) + { + var contextFactory = await InitializeAsync(); + + using var context = contextFactory.CreateContext(); + context.Database.AutoTransactionBehavior = autoTransactionBehavior; + + // Only way we can test this is by setting a session token that will fail the request if used.. + // Only way to do this for a write is to set an invalid session token.. + _sessionTokenStorage.SessionTokens = new Dictionary { { defaultContainer ? nameof(CosmosSessionTokenContext) : OtherContainerName, "invalidtoken" } }; + + if (defaultContainer) + { + context.Customers.Add(new Customer { Id = "1", PartitionKey = "1" }); + } + else + { + context.OtherContainerCustomers.Add(new OtherContainerCustomer { Id = "1", PartitionKey = "1" }); + } + + var ex = await Assert.ThrowsAsync(() => context.SaveChangesAsync()); + + Assert.Contains("The session token provided 'invalidtoken' is invalid.", ((CosmosException)ex.InnerException!).ResponseBody); + } + + [ConditionalTheory] + [InlineData(AutoTransactionBehavior.WhenNeeded, true)] + [InlineData(AutoTransactionBehavior.WhenNeeded, false)] + [InlineData(AutoTransactionBehavior.Never, false)] + [InlineData(AutoTransactionBehavior.Never, true)] + [InlineData(AutoTransactionBehavior.Always, false)] + [InlineData(AutoTransactionBehavior.Always, true)] + public virtual async Task Update_uses_session_token(AutoTransactionBehavior autoTransactionBehavior, bool defaultContainer) + { + var contextFactory = await InitializeAsync(); + + using var context = contextFactory.CreateContext(); + context.Database.AutoTransactionBehavior = autoTransactionBehavior; + + var sessionTokens = context.Database.GetSessionTokens(); + // Only way we can test this is by setting a session token that will fail the request if used.. + // Only way to do this for a write is to set an invalid session token.. + _sessionTokenStorage.SessionTokens = new Dictionary { { defaultContainer ? nameof(CosmosSessionTokenContext) : OtherContainerName, "invalidtoken" } }; + + if (defaultContainer) + { + context.Customers.Update(new Customer { Id = "1", PartitionKey = "1" }); + } + else + { + context.OtherContainerCustomers.Update(new OtherContainerCustomer { Id = "1", PartitionKey = "1" }); + } + + var ex = await Assert.ThrowsAsync(() => context.SaveChangesAsync()); + + Assert.Contains("The session token provided 'invalidtoken' is invalid.", ((CosmosException)ex.InnerException!).ResponseBody); + } + + [ConditionalTheory] + [InlineData(AutoTransactionBehavior.WhenNeeded, true)] + [InlineData(AutoTransactionBehavior.WhenNeeded, false)] + [InlineData(AutoTransactionBehavior.Never, false)] + [InlineData(AutoTransactionBehavior.Never, true)] + [InlineData(AutoTransactionBehavior.Always, false)] + [InlineData(AutoTransactionBehavior.Always, true)] + public virtual async Task Delete_uses_session_token(AutoTransactionBehavior autoTransactionBehavior, bool defaultContainer) + { + var contextFactory = await InitializeAsync(); + + using var context = contextFactory.CreateContext(); + context.Database.AutoTransactionBehavior = autoTransactionBehavior; + + var sessionTokens = context.Database.GetSessionTokens(); + // Only way we can test this is by setting a session token that will fail the request if used.. + // Only way to do this for a write is to set an invalid session token.. + _sessionTokenStorage.SessionTokens = new Dictionary { { defaultContainer ? nameof(CosmosSessionTokenContext) : OtherContainerName, "invalidtoken" } }; + + if (defaultContainer) + { + context.Customers.Remove(new Customer { Id = "1", PartitionKey = "1" }); + } + else + { + context.OtherContainerCustomers.Remove(new OtherContainerCustomer { Id = "1", PartitionKey = "1" }); + } + + var ex = await Assert.ThrowsAsync(() => context.SaveChangesAsync()); + + Assert.Contains("The session token provided 'invalidtoken' is invalid.", ((CosmosException)ex.InnerException!).ResponseBody); + } + + public class CosmosSessionTokenContext(DbContextOptions options) : PoolableDbContext(options) + { + public DbSet Customers { get; set; } = null!; + public DbSet OtherContainerCustomers { get; set; } = null!; + + protected override void OnModelCreating(ModelBuilder builder) + { + builder.Entity( + b => + { + b.HasKey(c => c.Id); + b.Property(c => c.ETag).IsETagConcurrency(); + b.OwnsMany(x => x.Children); + b.HasPartitionKey(c => c.PartitionKey); + }); + + builder.Entity( + b => + { + b.HasKey(c => c.Id); + b.HasPartitionKey(c => c.PartitionKey); + b.ToContainer(OtherContainerName); + }); + } + } + + public class Customer + { + public string? Id { get; set; } + + public string? Name { get; set; } + + public string? ETag { get; set; } + + public string? PartitionKey { get; set; } + + public ICollection Children { get; } = new HashSet(); + } + + public class DummyChild + { + public string? Id { get; init; } + } + + public class OtherContainerCustomer + { + public string? Id { get; set; } + + public string? Name { get; set; } + + public string? PartitionKey { get; set; } + } +} diff --git a/test/EFCore.Cosmos.FunctionalTests/CosmosSessionTokensTestBase.cs b/test/EFCore.Cosmos.FunctionalTests/CosmosSessionTokensTestBase.cs deleted file mode 100644 index 7937df854ac..00000000000 --- a/test/EFCore.Cosmos.FunctionalTests/CosmosSessionTokensTestBase.cs +++ /dev/null @@ -1,1344 +0,0 @@ -// Licensed to the .NET Foundation under one or more agreements. -// The .NET Foundation licenses this file to you under the MIT license. - -using Microsoft.Azure.Cosmos; -using Microsoft.EntityFrameworkCore.Cosmos.Infrastructure; -using Microsoft.EntityFrameworkCore.Cosmos.Internal; - -namespace Microsoft.EntityFrameworkCore; - -public abstract class CosmosSessionTokensTestBase(NonSharedFixture fixture) : NonSharedModelTestBase(fixture), IClassFixture -{ - protected const string OtherContainerName = "Other"; - - protected override ITestStoreFactory TestStoreFactory - => CosmosTestStoreFactory.Instance; - - protected override string StoreName => nameof(CosmosSessionTokensTestBase); - - protected abstract SessionTokenManagementMode Mode { get; } - - protected override DbContextOptionsBuilder AddOptions(DbContextOptionsBuilder builder) => base.AddOptions(builder).ConfigureWarnings(x => x.Ignore(CosmosEventId.SyncNotSupported)); - - protected override TestStore CreateTestStore() => CosmosTestStore.Create(StoreName, (c) => c.SessionTokenManagementMode(Mode)); - - public class CosmosSessionTokenContext(DbContextOptions options) : PoolableDbContext(options) - { - public DbSet Customers { get; set; } = null!; - public DbSet OtherContainerCustomers { get; set; } = null!; - - protected override void OnModelCreating(ModelBuilder builder) - { - builder.Entity( - b => - { - b.HasKey(c => c.Id); - b.Property(c => c.ETag).IsETagConcurrency(); - b.OwnsMany(x => x.Children); - b.HasPartitionKey(c => c.PartitionKey); - }); - - builder.Entity( - b => - { - b.HasKey(c => c.Id); - b.HasPartitionKey(c => c.PartitionKey); - b.ToContainer(OtherContainerName); - }); - } - } - - public class Customer - { - public string? Id { get; set; } - - public string? Name { get; set; } - - public string? ETag { get; set; } - - public string? PartitionKey { get; set; } - - public ICollection Children { get; } = new HashSet(); - } - - public class DummyChild - { - public string? Id { get; init; } - } - - public class OtherContainerCustomer - { - public string? Id { get; set; } - - public string? Name { get; set; } - - public string? PartitionKey { get; set; } - } -} - -public abstract class CosmosSessionTokensNonFullyAutomaticTestBase(NonSharedFixture fixture) : CosmosSessionTokensTestBase(fixture) -{ - [ConditionalFact] - public virtual async Task AppendSessionTokens_throws_for_non_existent_container() - { - var contextFactory = await InitializeAsync(); - - using var context = contextFactory.CreateContext(); - var exception = Assert.Throws(() => context.Database.AppendSessionTokens(new Dictionary { { OtherContainerName, "0:-1#231"}, { "Not the container name", "0:-1#231" } })); - Assert.Equal(CosmosStrings.ContainerNameDoesNotExist("Not the container name"), exception.Message); - } - - [ConditionalFact] - public virtual async Task AppendSessionToken_no_token_sets_token() - { - var contextFactory = await InitializeAsync(); - - using var context = contextFactory.CreateContext(); - IReadOnlyDictionary sessionTokens; - context.Database.AppendSessionToken("0:-1#231"); - sessionTokens = context.Database.GetSessionTokens(); - - Assert.Equal("0:-1#231", sessionTokens[nameof(CosmosSessionTokenContext)]); - Assert.Equal(nameof(CosmosSessionTokenContext), sessionTokens.First().Key); - Assert.Equal(sessionTokens[nameof(CosmosSessionTokenContext)], sessionTokens.First().Value); - } - - [ConditionalFact] - public virtual async Task AppendSessionTokens_no_tokens_sets_tokens() - { - var contextFactory = await InitializeAsync(); - - using var context = contextFactory.CreateContext(); - IReadOnlyDictionary sessionTokens; - - context.Database.AppendSessionTokens(new Dictionary { { OtherContainerName, "0:-1#123" }, { nameof(CosmosSessionTokenContext), "0:-1#231" } }); - sessionTokens = context.Database.GetSessionTokens(); - - Assert.Equal("0:-1#123", sessionTokens[OtherContainerName]); - Assert.Equal("0:-1#231", sessionTokens[nameof(CosmosSessionTokenContext)]); - } - - [ConditionalFact] - public virtual async Task AppendSessionToken_append_token_already_present_does_not_add_token() - { - var contextFactory = await InitializeAsync(); - - using var context = contextFactory.CreateContext(); - context.Add(new Customer { Id = "1", PartitionKey = "1" }); - - await context.SaveChangesAsync(); - - var initialToken = context.Database.GetSessionToken(); - Assert.False(string.IsNullOrWhiteSpace(initialToken)); - context.Database.AppendSessionToken(initialToken); - - var updatedToken = context.Database.GetSessionToken(); - - Assert.Equal(initialToken, updatedToken); - } - - [ConditionalFact] - public virtual async Task AppendSessionTokens_append_token_already_present_does_not_add_token() - { - var contextFactory = await InitializeAsync(); - - using var context = contextFactory.CreateContext(); - context.Add(new Customer { Id = "1", PartitionKey = "1" }); - context.Add(new OtherContainerCustomer { Id = "1", PartitionKey = "1" }); - - await context.SaveChangesAsync(); - - var initialTokens = context.Database.GetSessionTokens(); - context.Database.AppendSessionTokens(initialTokens!); - - var updatedTokens = context.Database.GetSessionTokens(); - - foreach (var pair in updatedTokens) - { - Assert.Equal(initialTokens[pair.Key], pair.Value); - } - } - - [ConditionalFact] - public virtual async Task AppendSessionToken_multiple_tokens_splits_tokens() - { - var contextFactory = await InitializeAsync(); - - using var context = contextFactory.CreateContext(); - - var sessionTokens = context.Database.GetSessionTokens(); - var newToken = "0:-1#123,1:-1#456"; - var appendix = "0:-1#123"; - context.Database.AppendSessionToken(newToken); - context.Database.AppendSessionToken(appendix); - - var updatedToken = context.Database.GetSessionToken(); - - Assert.Equal(newToken, updatedToken); - } - - [ConditionalFact] - public virtual async Task AppendSessionTokens_multiple_tokens_splits_tokens() - { - var contextFactory = await InitializeAsync(); - - using var context = contextFactory.CreateContext(); - - var sessionTokens = context.Database.GetSessionTokens(); - var newToken = "0:-1#123,1:-1#456"; - var appendix = "0:-1#123"; - - context.Database.AppendSessionTokens(new Dictionary { { OtherContainerName, newToken }, { nameof(CosmosSessionTokenContext), newToken } }); - context.Database.AppendSessionTokens(new Dictionary { { OtherContainerName, appendix }, { nameof(CosmosSessionTokenContext), appendix } }); - - var updatedTokens = context.Database.GetSessionTokens(); - - foreach (var pair in updatedTokens) - { - Assert.Equal(newToken, pair.Value); - } - } - - [ConditionalFact] - public virtual async Task UseSessionToken_sets_tokens() - { - var contextFactory = await InitializeAsync(); - - using var context = contextFactory.CreateContext(); - - var sessionTokens = context.Database.GetSessionTokens(); - var newToken = "0:-1#123,1:-1#456"; - var overwrite = "0:-1#123"; - - context.Database.UseSessionToken(newToken); - context.Database.UseSessionToken(overwrite); - - var updatedToken = context.Database.GetSessionToken(); - Assert.Equal(overwrite, updatedToken); - } - - [ConditionalFact] - public virtual async Task UseSessionToken_multiple_tokens_splits_tokens() - { - var contextFactory = await InitializeAsync(); - - using var context = contextFactory.CreateContext(); - - var sessionTokens = context.Database.GetSessionTokens(); - var newToken = "0:-1#123,1:-1#456"; - var appendix = "0:-1#123"; - - context.Database.UseSessionToken(newToken); - context.Database.AppendSessionToken(appendix); - - var updatedToken = context.Database.GetSessionToken(); - Assert.Equal(newToken, updatedToken); - } - - [ConditionalFact] - public virtual async Task UseSessionTokens_sets_tokens() - { - var contextFactory = await InitializeAsync(); - - using var context = contextFactory.CreateContext(); - - var sessionTokens = context.Database.GetSessionTokens(); - var newToken = "0:-1#123,1:-1#456"; - var overwrite = "0:-1#123"; - - context.Database.UseSessionTokens(new Dictionary { { OtherContainerName, newToken }, { nameof(CosmosSessionTokenContext), newToken } }); - context.Database.UseSessionTokens(new Dictionary { { OtherContainerName, overwrite }, { nameof(CosmosSessionTokenContext), overwrite } }); - - var updatedTokens = context.Database.GetSessionTokens(); - - foreach (var pair in updatedTokens) - { - Assert.Equal(overwrite, pair.Value); - } - } - - [ConditionalFact] - public virtual async Task UseSessionTokens_multiple_tokens_splits_tokens() - { - var contextFactory = await InitializeAsync(); - - using var context = contextFactory.CreateContext(); - - var sessionTokens = context.Database.GetSessionTokens(); - var newToken = "0:-1#123,1:-1#456"; - var appendix = "0:-1#123"; - - context.Database.UseSessionTokens(new Dictionary { { OtherContainerName, newToken }, { nameof(CosmosSessionTokenContext), newToken } }); - context.Database.AppendSessionTokens(new Dictionary { { OtherContainerName, appendix }, { nameof(CosmosSessionTokenContext), appendix } }); - - var updatedTokens = context.Database.GetSessionTokens(); - - foreach (var pair in updatedTokens) - { - Assert.Equal(newToken, pair.Value); - } - } - - [ConditionalFact] - public virtual async Task GetSessionTokens_no_token_returns_empty() - { - var contextFactory = await InitializeAsync(); - using var context = contextFactory.CreateContext(); - var sessionTokens = context.Database.GetSessionTokens(); - Assert.Equal(0, sessionTokens.Count); - } - - [ConditionalFact] - public virtual async Task GetSessionToken_no_token_returns_null() - { - var contextFactory = await InitializeAsync(); - using var context = contextFactory.CreateContext(); - var sessionToken = context.Database.GetSessionToken(); - Assert.Null(sessionToken); - } - - [ConditionalTheory, InlineData(true), InlineData(false)] - public virtual async Task Query_uses_session_token(bool async) - { - var contextFactory = await InitializeAsync(); - using var context = contextFactory.CreateContext(); - - context.Customers.Add(new Customer { Id = "1", PartitionKey = "1" }); - context.OtherContainerCustomers.Add(new OtherContainerCustomer { Id = "1", PartitionKey = "1" }); - - if (async) - { - await context.SaveChangesAsync(); - } - else - { - context.SaveChanges(); - } - - var sessionTokens = context.Database.GetSessionTokens(); - - // Only way we can test this is by setting a session token that will fail the request if used.. - var newTokens = sessionTokens.ToDictionary(x => x.Key, x => "invalidtoken"); - - context.Database.AppendSessionTokens(newTokens); - - CosmosException ex1; - CosmosException ex2; - - if (async) - { - ex1 = await Assert.ThrowsAsync(() => context.Customers.ToListAsync()); - ex2 = await Assert.ThrowsAsync(() => context.OtherContainerCustomers.ToListAsync()); - } - else - { - ex1 = Assert.Throws(() => context.Customers.ToList()); - ex2 = Assert.Throws(() => context.OtherContainerCustomers.ToList()); - } - - Assert.Contains("The session token provided 'invalidtoken' is invalid", ex1.ResponseBody); - Assert.Contains("The session token provided 'invalidtoken' is invalid", ex2.ResponseBody); - } - - [ConditionalTheory, InlineData(true), InlineData(false)] - public virtual async Task Query_on_new_context_does_not_use_same_session_token(bool async) - { - var contextFactory = await InitializeAsync(); - using var context = contextFactory.CreateContext(); - - context.Customers.Add(new Customer { Id = "1", PartitionKey = "1" }); - context.OtherContainerCustomers.Add(new OtherContainerCustomer { Id = "1", PartitionKey = "1" }); - - if (async) - { - await context.SaveChangesAsync(); - } - else - { - context.SaveChanges(); - } - - var sessionTokens = context.Database.GetSessionTokens(); - - // Only way we can test this is by setting a session token that will fail the request if used.. - var newTokens = sessionTokens.ToDictionary(x => x.Key, x => "invalidtoken"); - - context.Database.AppendSessionTokens(newTokens); - - if (async) - { - await Assert.ThrowsAsync(() => context.Customers.ToListAsync()); - await Assert.ThrowsAsync(() => context.OtherContainerCustomers.ToListAsync()); - } - else - { - Assert.Throws(() => context.Customers.ToList()); - Assert.Throws(() => context.OtherContainerCustomers.ToList()); - } - - using var newContext = contextFactory.CreateContext(); - Assert.NotSame(context, newContext); - if (async) - { - await newContext.Customers.ToListAsync(); - await newContext.OtherContainerCustomers.ToListAsync(); - } - else - { - newContext.Customers.ToList(); - newContext.OtherContainerCustomers.ToList(); - } - } - - [ConditionalTheory, InlineData(true), InlineData(false)] - public virtual async Task Query_on_same_newly_pooled_context_does_not_use_same_session_token(bool async) - { - var contextFactory = await InitializeAsync(); - DbContext contextCopy; - using (var context = contextFactory.CreateContext()) - { - contextCopy = context; - context.Customers.Add(new Customer { Id = "1", PartitionKey = "1" }); - context.OtherContainerCustomers.Add(new OtherContainerCustomer { Id = "1", PartitionKey = "1" }); - - if (async) - { - await context.SaveChangesAsync(); - } - else - { - context.SaveChanges(); - } - - var sessionTokens = context.Database.GetSessionTokens(); - - // Only way we can test this is by setting a session token that will fail the request if used.. - var newTokens = sessionTokens.ToDictionary(x => x.Key, x => "invalidtoken"); - - context.Database.AppendSessionTokens(newTokens); - - if (async) - { - await Assert.ThrowsAsync(() => context.Customers.ToListAsync()); - await Assert.ThrowsAsync(() => context.OtherContainerCustomers.ToListAsync()); - } - else - { - Assert.Throws(() => context.Customers.ToList()); - Assert.Throws(() => context.OtherContainerCustomers.ToList()); - } - } - - using var newContext = contextFactory.CreateContext(); - Assert.Same(newContext, contextCopy); - await newContext.Customers.ToListAsync(); - await newContext.OtherContainerCustomers.ToListAsync(); - } - - [ConditionalTheory, InlineData(true), InlineData(false)] - public virtual async Task PagingQuery_uses_session_token(bool async) - { - var contextFactory = await InitializeAsync(); - using var context = contextFactory.CreateContext(); - - context.Customers.Add(new Customer { Id = "1", PartitionKey = "1" }); - context.OtherContainerCustomers.Add(new OtherContainerCustomer { Id = "1", PartitionKey = "1" }); - - if (async) - { - await context.SaveChangesAsync(); - } - else - { - context.SaveChanges(); - } - - var sessionTokens = context.Database.GetSessionTokens(); - - // Only way we can test this is by setting a session token that will fail the request if used.. - var newTokens = sessionTokens.ToDictionary(x => x.Key, x => "invalidtoken"); - - context.Database.AppendSessionTokens(newTokens); - - var ex1 = await Assert.ThrowsAsync(() => context.Customers.ToPageAsync(1, null)); - var ex2 = await Assert.ThrowsAsync(() => context.OtherContainerCustomers.ToPageAsync(1, null)); - - Assert.Contains("The session token provided 'invalidtoken' is invalid", ex1.ResponseBody); - Assert.Contains("The session token provided 'invalidtoken' is invalid", ex2.ResponseBody); - - } - - [ConditionalTheory, InlineData(true), InlineData(false)] - public virtual async Task Shaped_query_uses_session_token(bool async) - { - var contextFactory = await InitializeAsync(); - using var context = contextFactory.CreateContext(); - - context.Customers.Add(new Customer { Id = "1", PartitionKey = "1" }); - context.OtherContainerCustomers.Add(new OtherContainerCustomer { Id = "1", PartitionKey = "1" }); - - if (async) - { - await context.SaveChangesAsync(); - } - else - { - context.SaveChanges(); - } - - var sessionTokens = context.Database.GetSessionTokens(); - - // Only way we can test this is by setting a session token that will fail the request if used.. - var newTokens = sessionTokens.ToDictionary(x => x.Key, x => "invalidtoken"); - - context.Database.AppendSessionTokens(newTokens); - - CosmosException ex1; - CosmosException ex2; - if (async) - { - ex1 = await Assert.ThrowsAsync(() => context.Customers.Select(x => new { x.Id, x.PartitionKey }).ToListAsync()); - ex2 = await Assert.ThrowsAsync(() => context.OtherContainerCustomers.Select(x => new { x.Id, x.PartitionKey }).ToListAsync()); - } - else - { - ex1 = Assert.Throws(() => context.Customers.Select(x => new { x.Id, x.PartitionKey }).ToList()); - ex2 = Assert.Throws(() => context.OtherContainerCustomers.Select(x => new { x.Id, x.PartitionKey }).ToList()); - } - - Assert.Contains("The session token provided 'invalidtoken' is invalid", ex1.ResponseBody); - Assert.Contains("The session token provided 'invalidtoken' is invalid", ex2.ResponseBody); - - } - - [ConditionalTheory, InlineData(true), InlineData(false)] - public virtual async Task Read_item_uses_session_token(bool async) - { - var contextFactory = await InitializeAsync(); - using var context = contextFactory.CreateContext(); - - context.Customers.Add(new Customer { Id = "1", PartitionKey = "1" }); - context.OtherContainerCustomers.Add(new OtherContainerCustomer { Id = "1", PartitionKey = "1" }); - - if (async) - { - await context.SaveChangesAsync(); - } - else - { - context.SaveChanges(); - } - - var sessionTokens = context.Database.GetSessionTokens(); - - // Only way we can test this is by setting a session token that will fail the request if used.. - var newTokens = sessionTokens.ToDictionary(x => x.Key, x => "invalidtoken"); - - context.Database.AppendSessionTokens(newTokens); - - CosmosException ex1; - CosmosException ex2; - if (async) - { - ex1 = await Assert.ThrowsAsync(() => context.Customers.FirstOrDefaultAsync(x => x.Id == "1" && x.PartitionKey == "1")); - ex2 = await Assert.ThrowsAsync(() => context.OtherContainerCustomers.FirstOrDefaultAsync(x => x.Id == "1" && x.PartitionKey == "1")); - } - else - { - ex1 = Assert.Throws(() => context.Customers.FirstOrDefault(x => x.Id == "1" && x.PartitionKey == "1")); - ex2 = Assert.Throws(() => context.OtherContainerCustomers.FirstOrDefault(x => x.Id == "1" && x.PartitionKey == "1")); - } - - Assert.Contains("The session token provided 'invalidtoken' is invalid", ex1.ResponseBody); - Assert.Contains("The session token provided 'invalidtoken' is invalid", ex2.ResponseBody); - - } - - [ConditionalTheory, InlineData(true), InlineData(false)] - public virtual async Task Query_sets_session_token(bool async) - { - var contextFactory = await InitializeAsync(); - using var context = contextFactory.CreateContext(); - - if (async) - { - await context.Customers.ToListAsync(); - await context.OtherContainerCustomers.ToListAsync(); - } - else - { - _ = context.Customers.ToList(); - _ = context.OtherContainerCustomers.ToList(); - } - - var sessionTokens = context.Database.GetSessionTokens(); - Assert.False(string.IsNullOrWhiteSpace(sessionTokens.First().Value)); - Assert.False(string.IsNullOrWhiteSpace(sessionTokens.Last().Value)); - } - - [ConditionalTheory, InlineData(true), InlineData(false)] - public virtual async Task Query_appends_session_token(bool async) - { - var contextFactory = await InitializeAsync(); - using var context = contextFactory.CreateContext(); - - if (async) - { - await context.Customers.ToListAsync(); - await context.OtherContainerCustomers.ToListAsync(); - } - else - { - _ = context.Customers.ToList(); - _ = context.OtherContainerCustomers.ToList(); - } - - var initialTokens = context.Database.GetSessionTokens(); - - using var otherContext = contextFactory.CreateContext(); - otherContext.Customers.Add(new Customer { Id = "1", PartitionKey = "1" }); - otherContext.OtherContainerCustomers.Add(new OtherContainerCustomer { Id = "1", PartitionKey = "1" }); - - if (async) - { - await otherContext.SaveChangesAsync(); - } - else - { - otherContext.SaveChanges(); - } - - var otherTokens = otherContext.Database.GetSessionTokens(); - - if (async) - { - await context.Customers.ToListAsync(); - await context.OtherContainerCustomers.ToListAsync(); - } - else - { - _ = context.Customers.ToList(); - _ = context.OtherContainerCustomers.ToList(); - } - - var sessionTokens = context.Database.GetSessionTokens(); - foreach (var token in sessionTokens) - { - var initialToken = initialTokens[token.Key]; - var otherToken = otherTokens[token.Key]; - Assert.Equal(initialToken + "," + otherToken, token.Value); - } - } - - [ConditionalTheory, InlineData(true), InlineData(false)] - public virtual async Task Query_same_session_does_not_append_session_token(bool async) - { - var contextFactory = await InitializeAsync(); - using var context = contextFactory.CreateContext(); - - if (async) - { - await context.Customers.ToListAsync(); - await context.OtherContainerCustomers.ToListAsync(); - } - else - { - _ = context.Customers.ToList(); - _ = context.OtherContainerCustomers.ToList(); - } - - var initialTokens = context.Database.GetSessionTokens(); - - if (async) - { - await context.Customers.ToListAsync(); - await context.OtherContainerCustomers.ToListAsync(); - } - else - { - _ = context.Customers.ToList(); - _ = context.OtherContainerCustomers.ToList(); - } - - var sessionTokens = context.Database.GetSessionTokens(); - foreach (var token in sessionTokens) - { - var initialToken = initialTokens[token.Key]; - Assert.Equal(initialToken, token.Value); - } - } - - [ConditionalFact] - public virtual async Task PagingQuery_appends_session_token() - { - var contextFactory = await InitializeAsync(); - using var context = contextFactory.CreateContext(); - - await context.Customers.ToPageAsync(1, null); - await context.OtherContainerCustomers.ToPageAsync(1, null); - - var initialTokens = context.Database.GetSessionTokens(); - - using var otherContext = contextFactory.CreateContext(); - otherContext.Customers.Add(new Customer { Id = "1", PartitionKey = "1" }); - otherContext.OtherContainerCustomers.Add(new OtherContainerCustomer { Id = "1", PartitionKey = "1" }); - - await otherContext.SaveChangesAsync(); - otherContext.SaveChanges(); - - var otherTokens = otherContext.Database.GetSessionTokens(); - - await context.Customers.ToPageAsync(1, null); - await context.OtherContainerCustomers.ToPageAsync(1, null); - - var sessionTokens = context.Database.GetSessionTokens(); - foreach (var token in sessionTokens) - { - var initialToken = initialTokens[token.Key]; - var otherToken = otherTokens[token.Key]; - Assert.Equal(initialToken + "," + otherToken, token.Value); - } - } - - [ConditionalTheory, InlineData(true), InlineData(false)] - public virtual async Task Read_item_appends_session_token(bool async) - { - var contextFactory = await InitializeAsync(); - using var context = contextFactory.CreateContext(); - - if (async) - { - await context.Customers.FirstOrDefaultAsync(x => x.Id == "1" && x.PartitionKey == "1"); - await context.OtherContainerCustomers.FirstOrDefaultAsync(x => x.Id == "1" && x.PartitionKey == "1"); - } - else - { - _ = context.Customers.ToList(); - _ = context.OtherContainerCustomers.ToList(); - } - - var initialTokens = context.Database.GetSessionTokens(); - - using var otherContext = contextFactory.CreateContext(); - otherContext.Customers.Add(new Customer { Id = "1", PartitionKey = "1" }); - otherContext.OtherContainerCustomers.Add(new OtherContainerCustomer { Id = "1", PartitionKey = "1" }); - - if (async) - { - await otherContext.SaveChangesAsync(); - } - else - { - otherContext.SaveChanges(); - } - - var otherTokens = otherContext.Database.GetSessionTokens(); - - if (async) - { - await context.Customers.FirstOrDefaultAsync(x => x.Id == "1" && x.PartitionKey == "1"); - await context.OtherContainerCustomers.FirstOrDefaultAsync(x => x.Id == "1" && x.PartitionKey == "1"); - } - else - { - _ = context.Customers.ToList(); - _ = context.OtherContainerCustomers.ToList(); - } - - var sessionTokens = context.Database.GetSessionTokens(); - foreach (var token in sessionTokens) - { - var initialToken = initialTokens[token.Key]; - var otherToken = otherTokens[token.Key]; - Assert.Equal(initialToken + "," + otherToken, token.Value); - } - } - - [ConditionalTheory, InlineData(true), InlineData(false)] - public virtual async Task Read_item_enumerable_sets_session_token(bool async) - { - var contextFactory = await InitializeAsync(); - using var context = contextFactory.CreateContext(); - - if (async) - { - await context.Customers.Where(x => x.Id == "1" && x.PartitionKey == "1").ToListAsync(); - await context.OtherContainerCustomers.Where(x => x.Id == "1" && x.PartitionKey == "1").ToListAsync(); - } - else - { - _ = context.Customers.ToList(); - _ = context.OtherContainerCustomers.ToList(); - } - - var initialTokens = context.Database.GetSessionTokens(); - - using var otherContext = contextFactory.CreateContext(); - otherContext.Customers.Add(new Customer { Id = "1", PartitionKey = "1" }); - otherContext.OtherContainerCustomers.Add(new OtherContainerCustomer { Id = "1", PartitionKey = "1" }); - - if (async) - { - await otherContext.SaveChangesAsync(); - } - else - { - otherContext.SaveChanges(); - } - - var otherTokens = otherContext.Database.GetSessionTokens(); - - if (async) - { - await context.Customers.Where(x => x.Id == "1" && x.PartitionKey == "1").ToListAsync(); - await context.OtherContainerCustomers.Where(x => x.Id == "1" && x.PartitionKey == "1").ToListAsync(); - } - else - { - _ = context.Customers.ToList(); - _ = context.OtherContainerCustomers.ToList(); - } - - var sessionTokens = context.Database.GetSessionTokens(); - foreach (var token in sessionTokens) - { - var initialToken = initialTokens[token.Key]; - var otherToken = otherTokens[token.Key]; - Assert.Equal(initialToken + "," + otherToken, token.Value); - } - } - - [ConditionalTheory] - [InlineData(true)] - [InlineData(false)] - public virtual async Task Add_AutoTransactionBehavior_never_sets_session_token(bool async) - { - var contextFactory = await InitializeAsync(); - - using var context = contextFactory.CreateContext(); - context.Database.AutoTransactionBehavior = AutoTransactionBehavior.Never; - context.Customers.Add(new Customer { Id = "1", PartitionKey = "1" }); - context.OtherContainerCustomers.Add(new OtherContainerCustomer { Id = "1", PartitionKey = "1" }); - - if (async) - { - await context.SaveChangesAsync(); - } - else - { - context.SaveChanges(); - } - - var sessionTokens = context.Database.GetSessionTokens(); - foreach (var sessionToken in sessionTokens.Values) - { - Assert.False(string.IsNullOrWhiteSpace(sessionToken)); - } - } - - [ConditionalTheory] - [InlineData(true, true)] - [InlineData(true, false)] - [InlineData(false, true)] - [InlineData(false, false)] - public virtual async Task Add_AutoTransactionBehavior_always_sets_session_token(bool async, bool defaultContainer) - { - var contextFactory = await InitializeAsync(); - - using var context = contextFactory.CreateContext(); - context.Database.AutoTransactionBehavior = AutoTransactionBehavior.Always; - if (defaultContainer) - { - context.Customers.Add(new Customer { Id = "1", PartitionKey = "1" }); - } - else - { - context.OtherContainerCustomers.Add(new OtherContainerCustomer { Id = "1", PartitionKey = "1" }); - } - - if (async) - { - await context.SaveChangesAsync(); - } - else - { - context.SaveChanges(); - } - - var sessionTokens = context.Database.GetSessionTokens(); - var sessionToken = defaultContainer ? sessionTokens[nameof(CosmosSessionTokenContext)]! : sessionTokens[OtherContainerName]!; - Assert.False(string.IsNullOrWhiteSpace(sessionToken)); - } - - [ConditionalTheory] - [InlineData(true)] - [InlineData(false)] - public virtual async Task Add_never_merges_session_token(bool async) - { - var contextFactory = await InitializeAsync(); - - using var context = contextFactory.CreateContext(); - context.Database.AutoTransactionBehavior = AutoTransactionBehavior.Never; - - context.Customers.Add(new Customer { Id = "1", PartitionKey = "1" }); - context.OtherContainerCustomers.Add(new OtherContainerCustomer { Id = "1", PartitionKey = "1" }); - - if (async) - { - await context.SaveChangesAsync(); - } - else - { - context.SaveChanges(); - } - - var initialTokens = context.Database.GetSessionTokens(); - - context.Customers.Add(new Customer { Id = "2", PartitionKey = "1" }); - context.OtherContainerCustomers.Add(new OtherContainerCustomer { Id = "2", PartitionKey = "1" }); - - if (async) - { - await context.SaveChangesAsync(); - } - else - { - context.SaveChanges(); - } - - var sessionTokens = context.Database.GetSessionTokens(); - foreach (var sessionToken in sessionTokens) - { - var initialToken = initialTokens[sessionToken.Key]; - Assert.NotEqual(sessionToken.Value, initialToken); - Assert.StartsWith(initialToken + ",", sessionToken.Value); - Assert.False(string.IsNullOrWhiteSpace(sessionToken.Value!.Substring(sessionToken.Value.IndexOf(",") + 1))); - } - } - - [ConditionalTheory] - [InlineData(true, true)] - [InlineData(true, false)] - [InlineData(false, true)] - [InlineData(false, false)] - public virtual async Task Add_always_merges_session_token(bool async, bool defaultContainer) - { - var contextFactory = await InitializeAsync(); - - using var context = contextFactory.CreateContext(); - context.Database.AutoTransactionBehavior = AutoTransactionBehavior.Always; - - if (defaultContainer) - { - context.Customers.Add(new Customer { Id = "1", PartitionKey = "1" }); - } - else - { - context.OtherContainerCustomers.Add(new OtherContainerCustomer { Id = "1", PartitionKey = "1" }); - } - - if (async) - { - await context.SaveChangesAsync(); - } - else - { - context.SaveChanges(); - } - - var initialTokens = context.Database.GetSessionTokens(); - - if (defaultContainer) - { - context.Customers.Add(new Customer { Id = "2", PartitionKey = "1" }); - } - else - { - context.OtherContainerCustomers.Add(new OtherContainerCustomer { Id = "2", PartitionKey = "1" }); - } - - if (async) - { - await context.SaveChangesAsync(); - } - else - { - context.SaveChanges(); - } - - var sessionTokens = context.Database.GetSessionTokens(); - var sessionToken = defaultContainer ? sessionTokens[nameof(CosmosSessionTokenContext)]! : sessionTokens[OtherContainerName]!; - Assert.False(string.IsNullOrWhiteSpace(sessionToken)); - } - - [ConditionalTheory] - [InlineData(true)] - [InlineData(false)] - public virtual async Task Delete_never_merges_session_token(bool async) - { - var contextFactory = await InitializeAsync(); - - using var context = contextFactory.CreateContext(); - context.Database.AutoTransactionBehavior = AutoTransactionBehavior.Never; - - var customer = new Customer { Id = "1", PartitionKey = "1" }; - var otherContainerCustomer = new OtherContainerCustomer { Id = "1", PartitionKey = "1" }; - context.Customers.Add(customer); - context.OtherContainerCustomers.Add(otherContainerCustomer); - - if (async) - { - await context.SaveChangesAsync(); - } - else - { - context.SaveChanges(); - } - - var initialTokens = context.Database.GetSessionTokens(); - - context.Customers.Remove(customer); - context.OtherContainerCustomers.Remove(otherContainerCustomer); - - if (async) - { - await context.SaveChangesAsync(); - } - else - { - context.SaveChanges(); - } - - var sessionTokens = context.Database.GetSessionTokens(); - foreach (var sessionToken in sessionTokens) - { - var initialToken = initialTokens[sessionToken.Key]; - Assert.NotEqual(sessionToken.Value, initialToken); - Assert.StartsWith(initialToken + ",", sessionToken.Value); - Assert.False(string.IsNullOrWhiteSpace(sessionToken.Value!.Substring(sessionToken.Value.IndexOf(",") + 1))); - } - } - - [ConditionalTheory] - [InlineData(true, true)] - [InlineData(true, false)] - [InlineData(false, true)] - [InlineData(false, false)] - public virtual async Task Delete_always_merges_session_token(bool async, bool defaultContainer) - { - var contextFactory = await InitializeAsync(); - - using var context = contextFactory.CreateContext(); - context.Database.AutoTransactionBehavior = AutoTransactionBehavior.Always; - - if (defaultContainer) - { - context.Customers.Add(new Customer { Id = "1", PartitionKey = "1" }); - } - else - { - context.OtherContainerCustomers.Add(new OtherContainerCustomer { Id = "1", PartitionKey = "1" }); - } - - if (async) - { - await context.SaveChangesAsync(); - } - else - { - context.SaveChanges(); - } - - context.ChangeTracker.Clear(); - var initialTokens = context.Database.GetSessionTokens(); - - if (defaultContainer) - { - context.Customers.Remove(new Customer { Id = "1", PartitionKey = "1" }); - } - else - { - context.OtherContainerCustomers.Remove(new OtherContainerCustomer { Id = "1", PartitionKey = "1" }); - } - - if (async) - { - await context.SaveChangesAsync(); - } - else - { - context.SaveChanges(); - } - - var sessionTokens = context.Database.GetSessionTokens(); - var sessionToken = defaultContainer ? sessionTokens[nameof(CosmosSessionTokenContext)]! : sessionTokens[OtherContainerName]!; - Assert.False(string.IsNullOrWhiteSpace(sessionToken)); - } - - [ConditionalTheory] - [InlineData(true)] - [InlineData(false)] - public virtual async Task Update_never_merges_session_token(bool async) - { - var contextFactory = await InitializeAsync(); - - using var context = contextFactory.CreateContext(); - context.Database.AutoTransactionBehavior = AutoTransactionBehavior.Never; - - var customer = new Customer { Id = "1", PartitionKey = "1" }; - var otherContainerCustomer = new OtherContainerCustomer { Id = "1", PartitionKey = "1" }; - context.Customers.Add(customer); - context.OtherContainerCustomers.Add(otherContainerCustomer); - - if (async) - { - await context.SaveChangesAsync(); - } - else - { - context.SaveChanges(); - } - - var initialTokens = context.Database.GetSessionTokens(); - - customer.Name = "updated"; - otherContainerCustomer.Name = "updated"; - - if (async) - { - await context.SaveChangesAsync(); - } - else - { - context.SaveChanges(); - } - - var sessionTokens = context.Database.GetSessionTokens(); - foreach (var sessionToken in sessionTokens) - { - var initialToken = initialTokens[sessionToken.Key]; - Assert.NotEqual(sessionToken.Value, initialToken); - Assert.StartsWith(initialToken + ",", sessionToken.Value); - Assert.False(string.IsNullOrWhiteSpace(sessionToken.Value!.Substring(sessionToken.Value.IndexOf(",") + 1))); - } - } - - [ConditionalTheory] - [InlineData(true, true)] - [InlineData(true, false)] - [InlineData(false, true)] - [InlineData(false, false)] - public virtual async Task Update_always_merges_session_token(bool async, bool defaultContainer) - { - var contextFactory = await InitializeAsync(); - - using var context = contextFactory.CreateContext(); - context.Database.AutoTransactionBehavior = AutoTransactionBehavior.Always; - - if (defaultContainer) - { - context.Customers.Add(new Customer { Id = "1", PartitionKey = "1" }); - } - else - { - context.OtherContainerCustomers.Add(new OtherContainerCustomer { Id = "1", PartitionKey = "1" }); - } - - if (async) - { - await context.SaveChangesAsync(); - } - else - { - context.SaveChanges(); - } - - context.ChangeTracker.Clear(); - var initialTokens = context.Database.GetSessionTokens(); - - if (defaultContainer) - { - context.Customers.Update(new Customer { Id = "1", Name = "updated", PartitionKey = "1" }); - } - else - { - context.OtherContainerCustomers.Update(new OtherContainerCustomer { Id = "1", Name = "updated", PartitionKey = "1" }); - } - - if (async) - { - await context.SaveChangesAsync(); - } - else - { - context.SaveChanges(); - } - - var sessionTokens = context.Database.GetSessionTokens(); - var sessionToken = defaultContainer ? sessionTokens[nameof(CosmosSessionTokenContext)]! : sessionTokens[OtherContainerName]!; - Assert.False(string.IsNullOrWhiteSpace(sessionToken)); - } - - [ConditionalTheory] - [InlineData(AutoTransactionBehavior.WhenNeeded, true, true)] - [InlineData(AutoTransactionBehavior.WhenNeeded, true, false)] - [InlineData(AutoTransactionBehavior.WhenNeeded, false, true)] - [InlineData(AutoTransactionBehavior.WhenNeeded, false, false)] - [InlineData(AutoTransactionBehavior.Never, true, true)] - [InlineData(AutoTransactionBehavior.Never, true, false)] - [InlineData(AutoTransactionBehavior.Never, false, true)] - [InlineData(AutoTransactionBehavior.Never, false, false)] - [InlineData(AutoTransactionBehavior.Always, true, true)] - [InlineData(AutoTransactionBehavior.Always, true, false)] - [InlineData(AutoTransactionBehavior.Always, false, true)] - [InlineData(AutoTransactionBehavior.Always, false, false)] - public virtual async Task Add_uses_session_token(AutoTransactionBehavior autoTransactionBehavior, bool async, bool defaultContainer) - { - var contextFactory = await InitializeAsync(); - - using var context = contextFactory.CreateContext(); - context.Database.AutoTransactionBehavior = autoTransactionBehavior; - - var sessionTokens = context.Database.GetSessionTokens(); - // Only way we can test this is by setting a session token that will fail the request if used.. - // Only way to do this for a write is to set an invalid session token.. - - if (defaultContainer) - { - context.Database.AppendSessionToken("invalidtoken"); - } - else - { - context.Database.AppendSessionTokens(new Dictionary { { OtherContainerName, "invalidtoken" } }); - } - - if (defaultContainer) - { - context.Customers.Add(new Customer { Id = "1", PartitionKey = "1" }); - } - else - { - context.OtherContainerCustomers.Add(new OtherContainerCustomer { Id = "1", PartitionKey = "1" }); - } - - DbUpdateException ex; - if (async) - { - ex = await Assert.ThrowsAsync(() => context.SaveChangesAsync()); - } - else - { - ex = Assert.Throws(() => context.SaveChanges()); - } - - Assert.Contains("The session token provided 'invalidtoken' is invalid.", ((CosmosException)ex.InnerException!).ResponseBody); - } - - [ConditionalTheory] - [InlineData(AutoTransactionBehavior.WhenNeeded, true, true)] - [InlineData(AutoTransactionBehavior.WhenNeeded, true, false)] - [InlineData(AutoTransactionBehavior.WhenNeeded, false, true)] - [InlineData(AutoTransactionBehavior.WhenNeeded, false, false)] - [InlineData(AutoTransactionBehavior.Never, true, true)] - [InlineData(AutoTransactionBehavior.Never, true, false)] - [InlineData(AutoTransactionBehavior.Never, false, true)] - [InlineData(AutoTransactionBehavior.Never, false, false)] - [InlineData(AutoTransactionBehavior.Always, true, true)] - [InlineData(AutoTransactionBehavior.Always, true, false)] - [InlineData(AutoTransactionBehavior.Always, false, true)] - [InlineData(AutoTransactionBehavior.Always, false, false)] - public virtual async Task Update_uses_session_token(AutoTransactionBehavior autoTransactionBehavior, bool async, bool defaultContainer) - { - var contextFactory = await InitializeAsync(); - - using var context = contextFactory.CreateContext(); - context.Database.AutoTransactionBehavior = autoTransactionBehavior; - - var sessionTokens = context.Database.GetSessionTokens(); - // Only way we can test this is by setting a session token that will fail the request if used.. - // Only way to do this for a write is to set an invalid session token.. - if (defaultContainer) - { - context.Database.AppendSessionToken("invalidtoken"); - } - else - { - context.Database.AppendSessionTokens(new Dictionary { { OtherContainerName, "invalidtoken" } }); - } - - if (defaultContainer) - { - context.Customers.Update(new Customer { Id = "1", PartitionKey = "1" }); - } - else - { - context.OtherContainerCustomers.Update(new OtherContainerCustomer { Id = "1", PartitionKey = "1" }); - } - - DbUpdateException ex; - if (async) - { - ex = await Assert.ThrowsAsync(() => context.SaveChangesAsync()); - } - else - { - ex = Assert.Throws(() => context.SaveChanges()); - } - - Assert.Contains("The session token provided 'invalidtoken' is invalid.", ((CosmosException)ex.InnerException!).ResponseBody); - } - - [ConditionalTheory] - [InlineData(AutoTransactionBehavior.WhenNeeded, true, true)] - [InlineData(AutoTransactionBehavior.WhenNeeded, true, false)] - [InlineData(AutoTransactionBehavior.WhenNeeded, false, true)] - [InlineData(AutoTransactionBehavior.WhenNeeded, false, false)] - [InlineData(AutoTransactionBehavior.Never, true, true)] - [InlineData(AutoTransactionBehavior.Never, true, false)] - [InlineData(AutoTransactionBehavior.Never, false, true)] - [InlineData(AutoTransactionBehavior.Never, false, false)] - [InlineData(AutoTransactionBehavior.Always, true, true)] - [InlineData(AutoTransactionBehavior.Always, true, false)] - [InlineData(AutoTransactionBehavior.Always, false, true)] - [InlineData(AutoTransactionBehavior.Always, false, false)] - public virtual async Task Delete_uses_session_token(AutoTransactionBehavior autoTransactionBehavior, bool async, bool defaultContainer) - { - var contextFactory = await InitializeAsync(); - - using var context = contextFactory.CreateContext(); - context.Database.AutoTransactionBehavior = autoTransactionBehavior; - - var sessionTokens = context.Database.GetSessionTokens(); - // Only way we can test this is by setting a session token that will fail the request if used.. - // Only way to do this for a write is to set an invalid session token.. - if (defaultContainer) - { - context.Database.AppendSessionToken("invalidtoken"); - } - else - { - context.Database.AppendSessionTokens(new Dictionary { { OtherContainerName, "invalidtoken" } }); - } - - if (defaultContainer) - { - context.Customers.Remove(new Customer { Id = "1", PartitionKey = "1" }); - } - else - { - context.OtherContainerCustomers.Remove(new OtherContainerCustomer { Id = "1", PartitionKey = "1" }); - } - - DbUpdateException ex; - if (async) - { - ex = await Assert.ThrowsAsync(() => context.SaveChangesAsync()); - } - else - { - ex = Assert.Throws(() => context.SaveChanges()); - } - - Assert.Contains("The session token provided 'invalidtoken' is invalid.", ((CosmosException)ex.InnerException!).ResponseBody); - } -} diff --git a/test/EFCore.Cosmos.Tests/Extensions/CosmosDbContextOptionsExtensionsTests.cs b/test/EFCore.Cosmos.Tests/Extensions/CosmosDbContextOptionsExtensionsTests.cs index 50b04a60a93..a7c8db0ea68 100644 --- a/test/EFCore.Cosmos.Tests/Extensions/CosmosDbContextOptionsExtensionsTests.cs +++ b/test/EFCore.Cosmos.Tests/Extensions/CosmosDbContextOptionsExtensionsTests.cs @@ -65,7 +65,7 @@ public void Can_create_options_with_valid_values() Test(o => o.MaxTcpConnectionsPerEndpoint(3), o => Assert.Equal(3, o.MaxTcpConnectionsPerEndpoint)); Test(o => o.LimitToEndpoint(), o => Assert.True(o.LimitToEndpoint)); Test(o => o.ContentResponseOnWriteEnabled(), o => Assert.True(o.EnableContentResponseOnWrite)); - Test(o => o.ManualSessionTokenManagementEnabled(), o => Assert.True(o.EnableManualSessionTokenManagement)); + Test(o => o.SessionTokenManagementMode(Cosmos.Infrastructure.SessionTokenManagementMode.EnforcedManual), o => Assert.Equal(Cosmos.Infrastructure.SessionTokenManagementMode.EnforcedManual, o.SessionTokenManagementMode)); var webProxy = new WebProxy(); Test(o => o.WebProxy(webProxy), o => Assert.Same(webProxy, o.WebProxy)); diff --git a/test/EFCore.Cosmos.Tests/Storage/Internal/SessionTokenStorageTest.cs b/test/EFCore.Cosmos.Tests/Storage/Internal/SessionTokenStorageTest.cs new file mode 100644 index 00000000000..ccdd4e80df2 --- /dev/null +++ b/test/EFCore.Cosmos.Tests/Storage/Internal/SessionTokenStorageTest.cs @@ -0,0 +1,785 @@ +// Licensed to the .NET Foundation under one or more agreements. +// The .NET Foundation licenses this file to you under the MIT license. + +#nullable enable + +using Microsoft.EntityFrameworkCore.Cosmos.Infrastructure; + +namespace Microsoft.EntityFrameworkCore.Cosmos.Storage.Internal; + +public class SessionTokenStorageTest +{ + private readonly string _defaultContainerName = "default"; + private readonly HashSet _containerNames = new(["default", "other"]); + + + [ConditionalTheory] + [InlineData(SessionTokenManagementMode.FullyAutomatic)] + [InlineData(SessionTokenManagementMode.SemiAutomatic)] + [InlineData(SessionTokenManagementMode.Manual)] + [InlineData(SessionTokenManagementMode.EnforcedManual)] + public virtual void TrackSessionToken_WhenContainerNameIsNull_ThrowsArgumentNullException(SessionTokenManagementMode mode) + { + var storage = CreateStorage(mode); + Assert.Throws(() => storage.TrackSessionToken(null!, "A")); + } + + [ConditionalTheory] + [InlineData(SessionTokenManagementMode.FullyAutomatic)] + [InlineData(SessionTokenManagementMode.SemiAutomatic)] + [InlineData(SessionTokenManagementMode.Manual)] + [InlineData(SessionTokenManagementMode.EnforcedManual)] + public virtual void TrackSessionToken_WhenContainerNameIsWhitespace_ThrowsArgumentNullException(SessionTokenManagementMode mode) + { + var storage = CreateStorage(mode); + Assert.Throws(() => storage.TrackSessionToken(" ", "A")); + Assert.Throws(() => storage.TrackSessionToken("", "A")); + } + + [ConditionalTheory] + [InlineData(SessionTokenManagementMode.FullyAutomatic)] + [InlineData(SessionTokenManagementMode.SemiAutomatic)] + [InlineData(SessionTokenManagementMode.Manual)] + [InlineData(SessionTokenManagementMode.EnforcedManual)] + public virtual void TrackSessionToken_WhenTokenIsNull_ThrowsArgumentNullException(SessionTokenManagementMode mode) + { + var storage = CreateStorage(mode); + Assert.Throws(() => storage.TrackSessionToken(_defaultContainerName, null!)); + } + + [ConditionalTheory] + [InlineData(SessionTokenManagementMode.FullyAutomatic)] + [InlineData(SessionTokenManagementMode.SemiAutomatic)] + [InlineData(SessionTokenManagementMode.Manual)] + [InlineData(SessionTokenManagementMode.EnforcedManual)] + public virtual void TrackSessionToken_WhenTokenIsWhitespace_ThrowsArgumentNullException(SessionTokenManagementMode mode) + { + var storage = CreateStorage(mode); + Assert.Throws(() => storage.TrackSessionToken(_defaultContainerName, " ")); + Assert.Throws(() => storage.TrackSessionToken(_defaultContainerName, "")); + } + + [ConditionalTheory] + [InlineData(SessionTokenManagementMode.SemiAutomatic)] + [InlineData(SessionTokenManagementMode.Manual)] + [InlineData(SessionTokenManagementMode.EnforcedManual)] + public virtual void SetSessionTokens_WhenContainerNameIsUnknown_ThrowsInvalidOperationException(SessionTokenManagementMode mode) + { + var storage = CreateStorage(mode); + var ex = Assert.Throws(() => + storage.SetSessionTokens(new Dictionary { { "bad", "A" } })); + Assert.Equal("invalid container name", ex.Message); + } + + [ConditionalTheory] + [InlineData(SessionTokenManagementMode.FullyAutomatic)] + [InlineData(SessionTokenManagementMode.SemiAutomatic)] + [InlineData(SessionTokenManagementMode.Manual)] + [InlineData(SessionTokenManagementMode.EnforcedManual)] + public virtual void AppendDefaultContainerSessionToken_WhenTokenIsNull_ThrowsArgumentNullException(SessionTokenManagementMode mode) + { + var storage = CreateStorage(mode); + Assert.Throws(() => storage.AppendDefaultContainerSessionToken(null!)); + } + + [ConditionalTheory] + [InlineData(SessionTokenManagementMode.FullyAutomatic)] + [InlineData(SessionTokenManagementMode.SemiAutomatic)] + [InlineData(SessionTokenManagementMode.Manual)] + [InlineData(SessionTokenManagementMode.EnforcedManual)] + public virtual void AppendDefaultContainerSessionToken_WhenTokenIsWhitespace_ThrowsArgumentException(SessionTokenManagementMode mode) + { + var storage = CreateStorage(mode); + Assert.Throws(() => storage.AppendDefaultContainerSessionToken(" ")); + Assert.Throws(() => storage.AppendDefaultContainerSessionToken("")); + } + + [ConditionalTheory] + [InlineData(SessionTokenManagementMode.SemiAutomatic)] + [InlineData(SessionTokenManagementMode.Manual)] + [InlineData(SessionTokenManagementMode.EnforcedManual)] + public virtual void AppendSessionTokens_WhenContainerNameIsUnknown_ThrowsInvalidOperationException(SessionTokenManagementMode mode) + { + var storage = CreateStorage(mode); + var ex = Assert.Throws(() => + storage.AppendSessionTokens(new Dictionary { { "bad", "A" } })); + Assert.Equal("invalid container name", ex.Message); + } + + [ConditionalTheory] + [InlineData(SessionTokenManagementMode.SemiAutomatic)] + [InlineData(SessionTokenManagementMode.Manual)] + [InlineData(SessionTokenManagementMode.EnforcedManual)] + public virtual void GetSessionToken_WhenContainerNameIsNull_ThrowsArgumentNullException(SessionTokenManagementMode mode) + { + var storage = CreateStorage(mode); + Assert.Throws(() => storage.GetSessionToken(null!)); + } + + [ConditionalTheory] + [InlineData(SessionTokenManagementMode.SemiAutomatic)] + [InlineData(SessionTokenManagementMode.Manual)] + [InlineData(SessionTokenManagementMode.EnforcedManual)] + public virtual void GetSessionToken_WhenContainerNameIsWhitespace_ThrowsArgumentNullException(SessionTokenManagementMode mode) + { + var storage = CreateStorage(mode); + Assert.Throws(() => storage.GetSessionToken(" ")); + Assert.Throws(() => storage.GetSessionToken("")); + } + + // ================================================================ + // FUNCTIONAL TESTS - SET AND RETRIEVE + // ================================================================ + + [ConditionalTheory] + [InlineData(SessionTokenManagementMode.SemiAutomatic)] + [InlineData(SessionTokenManagementMode.Manual)] + [InlineData(SessionTokenManagementMode.EnforcedManual)] + public virtual void SetSessionTokens_WhenSettingSingleToken_CanRetrieveFromAllMethods(SessionTokenManagementMode mode) + { + var storage = CreateStorage(mode); + + storage.SetSessionTokens(new Dictionary { { _defaultContainerName, "A" } }); + + var all = storage.GetTrackedTokens(); + Assert.Equal("A", all[_defaultContainerName]); + Assert.Null(all["other"]); + Assert.Equal("A", storage.GetSessionToken(_defaultContainerName)); + Assert.Equal("A", storage.GetDefaultContainerTrackedToken()); + } + + [ConditionalTheory] + [InlineData(SessionTokenManagementMode.SemiAutomatic)] + [InlineData(SessionTokenManagementMode.Manual)] + [InlineData(SessionTokenManagementMode.EnforcedManual)] + public virtual void SetSessionTokens_WhenSettingMultipleContainers_AllContainersAreSetCorrectly(SessionTokenManagementMode mode) + { + var storage = CreateStorage(mode); + + storage.SetSessionTokens(new Dictionary + { + { _defaultContainerName, "Token1" }, + { "other", "Token2" } + }); + + var all = storage.GetTrackedTokens(); + Assert.Equal("Token1", all[_defaultContainerName]); + Assert.Equal("Token2", all["other"]); + Assert.Equal("Token1", storage.GetSessionToken(_defaultContainerName)); + Assert.Equal("Token2", storage.GetSessionToken("other")); + } + + [ConditionalTheory] + [InlineData(SessionTokenManagementMode.SemiAutomatic)] + [InlineData(SessionTokenManagementMode.Manual)] + [InlineData(SessionTokenManagementMode.EnforcedManual)] + public virtual void SetSessionTokens_WhenSettingNullValue_ContainerIsCleared(SessionTokenManagementMode mode) + { + var storage = CreateStorage(mode); + storage.SetSessionTokens(new Dictionary { { _defaultContainerName, "A" } }); + storage.SetSessionTokens(new Dictionary { { _defaultContainerName, null } }); + + var token = storage.GetDefaultContainerTrackedToken(); + Assert.Null(token); + } + + [ConditionalTheory] + [InlineData(SessionTokenManagementMode.SemiAutomatic)] + [InlineData(SessionTokenManagementMode.Manual)] + [InlineData(SessionTokenManagementMode.EnforcedManual)] + public virtual void SetSessionTokens_WhenSettingEmptyDictionary_NoExceptionThrown(SessionTokenManagementMode mode) + { + var storage = CreateStorage(mode); + storage.SetSessionTokens(new Dictionary()); + + var all = storage.GetTrackedTokens(); + Assert.Null(all[_defaultContainerName]); + Assert.Null(all["other"]); + } + + [ConditionalTheory] + [InlineData(SessionTokenManagementMode.SemiAutomatic)] + [InlineData(SessionTokenManagementMode.Manual)] + [InlineData(SessionTokenManagementMode.EnforcedManual)] + public virtual void SetSessionTokens_WhenPartiallyUpdatingContainers_OnlySpecifiedContainersAreUpdated(SessionTokenManagementMode mode) + { + var storage = CreateStorage(mode); + storage.SetSessionTokens(new Dictionary + { + { _defaultContainerName, "A" }, + { "other", "B" } + }); + storage.SetSessionTokens(new Dictionary { { _defaultContainerName, "C" } }); + + var all = storage.GetTrackedTokens(); + Assert.Equal("C", all[_defaultContainerName]); + Assert.Equal("B", all["other"]); + } + + // ================================================================ + // FUNCTIONAL TESTS - APPEND OPERATIONS + // ================================================================ + + [ConditionalTheory] + [InlineData(SessionTokenManagementMode.SemiAutomatic)] + [InlineData(SessionTokenManagementMode.Manual)] + [InlineData(SessionTokenManagementMode.EnforcedManual)] + public virtual void AppendSessionTokens_WhenAppendingOverlappingTokens_MergesUniquely(SessionTokenManagementMode mode) + { + var storage = CreateStorage(mode); + + storage.AppendSessionTokens(new Dictionary { { _defaultContainerName, "A,B" } }); + storage.AppendSessionTokens(new Dictionary { { _defaultContainerName, "B,C" } }); + + var token = storage.GetDefaultContainerTrackedToken(); + Assert.Equal("A,B,C", token); + } + + [ConditionalTheory] + [InlineData(SessionTokenManagementMode.SemiAutomatic)] + [InlineData(SessionTokenManagementMode.Manual)] + [InlineData(SessionTokenManagementMode.EnforcedManual)] + public virtual void AppendSessionTokens_WhenAppendingToMultipleContainers_AllContainersAreUpdated(SessionTokenManagementMode mode) + { + var storage = CreateStorage(mode); + + storage.AppendSessionTokens(new Dictionary + { + { _defaultContainerName, "A" }, + { "other", "B" } + }); + storage.AppendSessionTokens(new Dictionary + { + { _defaultContainerName, "C" }, + { "other", "D" } + }); + + var all = storage.GetTrackedTokens(); + Assert.Equal("A,C", all[_defaultContainerName]); + Assert.Equal("B,D", all["other"]); + } + + [ConditionalTheory] + [InlineData(SessionTokenManagementMode.SemiAutomatic)] + [InlineData(SessionTokenManagementMode.Manual)] + [InlineData(SessionTokenManagementMode.EnforcedManual)] + public virtual void AppendSessionTokens_WhenAppendingEmptyDictionary_NoExceptionThrown(SessionTokenManagementMode mode) + { + var storage = CreateStorage(mode); + storage.AppendSessionTokens(new Dictionary()); + + var all = storage.GetTrackedTokens(); + Assert.Null(all[_defaultContainerName]); + Assert.Null(all["other"]); + } + + [ConditionalTheory] + [InlineData(SessionTokenManagementMode.SemiAutomatic)] + [InlineData(SessionTokenManagementMode.Manual)] + [InlineData(SessionTokenManagementMode.EnforcedManual)] + public virtual void AppendDefaultContainerSessionToken_WhenAppendingMultipleTokens_AccumulatesUniquely(SessionTokenManagementMode mode) + { + var storage = CreateStorage(mode); + + storage.AppendDefaultContainerSessionToken("A"); + storage.AppendDefaultContainerSessionToken("B"); + storage.AppendDefaultContainerSessionToken("B"); + + var token = storage.GetDefaultContainerTrackedToken(); + Assert.Equal("A,B", token); + } + + // ================================================================ + // FUNCTIONAL TESTS - SET OPERATIONS + // ================================================================ + + [ConditionalTheory] + [InlineData(SessionTokenManagementMode.SemiAutomatic)] + [InlineData(SessionTokenManagementMode.Manual)] + [InlineData(SessionTokenManagementMode.EnforcedManual)] + public virtual void SetDefaultContainerSessionToken_WhenReplacingExistingTokens_ReplacesAllPreviousTokens(SessionTokenManagementMode mode) + { + var storage = CreateStorage(mode); + storage.AppendDefaultContainerSessionToken("A"); + storage.AppendDefaultContainerSessionToken("B"); + storage.SetDefaultContainerSessionToken("XYZ"); + + var token = storage.GetDefaultContainerTrackedToken(); + Assert.Equal("XYZ", token); + } + + [ConditionalTheory] + [InlineData(SessionTokenManagementMode.SemiAutomatic)] + [InlineData(SessionTokenManagementMode.Manual)] + [InlineData(SessionTokenManagementMode.EnforcedManual)] + public virtual void SetDefaultContainerSessionToken_WhenSettingNull_ClearsContainerToken(SessionTokenManagementMode mode) + { + var storage = CreateStorage(mode); + storage.SetDefaultContainerSessionToken("A"); + storage.SetDefaultContainerSessionToken(null); + + var token = storage.GetDefaultContainerTrackedToken(); + Assert.Null(token); + } + + [ConditionalTheory] + [InlineData(SessionTokenManagementMode.SemiAutomatic)] + [InlineData(SessionTokenManagementMode.Manual)] + [InlineData(SessionTokenManagementMode.EnforcedManual)] + public virtual void SetDefaultContainerSessionToken_WhenSettingEmptyString_StoresEmptyString(SessionTokenManagementMode mode) + { + var storage = CreateStorage(mode); + storage.SetDefaultContainerSessionToken(""); + var token = storage.GetDefaultContainerTrackedToken(); + Assert.Equal("", token); + } + + // ================================================================ + // FUNCTIONAL TESTS - TRACK OPERATIONS + // ================================================================ + + [ConditionalTheory] + [InlineData(SessionTokenManagementMode.SemiAutomatic)] + [InlineData(SessionTokenManagementMode.Manual)] + public virtual void TrackSessionToken_WhenTrackingMultipleTokens_AppendsUniquely(SessionTokenManagementMode mode) + { + var storage = CreateStorage(mode); + + storage.TrackSessionToken(_defaultContainerName, "A"); + storage.TrackSessionToken(_defaultContainerName, "B"); + storage.TrackSessionToken(_defaultContainerName, "A"); + + var token = storage.GetSessionToken(_defaultContainerName); + Assert.Equal("A,B", token); + } + + [ConditionalTheory] + [InlineData(SessionTokenManagementMode.SemiAutomatic)] + [InlineData(SessionTokenManagementMode.Manual)] + public virtual void TrackSessionToken_WhenTrackingToDifferentContainers_ContainersAreIndependent(SessionTokenManagementMode mode) + { + var storage = CreateStorage(mode); + + storage.TrackSessionToken(_defaultContainerName, "A"); + storage.TrackSessionToken("other", "B"); + storage.TrackSessionToken(_defaultContainerName, "C"); + + Assert.Equal("A,C", storage.GetSessionToken(_defaultContainerName)); + Assert.Equal("B", storage.GetSessionToken("other")); + } + + [ConditionalTheory] + [InlineData(SessionTokenManagementMode.SemiAutomatic)] + [InlineData(SessionTokenManagementMode.Manual)] + public virtual void TrackSessionToken_WhenTrackingCommaSeparatedToken_ParsesAndMergesCorrectly(SessionTokenManagementMode mode) + { + var storage = CreateStorage(mode); + + storage.TrackSessionToken(_defaultContainerName, "A,B"); + storage.TrackSessionToken(_defaultContainerName, "C"); + storage.TrackSessionToken(_defaultContainerName, "B,D"); + + var token = storage.GetSessionToken(_defaultContainerName); + Assert.Equal("A,B,C,D", token); + } + + // ================================================================ + // FUNCTIONAL TESTS - CLEAR OPERATIONS + // ================================================================ + + [ConditionalTheory] + [InlineData(SessionTokenManagementMode.SemiAutomatic)] + [InlineData(SessionTokenManagementMode.Manual)] + [InlineData(SessionTokenManagementMode.EnforcedManual)] + public virtual void Clear_WhenClearingAllTokens_ResetsAllContainersToNull(SessionTokenManagementMode mode) + { + var storage = CreateStorage(mode); + + storage.AppendSessionTokens(new Dictionary { { _defaultContainerName, "A" }, { "other", "B" } }); + storage.Clear(); + + var all = storage.GetTrackedTokens(); + Assert.Null(all[_defaultContainerName]); + Assert.Null(all["other"]); + } + + [ConditionalTheory] + [InlineData(SessionTokenManagementMode.SemiAutomatic)] + [InlineData(SessionTokenManagementMode.Manual)] + [InlineData(SessionTokenManagementMode.EnforcedManual)] + public virtual void Clear_WhenClearing_ContainersStillExistInTrackedTokens(SessionTokenManagementMode mode) + { + var storage = CreateStorage(mode); + storage.AppendSessionTokens(new Dictionary { { _defaultContainerName, "A" } }); + storage.Clear(); + + var tracked = storage.GetTrackedTokens(); + Assert.Contains(_defaultContainerName, tracked.Keys); + Assert.Contains("other", tracked.Keys); + } + + [ConditionalTheory] + [InlineData(SessionTokenManagementMode.SemiAutomatic)] + [InlineData(SessionTokenManagementMode.Manual)] + [InlineData(SessionTokenManagementMode.EnforcedManual)] + public virtual void Clear_WhenClearing_CanSetNewTokensAfterClear(SessionTokenManagementMode mode) + { + var storage = CreateStorage(mode); + storage.AppendDefaultContainerSessionToken("A"); + storage.Clear(); + storage.AppendDefaultContainerSessionToken("B"); + + var token = storage.GetDefaultContainerTrackedToken(); + Assert.Equal("B", token); + } + + // ================================================================ + // FULLY AUTOMATIC MODE TESTS + // ================================================================ + + [ConditionalFact] + public virtual void FullyAutomatic_WhenCallingSetSessionTokens_ThrowsInvalidOperationException() + { + var storage = CreateStorage(SessionTokenManagementMode.FullyAutomatic); + var ex = Assert.Throws(() => + storage.SetSessionTokens(new Dictionary())); + Assert.Equal("Can't use session tokens with FullyAutomatic", ex.Message); + } + + [ConditionalFact] + public virtual void FullyAutomatic_WhenCallingGetTrackedTokens_ThrowsInvalidOperationException() + { + var storage = CreateStorage(SessionTokenManagementMode.FullyAutomatic); + var ex = Assert.Throws(() => storage.GetTrackedTokens()); + Assert.Equal("Can't use session tokens with FullyAutomatic", ex.Message); + } + + [ConditionalFact] + public virtual void FullyAutomatic_WhenCallingAppendSessionTokens_ThrowsInvalidOperationException() + { + var storage = CreateStorage(SessionTokenManagementMode.FullyAutomatic); + var ex = Assert.Throws(() => + storage.AppendSessionTokens(new Dictionary())); + Assert.Equal("Can't use session tokens with FullyAutomatic", ex.Message); + } + + [ConditionalFact] + public virtual void FullyAutomatic_WhenCallingSetDefaultContainerSessionToken_ThrowsInvalidOperationException() + { + var storage = CreateStorage(SessionTokenManagementMode.FullyAutomatic); + var ex = Assert.Throws(() => + storage.SetDefaultContainerSessionToken(null)); + Assert.Equal("Can't use session tokens with FullyAutomatic", ex.Message); + } + + [ConditionalFact] + public virtual void FullyAutomatic_WhenCallingAppendDefaultContainerSessionToken_ThrowsInvalidOperationException() + { + var storage = CreateStorage(SessionTokenManagementMode.FullyAutomatic); + var ex = Assert.Throws(() => + storage.AppendDefaultContainerSessionToken("A")); + Assert.Equal("Can't use session tokens with FullyAutomatic", ex.Message); + } + + [ConditionalFact] + public virtual void FullyAutomatic_WhenCallingGetDefaultContainerTrackedToken_ThrowsInvalidOperationException() + { + var storage = CreateStorage(SessionTokenManagementMode.FullyAutomatic); + var ex = Assert.Throws(() => + storage.GetDefaultContainerTrackedToken()); + Assert.Equal("Can't use session tokens with FullyAutomatic", ex.Message); + } + + [ConditionalFact] + public virtual void FullyAutomatic_WhenTrackingToken_AlwaysReturnsNull() + { + var storage = CreateStorage(SessionTokenManagementMode.FullyAutomatic); + storage.TrackSessionToken(_defaultContainerName, "A"); + Assert.Null(storage.GetSessionToken(_defaultContainerName)); + } + + [ConditionalFact] + public virtual void FullyAutomatic_WhenTrackingMultipleTokens_AlwaysReturnsNull() + { + var storage = CreateStorage(SessionTokenManagementMode.FullyAutomatic); + storage.TrackSessionToken(_defaultContainerName, "A"); + storage.TrackSessionToken(_defaultContainerName, "B"); + storage.TrackSessionToken("other", "C"); + Assert.Null(storage.GetSessionToken(_defaultContainerName)); + Assert.Null(storage.GetSessionToken("other")); + } + + // ================================================================ + // ENFORCED MANUAL MODE TESTS + // ================================================================ + + [ConditionalFact] + public virtual void EnforcedManual_WhenGettingTokenBeforeSet_ThrowsInvalidOperationException() + { + var storage = CreateStorage(SessionTokenManagementMode.EnforcedManual); + var ex = Assert.Throws(() => + storage.GetSessionToken(_defaultContainerName)); + Assert.Contains("No session token set for container while EnforcedManual", ex.Message); + } + + [ConditionalFact] + public virtual void EnforcedManual_WhenGettingTokenAfterSet_ReturnsToken() + { + var storage = CreateStorage(SessionTokenManagementMode.EnforcedManual); + storage.SetDefaultContainerSessionToken("A"); + var token = storage.GetSessionToken(_defaultContainerName); + Assert.Equal("A", token); + } + + [ConditionalFact] + public virtual void EnforcedManual_WhenGettingTokenAfterClear_ThrowsInvalidOperationException() + { + var storage = CreateStorage(SessionTokenManagementMode.EnforcedManual); + storage.SetDefaultContainerSessionToken("A"); + storage.Clear(); + var ex = Assert.Throws(() => + storage.GetSessionToken(_defaultContainerName)); + Assert.Contains("No session token set for container while EnforcedManual", ex.Message); + } + + [ConditionalFact] + public virtual void EnforcedManual_WhenGettingTokenAfterSetThenClearThenSet_ReturnsNewToken() + { + var storage = CreateStorage(SessionTokenManagementMode.EnforcedManual); + storage.SetDefaultContainerSessionToken("A"); + storage.Clear(); + storage.SetDefaultContainerSessionToken("B"); + var token = storage.GetSessionToken(_defaultContainerName); + Assert.Equal("B", token); + } + + [ConditionalFact] + public virtual void EnforcedManual_WhenSettingMultipleContainers_AllContainersCanBeRetrieved() + { + var storage = CreateStorage(SessionTokenManagementMode.EnforcedManual); + storage.SetSessionTokens(new Dictionary + { + { _defaultContainerName, "A" }, + { "other", "B" } + }); + + Assert.Equal("A", storage.GetSessionToken(_defaultContainerName)); + Assert.Equal("B", storage.GetSessionToken("other")); + } + + [ConditionalFact] + public virtual void EnforcedManual_WhenOneContainerNotSet_ThrowsForThatContainerOnly() + { + var storage = CreateStorage(SessionTokenManagementMode.EnforcedManual); + storage.SetSessionTokens(new Dictionary { { _defaultContainerName, "A" } }); + + Assert.Equal("A", storage.GetSessionToken(_defaultContainerName)); + var ex = Assert.Throws(() => storage.GetSessionToken("other")); + Assert.Contains("No session token set for container while EnforcedManual", ex.Message); + } + + [ConditionalFact] + public virtual void EnforcedManual_WhenTrackingToken_ThrowsWhenRetrieving() + { + var storage = CreateStorage(SessionTokenManagementMode.EnforcedManual); + storage.TrackSessionToken(_defaultContainerName, "A"); + var ex = Assert.Throws(() => storage.GetSessionToken(_defaultContainerName)); + Assert.Contains("No session token set for container while EnforcedManual", ex.Message); + } + + [ConditionalFact] + public virtual void EnforcedManual_WhenAppendingToken_CanRetrieveToken() + { + var storage = CreateStorage(SessionTokenManagementMode.EnforcedManual); + storage.AppendDefaultContainerSessionToken("A"); + var token = storage.GetSessionToken(_defaultContainerName); + Assert.Equal("A", token); + } + + // ================================================================ + // INITIALIZATION AND CONTAINER MANAGEMENT TESTS + // ================================================================ + + [ConditionalFact] + public virtual void Constructor_WhenInitializing_AllContainersAreInitialized() + { + var storage = CreateStorage(SessionTokenManagementMode.Manual); + + var tracked = storage.GetTrackedTokens(); + Assert.True(tracked.ContainsKey(_defaultContainerName)); + Assert.True(tracked.ContainsKey("other")); + Assert.Equal(2, tracked.Count); + } + + [ConditionalFact] + public virtual void Constructor_WhenDefaultContainerNotInContainerNames_ThrowsException() + { + var containers = new HashSet(["other"]); + Assert.True(!containers.Contains("default")); + Assert.ThrowsAny(() => + { + _ = new SessionTokenStorage("default", containers, SessionTokenManagementMode.Manual); + }); + } + + [ConditionalFact] + public virtual void Constructor_WhenInitializing_AllContainersStartWithNullTokens() + { + var storage = CreateStorage(SessionTokenManagementMode.Manual); + + var tracked = storage.GetTrackedTokens(); + Assert.Null(tracked[_defaultContainerName]); + Assert.Null(tracked["other"]); + } + + // ================================================================ + // GETTRACKEDTOKENS TESTS + // ================================================================ + + [ConditionalFact] + public virtual void GetTrackedTokens_WhenCalled_ReturnsSnapshotNotLiveReference() + { + var storage = CreateStorage(SessionTokenManagementMode.Manual); + var snapshot = storage.GetTrackedTokens(); + + storage.AppendDefaultContainerSessionToken("A"); + var snapshot2 = storage.GetTrackedTokens(); + + Assert.NotSame(snapshot, snapshot2); + Assert.Null(snapshot[_defaultContainerName]); + Assert.Equal("A", snapshot2[_defaultContainerName]); + } + + [ConditionalFact] + public virtual void GetTrackedTokens_WhenModifyingReturnedDictionary_DoesNotAffectStorage() + { + var storage = CreateStorage(SessionTokenManagementMode.Manual); + storage.AppendDefaultContainerSessionToken("A"); + var tracked = storage.GetTrackedTokens(); + + // This should not compile or should not affect storage + // The returned dictionary is read-only, so we can't modify it + var token = storage.GetDefaultContainerTrackedToken(); + Assert.Equal("A", token); + } + + // ================================================================ + // COMPOSITE TOKEN TESTS + // ================================================================ + + [ConditionalFact] + public virtual void CompositeSessionToken_WhenAppendingDuplicateTokens_RemovesDuplicates() + { + var storage = CreateStorage(SessionTokenManagementMode.Manual); + storage.AppendDefaultContainerSessionToken("A,A,B"); + var token = storage.GetDefaultContainerTrackedToken(); + Assert.Equal("A,B", token); + } + + [ConditionalFact] + public virtual void CompositeSessionToken_WhenAppendingDuplicateTokensInSeparateCalls_RemovesDuplicates() + { + var storage = CreateStorage(SessionTokenManagementMode.Manual); + storage.AppendDefaultContainerSessionToken("A"); + storage.AppendDefaultContainerSessionToken("A"); + storage.AppendDefaultContainerSessionToken("B"); + storage.AppendDefaultContainerSessionToken("A"); + var token = storage.GetDefaultContainerTrackedToken(); + Assert.Equal("A,B", token); + } + + [ConditionalFact] + public virtual void CompositeSessionToken_WhenSettingCommaSeparatedTokens_StoresAllTokens() + { + var storage = CreateStorage(SessionTokenManagementMode.Manual); + storage.SetDefaultContainerSessionToken("A,B,C"); + var token = storage.GetDefaultContainerTrackedToken(); + Assert.Equal("A,B,C", token); + } + + [ConditionalFact] + public virtual void CompositeSessionToken_WhenAppendingCommaSeparatedTokens_ParsesAndMergesCorrectly() + { + var storage = CreateStorage(SessionTokenManagementMode.Manual); + storage.AppendDefaultContainerSessionToken("A,B,C"); + var token = storage.GetDefaultContainerTrackedToken(); + Assert.Equal("A,B,C", token); + } + + // ================================================================ + // MULTI-CONTAINER OPERATIONS TESTS + // ================================================================ + + [ConditionalFact] + public virtual void TrackSessionToken_WhenTrackingToNonDefaultContainer_MergesCorrectly() + { + var storage = CreateStorage(SessionTokenManagementMode.Manual); + storage.TrackSessionToken("other", "A"); + storage.TrackSessionToken("other", "B"); + Assert.Equal("A,B", storage.GetSessionToken("other")); + } + + // ================================================================ + // NULL AND EMPTY VALUE TESTS + // ================================================================ + + [ConditionalFact] + public virtual void GetDefaultContainerTrackedToken_WhenNoTokenSet_ReturnsNull() + { + var storage = CreateStorage(SessionTokenManagementMode.Manual); + Assert.Null(storage.GetDefaultContainerTrackedToken()); + } + + [ConditionalFact] + public virtual void GetDefaultContainerTrackedToken_AfterClear_ReturnsNull() + { + var storage = CreateStorage(SessionTokenManagementMode.Manual); + storage.AppendDefaultContainerSessionToken("A"); + storage.Clear(); + Assert.Null(storage.GetDefaultContainerTrackedToken()); + } + + [ConditionalFact] + public virtual void GetSessionToken_AfterClear_ReturnsNull() + { + var storage = CreateStorage(SessionTokenManagementMode.Manual); + storage.AppendDefaultContainerSessionToken("A"); + storage.Clear(); + var token = storage.GetSessionToken(_defaultContainerName); + Assert.Null(token); + } + + // ================================================================ + // SEMI-AUTOMATIC MODE SPECIFIC TESTS + // ================================================================ + + [ConditionalFact] + public virtual void SemiAutomatic_WhenTrackingToken_CanRetrieveToken() + { + var storage = CreateStorage(SessionTokenManagementMode.SemiAutomatic); + storage.TrackSessionToken(_defaultContainerName, "A"); + var token = storage.GetSessionToken(_defaultContainerName); + Assert.Equal("A", token); + } + + [ConditionalFact] + public virtual void SemiAutomatic_WhenGettingTokenBeforeSet_ReturnsNull() + { + var storage = CreateStorage(SessionTokenManagementMode.SemiAutomatic); + var token = storage.GetSessionToken(_defaultContainerName); + Assert.Null(token); + } + + [ConditionalFact] + public virtual void SemiAutomatic_WhenGettingTokenAfterClear_ReturnsNull() + { + var storage = CreateStorage(SessionTokenManagementMode.SemiAutomatic); + storage.SetDefaultContainerSessionToken("A"); + storage.Clear(); + var token = storage.GetSessionToken(_defaultContainerName); + Assert.Null(token); + } + + private SessionTokenStorage CreateStorage(SessionTokenManagementMode mode) + => new(_defaultContainerName, _containerNames, mode); +} From 7f4105b745e0934e336d55a7d2fc9e70b9e0ec23 Mon Sep 17 00:00:00 2001 From: JoasE <32096708+JoasE@users.noreply.github.com> Date: Thu, 6 Nov 2025 22:35:41 +0100 Subject: [PATCH 32/37] Use cosmos strings --- .../Properties/CosmosStrings.Designer.cs | 8 ++++++ .../Properties/CosmosStrings.resx | 4 +++ .../Storage/Internal/SessionTokenStorage.cs | 9 ++++--- .../Internal/SessionTokenStorageTest.cs | 25 ++++++++++--------- 4 files changed, 30 insertions(+), 16 deletions(-) diff --git a/src/EFCore.Cosmos/Properties/CosmosStrings.Designer.cs b/src/EFCore.Cosmos/Properties/CosmosStrings.Designer.cs index 3c5f5e3c41f..207ff19e3a8 100644 --- a/src/EFCore.Cosmos/Properties/CosmosStrings.Designer.cs +++ b/src/EFCore.Cosmos/Properties/CosmosStrings.Designer.cs @@ -275,6 +275,14 @@ public static string LimitOffsetNotSupportedInSubqueries public static string MissingOrderingInSelectExpression => GetString("MissingOrderingInSelectExpression"); + /// + /// No session token has been set for container: {container}. While using EnforceManual you must always set a session token for any container used. + /// + public static string MissingSessionTokenEnforceManual(object? container) + => string.Format( + GetString("MissingSessionTokenEnforceManual", nameof(container)), + container); + /// /// Root entity type '{entityType1}' is referenced by the query, but '{entityType2}' is already being referenced. A query can only reference a single root entity type. /// diff --git a/src/EFCore.Cosmos/Properties/CosmosStrings.resx b/src/EFCore.Cosmos/Properties/CosmosStrings.resx index 9622d0028bc..141fecfc397 100644 --- a/src/EFCore.Cosmos/Properties/CosmosStrings.resx +++ b/src/EFCore.Cosmos/Properties/CosmosStrings.resx @@ -264,6 +264,10 @@ 'Reverse' could not be translated to the server because there is no ordering on the server side. + + No session token has been set for container: {container}. While using EnforceManual you must always set a session token for any container used. + string + Root entity type '{entityType1}' is referenced by the query, but '{entityType2}' is already being referenced. A query can only reference a single root entity type. diff --git a/src/EFCore.Cosmos/Storage/Internal/SessionTokenStorage.cs b/src/EFCore.Cosmos/Storage/Internal/SessionTokenStorage.cs index 42aff6f34ac..87c2513b43c 100644 --- a/src/EFCore.Cosmos/Storage/Internal/SessionTokenStorage.cs +++ b/src/EFCore.Cosmos/Storage/Internal/SessionTokenStorage.cs @@ -2,6 +2,7 @@ // The .NET Foundation licenses this file to you under the MIT license. using Microsoft.EntityFrameworkCore.Cosmos.Infrastructure; +using Microsoft.EntityFrameworkCore.Cosmos.Internal; namespace Microsoft.EntityFrameworkCore.Cosmos.Storage.Internal; @@ -47,7 +48,7 @@ public virtual void SetSessionTokens(IReadOnlyDictionary sessio { if (!_containerNames.Contains(sessionToken.Key)) { - throw new InvalidOperationException("invalid container name"); + throw new InvalidOperationException(CosmosStrings.ContainerNameDoesNotExist(sessionToken.Key)); } _containerSessionTokens[sessionToken.Key] = new CompositeSessionToken(sessionToken.Value, true); @@ -67,7 +68,7 @@ public virtual void AppendSessionTokens(IReadOnlyDictionary sess { if (!_containerNames.Contains(sessionToken.Key)) { - throw new InvalidOperationException("invalid container name"); + throw new InvalidOperationException(CosmosStrings.ContainerNameDoesNotExist("bad")); } _containerSessionTokens[sessionToken.Key].Add(sessionToken.Value, true); @@ -142,7 +143,7 @@ public virtual void SetDefaultContainerSessionToken(string? sessionToken) if (!sessionToken.IsSet && _mode == SessionTokenManagementMode.EnforcedManual) { - throw new InvalidOperationException("No session token set for container while EnforcedManual"); + throw new InvalidOperationException(CosmosStrings.MissingSessionTokenEnforceManual(containerName)); } return sessionToken.ConvertToString(); @@ -185,7 +186,7 @@ private void CheckMode() { if (_mode == SessionTokenManagementMode.FullyAutomatic) { - throw new InvalidOperationException("Can't use session tokens with FullyAutomatic"); + throw new InvalidOperationException(CosmosStrings.EnableManualSessionTokenManagement); } } diff --git a/test/EFCore.Cosmos.Tests/Storage/Internal/SessionTokenStorageTest.cs b/test/EFCore.Cosmos.Tests/Storage/Internal/SessionTokenStorageTest.cs index ccdd4e80df2..85ebbf1362b 100644 --- a/test/EFCore.Cosmos.Tests/Storage/Internal/SessionTokenStorageTest.cs +++ b/test/EFCore.Cosmos.Tests/Storage/Internal/SessionTokenStorageTest.cs @@ -4,6 +4,7 @@ #nullable enable using Microsoft.EntityFrameworkCore.Cosmos.Infrastructure; +using Microsoft.EntityFrameworkCore.Cosmos.Internal; namespace Microsoft.EntityFrameworkCore.Cosmos.Storage.Internal; @@ -68,7 +69,7 @@ public virtual void SetSessionTokens_WhenContainerNameIsUnknown_ThrowsInvalidOpe var storage = CreateStorage(mode); var ex = Assert.Throws(() => storage.SetSessionTokens(new Dictionary { { "bad", "A" } })); - Assert.Equal("invalid container name", ex.Message); + Assert.Equal(CosmosStrings.ContainerNameDoesNotExist("bad"), ex.Message); } [ConditionalTheory] @@ -103,7 +104,7 @@ public virtual void AppendSessionTokens_WhenContainerNameIsUnknown_ThrowsInvalid var storage = CreateStorage(mode); var ex = Assert.Throws(() => storage.AppendSessionTokens(new Dictionary { { "bad", "A" } })); - Assert.Equal("invalid container name", ex.Message); + Assert.Equal(CosmosStrings.ContainerNameDoesNotExist("bad"), ex.Message); } [ConditionalTheory] @@ -443,7 +444,7 @@ public virtual void FullyAutomatic_WhenCallingSetSessionTokens_ThrowsInvalidOper var storage = CreateStorage(SessionTokenManagementMode.FullyAutomatic); var ex = Assert.Throws(() => storage.SetSessionTokens(new Dictionary())); - Assert.Equal("Can't use session tokens with FullyAutomatic", ex.Message); + Assert.Equal(CosmosStrings.EnableManualSessionTokenManagement, ex.Message); } [ConditionalFact] @@ -451,7 +452,7 @@ public virtual void FullyAutomatic_WhenCallingGetTrackedTokens_ThrowsInvalidOper { var storage = CreateStorage(SessionTokenManagementMode.FullyAutomatic); var ex = Assert.Throws(() => storage.GetTrackedTokens()); - Assert.Equal("Can't use session tokens with FullyAutomatic", ex.Message); + Assert.Equal(CosmosStrings.EnableManualSessionTokenManagement, ex.Message); } [ConditionalFact] @@ -460,7 +461,7 @@ public virtual void FullyAutomatic_WhenCallingAppendSessionTokens_ThrowsInvalidO var storage = CreateStorage(SessionTokenManagementMode.FullyAutomatic); var ex = Assert.Throws(() => storage.AppendSessionTokens(new Dictionary())); - Assert.Equal("Can't use session tokens with FullyAutomatic", ex.Message); + Assert.Equal(CosmosStrings.EnableManualSessionTokenManagement, ex.Message); } [ConditionalFact] @@ -469,7 +470,7 @@ public virtual void FullyAutomatic_WhenCallingSetDefaultContainerSessionToken_Th var storage = CreateStorage(SessionTokenManagementMode.FullyAutomatic); var ex = Assert.Throws(() => storage.SetDefaultContainerSessionToken(null)); - Assert.Equal("Can't use session tokens with FullyAutomatic", ex.Message); + Assert.Equal(CosmosStrings.EnableManualSessionTokenManagement, ex.Message); } [ConditionalFact] @@ -478,7 +479,7 @@ public virtual void FullyAutomatic_WhenCallingAppendDefaultContainerSessionToken var storage = CreateStorage(SessionTokenManagementMode.FullyAutomatic); var ex = Assert.Throws(() => storage.AppendDefaultContainerSessionToken("A")); - Assert.Equal("Can't use session tokens with FullyAutomatic", ex.Message); + Assert.Equal(CosmosStrings.EnableManualSessionTokenManagement, ex.Message); } [ConditionalFact] @@ -487,7 +488,7 @@ public virtual void FullyAutomatic_WhenCallingGetDefaultContainerTrackedToken_Th var storage = CreateStorage(SessionTokenManagementMode.FullyAutomatic); var ex = Assert.Throws(() => storage.GetDefaultContainerTrackedToken()); - Assert.Equal("Can't use session tokens with FullyAutomatic", ex.Message); + Assert.Equal(CosmosStrings.EnableManualSessionTokenManagement, ex.Message); } [ConditionalFact] @@ -519,7 +520,7 @@ public virtual void EnforcedManual_WhenGettingTokenBeforeSet_ThrowsInvalidOperat var storage = CreateStorage(SessionTokenManagementMode.EnforcedManual); var ex = Assert.Throws(() => storage.GetSessionToken(_defaultContainerName)); - Assert.Contains("No session token set for container while EnforcedManual", ex.Message); + Assert.Contains(CosmosStrings.MissingSessionTokenEnforceManual(_defaultContainerName), ex.Message); } [ConditionalFact] @@ -539,7 +540,7 @@ public virtual void EnforcedManual_WhenGettingTokenAfterClear_ThrowsInvalidOpera storage.Clear(); var ex = Assert.Throws(() => storage.GetSessionToken(_defaultContainerName)); - Assert.Contains("No session token set for container while EnforcedManual", ex.Message); + Assert.Contains(CosmosStrings.MissingSessionTokenEnforceManual(_defaultContainerName), ex.Message); } [ConditionalFact] @@ -575,7 +576,7 @@ public virtual void EnforcedManual_WhenOneContainerNotSet_ThrowsForThatContainer Assert.Equal("A", storage.GetSessionToken(_defaultContainerName)); var ex = Assert.Throws(() => storage.GetSessionToken("other")); - Assert.Contains("No session token set for container while EnforcedManual", ex.Message); + Assert.Contains(CosmosStrings.MissingSessionTokenEnforceManual("other"), ex.Message); } [ConditionalFact] @@ -584,7 +585,7 @@ public virtual void EnforcedManual_WhenTrackingToken_ThrowsWhenRetrieving() var storage = CreateStorage(SessionTokenManagementMode.EnforcedManual); storage.TrackSessionToken(_defaultContainerName, "A"); var ex = Assert.Throws(() => storage.GetSessionToken(_defaultContainerName)); - Assert.Contains("No session token set for container while EnforcedManual", ex.Message); + Assert.Contains(CosmosStrings.MissingSessionTokenEnforceManual(_defaultContainerName), ex.Message); } [ConditionalFact] From 20c2f1eefcf1dc1d2c1e638f6cc016b5e374bf4c Mon Sep 17 00:00:00 2001 From: JoasE <32096708+JoasE@users.noreply.github.com> Date: Mon, 10 Nov 2025 12:58:17 +0100 Subject: [PATCH 33/37] Add tests and fix cases --- .../CosmosServiceCollectionExtensions.cs | 2 +- .../Storage/Internal/CosmosDatabaseWrapper.cs | 2 +- .../Internal/ISessionTokenStorageFactory.cs | 2 +- .../Storage/Internal/SessionTokenStorage.cs | 30 +- .../Internal/SessionTokenStorageFactory.cs | 17 +- .../CosmosSessionTokensTest.cs | 130 +- .../TestUtilities}/NullSessionTokenStorage.cs | 13 +- .../Internal/SessionTokenStorageTest.cs | 1055 +++++++++++------ 8 files changed, 813 insertions(+), 438 deletions(-) rename {src/EFCore.Cosmos/Storage/Internal => test/EFCore.Cosmos.FunctionalTests/TestUtilities}/NullSessionTokenStorage.cs (67%) diff --git a/src/EFCore.Cosmos/Extensions/CosmosServiceCollectionExtensions.cs b/src/EFCore.Cosmos/Extensions/CosmosServiceCollectionExtensions.cs index f560d9b4ae6..09018b0f2a5 100644 --- a/src/EFCore.Cosmos/Extensions/CosmosServiceCollectionExtensions.cs +++ b/src/EFCore.Cosmos/Extensions/CosmosServiceCollectionExtensions.cs @@ -123,7 +123,7 @@ public static IServiceCollection AddEntityFrameworkCosmos(this IServiceCollectio .TryAddScoped() .TryAddScoped() .TryAddScoped() - .TryAddScoped()); + .TryAddSingleton()); builder.TryAddCoreServices(); diff --git a/src/EFCore.Cosmos/Storage/Internal/CosmosDatabaseWrapper.cs b/src/EFCore.Cosmos/Storage/Internal/CosmosDatabaseWrapper.cs index f9d742d9c78..491f70dbdae 100644 --- a/src/EFCore.Cosmos/Storage/Internal/CosmosDatabaseWrapper.cs +++ b/src/EFCore.Cosmos/Storage/Internal/CosmosDatabaseWrapper.cs @@ -44,7 +44,7 @@ public CosmosDatabaseWrapper( { _currentDbContext = currentDbContext; _cosmosClient = cosmosClient; - SessionTokenStorage = sessionTokenStorageFactory.Create(); + SessionTokenStorage = sessionTokenStorageFactory.Create(currentDbContext.Context); if (loggingOptions.IsSensitiveDataLoggingEnabled) { diff --git a/src/EFCore.Cosmos/Storage/Internal/ISessionTokenStorageFactory.cs b/src/EFCore.Cosmos/Storage/Internal/ISessionTokenStorageFactory.cs index 25311d88d48..91abc884a01 100644 --- a/src/EFCore.Cosmos/Storage/Internal/ISessionTokenStorageFactory.cs +++ b/src/EFCore.Cosmos/Storage/Internal/ISessionTokenStorageFactory.cs @@ -17,5 +17,5 @@ public interface ISessionTokenStorageFactory /// 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. /// - public ISessionTokenStorage Create(); + public ISessionTokenStorage Create(DbContext dbContext); } diff --git a/src/EFCore.Cosmos/Storage/Internal/SessionTokenStorage.cs b/src/EFCore.Cosmos/Storage/Internal/SessionTokenStorage.cs index 87c2513b43c..3c0a903c143 100644 --- a/src/EFCore.Cosmos/Storage/Internal/SessionTokenStorage.cs +++ b/src/EFCore.Cosmos/Storage/Internal/SessionTokenStorage.cs @@ -141,9 +141,17 @@ public virtual void SetDefaultContainerSessionToken(string? sessionToken) var sessionToken = _containerSessionTokens[containerName]; - if (!sessionToken.IsSet && _mode == SessionTokenManagementMode.EnforcedManual) + if (!sessionToken.IsSet) { - throw new InvalidOperationException(CosmosStrings.MissingSessionTokenEnforceManual(containerName)); + if (_mode == SessionTokenManagementMode.EnforcedManual) + { + throw new InvalidOperationException(CosmosStrings.MissingSessionTokenEnforceManual(containerName)); + } + + if (_mode == SessionTokenManagementMode.SemiAutomatic) + { + return null; + } } return sessionToken.ConvertToString(); @@ -215,13 +223,18 @@ public void Add(string token, bool isSet = false) { IsSet = IsSet || isSet; - if (token == null) + if (token == string.Empty && _tokens.Count == 0) { - return; + _string = ""; } foreach (var tokenPart in token.Split(',')) { + if (string.IsNullOrEmpty(tokenPart)) + { + continue; + } + if (_tokens.Add(tokenPart)) { _isChanged = true; @@ -234,7 +247,14 @@ public void Add(string token, bool isSet = false) if (_isChanged) { _isChanged = false; - _string = IsSet && _tokens.Count == 0 ? null : string.Join(",", _tokens); + if (_tokens.Count == 0) + { + _string = null; + } + else + { + _string = string.Join(",", _tokens); + } } return _string; diff --git a/src/EFCore.Cosmos/Storage/Internal/SessionTokenStorageFactory.cs b/src/EFCore.Cosmos/Storage/Internal/SessionTokenStorageFactory.cs index 4db1d15de7e..fd0aa7e5c24 100644 --- a/src/EFCore.Cosmos/Storage/Internal/SessionTokenStorageFactory.cs +++ b/src/EFCore.Cosmos/Storage/Internal/SessionTokenStorageFactory.cs @@ -15,8 +15,8 @@ namespace Microsoft.EntityFrameworkCore.Cosmos.Storage.Internal; /// public class SessionTokenStorageFactory : ISessionTokenStorageFactory { - private readonly string _defaultContainerName; - private readonly HashSet _containerNames; + private string? _defaultContainerName; + private HashSet? _containerNames; private readonly SessionTokenManagementMode _mode; /// @@ -25,10 +25,8 @@ public class SessionTokenStorageFactory : ISessionTokenStorageFactory /// 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. /// - public SessionTokenStorageFactory(ICurrentDbContext currentDbContext, ICosmosSingletonOptions options) + public SessionTokenStorageFactory(ICosmosSingletonOptions options) { - _defaultContainerName = (string)currentDbContext.Context.Model.GetAnnotation(CosmosAnnotationNames.ContainerName).Value!; - _containerNames = new HashSet([_defaultContainerName, ..GetContainerNames(currentDbContext.Context.Model)]); _mode = options.SessionTokenManagementMode; } @@ -38,10 +36,11 @@ public SessionTokenStorageFactory(ICurrentDbContext currentDbContext, ICosmosSin /// 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. /// - public ISessionTokenStorage Create() - => _mode == SessionTokenManagementMode.FullyAutomatic ? - new NullSessionTokenStorage() : - new SessionTokenStorage(_defaultContainerName, _containerNames, _mode); + public ISessionTokenStorage Create(DbContext dbContext) + => new SessionTokenStorage( + _defaultContainerName ??= (string)dbContext.Model.GetAnnotation(CosmosAnnotationNames.ContainerName).Value!, + _containerNames ??= new HashSet([_defaultContainerName, ..GetContainerNames(dbContext.Model)]), + _mode); private static IEnumerable GetContainerNames(IModel model) diff --git a/test/EFCore.Cosmos.FunctionalTests/CosmosSessionTokensTest.cs b/test/EFCore.Cosmos.FunctionalTests/CosmosSessionTokensTest.cs index e4a98ececba..6bf0f9a7dea 100644 --- a/test/EFCore.Cosmos.FunctionalTests/CosmosSessionTokensTest.cs +++ b/test/EFCore.Cosmos.FunctionalTests/CosmosSessionTokensTest.cs @@ -16,44 +16,54 @@ protected override ITestStoreFactory TestStoreFactory protected override string StoreName => nameof(CosmosSessionTokensTest); + private bool _mock = true; protected override IServiceCollection AddServices(IServiceCollection serviceCollection) - => base.AddServices(serviceCollection).Replace(ServiceDescriptor.Singleton()); + { + var services = base.AddServices(serviceCollection); + + return _mock + ? services.Replace(ServiceDescriptor.Singleton()) + : services; + } + + protected override TestStore CreateTestStore() => CosmosTestStore.Create(StoreName, (cfg) => cfg.SessionTokenManagementMode(Cosmos.Infrastructure.SessionTokenManagementMode.SemiAutomatic)); private static TestSessionTokenStorage _sessionTokenStorage = null!; - private class TestSessionTokenStorageFactory : ISessionTokenStorageFactory + [ConditionalFact] + public virtual async Task Can_use_session_tokens_no_mock() { - public ISessionTokenStorage Create() - => _sessionTokenStorage = new(); - } + _mock = false; + var contextFactory = await InitializeAsync(); + using var context = contextFactory.CreateContext(); - private class TestSessionTokenStorage : ISessionTokenStorage - { - public Dictionary SessionTokens { get; set; } = new() { { nameof(CosmosSessionTokenContext), null }, { OtherContainerName, null } }; + context.Customers.Add(new Customer { Id = "1", PartitionKey = "1" }); + context.OtherContainerCustomers.Add(new OtherContainerCustomer { Id = "1", PartitionKey = "1" }); - public List AppendDefaultContainerSessionTokenCalls { get; set; } = new(); - public List> AppendSessionTokensCalls { get; set; } = new(); - public List SetDefaultContainerSessionTokenCalls { get; set; } = new(); + await context.SaveChangesAsync(); - public List> SetSessionTokensCalls { get; set; } = new(); - public List<(string containerName, string sessionToken)> TrackSessionTokenCalls { get; set; } = new(); - public bool ClearCalled { get; set; } + var sessionTokens = context.Database.GetSessionTokens(); + Assert.NotNull(sessionTokens[nameof(CosmosSessionTokenContext)]); + Assert.NotNull(sessionTokens[OtherContainerName]); + // Only way we can test this is by setting a session token that will fail the request if used.. + // This will take a couple of seconds to fail + var newTokens = sessionTokens.ToDictionary(x => x.Key, x => x.Value!.Substring(0, x.Value.IndexOf('#') + 1) + int.MaxValue); + context.Database.UseSessionTokens(newTokens!); - public void AppendDefaultContainerSessionToken(string sessionToken) => AppendDefaultContainerSessionTokenCalls.Add(sessionToken); + var exes = new List() + { + await Assert.ThrowsAsync(() => context.Customers.ToListAsync()), + await Assert.ThrowsAsync(() => context.OtherContainerCustomers.ToListAsync()) + }; - public void AppendSessionTokens(IReadOnlyDictionary sessionTokens) => AppendSessionTokensCalls.Add(sessionTokens); - public void Clear() => ClearCalled = true; - public string? GetDefaultContainerTrackedToken() => SessionTokens.FirstOrDefault().Value; - public string? GetSessionToken(string containerName) => SessionTokens[containerName]; - public IReadOnlyDictionary GetTrackedTokens() => SessionTokens; - public void SetDefaultContainerSessionToken(string sessionToken) => SetDefaultContainerSessionTokenCalls.Add(sessionToken); - public void SetSessionTokens(IReadOnlyDictionary sessionTokens) => SetSessionTokensCalls.Add(sessionTokens); - public void TrackSessionToken(string containerName, string sessionToken) => TrackSessionTokenCalls.Add((containerName, sessionToken)); + foreach (var ex in exes) + { + Assert.Contains("The read session is not available for the input session token.", ex.ResponseBody); + } } - [ConditionalFact] public virtual async Task AppendSessionToken_uses_AppendDefaultContainerSessionToken() { @@ -133,32 +143,57 @@ await Assert.ThrowsAsync(() => context.OtherContainerCustomers. [ConditionalFact] public virtual async Task New_context_does_not_use_same_SessionTokenStorage() { + _mock = false; var contextFactory = await InitializeAsync(); using var context = contextFactory.CreateContext(); - - var oldSessionTokenStorage = _sessionTokenStorage; + context.Database.UseSessionToken("A"); using var newContext = contextFactory.CreateContext(); Assert.NotSame(context, newContext); - Assert.NotSame(oldSessionTokenStorage, _sessionTokenStorage); + Assert.Null(newContext.Database.GetSessionToken()); + Assert.Equal("A", context.Database.GetSessionToken()); + Assert.NotSame(((CosmosDatabaseWrapper)context.GetService()).SessionTokenStorage, ((CosmosDatabaseWrapper)newContext.GetService()).SessionTokenStorage); } [ConditionalFact] - public virtual async Task Pooled_context_clears_SessionTokenStorage() + public virtual async Task Pooled_context_uses_same_SessionTokenStorage() { + _mock = false; + var contextFactory = await InitializeAsync(); DbContext contextCopy; ISessionTokenStorage sessionTokenStorageCopy; using (var context = contextFactory.CreateContext()) { contextCopy = context; - sessionTokenStorageCopy = _sessionTokenStorage; - Assert.False(_sessionTokenStorage.ClearCalled); + context.Database.UseSessionToken("A"); + sessionTokenStorageCopy = ((CosmosDatabaseWrapper)context.GetService()).SessionTokenStorage; } using var newContext = contextFactory.CreateContext(); + + Assert.Same(newContext, contextCopy); + Assert.Same(sessionTokenStorageCopy, ((CosmosDatabaseWrapper)newContext.GetService()).SessionTokenStorage); + Assert.Null(newContext.Database.GetSessionToken()); + } + + [ConditionalFact] + public virtual async Task Pooled_context_clears_SessionTokenStorage() + { + var contextFactory = await InitializeAsync(); + DbContext contextCopy; + ISessionTokenStorage sessionTokenStorageCopy; + using (var context = contextFactory.CreateContext()) + { + contextCopy = context; + sessionTokenStorageCopy = ((CosmosDatabaseWrapper)context.GetService()).SessionTokenStorage; + _sessionTokenStorage.ClearCalled = false; + } + + using var newContext = contextFactory.CreateContext(); + Assert.Same(newContext, contextCopy); - Assert.Same(_sessionTokenStorage, sessionTokenStorageCopy); + Assert.Same(sessionTokenStorageCopy, ((CosmosDatabaseWrapper)newContext.GetService()).SessionTokenStorage); Assert.True(_sessionTokenStorage.ClearCalled); } @@ -644,6 +679,39 @@ public virtual async Task Delete_uses_session_token(AutoTransactionBehavior auto Assert.Contains("The session token provided 'invalidtoken' is invalid.", ((CosmosException)ex.InnerException!).ResponseBody); } + private class TestSessionTokenStorageFactory : ISessionTokenStorageFactory + { + public ISessionTokenStorage Create(DbContext _) + => _sessionTokenStorage = new(); + } + + private class TestSessionTokenStorage : ISessionTokenStorage + { + public Dictionary SessionTokens { get; set; } = new() { { nameof(CosmosSessionTokenContext), null }, { OtherContainerName, null } }; + + public List AppendDefaultContainerSessionTokenCalls { get; set; } = new(); + public List> AppendSessionTokensCalls { get; set; } = new(); + public List SetDefaultContainerSessionTokenCalls { get; set; } = new(); + + public List> SetSessionTokensCalls { get; set; } = new(); + public List<(string containerName, string sessionToken)> TrackSessionTokenCalls { get; set; } = new(); + public bool ClearCalled { get; set; } + + + + public void AppendDefaultContainerSessionToken(string sessionToken) => AppendDefaultContainerSessionTokenCalls.Add(sessionToken); + + public void AppendSessionTokens(IReadOnlyDictionary sessionTokens) => AppendSessionTokensCalls.Add(sessionTokens); + public void Clear() => ClearCalled = true; + public string? GetDefaultContainerTrackedToken() => SessionTokens.FirstOrDefault().Value; + public string? GetSessionToken(string containerName) => SessionTokens[containerName]; + public IReadOnlyDictionary GetTrackedTokens() => SessionTokens; + public void SetDefaultContainerSessionToken(string sessionToken) => SetDefaultContainerSessionTokenCalls.Add(sessionToken); + public void SetSessionTokens(IReadOnlyDictionary sessionTokens) => SetSessionTokensCalls.Add(sessionTokens); + public void TrackSessionToken(string containerName, string sessionToken) => TrackSessionTokenCalls.Add((containerName, sessionToken)); + } + + public class CosmosSessionTokenContext(DbContextOptions options) : PoolableDbContext(options) { public DbSet Customers { get; set; } = null!; diff --git a/src/EFCore.Cosmos/Storage/Internal/NullSessionTokenStorage.cs b/test/EFCore.Cosmos.FunctionalTests/TestUtilities/NullSessionTokenStorage.cs similarity index 67% rename from src/EFCore.Cosmos/Storage/Internal/NullSessionTokenStorage.cs rename to test/EFCore.Cosmos.FunctionalTests/TestUtilities/NullSessionTokenStorage.cs index 1fb63959e8c..9aec9cdbe9f 100644 --- a/src/EFCore.Cosmos/Storage/Internal/NullSessionTokenStorage.cs +++ b/test/EFCore.Cosmos.FunctionalTests/TestUtilities/NullSessionTokenStorage.cs @@ -2,14 +2,11 @@ // The .NET Foundation licenses this file to you under the MIT license. -namespace Microsoft.EntityFrameworkCore.Cosmos.Storage.Internal; - -/// -/// 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. -/// +using Microsoft.EntityFrameworkCore.Cosmos.Storage.Internal; + +namespace Microsoft.EntityFrameworkCore.TestUtilities; + +/// public class NullSessionTokenStorage : ISessionTokenStorage { /// diff --git a/test/EFCore.Cosmos.Tests/Storage/Internal/SessionTokenStorageTest.cs b/test/EFCore.Cosmos.Tests/Storage/Internal/SessionTokenStorageTest.cs index 85ebbf1362b..803c6511614 100644 --- a/test/EFCore.Cosmos.Tests/Storage/Internal/SessionTokenStorageTest.cs +++ b/test/EFCore.Cosmos.Tests/Storage/Internal/SessionTokenStorageTest.cs @@ -11,636 +11,798 @@ namespace Microsoft.EntityFrameworkCore.Cosmos.Storage.Internal; public class SessionTokenStorageTest { private readonly string _defaultContainerName = "default"; + private readonly string _otherContainerName = "other"; private readonly HashSet _containerNames = new(["default", "other"]); - + // ================================================================ + // FUNCTIONAL TESTS - SET AND RETRIEVE + // ================================================================ + [ConditionalTheory] - [InlineData(SessionTokenManagementMode.FullyAutomatic)] [InlineData(SessionTokenManagementMode.SemiAutomatic)] [InlineData(SessionTokenManagementMode.Manual)] [InlineData(SessionTokenManagementMode.EnforcedManual)] - public virtual void TrackSessionToken_WhenContainerNameIsNull_ThrowsArgumentNullException(SessionTokenManagementMode mode) + public virtual void SetSessionTokens_SetSingle_Default(SessionTokenManagementMode mode) { var storage = CreateStorage(mode); - Assert.Throws(() => storage.TrackSessionToken(null!, "A")); + + storage.SetSessionTokens(new Dictionary { { _defaultContainerName, "A" } }); + + AssertDefault(storage, "A"); + if (mode != SessionTokenManagementMode.EnforcedManual) + { + AssertOther(storage, null); + } } [ConditionalTheory] - [InlineData(SessionTokenManagementMode.FullyAutomatic)] [InlineData(SessionTokenManagementMode.SemiAutomatic)] [InlineData(SessionTokenManagementMode.Manual)] [InlineData(SessionTokenManagementMode.EnforcedManual)] - public virtual void TrackSessionToken_WhenContainerNameIsWhitespace_ThrowsArgumentNullException(SessionTokenManagementMode mode) + public virtual void SetSessionTokens_SetSingle_Other(SessionTokenManagementMode mode) { var storage = CreateStorage(mode); - Assert.Throws(() => storage.TrackSessionToken(" ", "A")); - Assert.Throws(() => storage.TrackSessionToken("", "A")); + + storage.SetSessionTokens(new Dictionary { { _otherContainerName, "A" } }); + + if (mode != SessionTokenManagementMode.EnforcedManual) + { + AssertDefault(storage, null); + } + AssertOther(storage, "A"); } [ConditionalTheory] - [InlineData(SessionTokenManagementMode.FullyAutomatic)] [InlineData(SessionTokenManagementMode.SemiAutomatic)] [InlineData(SessionTokenManagementMode.Manual)] [InlineData(SessionTokenManagementMode.EnforcedManual)] - public virtual void TrackSessionToken_WhenTokenIsNull_ThrowsArgumentNullException(SessionTokenManagementMode mode) + public virtual void SetSessionTokens_Multiple(SessionTokenManagementMode mode) { var storage = CreateStorage(mode); - Assert.Throws(() => storage.TrackSessionToken(_defaultContainerName, null!)); + + storage.SetSessionTokens(new Dictionary + { + { _defaultContainerName, "Token1" }, + { _otherContainerName, "Token2" } + }); + + AssertDefault(storage, "Token1"); + AssertOther(storage, "Token2"); } [ConditionalTheory] - [InlineData(SessionTokenManagementMode.FullyAutomatic)] [InlineData(SessionTokenManagementMode.SemiAutomatic)] [InlineData(SessionTokenManagementMode.Manual)] [InlineData(SessionTokenManagementMode.EnforcedManual)] - public virtual void TrackSessionToken_WhenTokenIsWhitespace_ThrowsArgumentNullException(SessionTokenManagementMode mode) + public virtual void SetSessionTokens_OverwritesSet(SessionTokenManagementMode mode) { var storage = CreateStorage(mode); - Assert.Throws(() => storage.TrackSessionToken(_defaultContainerName, " ")); - Assert.Throws(() => storage.TrackSessionToken(_defaultContainerName, "")); + + storage.SetSessionTokens(new Dictionary + { + { _defaultContainerName, "A" }, + { _otherContainerName, "B" } + }); + storage.SetSessionTokens(new Dictionary + { + { _defaultContainerName, "" }, + { _otherContainerName, "" } + }); + + AssertDefault(storage, ""); + AssertOther(storage, ""); } [ConditionalTheory] [InlineData(SessionTokenManagementMode.SemiAutomatic)] [InlineData(SessionTokenManagementMode.Manual)] [InlineData(SessionTokenManagementMode.EnforcedManual)] - public virtual void SetSessionTokens_WhenContainerNameIsUnknown_ThrowsInvalidOperationException(SessionTokenManagementMode mode) + public virtual void SetSessionTokens_OverwritesTracked(SessionTokenManagementMode mode) { var storage = CreateStorage(mode); - var ex = Assert.Throws(() => - storage.SetSessionTokens(new Dictionary { { "bad", "A" } })); - Assert.Equal(CosmosStrings.ContainerNameDoesNotExist("bad"), ex.Message); + storage.TrackSessionToken(_defaultContainerName, "Token1"); + storage.TrackSessionToken(_otherContainerName, "Token2"); + storage.SetSessionTokens(new Dictionary + { + { _defaultContainerName, "A" }, + { _otherContainerName, "B" } + }); + + AssertDefault(storage, "A"); + AssertOther(storage, "B"); } [ConditionalTheory] - [InlineData(SessionTokenManagementMode.FullyAutomatic)] [InlineData(SessionTokenManagementMode.SemiAutomatic)] [InlineData(SessionTokenManagementMode.Manual)] [InlineData(SessionTokenManagementMode.EnforcedManual)] - public virtual void AppendDefaultContainerSessionToken_WhenTokenIsNull_ThrowsArgumentNullException(SessionTokenManagementMode mode) + public virtual void SetSessionTokens_SingleContainer_OverwritesOnlySingleContainer(SessionTokenManagementMode mode) { var storage = CreateStorage(mode); - Assert.Throws(() => storage.AppendDefaultContainerSessionToken(null!)); + storage.SetSessionTokens(new Dictionary + { + { _defaultContainerName, "A" }, + { _otherContainerName, "B" } + }); + storage.SetSessionTokens(new Dictionary { { _defaultContainerName, "C" } }); + + AssertDefault(storage, "C"); + AssertOther(storage, "B"); } [ConditionalTheory] - [InlineData(SessionTokenManagementMode.FullyAutomatic)] [InlineData(SessionTokenManagementMode.SemiAutomatic)] [InlineData(SessionTokenManagementMode.Manual)] [InlineData(SessionTokenManagementMode.EnforcedManual)] - public virtual void AppendDefaultContainerSessionToken_WhenTokenIsWhitespace_ThrowsArgumentException(SessionTokenManagementMode mode) + public virtual void SetSessionTokens_Null_SetsNull(SessionTokenManagementMode mode) { var storage = CreateStorage(mode); - Assert.Throws(() => storage.AppendDefaultContainerSessionToken(" ")); - Assert.Throws(() => storage.AppendDefaultContainerSessionToken("")); + storage.SetSessionTokens(new Dictionary { { _defaultContainerName, "A" }, { _otherContainerName, "B" } }); + storage.SetSessionTokens(new Dictionary { { _defaultContainerName, null } }); + + AssertDefault(storage, null); + AssertOther(storage, "B"); } [ConditionalTheory] [InlineData(SessionTokenManagementMode.SemiAutomatic)] [InlineData(SessionTokenManagementMode.Manual)] [InlineData(SessionTokenManagementMode.EnforcedManual)] - public virtual void AppendSessionTokens_WhenContainerNameIsUnknown_ThrowsInvalidOperationException(SessionTokenManagementMode mode) + public virtual void SetDefaultContainerSessionToken_SetsToken(SessionTokenManagementMode mode) { var storage = CreateStorage(mode); - var ex = Assert.Throws(() => - storage.AppendSessionTokens(new Dictionary { { "bad", "A" } })); - Assert.Equal(CosmosStrings.ContainerNameDoesNotExist("bad"), ex.Message); + storage.SetDefaultContainerSessionToken("A"); + + AssertDefault(storage, "A"); + + if (mode != SessionTokenManagementMode.EnforcedManual) + { + AssertOther(storage, null); + } } [ConditionalTheory] [InlineData(SessionTokenManagementMode.SemiAutomatic)] [InlineData(SessionTokenManagementMode.Manual)] [InlineData(SessionTokenManagementMode.EnforcedManual)] - public virtual void GetSessionToken_WhenContainerNameIsNull_ThrowsArgumentNullException(SessionTokenManagementMode mode) + public virtual void SetDefaultContainerSessionToken_OverwritesSet(SessionTokenManagementMode mode) { var storage = CreateStorage(mode); - Assert.Throws(() => storage.GetSessionToken(null!)); + storage.SetDefaultContainerSessionToken("A"); + storage.SetDefaultContainerSessionToken("B"); + + AssertDefault(storage, "B"); + + if (mode != SessionTokenManagementMode.EnforcedManual) + { + AssertOther(storage, null); + } } [ConditionalTheory] [InlineData(SessionTokenManagementMode.SemiAutomatic)] [InlineData(SessionTokenManagementMode.Manual)] [InlineData(SessionTokenManagementMode.EnforcedManual)] - public virtual void GetSessionToken_WhenContainerNameIsWhitespace_ThrowsArgumentNullException(SessionTokenManagementMode mode) + public virtual void SetDefaultContainerSessionToken_OverwritesTracked(SessionTokenManagementMode mode) { var storage = CreateStorage(mode); - Assert.Throws(() => storage.GetSessionToken(" ")); - Assert.Throws(() => storage.GetSessionToken("")); - } + storage.TrackSessionToken(_defaultContainerName, "A"); + storage.SetDefaultContainerSessionToken("B"); - // ================================================================ - // FUNCTIONAL TESTS - SET AND RETRIEVE - // ================================================================ + AssertDefault(storage, "B"); + if (mode != SessionTokenManagementMode.EnforcedManual) + { + AssertOther(storage, null); + } + } [ConditionalTheory] [InlineData(SessionTokenManagementMode.SemiAutomatic)] [InlineData(SessionTokenManagementMode.Manual)] [InlineData(SessionTokenManagementMode.EnforcedManual)] - public virtual void SetSessionTokens_WhenSettingSingleToken_CanRetrieveFromAllMethods(SessionTokenManagementMode mode) + public virtual void AppendDefaultContainerSessionToken_NoPreviousToken_SetsToken(SessionTokenManagementMode mode) { var storage = CreateStorage(mode); + storage.AppendDefaultContainerSessionToken("A"); - storage.SetSessionTokens(new Dictionary { { _defaultContainerName, "A" } }); + AssertDefault(storage, "A"); - var all = storage.GetTrackedTokens(); - Assert.Equal("A", all[_defaultContainerName]); - Assert.Null(all["other"]); - Assert.Equal("A", storage.GetSessionToken(_defaultContainerName)); - Assert.Equal("A", storage.GetDefaultContainerTrackedToken()); + if (mode != SessionTokenManagementMode.EnforcedManual) + { + AssertOther(storage, null); + } } [ConditionalTheory] [InlineData(SessionTokenManagementMode.SemiAutomatic)] [InlineData(SessionTokenManagementMode.Manual)] [InlineData(SessionTokenManagementMode.EnforcedManual)] - public virtual void SetSessionTokens_WhenSettingMultipleContainers_AllContainersAreSetCorrectly(SessionTokenManagementMode mode) + public virtual void AppendDefaultContainerSessionToken_PreviousSetToken_AppendsToken(SessionTokenManagementMode mode) { var storage = CreateStorage(mode); + storage.SetDefaultContainerSessionToken("A"); + storage.AppendDefaultContainerSessionToken("B"); - storage.SetSessionTokens(new Dictionary + AssertDefault(storage, "A,B"); + + if (mode != SessionTokenManagementMode.EnforcedManual) { - { _defaultContainerName, "Token1" }, - { "other", "Token2" } - }); + AssertOther(storage, null); + } + } + + [ConditionalTheory] + [InlineData(SessionTokenManagementMode.SemiAutomatic)] + [InlineData(SessionTokenManagementMode.Manual)] + [InlineData(SessionTokenManagementMode.EnforcedManual)] + public virtual void AppendDefaultContainerSessionToken_PreviousSetToken_Duplicate_DoesNotAppendToken(SessionTokenManagementMode mode) + { + var storage = CreateStorage(mode); + storage.SetDefaultContainerSessionToken("A"); + storage.AppendDefaultContainerSessionToken("A"); - var all = storage.GetTrackedTokens(); - Assert.Equal("Token1", all[_defaultContainerName]); - Assert.Equal("Token2", all["other"]); - Assert.Equal("Token1", storage.GetSessionToken(_defaultContainerName)); - Assert.Equal("Token2", storage.GetSessionToken("other")); + AssertDefault(storage, "A"); + if (mode != SessionTokenManagementMode.EnforcedManual) + { + AssertOther(storage, null); + } } [ConditionalTheory] [InlineData(SessionTokenManagementMode.SemiAutomatic)] [InlineData(SessionTokenManagementMode.Manual)] [InlineData(SessionTokenManagementMode.EnforcedManual)] - public virtual void SetSessionTokens_WhenSettingNullValue_ContainerIsCleared(SessionTokenManagementMode mode) + public virtual void AppendDefaultContainerSessionToken_PreviousTrackedToken_AppendsToken(SessionTokenManagementMode mode) { var storage = CreateStorage(mode); - storage.SetSessionTokens(new Dictionary { { _defaultContainerName, "A" } }); - storage.SetSessionTokens(new Dictionary { { _defaultContainerName, null } }); + storage.TrackSessionToken(_defaultContainerName, "A"); + storage.AppendDefaultContainerSessionToken("B"); - var token = storage.GetDefaultContainerTrackedToken(); - Assert.Null(token); + AssertDefault(storage, "A,B"); + if (mode != SessionTokenManagementMode.EnforcedManual) + { + AssertOther(storage, null); + } } [ConditionalTheory] [InlineData(SessionTokenManagementMode.SemiAutomatic)] [InlineData(SessionTokenManagementMode.Manual)] [InlineData(SessionTokenManagementMode.EnforcedManual)] - public virtual void SetSessionTokens_WhenSettingEmptyDictionary_NoExceptionThrown(SessionTokenManagementMode mode) + public virtual void AppendDefaultContainerSessionToken_PreviousTrackedToken_Duplicate_DoesNotAppendToken(SessionTokenManagementMode mode) { var storage = CreateStorage(mode); - storage.SetSessionTokens(new Dictionary()); + storage.TrackSessionToken(_defaultContainerName, "A"); + storage.AppendDefaultContainerSessionToken("A"); - var all = storage.GetTrackedTokens(); - Assert.Null(all[_defaultContainerName]); - Assert.Null(all["other"]); + AssertDefault(storage, "A"); + if (mode != SessionTokenManagementMode.EnforcedManual) + { + AssertOther(storage, null); + } } [ConditionalTheory] [InlineData(SessionTokenManagementMode.SemiAutomatic)] [InlineData(SessionTokenManagementMode.Manual)] [InlineData(SessionTokenManagementMode.EnforcedManual)] - public virtual void SetSessionTokens_WhenPartiallyUpdatingContainers_OnlySpecifiedContainersAreUpdated(SessionTokenManagementMode mode) + public virtual void AppendSessionTokens_MultipleContainers_NoPreviousTokens_SetsTokens(SessionTokenManagementMode mode) { var storage = CreateStorage(mode); - storage.SetSessionTokens(new Dictionary + + storage.AppendSessionTokens(new Dictionary { { _defaultContainerName, "A" }, - { "other", "B" } + { _otherContainerName, "B" } }); - storage.SetSessionTokens(new Dictionary { { _defaultContainerName, "C" } }); - var all = storage.GetTrackedTokens(); - Assert.Equal("C", all[_defaultContainerName]); - Assert.Equal("B", all["other"]); + AssertDefault(storage, "A"); + AssertOther(storage, "B"); } - // ================================================================ - // FUNCTIONAL TESTS - APPEND OPERATIONS - // ================================================================ - [ConditionalTheory] [InlineData(SessionTokenManagementMode.SemiAutomatic)] [InlineData(SessionTokenManagementMode.Manual)] [InlineData(SessionTokenManagementMode.EnforcedManual)] - public virtual void AppendSessionTokens_WhenAppendingOverlappingTokens_MergesUniquely(SessionTokenManagementMode mode) + public virtual void AppendSessionTokens_SingleContainer_NoPreviousTokens_SetsTokens(SessionTokenManagementMode mode) { var storage = CreateStorage(mode); - storage.AppendSessionTokens(new Dictionary { { _defaultContainerName, "A,B" } }); - storage.AppendSessionTokens(new Dictionary { { _defaultContainerName, "B,C" } }); + storage.AppendSessionTokens(new Dictionary + { + { _otherContainerName, "B" } + }); - var token = storage.GetDefaultContainerTrackedToken(); - Assert.Equal("A,B,C", token); + if (mode != SessionTokenManagementMode.EnforcedManual) + { + AssertDefault(storage, null); + } + AssertOther(storage, "B"); } [ConditionalTheory] [InlineData(SessionTokenManagementMode.SemiAutomatic)] [InlineData(SessionTokenManagementMode.Manual)] [InlineData(SessionTokenManagementMode.EnforcedManual)] - public virtual void AppendSessionTokens_WhenAppendingToMultipleContainers_AllContainersAreUpdated(SessionTokenManagementMode mode) + public virtual void AppendSessionTokens_PreviousSetTokens_AppendsTokens(SessionTokenManagementMode mode) { var storage = CreateStorage(mode); - - storage.AppendSessionTokens(new Dictionary + storage.SetSessionTokens(new Dictionary { { _defaultContainerName, "A" }, - { "other", "B" } + { _otherContainerName, "B" } }); storage.AppendSessionTokens(new Dictionary { { _defaultContainerName, "C" }, - { "other", "D" } + { _otherContainerName, "D" } }); - var all = storage.GetTrackedTokens(); - Assert.Equal("A,C", all[_defaultContainerName]); - Assert.Equal("B,D", all["other"]); + AssertDefault(storage, "A,C"); + AssertOther(storage, "B,D"); } [ConditionalTheory] [InlineData(SessionTokenManagementMode.SemiAutomatic)] [InlineData(SessionTokenManagementMode.Manual)] [InlineData(SessionTokenManagementMode.EnforcedManual)] - public virtual void AppendSessionTokens_WhenAppendingEmptyDictionary_NoExceptionThrown(SessionTokenManagementMode mode) + public virtual void AppendSessionTokens_PreviousSetToken_AppendsAndSetsTokens(SessionTokenManagementMode mode) { var storage = CreateStorage(mode); - storage.AppendSessionTokens(new Dictionary()); + storage.SetSessionTokens(new Dictionary + { + { _otherContainerName, "B" } + }); + storage.AppendSessionTokens(new Dictionary + { + { _defaultContainerName, "C" }, + { _otherContainerName, "D" } + }); - var all = storage.GetTrackedTokens(); - Assert.Null(all[_defaultContainerName]); - Assert.Null(all["other"]); + AssertDefault(storage, "C"); + AssertOther(storage, "B,D"); } [ConditionalTheory] [InlineData(SessionTokenManagementMode.SemiAutomatic)] [InlineData(SessionTokenManagementMode.Manual)] [InlineData(SessionTokenManagementMode.EnforcedManual)] - public virtual void AppendDefaultContainerSessionToken_WhenAppendingMultipleTokens_AccumulatesUniquely(SessionTokenManagementMode mode) + public virtual void AppendSessionTokens_PreviousTrackedToken_AppendsAndSetsTokens(SessionTokenManagementMode mode) { var storage = CreateStorage(mode); + storage.TrackSessionToken(_otherContainerName, "B"); + storage.AppendSessionTokens(new Dictionary + { + { _defaultContainerName, "C" }, + { _otherContainerName, "D" } + }); - storage.AppendDefaultContainerSessionToken("A"); - storage.AppendDefaultContainerSessionToken("B"); - storage.AppendDefaultContainerSessionToken("B"); - - var token = storage.GetDefaultContainerTrackedToken(); - Assert.Equal("A,B", token); + AssertDefault(storage, "C"); + AssertOther(storage, "B,D"); } - // ================================================================ - // FUNCTIONAL TESTS - SET OPERATIONS - // ================================================================ - [ConditionalTheory] [InlineData(SessionTokenManagementMode.SemiAutomatic)] [InlineData(SessionTokenManagementMode.Manual)] [InlineData(SessionTokenManagementMode.EnforcedManual)] - public virtual void SetDefaultContainerSessionToken_WhenReplacingExistingTokens_ReplacesAllPreviousTokens(SessionTokenManagementMode mode) + public virtual void SetDefaultContainerSessionToken_RemovesDuplicates(SessionTokenManagementMode mode) { var storage = CreateStorage(mode); - storage.AppendDefaultContainerSessionToken("A"); - storage.AppendDefaultContainerSessionToken("B"); - storage.SetDefaultContainerSessionToken("XYZ"); - - var token = storage.GetDefaultContainerTrackedToken(); - Assert.Equal("XYZ", token); + storage.SetDefaultContainerSessionToken("A,A,B"); + AssertDefault(storage, "A,B"); } [ConditionalTheory] [InlineData(SessionTokenManagementMode.SemiAutomatic)] [InlineData(SessionTokenManagementMode.Manual)] [InlineData(SessionTokenManagementMode.EnforcedManual)] - public virtual void SetDefaultContainerSessionToken_WhenSettingNull_ClearsContainerToken(SessionTokenManagementMode mode) + public virtual void AppendDefaultContainerSessionToken_RemovesDuplicates(SessionTokenManagementMode mode) { var storage = CreateStorage(mode); - storage.SetDefaultContainerSessionToken("A"); - storage.SetDefaultContainerSessionToken(null); - - var token = storage.GetDefaultContainerTrackedToken(); - Assert.Null(token); + storage.AppendDefaultContainerSessionToken("A,A,B"); + AssertDefault(storage, "A,B"); } [ConditionalTheory] [InlineData(SessionTokenManagementMode.SemiAutomatic)] [InlineData(SessionTokenManagementMode.Manual)] [InlineData(SessionTokenManagementMode.EnforcedManual)] - public virtual void SetDefaultContainerSessionToken_WhenSettingEmptyString_StoresEmptyString(SessionTokenManagementMode mode) + public virtual void SetSessionTokens_RemovesDuplicates(SessionTokenManagementMode mode) { var storage = CreateStorage(mode); - storage.SetDefaultContainerSessionToken(""); - var token = storage.GetDefaultContainerTrackedToken(); - Assert.Equal("", token); + storage.SetSessionTokens(new Dictionary { { _defaultContainerName, "A,B,A" }, { _otherContainerName, "B,C,B" } }); + AssertDefault(storage, "A,B"); + AssertOther(storage, "B,C"); } - // ================================================================ - // FUNCTIONAL TESTS - TRACK OPERATIONS - // ================================================================ - [ConditionalTheory] [InlineData(SessionTokenManagementMode.SemiAutomatic)] [InlineData(SessionTokenManagementMode.Manual)] - public virtual void TrackSessionToken_WhenTrackingMultipleTokens_AppendsUniquely(SessionTokenManagementMode mode) + [InlineData(SessionTokenManagementMode.EnforcedManual)] + public virtual void AppendSessionTokens_RemovesDuplicates(SessionTokenManagementMode mode) { var storage = CreateStorage(mode); - - storage.TrackSessionToken(_defaultContainerName, "A"); - storage.TrackSessionToken(_defaultContainerName, "B"); - storage.TrackSessionToken(_defaultContainerName, "A"); - - var token = storage.GetSessionToken(_defaultContainerName); - Assert.Equal("A,B", token); + storage.AppendSessionTokens(new Dictionary { { _defaultContainerName, "A,B,A" }, { _otherContainerName, "B,C,B" } }); + AssertDefault(storage, "A,B"); + AssertOther(storage, "B,C"); } [ConditionalTheory] [InlineData(SessionTokenManagementMode.SemiAutomatic)] [InlineData(SessionTokenManagementMode.Manual)] - public virtual void TrackSessionToken_WhenTrackingToDifferentContainers_ContainersAreIndependent(SessionTokenManagementMode mode) + [InlineData(SessionTokenManagementMode.EnforcedManual)] + public virtual void AppendSessionTokens_PreviouslySetTokens_RemovesDuplicates(SessionTokenManagementMode mode) { var storage = CreateStorage(mode); - - storage.TrackSessionToken(_defaultContainerName, "A"); - storage.TrackSessionToken("other", "B"); - storage.TrackSessionToken(_defaultContainerName, "C"); - - Assert.Equal("A,C", storage.GetSessionToken(_defaultContainerName)); - Assert.Equal("B", storage.GetSessionToken("other")); + storage.SetSessionTokens(new Dictionary { { _defaultContainerName, "A,C,E" }, { _otherContainerName, "J,K,L" } }); + storage.AppendSessionTokens(new Dictionary { { _defaultContainerName, "A,B,B" }, { _otherContainerName, "K,A,A" } }); + AssertDefault(storage, "A,C,E,B"); + AssertOther(storage, "J,K,L,A"); } [ConditionalTheory] [InlineData(SessionTokenManagementMode.SemiAutomatic)] [InlineData(SessionTokenManagementMode.Manual)] - public virtual void TrackSessionToken_WhenTrackingCommaSeparatedToken_ParsesAndMergesCorrectly(SessionTokenManagementMode mode) + [InlineData(SessionTokenManagementMode.EnforcedManual)] + public virtual void AppendSessionTokens_EmptyStrings_DoesNotAppend(SessionTokenManagementMode mode) { var storage = CreateStorage(mode); - - storage.TrackSessionToken(_defaultContainerName, "A,B"); - storage.TrackSessionToken(_defaultContainerName, "C"); - storage.TrackSessionToken(_defaultContainerName, "B,D"); - - var token = storage.GetSessionToken(_defaultContainerName); - Assert.Equal("A,B,C,D", token); + storage.SetSessionTokens(new Dictionary { { _defaultContainerName, "A,C" }, { _otherContainerName, "J,K" } }); + storage.AppendSessionTokens(new Dictionary { { _defaultContainerName, "" }, { _otherContainerName, "" } }); + AssertDefault(storage, "A,C"); + AssertOther(storage, "J,K"); } - // ================================================================ - // FUNCTIONAL TESTS - CLEAR OPERATIONS - // ================================================================ - [ConditionalTheory] [InlineData(SessionTokenManagementMode.SemiAutomatic)] [InlineData(SessionTokenManagementMode.Manual)] [InlineData(SessionTokenManagementMode.EnforcedManual)] - public virtual void Clear_WhenClearingAllTokens_ResetsAllContainersToNull(SessionTokenManagementMode mode) + public virtual void AppendSessionTokens_NoPreviousTokens_EmptyStrings_Sets(SessionTokenManagementMode mode) { var storage = CreateStorage(mode); - - storage.AppendSessionTokens(new Dictionary { { _defaultContainerName, "A" }, { "other", "B" } }); - storage.Clear(); - - var all = storage.GetTrackedTokens(); - Assert.Null(all[_defaultContainerName]); - Assert.Null(all["other"]); + storage.AppendSessionTokens(new Dictionary { { _defaultContainerName, "" }, { _otherContainerName, "" } }); + AssertDefault(storage, ""); + AssertOther(storage, ""); } + // ================================================================ + // TRACK + // ================================================================ + [ConditionalTheory] [InlineData(SessionTokenManagementMode.SemiAutomatic)] [InlineData(SessionTokenManagementMode.Manual)] [InlineData(SessionTokenManagementMode.EnforcedManual)] - public virtual void Clear_WhenClearing_ContainersStillExistInTrackedTokens(SessionTokenManagementMode mode) + public virtual void TrackSessionToken_SetsToken(SessionTokenManagementMode mode) { var storage = CreateStorage(mode); - storage.AppendSessionTokens(new Dictionary { { _defaultContainerName, "A" } }); - storage.Clear(); - var tracked = storage.GetTrackedTokens(); - Assert.Contains(_defaultContainerName, tracked.Keys); - Assert.Contains("other", tracked.Keys); + storage.TrackSessionToken(_defaultContainerName, "A"); + storage.TrackSessionToken(_otherContainerName, "A"); + + AssertDefaultTracked(storage, "A"); + AssertOtherTracked(storage, "A"); + + if (mode == SessionTokenManagementMode.Manual) + { + AssertDefaultUsed(storage, "A"); + AssertOtherUsed(storage, "A"); + } + else if (mode != SessionTokenManagementMode.EnforcedManual) + { + + AssertDefaultUsed(storage, null); + AssertOtherUsed(storage, null); + } } [ConditionalTheory] [InlineData(SessionTokenManagementMode.SemiAutomatic)] [InlineData(SessionTokenManagementMode.Manual)] [InlineData(SessionTokenManagementMode.EnforcedManual)] - public virtual void Clear_WhenClearing_CanSetNewTokensAfterClear(SessionTokenManagementMode mode) + public virtual void TrackSessionToken_Appends(SessionTokenManagementMode mode) { var storage = CreateStorage(mode); - storage.AppendDefaultContainerSessionToken("A"); - storage.Clear(); - storage.AppendDefaultContainerSessionToken("B"); - var token = storage.GetDefaultContainerTrackedToken(); - Assert.Equal("B", token); + storage.TrackSessionToken(_defaultContainerName, "A"); + storage.TrackSessionToken(_defaultContainerName, "B"); + storage.TrackSessionToken(_defaultContainerName, "A"); + + storage.TrackSessionToken(_otherContainerName, "A"); + storage.TrackSessionToken(_otherContainerName, "C"); + storage.TrackSessionToken(_otherContainerName, "A"); + + AssertDefaultTracked(storage, "A,B"); + AssertOtherTracked(storage, "A,C"); + + if (mode == SessionTokenManagementMode.Manual) + { + AssertDefaultUsed(storage, "A,B"); + AssertOtherUsed(storage, "A,C"); + } + else if(mode != SessionTokenManagementMode.EnforcedManual) + { + AssertDefaultUsed(storage, null); + AssertOtherUsed(storage, null); + } } // ================================================================ - // FULLY AUTOMATIC MODE TESTS + // ENFORCED MANUAL MODE TESTS // ================================================================ [ConditionalFact] - public virtual void FullyAutomatic_WhenCallingSetSessionTokens_ThrowsInvalidOperationException() + public virtual void EnforcedManual_WhenGettingTokenBeforeSet_ThrowsInvalidOperationException() { - var storage = CreateStorage(SessionTokenManagementMode.FullyAutomatic); + var storage = CreateStorage(SessionTokenManagementMode.EnforcedManual); var ex = Assert.Throws(() => - storage.SetSessionTokens(new Dictionary())); - Assert.Equal(CosmosStrings.EnableManualSessionTokenManagement, ex.Message); + storage.GetSessionToken(_defaultContainerName)); + Assert.Contains(CosmosStrings.MissingSessionTokenEnforceManual(_defaultContainerName), ex.Message); } [ConditionalFact] - public virtual void FullyAutomatic_WhenCallingGetTrackedTokens_ThrowsInvalidOperationException() + public virtual void EnforcedManual_WhenGettingTokenAfterClear_ThrowsInvalidOperationException() { - var storage = CreateStorage(SessionTokenManagementMode.FullyAutomatic); - var ex = Assert.Throws(() => storage.GetTrackedTokens()); - Assert.Equal(CosmosStrings.EnableManualSessionTokenManagement, ex.Message); + var storage = CreateStorage(SessionTokenManagementMode.EnforcedManual); + storage.SetDefaultContainerSessionToken("A"); + storage.Clear(); + var ex = Assert.Throws(() => + storage.GetSessionToken(_defaultContainerName)); + Assert.Contains(CosmosStrings.MissingSessionTokenEnforceManual(_defaultContainerName), ex.Message); } [ConditionalFact] - public virtual void FullyAutomatic_WhenCallingAppendSessionTokens_ThrowsInvalidOperationException() + public virtual void EnforcedManual_SetDefaultContainerSessionToken_SetsAndUses() { - var storage = CreateStorage(SessionTokenManagementMode.FullyAutomatic); + var storage = CreateStorage(SessionTokenManagementMode.EnforcedManual); + storage.SetDefaultContainerSessionToken("A"); + AssertDefault(storage, "A"); var ex = Assert.Throws(() => - storage.AppendSessionTokens(new Dictionary())); - Assert.Equal(CosmosStrings.EnableManualSessionTokenManagement, ex.Message); + storage.GetSessionToken(_otherContainerName)); + Assert.Contains(CosmosStrings.MissingSessionTokenEnforceManual(_otherContainerName), ex.Message); } [ConditionalFact] - public virtual void FullyAutomatic_WhenCallingSetDefaultContainerSessionToken_ThrowsInvalidOperationException() + public virtual void EnforcedManual_SetSessionTokens_SetsAndUses() { - var storage = CreateStorage(SessionTokenManagementMode.FullyAutomatic); - var ex = Assert.Throws(() => - storage.SetDefaultContainerSessionToken(null)); - Assert.Equal(CosmosStrings.EnableManualSessionTokenManagement, ex.Message); + var storage = CreateStorage(SessionTokenManagementMode.EnforcedManual); + storage.SetSessionTokens(new Dictionary + { + { _defaultContainerName, "A" }, + { _otherContainerName, "B" } + }); + + AssertDefault(storage, "A"); + AssertOther(storage, "B"); } [ConditionalFact] - public virtual void FullyAutomatic_WhenCallingAppendDefaultContainerSessionToken_ThrowsInvalidOperationException() + public virtual void EnforcedManual_WhenOneContainerNotSet_ThrowsForThatContainerOnly() { - var storage = CreateStorage(SessionTokenManagementMode.FullyAutomatic); - var ex = Assert.Throws(() => - storage.AppendDefaultContainerSessionToken("A")); - Assert.Equal(CosmosStrings.EnableManualSessionTokenManagement, ex.Message); + var storage = CreateStorage(SessionTokenManagementMode.EnforcedManual); + storage.SetSessionTokens(new Dictionary { { _defaultContainerName, "A" } }); + + Assert.Equal("A", storage.GetSessionToken(_defaultContainerName)); + var ex = Assert.Throws(() => storage.GetSessionToken(_otherContainerName)); + Assert.Contains(CosmosStrings.MissingSessionTokenEnforceManual(_otherContainerName), ex.Message); } + // ================================================================ + // SEMI-AUTOMATIC MODE SPECIFIC TESTS + // ================================================================ + [ConditionalFact] - public virtual void FullyAutomatic_WhenCallingGetDefaultContainerTrackedToken_ThrowsInvalidOperationException() + public virtual void SemiAutomatic_WhenTrackingToken_SetsButDoesntUseToken() { - var storage = CreateStorage(SessionTokenManagementMode.FullyAutomatic); - var ex = Assert.Throws(() => - storage.GetDefaultContainerTrackedToken()); - Assert.Equal(CosmosStrings.EnableManualSessionTokenManagement, ex.Message); - } - - [ConditionalFact] - public virtual void FullyAutomatic_WhenTrackingToken_AlwaysReturnsNull() - { - var storage = CreateStorage(SessionTokenManagementMode.FullyAutomatic); + var storage = CreateStorage(SessionTokenManagementMode.SemiAutomatic); storage.TrackSessionToken(_defaultContainerName, "A"); - Assert.Null(storage.GetSessionToken(_defaultContainerName)); + AssertDefaultTracked(storage, "A"); + AssertDefaultUsed(storage, null); } [ConditionalFact] - public virtual void FullyAutomatic_WhenTrackingMultipleTokens_AlwaysReturnsNull() + public virtual void SemiAutomatic_WhenSetToken_SetsAndUses() { - var storage = CreateStorage(SessionTokenManagementMode.FullyAutomatic); + var storage = CreateStorage(SessionTokenManagementMode.SemiAutomatic); storage.TrackSessionToken(_defaultContainerName, "A"); - storage.TrackSessionToken(_defaultContainerName, "B"); - storage.TrackSessionToken("other", "C"); - Assert.Null(storage.GetSessionToken(_defaultContainerName)); - Assert.Null(storage.GetSessionToken("other")); + storage.AppendDefaultContainerSessionToken("A"); + AssertDefault(storage, "A"); } // ================================================================ - // ENFORCED MANUAL MODE TESTS + // CLEAR // ================================================================ - [ConditionalFact] - public virtual void EnforcedManual_WhenGettingTokenBeforeSet_ThrowsInvalidOperationException() + [ConditionalTheory] + [InlineData(SessionTokenManagementMode.SemiAutomatic)] + [InlineData(SessionTokenManagementMode.Manual)] + [InlineData(SessionTokenManagementMode.EnforcedManual)] + public virtual void Clear_WhenClearing_ContainersStillExistInTrackedTokens(SessionTokenManagementMode mode) { - var storage = CreateStorage(SessionTokenManagementMode.EnforcedManual); - var ex = Assert.Throws(() => - storage.GetSessionToken(_defaultContainerName)); - Assert.Contains(CosmosStrings.MissingSessionTokenEnforceManual(_defaultContainerName), ex.Message); - } + var storage = CreateStorage(mode); + storage.Clear(); - [ConditionalFact] - public virtual void EnforcedManual_WhenGettingTokenAfterSet_ReturnsToken() - { - var storage = CreateStorage(SessionTokenManagementMode.EnforcedManual); - storage.SetDefaultContainerSessionToken("A"); - var token = storage.GetSessionToken(_defaultContainerName); - Assert.Equal("A", token); + var tokens = storage.GetTrackedTokens(); + Assert.Contains(_defaultContainerName, tokens.Keys); + Assert.Contains(_otherContainerName, tokens.Keys); } - [ConditionalFact] - public virtual void EnforcedManual_WhenGettingTokenAfterClear_ThrowsInvalidOperationException() + [ConditionalTheory] + [InlineData(SessionTokenManagementMode.SemiAutomatic)] + [InlineData(SessionTokenManagementMode.Manual)] + [InlineData(SessionTokenManagementMode.EnforcedManual)] + public virtual void Clear_WhenClearingSetTokens_ResetsAllContainersToNull(SessionTokenManagementMode mode) { - var storage = CreateStorage(SessionTokenManagementMode.EnforcedManual); - storage.SetDefaultContainerSessionToken("A"); + var storage = CreateStorage(mode); + + storage.AppendSessionTokens(new Dictionary { { _defaultContainerName, "A" }, { _otherContainerName, "B" } }); storage.Clear(); - var ex = Assert.Throws(() => - storage.GetSessionToken(_defaultContainerName)); - Assert.Contains(CosmosStrings.MissingSessionTokenEnforceManual(_defaultContainerName), ex.Message); + + var tokens = storage.GetTrackedTokens(); + Assert.Null(tokens[_defaultContainerName]); + Assert.Null(tokens[_otherContainerName]); + Assert.Null(storage.GetDefaultContainerTrackedToken()); + + if (mode != SessionTokenManagementMode.EnforcedManual) + { + Assert.Null(storage.GetSessionToken(_defaultContainerName)); + Assert.Null(storage.GetSessionToken(_otherContainerName)); + } } - [ConditionalFact] - public virtual void EnforcedManual_WhenGettingTokenAfterSetThenClearThenSet_ReturnsNewToken() + [ConditionalTheory] + [InlineData(SessionTokenManagementMode.SemiAutomatic)] + [InlineData(SessionTokenManagementMode.Manual)] + [InlineData(SessionTokenManagementMode.EnforcedManual)] + public virtual void Clear_WhenClearingTrackedTokens_ResetsAllContainersToNull(SessionTokenManagementMode mode) { - var storage = CreateStorage(SessionTokenManagementMode.EnforcedManual); - storage.SetDefaultContainerSessionToken("A"); + var storage = CreateStorage(mode); + + storage.TrackSessionToken(_defaultContainerName, "A"); + storage.TrackSessionToken(_otherContainerName, "B"); storage.Clear(); - storage.SetDefaultContainerSessionToken("B"); - var token = storage.GetSessionToken(_defaultContainerName); - Assert.Equal("B", token); - } - [ConditionalFact] - public virtual void EnforcedManual_WhenSettingMultipleContainers_AllContainersCanBeRetrieved() - { - var storage = CreateStorage(SessionTokenManagementMode.EnforcedManual); - storage.SetSessionTokens(new Dictionary - { - { _defaultContainerName, "A" }, - { "other", "B" } - }); + var tokens = storage.GetTrackedTokens(); + Assert.Null(tokens[_defaultContainerName]); + Assert.Null(tokens[_otherContainerName]); + Assert.Null(storage.GetDefaultContainerTrackedToken()); - Assert.Equal("A", storage.GetSessionToken(_defaultContainerName)); - Assert.Equal("B", storage.GetSessionToken("other")); + if (mode != SessionTokenManagementMode.EnforcedManual) + { + Assert.Null(storage.GetSessionToken(_defaultContainerName)); + Assert.Null(storage.GetSessionToken(_otherContainerName)); + } } - [ConditionalFact] - public virtual void EnforcedManual_WhenOneContainerNotSet_ThrowsForThatContainerOnly() + [ConditionalTheory] + [InlineData(SessionTokenManagementMode.SemiAutomatic)] + [InlineData(SessionTokenManagementMode.Manual)] + [InlineData(SessionTokenManagementMode.EnforcedManual)] + public virtual void Clear_WhenClearing_CanSetNewTokensAfterClear(SessionTokenManagementMode mode) { - var storage = CreateStorage(SessionTokenManagementMode.EnforcedManual); - storage.SetSessionTokens(new Dictionary { { _defaultContainerName, "A" } }); + var storage = CreateStorage(mode); + storage.SetSessionTokens(new Dictionary { { _defaultContainerName, "A" }, { _otherContainerName, "B" } }); - Assert.Equal("A", storage.GetSessionToken(_defaultContainerName)); - var ex = Assert.Throws(() => storage.GetSessionToken("other")); - Assert.Contains(CosmosStrings.MissingSessionTokenEnforceManual("other"), ex.Message); + storage.Clear(); + storage.SetSessionTokens(new Dictionary { { _defaultContainerName, "C" }, { _otherContainerName, "D" } }); + + var expectedDefault = "C"; + var expectedOther = "D"; + var defaultContainerTrackedToken = storage.GetDefaultContainerTrackedToken(); + var tokens = storage.GetTrackedTokens(); + + Assert.Equal(expectedDefault, defaultContainerTrackedToken); + + Assert.Equal(expectedDefault, tokens[_defaultContainerName]); + Assert.Equal(expectedOther, tokens[_otherContainerName]); + + if (mode != SessionTokenManagementMode.EnforcedManual) + { + Assert.Equal(expectedDefault, storage.GetSessionToken(_defaultContainerName)); + Assert.Equal(expectedOther, storage.GetSessionToken(_otherContainerName)); + } } - [ConditionalFact] - public virtual void EnforcedManual_WhenTrackingToken_ThrowsWhenRetrieving() + [ConditionalTheory] + [InlineData(SessionTokenManagementMode.SemiAutomatic)] + [InlineData(SessionTokenManagementMode.Manual)] + [InlineData(SessionTokenManagementMode.EnforcedManual)] + public virtual void Clear_WhenClearing_CanAppendNewTokensAfterClear(SessionTokenManagementMode mode) { - var storage = CreateStorage(SessionTokenManagementMode.EnforcedManual); - storage.TrackSessionToken(_defaultContainerName, "A"); - var ex = Assert.Throws(() => storage.GetSessionToken(_defaultContainerName)); - Assert.Contains(CosmosStrings.MissingSessionTokenEnforceManual(_defaultContainerName), ex.Message); + var storage = CreateStorage(mode); + storage.AppendSessionTokens(new Dictionary { { _defaultContainerName, "A" }, { _otherContainerName, "B" } }); + + storage.Clear(); + storage.AppendSessionTokens(new Dictionary { { _defaultContainerName, "C" }, { _otherContainerName, "D" } }); + + var expectedDefault = "C"; + var expectedOther = "D"; + var defaultContainerTrackedToken = storage.GetDefaultContainerTrackedToken(); + var tokens = storage.GetTrackedTokens(); + + Assert.Equal(expectedDefault, defaultContainerTrackedToken); + + Assert.Equal(expectedDefault, tokens[_defaultContainerName]); + Assert.Equal(expectedOther, tokens[_otherContainerName]); + + if (mode != SessionTokenManagementMode.EnforcedManual) + { + Assert.Equal(expectedDefault, storage.GetSessionToken(_defaultContainerName)); + Assert.Equal(expectedOther, storage.GetSessionToken(_otherContainerName)); + } } - [ConditionalFact] - public virtual void EnforcedManual_WhenAppendingToken_CanRetrieveToken() + [ConditionalTheory] + [InlineData(SessionTokenManagementMode.SemiAutomatic)] + [InlineData(SessionTokenManagementMode.Manual)] + [InlineData(SessionTokenManagementMode.EnforcedManual)] + public virtual void Clear_WhenClearing_CanTrackNewTokensAfterClear(SessionTokenManagementMode mode) { - var storage = CreateStorage(SessionTokenManagementMode.EnforcedManual); - storage.AppendDefaultContainerSessionToken("A"); - var token = storage.GetSessionToken(_defaultContainerName); - Assert.Equal("A", token); + var storage = CreateStorage(mode); + storage.AppendSessionTokens(new Dictionary { { _defaultContainerName, "A" }, { _otherContainerName, "B" } }); + + storage.Clear(); + } // ================================================================ // INITIALIZATION AND CONTAINER MANAGEMENT TESTS // ================================================================ - [ConditionalFact] - public virtual void Constructor_WhenInitializing_AllContainersAreInitialized() + [ConditionalTheory] + [InlineData(SessionTokenManagementMode.SemiAutomatic)] + [InlineData(SessionTokenManagementMode.Manual)] + [InlineData(SessionTokenManagementMode.EnforcedManual)] + public virtual void Constructor_AllContainersAreInitialized(SessionTokenManagementMode mode) { - var storage = CreateStorage(SessionTokenManagementMode.Manual); + var storage = CreateStorage(mode); - var tracked = storage.GetTrackedTokens(); - Assert.True(tracked.ContainsKey(_defaultContainerName)); - Assert.True(tracked.ContainsKey("other")); - Assert.Equal(2, tracked.Count); + var tokens = storage.GetTrackedTokens(); + Assert.Equal(2, tokens.Count); + Assert.True(tokens.ContainsKey(_defaultContainerName)); + Assert.True(tokens.ContainsKey(_otherContainerName)); } - [ConditionalFact] - public virtual void Constructor_WhenDefaultContainerNotInContainerNames_ThrowsException() + [ConditionalTheory] + [InlineData(SessionTokenManagementMode.SemiAutomatic)] + [InlineData(SessionTokenManagementMode.Manual)] + [InlineData(SessionTokenManagementMode.EnforcedManual)] + public virtual void Constructor_WhenDefaultContainerNotInContainerNames_ThrowsException(SessionTokenManagementMode mode) { - var containers = new HashSet(["other"]); + var containers = new HashSet([_otherContainerName]); Assert.True(!containers.Contains("default")); Assert.ThrowsAny(() => { - _ = new SessionTokenStorage("default", containers, SessionTokenManagementMode.Manual); + _ = new SessionTokenStorage("default", containers, mode); }); } - [ConditionalFact] - public virtual void Constructor_WhenInitializing_AllContainersStartWithNullTokens() + [ConditionalTheory] + [InlineData(SessionTokenManagementMode.SemiAutomatic)] + [InlineData(SessionTokenManagementMode.Manual)] + [InlineData(SessionTokenManagementMode.EnforcedManual)] + public virtual void Constructor_WhenInitializing_AllContainersStartWithNullTokens(SessionTokenManagementMode mode) { - var storage = CreateStorage(SessionTokenManagementMode.Manual); + var storage = CreateStorage(mode); - var tracked = storage.GetTrackedTokens(); - Assert.Null(tracked[_defaultContainerName]); - Assert.Null(tracked["other"]); + var tokens = storage.GetTrackedTokens(); + Assert.Null(tokens[_defaultContainerName]); + Assert.Null(tokens[_otherContainerName]); + Assert.Null(storage.GetDefaultContainerTrackedToken()); + if (mode != SessionTokenManagementMode.EnforcedManual) + { + Assert.Null(storage.GetSessionToken(_defaultContainerName)); + Assert.Null(storage.GetSessionToken(_otherContainerName)); + } } - // ================================================================ - // GETTRACKEDTOKENS TESTS - // ================================================================ - - [ConditionalFact] - public virtual void GetTrackedTokens_WhenCalled_ReturnsSnapshotNotLiveReference() + [ConditionalTheory] + [InlineData(SessionTokenManagementMode.SemiAutomatic)] + [InlineData(SessionTokenManagementMode.Manual)] + [InlineData(SessionTokenManagementMode.EnforcedManual)] + public virtual void GetTrackedTokens_WhenCalled_ReturnsSnapshotNotLiveReference(SessionTokenManagementMode mode) { - var storage = CreateStorage(SessionTokenManagementMode.Manual); + var storage = CreateStorage(mode); var snapshot = storage.GetTrackedTokens(); storage.AppendDefaultContainerSessionToken("A"); @@ -651,136 +813,265 @@ public virtual void GetTrackedTokens_WhenCalled_ReturnsSnapshotNotLiveReference( Assert.Equal("A", snapshot2[_defaultContainerName]); } + // ================================================================ + // FULLY AUTOMATIC MODE SPECIFIC TESTS + // ================================================================ + [ConditionalFact] - public virtual void GetTrackedTokens_WhenModifyingReturnedDictionary_DoesNotAffectStorage() + public virtual void FullyAutomatic_WhenCallingSetSessionTokens_ThrowsInvalidOperationException() { - var storage = CreateStorage(SessionTokenManagementMode.Manual); - storage.AppendDefaultContainerSessionToken("A"); - var tracked = storage.GetTrackedTokens(); - - // This should not compile or should not affect storage - // The returned dictionary is read-only, so we can't modify it - var token = storage.GetDefaultContainerTrackedToken(); - Assert.Equal("A", token); + var storage = CreateStorage(SessionTokenManagementMode.FullyAutomatic); + var ex = Assert.Throws(() => + storage.SetSessionTokens(new Dictionary())); + Assert.Equal(CosmosStrings.EnableManualSessionTokenManagement, ex.Message); } - // ================================================================ - // COMPOSITE TOKEN TESTS - // ================================================================ + [ConditionalFact] + public virtual void FullyAutomatic_WhenCallingGetTrackedTokens_ThrowsInvalidOperationException() + { + var storage = CreateStorage(SessionTokenManagementMode.FullyAutomatic); + var ex = Assert.Throws(() => storage.GetTrackedTokens()); + Assert.Equal(CosmosStrings.EnableManualSessionTokenManagement, ex.Message); + } [ConditionalFact] - public virtual void CompositeSessionToken_WhenAppendingDuplicateTokens_RemovesDuplicates() + public virtual void FullyAutomatic_WhenCallingAppendSessionTokens_ThrowsInvalidOperationException() { - var storage = CreateStorage(SessionTokenManagementMode.Manual); - storage.AppendDefaultContainerSessionToken("A,A,B"); - var token = storage.GetDefaultContainerTrackedToken(); - Assert.Equal("A,B", token); + var storage = CreateStorage(SessionTokenManagementMode.FullyAutomatic); + var ex = Assert.Throws(() => + storage.AppendSessionTokens(new Dictionary())); + Assert.Equal(CosmosStrings.EnableManualSessionTokenManagement, ex.Message); } [ConditionalFact] - public virtual void CompositeSessionToken_WhenAppendingDuplicateTokensInSeparateCalls_RemovesDuplicates() + public virtual void FullyAutomatic_WhenCallingSetDefaultContainerSessionToken_ThrowsInvalidOperationException() { - var storage = CreateStorage(SessionTokenManagementMode.Manual); - storage.AppendDefaultContainerSessionToken("A"); - storage.AppendDefaultContainerSessionToken("A"); - storage.AppendDefaultContainerSessionToken("B"); - storage.AppendDefaultContainerSessionToken("A"); - var token = storage.GetDefaultContainerTrackedToken(); - Assert.Equal("A,B", token); + var storage = CreateStorage(SessionTokenManagementMode.FullyAutomatic); + var ex = Assert.Throws(() => + storage.SetDefaultContainerSessionToken(null)); + Assert.Equal(CosmosStrings.EnableManualSessionTokenManagement, ex.Message); } [ConditionalFact] - public virtual void CompositeSessionToken_WhenSettingCommaSeparatedTokens_StoresAllTokens() + public virtual void FullyAutomatic_WhenCallingAppendDefaultContainerSessionToken_ThrowsInvalidOperationException() { - var storage = CreateStorage(SessionTokenManagementMode.Manual); - storage.SetDefaultContainerSessionToken("A,B,C"); - var token = storage.GetDefaultContainerTrackedToken(); - Assert.Equal("A,B,C", token); + var storage = CreateStorage(SessionTokenManagementMode.FullyAutomatic); + var ex = Assert.Throws(() => + storage.AppendDefaultContainerSessionToken("A")); + Assert.Equal(CosmosStrings.EnableManualSessionTokenManagement, ex.Message); } [ConditionalFact] - public virtual void CompositeSessionToken_WhenAppendingCommaSeparatedTokens_ParsesAndMergesCorrectly() + public virtual void FullyAutomatic_WhenCallingGetDefaultContainerTrackedToken_ThrowsInvalidOperationException() { - var storage = CreateStorage(SessionTokenManagementMode.Manual); - storage.AppendDefaultContainerSessionToken("A,B,C"); - var token = storage.GetDefaultContainerTrackedToken(); - Assert.Equal("A,B,C", token); + var storage = CreateStorage(SessionTokenManagementMode.FullyAutomatic); + var ex = Assert.Throws(() => + storage.GetDefaultContainerTrackedToken()); + Assert.Equal(CosmosStrings.EnableManualSessionTokenManagement, ex.Message); } - // ================================================================ - // MULTI-CONTAINER OPERATIONS TESTS - // ================================================================ + [ConditionalFact] + public virtual void FullyAutomatic_WhenTrackingToken_AlwaysReturnsNull() + { + var storage = CreateStorage(SessionTokenManagementMode.FullyAutomatic); + storage.TrackSessionToken(_defaultContainerName, "A"); + Assert.Null(storage.GetSessionToken(_defaultContainerName)); + } [ConditionalFact] - public virtual void TrackSessionToken_WhenTrackingToNonDefaultContainer_MergesCorrectly() + public virtual void FullyAutomatic_WhenTrackingMultipleTokens_AlwaysReturnsNull() { - var storage = CreateStorage(SessionTokenManagementMode.Manual); - storage.TrackSessionToken("other", "A"); - storage.TrackSessionToken("other", "B"); - Assert.Equal("A,B", storage.GetSessionToken("other")); + var storage = CreateStorage(SessionTokenManagementMode.FullyAutomatic); + storage.TrackSessionToken(_defaultContainerName, "A"); + storage.TrackSessionToken(_defaultContainerName, "B"); + storage.TrackSessionToken(_otherContainerName, "C"); + Assert.Null(storage.GetSessionToken(_defaultContainerName)); + Assert.Null(storage.GetSessionToken(_otherContainerName)); } // ================================================================ - // NULL AND EMPTY VALUE TESTS + // Argument exceptions // ================================================================ - [ConditionalFact] - public virtual void GetDefaultContainerTrackedToken_WhenNoTokenSet_ReturnsNull() + [ConditionalTheory] + [InlineData(SessionTokenManagementMode.FullyAutomatic)] + [InlineData(SessionTokenManagementMode.SemiAutomatic)] + [InlineData(SessionTokenManagementMode.Manual)] + [InlineData(SessionTokenManagementMode.EnforcedManual)] + public virtual void TrackSessionToken_WhenContainerNameIsNull_ThrowsArgumentNullException(SessionTokenManagementMode mode) { - var storage = CreateStorage(SessionTokenManagementMode.Manual); - Assert.Null(storage.GetDefaultContainerTrackedToken()); + var storage = CreateStorage(mode); + Assert.Throws(() => storage.TrackSessionToken(null!, "A")); } - [ConditionalFact] - public virtual void GetDefaultContainerTrackedToken_AfterClear_ReturnsNull() + [ConditionalTheory] + [InlineData(SessionTokenManagementMode.FullyAutomatic)] + [InlineData(SessionTokenManagementMode.SemiAutomatic)] + [InlineData(SessionTokenManagementMode.Manual)] + [InlineData(SessionTokenManagementMode.EnforcedManual)] + public virtual void TrackSessionToken_WhenContainerNameIsWhitespace_ThrowsArgumentNullException(SessionTokenManagementMode mode) { - var storage = CreateStorage(SessionTokenManagementMode.Manual); - storage.AppendDefaultContainerSessionToken("A"); - storage.Clear(); - Assert.Null(storage.GetDefaultContainerTrackedToken()); + var storage = CreateStorage(mode); + Assert.Throws(() => storage.TrackSessionToken(" ", "A")); + Assert.Throws(() => storage.TrackSessionToken("", "A")); } - [ConditionalFact] - public virtual void GetSessionToken_AfterClear_ReturnsNull() + [ConditionalTheory] + [InlineData(SessionTokenManagementMode.FullyAutomatic)] + [InlineData(SessionTokenManagementMode.SemiAutomatic)] + [InlineData(SessionTokenManagementMode.Manual)] + [InlineData(SessionTokenManagementMode.EnforcedManual)] + public virtual void TrackSessionToken_WhenTokenIsNull_ThrowsArgumentNullException(SessionTokenManagementMode mode) { - var storage = CreateStorage(SessionTokenManagementMode.Manual); - storage.AppendDefaultContainerSessionToken("A"); - storage.Clear(); - var token = storage.GetSessionToken(_defaultContainerName); - Assert.Null(token); + var storage = CreateStorage(mode); + Assert.Throws(() => storage.TrackSessionToken(_defaultContainerName, null!)); } - // ================================================================ - // SEMI-AUTOMATIC MODE SPECIFIC TESTS - // ================================================================ + [ConditionalTheory] + [InlineData(SessionTokenManagementMode.FullyAutomatic)] + [InlineData(SessionTokenManagementMode.SemiAutomatic)] + [InlineData(SessionTokenManagementMode.Manual)] + [InlineData(SessionTokenManagementMode.EnforcedManual)] + public virtual void TrackSessionToken_WhenTokenIsWhitespace_ThrowsArgumentNullException(SessionTokenManagementMode mode) + { + var storage = CreateStorage(mode); + Assert.Throws(() => storage.TrackSessionToken(_defaultContainerName, " ")); + Assert.Throws(() => storage.TrackSessionToken(_defaultContainerName, "")); + } - [ConditionalFact] - public virtual void SemiAutomatic_WhenTrackingToken_CanRetrieveToken() + [ConditionalTheory] + [InlineData(SessionTokenManagementMode.FullyAutomatic)] + [InlineData(SessionTokenManagementMode.SemiAutomatic)] + [InlineData(SessionTokenManagementMode.Manual)] + [InlineData(SessionTokenManagementMode.EnforcedManual)] + public virtual void AppendDefaultContainerSessionToken_WhenTokenIsNull_ThrowsArgumentNullException(SessionTokenManagementMode mode) { - var storage = CreateStorage(SessionTokenManagementMode.SemiAutomatic); - storage.TrackSessionToken(_defaultContainerName, "A"); - var token = storage.GetSessionToken(_defaultContainerName); - Assert.Equal("A", token); + var storage = CreateStorage(mode); + Assert.Throws(() => storage.AppendDefaultContainerSessionToken(null!)); } - [ConditionalFact] - public virtual void SemiAutomatic_WhenGettingTokenBeforeSet_ReturnsNull() + [ConditionalTheory] + [InlineData(SessionTokenManagementMode.FullyAutomatic)] + [InlineData(SessionTokenManagementMode.SemiAutomatic)] + [InlineData(SessionTokenManagementMode.Manual)] + [InlineData(SessionTokenManagementMode.EnforcedManual)] + public virtual void AppendDefaultContainerSessionToken_WhenTokenIsWhitespace_ThrowsArgumentException(SessionTokenManagementMode mode) { - var storage = CreateStorage(SessionTokenManagementMode.SemiAutomatic); - var token = storage.GetSessionToken(_defaultContainerName); - Assert.Null(token); + var storage = CreateStorage(mode); + Assert.Throws(() => storage.AppendDefaultContainerSessionToken(" ")); + Assert.Throws(() => storage.AppendDefaultContainerSessionToken("")); } - [ConditionalFact] - public virtual void SemiAutomatic_WhenGettingTokenAfterClear_ReturnsNull() + + [ConditionalTheory] + [InlineData(SessionTokenManagementMode.SemiAutomatic)] + [InlineData(SessionTokenManagementMode.Manual)] + [InlineData(SessionTokenManagementMode.EnforcedManual)] + public virtual void SetSessionTokens_WhenContainerNameIsUnknown_ThrowsInvalidOperationException(SessionTokenManagementMode mode) { - var storage = CreateStorage(SessionTokenManagementMode.SemiAutomatic); - storage.SetDefaultContainerSessionToken("A"); - storage.Clear(); - var token = storage.GetSessionToken(_defaultContainerName); - Assert.Null(token); + var storage = CreateStorage(mode); + var ex = Assert.Throws(() => + storage.SetSessionTokens(new Dictionary { { "bad", "A" } })); + Assert.Equal(CosmosStrings.ContainerNameDoesNotExist("bad"), ex.Message); + } + + [ConditionalTheory] + [InlineData(SessionTokenManagementMode.SemiAutomatic)] + [InlineData(SessionTokenManagementMode.Manual)] + [InlineData(SessionTokenManagementMode.EnforcedManual)] + public virtual void AppendSessionTokens_WhenContainerNameIsUnknown_ThrowsInvalidOperationException(SessionTokenManagementMode mode) + { + var storage = CreateStorage(mode); + var ex = Assert.Throws(() => + storage.AppendSessionTokens(new Dictionary { { "bad", "A" } })); + Assert.Equal(CosmosStrings.ContainerNameDoesNotExist("bad"), ex.Message); } + [ConditionalTheory] + [InlineData(SessionTokenManagementMode.SemiAutomatic)] + [InlineData(SessionTokenManagementMode.Manual)] + [InlineData(SessionTokenManagementMode.EnforcedManual)] + public virtual void GetSessionToken_WhenContainerNameIsNull_ThrowsArgumentNullException(SessionTokenManagementMode mode) + { + var storage = CreateStorage(mode); + Assert.Throws(() => storage.GetSessionToken(null!)); + } + + [ConditionalTheory] + [InlineData(SessionTokenManagementMode.SemiAutomatic)] + [InlineData(SessionTokenManagementMode.Manual)] + [InlineData(SessionTokenManagementMode.EnforcedManual)] + public virtual void GetSessionToken_WhenContainerNameIsWhitespace_ThrowsArgumentNullException(SessionTokenManagementMode mode) + { + var storage = CreateStorage(mode); + Assert.Throws(() => storage.GetSessionToken(" ")); + Assert.Throws(() => storage.GetSessionToken("")); + } + + private SessionTokenStorage CreateStorage(SessionTokenManagementMode mode) => new(_defaultContainerName, _containerNames, mode); + + private void AssertDefault(SessionTokenStorage storage, string? value) + { + AssertDefaultTracked(storage, value); + AssertDefaultUsed(storage, value); + } + + private void AssertDefaultUsed(SessionTokenStorage storage, string? value) + { + if (value == null) + { + Assert.Null(storage.GetSessionToken(_defaultContainerName)); + } + else + { + Assert.Equal(value, storage.GetSessionToken(_defaultContainerName)); + } + } + + private void AssertDefaultTracked(SessionTokenStorage storage, string? value) + { + if (value == null) + { + Assert.Null(storage.GetDefaultContainerTrackedToken()); + Assert.Null(storage.GetTrackedTokens()[_defaultContainerName]); + } + else + { + Assert.Equal(value, storage.GetDefaultContainerTrackedToken()); + Assert.Equal(value, storage.GetTrackedTokens()[_defaultContainerName]); + } + } + + private void AssertOther(SessionTokenStorage storage, string? value) + { + AssertOtherTracked(storage, value); + AssertOtherUsed(storage, value); + } + + private void AssertOtherUsed(SessionTokenStorage storage, string? value) + { + if (value == null) + { + Assert.Null(storage.GetSessionToken(_otherContainerName)); + } + else + { + Assert.Equal(value, storage.GetSessionToken(_otherContainerName)); + } + } + + private void AssertOtherTracked(SessionTokenStorage storage, string? value) + { + if (value == null) + { + Assert.Null(storage.GetTrackedTokens()[_otherContainerName]); + } + else + { + Assert.Equal(value, storage.GetTrackedTokens()[_otherContainerName]); + } + } + } From 7427a658baa6452b3f39d68558d04f53f6bdc9f1 Mon Sep 17 00:00:00 2001 From: JoasE <32096708+JoasE@users.noreply.github.com> Date: Mon, 10 Nov 2025 17:49:52 +0100 Subject: [PATCH 34/37] Do not use automatic session tokens for manual mode if not provided. --- .../SessionTokenManagementMode.cs | 2 + .../Storage/Internal/SessionTokenStorage.cs | 10 +- .../Internal/SessionTokenStorageTest.cs | 236 ++++++++++++++---- 3 files changed, 200 insertions(+), 48 deletions(-) diff --git a/src/EFCore.Cosmos/Infrastructure/SessionTokenManagementMode.cs b/src/EFCore.Cosmos/Infrastructure/SessionTokenManagementMode.cs index e0cb402f513..29f9ee26e81 100644 --- a/src/EFCore.Cosmos/Infrastructure/SessionTokenManagementMode.cs +++ b/src/EFCore.Cosmos/Infrastructure/SessionTokenManagementMode.cs @@ -23,12 +23,14 @@ public enum SessionTokenManagementMode /// /// Allows the usage of UseSessionTokens to overwrite the default Cosmos DB SDK automatic session token management by use of the UseSessionTokens method on a instance. + /// If UseSessionTokens 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 . /// SemiAutomatic, /// /// Fully overwrites the Cosmos DB SDK automatic session token management, and only uses session tokens specified via UseSessionTokens. + /// If UseSessionTokens 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 . /// Manual, diff --git a/src/EFCore.Cosmos/Storage/Internal/SessionTokenStorage.cs b/src/EFCore.Cosmos/Storage/Internal/SessionTokenStorage.cs index 3c0a903c143..843de89ba73 100644 --- a/src/EFCore.Cosmos/Storage/Internal/SessionTokenStorage.cs +++ b/src/EFCore.Cosmos/Storage/Internal/SessionTokenStorage.cs @@ -18,6 +18,7 @@ public class SessionTokenStorage : ISessionTokenStorage private readonly string _defaultContainerName; private readonly HashSet _containerNames; private readonly SessionTokenManagementMode _mode; + private readonly string? _defaultToken; /// /// This is an internal API that supports the Entity Framework Core infrastructure and not subject to @@ -31,8 +32,9 @@ public SessionTokenStorage(string defaultContainerName, HashSet containe _defaultContainerName = defaultContainerName; _containerNames = containerNames; _mode = mode; + _defaultToken = _mode == SessionTokenManagementMode.Manual || _mode == SessionTokenManagementMode.EnforcedManual ? "" : null; - _containerSessionTokens = containerNames.ToDictionary(x => x, x => new CompositeSessionToken()); + _containerSessionTokens = containerNames.ToDictionary(x => x, x => new CompositeSessionToken(_defaultToken)); } /// @@ -186,7 +188,7 @@ public virtual void Clear() { foreach (var key in _containerSessionTokens.Keys) { - _containerSessionTokens[key] = new CompositeSessionToken(); + _containerSessionTokens[key] = new CompositeSessionToken(_defaultToken); } } @@ -213,10 +215,6 @@ public CompositeSessionToken(string? token, bool isSet = false) IsSet = isSet; } - public CompositeSessionToken() - { - } - public bool IsSet { get; private set; } public void Add(string token, bool isSet = false) diff --git a/test/EFCore.Cosmos.Tests/Storage/Internal/SessionTokenStorageTest.cs b/test/EFCore.Cosmos.Tests/Storage/Internal/SessionTokenStorageTest.cs index 803c6511614..0287b91cfac 100644 --- a/test/EFCore.Cosmos.Tests/Storage/Internal/SessionTokenStorageTest.cs +++ b/test/EFCore.Cosmos.Tests/Storage/Internal/SessionTokenStorageTest.cs @@ -31,7 +31,14 @@ public virtual void SetSessionTokens_SetSingle_Default(SessionTokenManagementMod AssertDefault(storage, "A"); if (mode != SessionTokenManagementMode.EnforcedManual) { - AssertOther(storage, null); + if (mode == SessionTokenManagementMode.Manual) + { + AssertOther(storage, ""); + } + else + { + AssertOther(storage, null); + } } } @@ -47,7 +54,14 @@ public virtual void SetSessionTokens_SetSingle_Other(SessionTokenManagementMode if (mode != SessionTokenManagementMode.EnforcedManual) { - AssertDefault(storage, null); + if (mode == SessionTokenManagementMode.Manual) + { + AssertDefault(storage, ""); + } + else + { + AssertDefault(storage, null); + } } AssertOther(storage, "A"); } @@ -157,7 +171,14 @@ public virtual void SetDefaultContainerSessionToken_SetsToken(SessionTokenManage if (mode != SessionTokenManagementMode.EnforcedManual) { - AssertOther(storage, null); + if (mode == SessionTokenManagementMode.Manual) + { + AssertOther(storage, ""); + } + else + { + AssertOther(storage, null); + } } } @@ -175,7 +196,14 @@ public virtual void SetDefaultContainerSessionToken_OverwritesSet(SessionTokenMa if (mode != SessionTokenManagementMode.EnforcedManual) { - AssertOther(storage, null); + if (mode == SessionTokenManagementMode.Manual) + { + AssertOther(storage, ""); + } + else + { + AssertOther(storage, null); + } } } @@ -192,7 +220,14 @@ public virtual void SetDefaultContainerSessionToken_OverwritesTracked(SessionTok AssertDefault(storage, "B"); if (mode != SessionTokenManagementMode.EnforcedManual) { - AssertOther(storage, null); + if (mode == SessionTokenManagementMode.Manual) + { + AssertOther(storage, ""); + } + else + { + AssertOther(storage, null); + } } } @@ -209,7 +244,14 @@ public virtual void AppendDefaultContainerSessionToken_NoPreviousToken_SetsToken if (mode != SessionTokenManagementMode.EnforcedManual) { - AssertOther(storage, null); + if (mode == SessionTokenManagementMode.Manual) + { + AssertOther(storage, ""); + } + else + { + AssertOther(storage, null); + } } } @@ -227,7 +269,14 @@ public virtual void AppendDefaultContainerSessionToken_PreviousSetToken_AppendsT if (mode != SessionTokenManagementMode.EnforcedManual) { - AssertOther(storage, null); + if (mode == SessionTokenManagementMode.Manual) + { + AssertOther(storage, ""); + } + else + { + AssertOther(storage, null); + } } } @@ -244,7 +293,14 @@ public virtual void AppendDefaultContainerSessionToken_PreviousSetToken_Duplicat AssertDefault(storage, "A"); if (mode != SessionTokenManagementMode.EnforcedManual) { - AssertOther(storage, null); + if (mode == SessionTokenManagementMode.Manual) + { + AssertOther(storage, ""); + } + else + { + AssertOther(storage, null); + } } } @@ -261,7 +317,14 @@ public virtual void AppendDefaultContainerSessionToken_PreviousTrackedToken_Appe AssertDefault(storage, "A,B"); if (mode != SessionTokenManagementMode.EnforcedManual) { - AssertOther(storage, null); + if (mode == SessionTokenManagementMode.Manual) + { + AssertOther(storage, ""); + } + else + { + AssertOther(storage, null); + } } } @@ -278,7 +341,14 @@ public virtual void AppendDefaultContainerSessionToken_PreviousTrackedToken_Dupl AssertDefault(storage, "A"); if (mode != SessionTokenManagementMode.EnforcedManual) { - AssertOther(storage, null); + if (mode == SessionTokenManagementMode.Manual) + { + AssertOther(storage, ""); + } + else + { + AssertOther(storage, null); + } } } @@ -315,7 +385,14 @@ public virtual void AppendSessionTokens_SingleContainer_NoPreviousTokens_SetsTok if (mode != SessionTokenManagementMode.EnforcedManual) { - AssertDefault(storage, null); + if (mode == SessionTokenManagementMode.Manual) + { + AssertDefault(storage, ""); + } + else + { + AssertDefault(storage, null); + } } AssertOther(storage, "B"); } @@ -490,9 +567,16 @@ public virtual void TrackSessionToken_SetsToken(SessionTokenManagementMode mode) } else if (mode != SessionTokenManagementMode.EnforcedManual) { - - AssertDefaultUsed(storage, null); - AssertOtherUsed(storage, null); + if (mode == SessionTokenManagementMode.Manual) + { + AssertDefaultUsed(storage, ""); + AssertOtherUsed(storage, ""); + } + else + { + AssertDefaultUsed(storage, null); + AssertOtherUsed(storage, null); + } } } @@ -520,10 +604,18 @@ public virtual void TrackSessionToken_Appends(SessionTokenManagementMode mode) AssertDefaultUsed(storage, "A,B"); AssertOtherUsed(storage, "A,C"); } - else if(mode != SessionTokenManagementMode.EnforcedManual) + else if (mode != SessionTokenManagementMode.EnforcedManual) { - AssertDefaultUsed(storage, null); - AssertOtherUsed(storage, null); + if (mode == SessionTokenManagementMode.Manual) + { + AssertDefaultUsed(storage, ""); + AssertOtherUsed(storage, ""); + } + else + { + AssertDefaultUsed(storage, null); + AssertOtherUsed(storage, null); + } } } @@ -588,11 +680,11 @@ public virtual void EnforcedManual_WhenOneContainerNotSet_ThrowsForThatContainer } // ================================================================ - // SEMI-AUTOMATIC MODE SPECIFIC TESTS + // SEMI-AUTOMATIC SPECIFIC TESTS // ================================================================ [ConditionalFact] - public virtual void SemiAutomatic_WhenTrackingToken_SetsButDoesntUseToken() + public virtual void SemiAutomatic_WhenTrackingToken_SetsButDoesnotUseToken() { var storage = CreateStorage(SessionTokenManagementMode.SemiAutomatic); storage.TrackSessionToken(_defaultContainerName, "A"); @@ -609,6 +701,39 @@ public virtual void SemiAutomatic_WhenSetToken_SetsAndUses() AssertDefault(storage, "A"); } + // ================================================================ + // MANUAL SPECIFIC TESTS + // ================================================================ + + [ConditionalFact] + public virtual void Manual_Constructor_AllContainersHaveEmptyString() + { + var storage = CreateStorage(SessionTokenManagementMode.Manual); + + var tokens = storage.GetTrackedTokens(); + Assert.True(tokens[_defaultContainerName] == ""); + Assert.True(tokens[_otherContainerName] == ""); + Assert.True(storage.GetDefaultContainerTrackedToken() == ""); + Assert.True(storage.GetSessionToken(_defaultContainerName) == ""); + Assert.True(storage.GetSessionToken(_otherContainerName) == ""); + } + + [ConditionalFact] + public virtual void Manual_Clear_ResetsAllContainersToEmptyString() + { + var storage = CreateStorage(SessionTokenManagementMode.Manual); + + storage.AppendSessionTokens(new Dictionary { { _defaultContainerName, "A" }, { _otherContainerName, "B" } }); + storage.Clear(); + + var tokens = storage.GetTrackedTokens(); + Assert.True(tokens[_defaultContainerName] == ""); + Assert.True(tokens[_otherContainerName] == ""); + Assert.True(storage.GetDefaultContainerTrackedToken() == ""); + Assert.True(storage.GetSessionToken(_defaultContainerName) == ""); + Assert.True(storage.GetSessionToken(_otherContainerName) == ""); + } + // ================================================================ // CLEAR // ================================================================ @@ -631,20 +756,33 @@ public virtual void Clear_WhenClearing_ContainersStillExistInTrackedTokens(Sessi [InlineData(SessionTokenManagementMode.SemiAutomatic)] [InlineData(SessionTokenManagementMode.Manual)] [InlineData(SessionTokenManagementMode.EnforcedManual)] - public virtual void Clear_WhenClearingSetTokens_ResetsAllContainersToNull(SessionTokenManagementMode mode) + public virtual void Clear_WhenClearingSetTokens_ResetsAllContainers(SessionTokenManagementMode mode) { var storage = CreateStorage(mode); storage.AppendSessionTokens(new Dictionary { { _defaultContainerName, "A" }, { _otherContainerName, "B" } }); storage.Clear(); - + var tokens = storage.GetTrackedTokens(); - Assert.Null(tokens[_defaultContainerName]); - Assert.Null(tokens[_otherContainerName]); - Assert.Null(storage.GetDefaultContainerTrackedToken()); - if (mode != SessionTokenManagementMode.EnforcedManual) + if (mode == SessionTokenManagementMode.Manual || mode == SessionTokenManagementMode.EnforcedManual) { + Assert.True(tokens[_defaultContainerName] == ""); + Assert.True(tokens[_otherContainerName] == ""); + Assert.True(storage.GetDefaultContainerTrackedToken() == ""); + + if (mode != SessionTokenManagementMode.EnforcedManual) + { + Assert.True(storage.GetSessionToken(_defaultContainerName) == ""); + Assert.True(storage.GetSessionToken(_otherContainerName) == ""); + } + } + else + { + Assert.Null(tokens[_defaultContainerName]); + Assert.Null(tokens[_otherContainerName]); + Assert.Null(storage.GetDefaultContainerTrackedToken()); + Assert.Null(storage.GetSessionToken(_defaultContainerName)); Assert.Null(storage.GetSessionToken(_otherContainerName)); } @@ -654,7 +792,7 @@ public virtual void Clear_WhenClearingSetTokens_ResetsAllContainersToNull(Sessio [InlineData(SessionTokenManagementMode.SemiAutomatic)] [InlineData(SessionTokenManagementMode.Manual)] [InlineData(SessionTokenManagementMode.EnforcedManual)] - public virtual void Clear_WhenClearingTrackedTokens_ResetsAllContainersToNull(SessionTokenManagementMode mode) + public virtual void Clear_WhenClearingTrackedTokens_ResetsAllContainers(SessionTokenManagementMode mode) { var storage = CreateStorage(mode); @@ -663,12 +801,25 @@ public virtual void Clear_WhenClearingTrackedTokens_ResetsAllContainersToNull(Se storage.Clear(); var tokens = storage.GetTrackedTokens(); - Assert.Null(tokens[_defaultContainerName]); - Assert.Null(tokens[_otherContainerName]); - Assert.Null(storage.GetDefaultContainerTrackedToken()); - if (mode != SessionTokenManagementMode.EnforcedManual) + if (mode == SessionTokenManagementMode.Manual || mode == SessionTokenManagementMode.EnforcedManual) + { + Assert.True(tokens[_defaultContainerName] == ""); + Assert.True(tokens[_otherContainerName] == ""); + Assert.True(storage.GetDefaultContainerTrackedToken() == ""); + + if (mode != SessionTokenManagementMode.EnforcedManual) + { + Assert.True(storage.GetSessionToken(_defaultContainerName) == ""); + Assert.True(storage.GetSessionToken(_otherContainerName) == ""); + } + } + else { + Assert.Null(tokens[_defaultContainerName]); + Assert.Null(tokens[_otherContainerName]); + Assert.Null(storage.GetDefaultContainerTrackedToken()); + Assert.Null(storage.GetSessionToken(_defaultContainerName)); Assert.Null(storage.GetSessionToken(_otherContainerName)); } @@ -777,23 +928,17 @@ public virtual void Constructor_WhenDefaultContainerNotInContainerNames_ThrowsEx }); } - [ConditionalTheory] - [InlineData(SessionTokenManagementMode.SemiAutomatic)] - [InlineData(SessionTokenManagementMode.Manual)] - [InlineData(SessionTokenManagementMode.EnforcedManual)] - public virtual void Constructor_WhenInitializing_AllContainersStartWithNullTokens(SessionTokenManagementMode mode) + [ConditionalFact] + public virtual void Constructor_WhenInitializing_AllContainersStartWithNullTokens() { - var storage = CreateStorage(mode); + var storage = CreateStorage(SessionTokenManagementMode.SemiAutomatic); var tokens = storage.GetTrackedTokens(); Assert.Null(tokens[_defaultContainerName]); Assert.Null(tokens[_otherContainerName]); Assert.Null(storage.GetDefaultContainerTrackedToken()); - if (mode != SessionTokenManagementMode.EnforcedManual) - { - Assert.Null(storage.GetSessionToken(_defaultContainerName)); - Assert.Null(storage.GetSessionToken(_otherContainerName)); - } + Assert.Null(storage.GetSessionToken(_defaultContainerName)); + Assert.Null(storage.GetSessionToken(_otherContainerName)); } [ConditionalTheory] @@ -809,7 +954,14 @@ public virtual void GetTrackedTokens_WhenCalled_ReturnsSnapshotNotLiveReference( var snapshot2 = storage.GetTrackedTokens(); Assert.NotSame(snapshot, snapshot2); - Assert.Null(snapshot[_defaultContainerName]); + if (mode == SessionTokenManagementMode.Manual || mode == SessionTokenManagementMode.EnforcedManual) + { + Assert.True(snapshot[_defaultContainerName] == ""); + } + else + { + Assert.Null(snapshot[_defaultContainerName]); + } Assert.Equal("A", snapshot2[_defaultContainerName]); } From 8d46fedfbc39fde14a41d66e2d8ef3a5f288de8b Mon Sep 17 00:00:00 2001 From: JoasE <32096708+JoasE@users.noreply.github.com> Date: Mon, 10 Nov 2025 18:13:41 +0100 Subject: [PATCH 35/37] Cleanup --- .../Storage/Internal/SessionTokenStorage.cs | 9 +--- .../Internal/SessionTokenStorageTest.cs | 44 +++++++++++++++---- 2 files changed, 37 insertions(+), 16 deletions(-) diff --git a/src/EFCore.Cosmos/Storage/Internal/SessionTokenStorage.cs b/src/EFCore.Cosmos/Storage/Internal/SessionTokenStorage.cs index 843de89ba73..358f62525e3 100644 --- a/src/EFCore.Cosmos/Storage/Internal/SessionTokenStorage.cs +++ b/src/EFCore.Cosmos/Storage/Internal/SessionTokenStorage.cs @@ -245,14 +245,7 @@ public void Add(string token, bool isSet = false) if (_isChanged) { _isChanged = false; - if (_tokens.Count == 0) - { - _string = null; - } - else - { - _string = string.Join(",", _tokens); - } + _string = string.Join(",", _tokens); } return _string; diff --git a/test/EFCore.Cosmos.Tests/Storage/Internal/SessionTokenStorageTest.cs b/test/EFCore.Cosmos.Tests/Storage/Internal/SessionTokenStorageTest.cs index 0287b91cfac..fe9487af4a1 100644 --- a/test/EFCore.Cosmos.Tests/Storage/Internal/SessionTokenStorageTest.cs +++ b/test/EFCore.Cosmos.Tests/Storage/Internal/SessionTokenStorageTest.cs @@ -706,22 +706,42 @@ public virtual void SemiAutomatic_WhenSetToken_SetsAndUses() // ================================================================ [ConditionalFact] - public virtual void Manual_Constructor_AllContainersHaveEmptyString() + public virtual void Manual_TrackedToken_UsesToken() { var storage = CreateStorage(SessionTokenManagementMode.Manual); + storage.TrackSessionToken(_defaultContainerName, "A"); + storage.TrackSessionToken(_otherContainerName, "B"); + + Assert.True(storage.GetSessionToken(_defaultContainerName) == "A"); + Assert.True(storage.GetSessionToken(_otherContainerName) == "B"); + } + + [ConditionalTheory] + [InlineData(SessionTokenManagementMode.Manual)] + [InlineData(SessionTokenManagementMode.EnforcedManual)] + public virtual void Manual_Constructor_AllContainersHaveEmptyString(SessionTokenManagementMode mode) + { + var storage = CreateStorage(mode); + var tokens = storage.GetTrackedTokens(); Assert.True(tokens[_defaultContainerName] == ""); Assert.True(tokens[_otherContainerName] == ""); Assert.True(storage.GetDefaultContainerTrackedToken() == ""); - Assert.True(storage.GetSessionToken(_defaultContainerName) == ""); - Assert.True(storage.GetSessionToken(_otherContainerName) == ""); + + if (mode != SessionTokenManagementMode.EnforcedManual) + { + Assert.True(storage.GetSessionToken(_defaultContainerName) == ""); + Assert.True(storage.GetSessionToken(_otherContainerName) == ""); + } } - [ConditionalFact] - public virtual void Manual_Clear_ResetsAllContainersToEmptyString() + [ConditionalTheory] + [InlineData(SessionTokenManagementMode.Manual)] + [InlineData(SessionTokenManagementMode.EnforcedManual)] + public virtual void Manual_Clear_ResetsAllContainersToEmptyString(SessionTokenManagementMode mode) { - var storage = CreateStorage(SessionTokenManagementMode.Manual); + var storage = CreateStorage(mode); storage.AppendSessionTokens(new Dictionary { { _defaultContainerName, "A" }, { _otherContainerName, "B" } }); storage.Clear(); @@ -730,8 +750,11 @@ public virtual void Manual_Clear_ResetsAllContainersToEmptyString() Assert.True(tokens[_defaultContainerName] == ""); Assert.True(tokens[_otherContainerName] == ""); Assert.True(storage.GetDefaultContainerTrackedToken() == ""); - Assert.True(storage.GetSessionToken(_defaultContainerName) == ""); - Assert.True(storage.GetSessionToken(_otherContainerName) == ""); + if (mode != SessionTokenManagementMode.EnforcedManual) + { + Assert.True(storage.GetSessionToken(_defaultContainerName) == ""); + Assert.True(storage.GetSessionToken(_otherContainerName) == ""); + } } // ================================================================ @@ -894,6 +917,11 @@ public virtual void Clear_WhenClearing_CanTrackNewTokensAfterClear(SessionTokenM storage.Clear(); + storage.TrackSessionToken(_defaultContainerName, "C"); + storage.TrackSessionToken(_otherContainerName, "D"); + + AssertDefaultTracked(storage, "C"); + AssertOtherTracked(storage, "D"); } // ================================================================ From 2ea28cb2e9566e921b90110485b5ac51b23ff72c Mon Sep 17 00:00:00 2001 From: JoasE <32096708+JoasE@users.noreply.github.com> Date: Mon, 10 Nov 2025 19:39:30 +0100 Subject: [PATCH 36/37] Cleanup --- .../Infrastructure/SessionTokenManagementMode.cs | 12 ++++++------ .../Properties/CosmosStrings.Designer.cs | 2 +- src/EFCore.Cosmos/Properties/CosmosStrings.resx | 2 +- ...lingExpressionVisitor.PagingQueryingEnumerable.cs | 3 ++- ...yCompilingExpressionVisitor.QueryingEnumerable.cs | 3 --- ...ngExpressionVisitor.ReadItemQueryingEnumerable.cs | 2 +- 6 files changed, 11 insertions(+), 13 deletions(-) diff --git a/src/EFCore.Cosmos/Infrastructure/SessionTokenManagementMode.cs b/src/EFCore.Cosmos/Infrastructure/SessionTokenManagementMode.cs index 29f9ee26e81..f023127b05f 100644 --- a/src/EFCore.Cosmos/Infrastructure/SessionTokenManagementMode.cs +++ b/src/EFCore.Cosmos/Infrastructure/SessionTokenManagementMode.cs @@ -14,7 +14,7 @@ public enum SessionTokenManagementMode /// /// 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. UseSessionTokens and GetSessionTokens methods will throw when invoked. + /// EF will not track or parse session tokens returned from Cosmos DB. and 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. @@ -22,21 +22,21 @@ public enum SessionTokenManagementMode FullyAutomatic, /// - /// Allows the usage of UseSessionTokens to overwrite the default Cosmos DB SDK automatic session token management by use of the UseSessionTokens method on a instance. - /// If UseSessionTokens has not been invoked for an container, the default Cosmos DB SDK automatic session token management will be used. + /// Allows the usage of to overwrite the default Cosmos DB SDK automatic session token management by use of the method on a instance. + /// If 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 . /// SemiAutomatic, /// - /// Fully overwrites the Cosmos DB SDK automatic session token management, and only uses session tokens specified via UseSessionTokens. - /// If UseSessionTokens has not been invoked for an container, no session token will be used. + /// Fully overwrites the Cosmos DB SDK automatic session token management, and only uses session tokens specified via . + /// If 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 . /// Manual, /// - /// Same as , but will throw an exception if UseSessionTokens was not invoked before executong a read. + /// Same as , but will throw an exception if was not invoked before executong a read. /// EnforcedManual } diff --git a/src/EFCore.Cosmos/Properties/CosmosStrings.Designer.cs b/src/EFCore.Cosmos/Properties/CosmosStrings.Designer.cs index 6a19bc63043..eead7f1c210 100644 --- a/src/EFCore.Cosmos/Properties/CosmosStrings.Designer.cs +++ b/src/EFCore.Cosmos/Properties/CosmosStrings.Designer.cs @@ -166,7 +166,7 @@ public static string ElementWithValueConverter(object? propertyType, object? str propertyType, structuralType, property, elementType); /// - /// Enable manual session token management using 'options.ManualSessionTokenManagementEnabled' to use this method. + /// Enable manual session token management using 'options.SessionTokenManagementMode' to use this method. /// public static string EnableManualSessionTokenManagement => GetString("EnableManualSessionTokenManagement"); diff --git a/src/EFCore.Cosmos/Properties/CosmosStrings.resx b/src/EFCore.Cosmos/Properties/CosmosStrings.resx index 219557c6338..38b46cfe53d 100644 --- a/src/EFCore.Cosmos/Properties/CosmosStrings.resx +++ b/src/EFCore.Cosmos/Properties/CosmosStrings.resx @@ -176,7 +176,7 @@ 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. - Enable manual session token management using 'options.ManualSessionTokenManagementEnabled' to use this method. + Enable manual session token management using 'options.SessionTokenManagementMode' to use this method. The type of the etag property '{property}' on '{entityType}' is '{propertyType}'. All etag properties must be strings or have a string value converter. diff --git a/src/EFCore.Cosmos/Query/Internal/CosmosShapedQueryCompilingExpressionVisitor.PagingQueryingEnumerable.cs b/src/EFCore.Cosmos/Query/Internal/CosmosShapedQueryCompilingExpressionVisitor.PagingQueryingEnumerable.cs index 536a973d106..5687867d455 100644 --- a/src/EFCore.Cosmos/Query/Internal/CosmosShapedQueryCompilingExpressionVisitor.PagingQueryingEnumerable.cs +++ b/src/EFCore.Cosmos/Query/Internal/CosmosShapedQueryCompilingExpressionVisitor.PagingQueryingEnumerable.cs @@ -4,7 +4,6 @@ #nullable disable using Microsoft.EntityFrameworkCore.Cosmos.Diagnostics.Internal; -using Microsoft.EntityFrameworkCore.Cosmos.Storage; using Microsoft.EntityFrameworkCore.Cosmos.Storage.Internal; using Newtonsoft.Json.Linq; @@ -63,6 +62,7 @@ public PagingQueryingEnumerable( _threadSafetyChecksEnabled = threadSafetyChecksEnabled; _maxItemCountParameterName = maxItemCountParameterName; _continuationTokenParameterName = continuationTokenParameterName; + _responseContinuationTokenLimitInKbParameterName = responseContinuationTokenLimitInKbParameterName; _cosmosContainer = rootEntityType.GetContainer() ?? throw new UnreachableException("Root entity type without a Cosmos container."); @@ -95,6 +95,7 @@ private sealed class AsyncEnumerator : IAsyncEnumerator> private readonly CancellationToken _cancellationToken; private readonly IConcurrencyDetector _concurrencyDetector; private readonly IExceptionDetector _exceptionDetector; + private bool _hasExecuted; private bool _isDisposed; diff --git a/src/EFCore.Cosmos/Query/Internal/CosmosShapedQueryCompilingExpressionVisitor.QueryingEnumerable.cs b/src/EFCore.Cosmos/Query/Internal/CosmosShapedQueryCompilingExpressionVisitor.QueryingEnumerable.cs index 8c67856f4df..d2906182de9 100644 --- a/src/EFCore.Cosmos/Query/Internal/CosmosShapedQueryCompilingExpressionVisitor.QueryingEnumerable.cs +++ b/src/EFCore.Cosmos/Query/Internal/CosmosShapedQueryCompilingExpressionVisitor.QueryingEnumerable.cs @@ -5,7 +5,6 @@ using System.Collections; using System.Text; -using Microsoft.EntityFrameworkCore.Cosmos.Storage; using Microsoft.EntityFrameworkCore.Cosmos.Internal; using Microsoft.EntityFrameworkCore.Cosmos.Storage.Internal; using Newtonsoft.Json.Linq; @@ -115,7 +114,6 @@ private sealed class AsyncEnumerator : IAsyncEnumerator private readonly IConcurrencyDetector _concurrencyDetector; private readonly IExceptionDetector _exceptionDetector; - private IAsyncEnumerator _enumerator; public AsyncEnumerator(QueryingEnumerable queryingEnumerable, CancellationToken cancellationToken) @@ -129,7 +127,6 @@ public AsyncEnumerator(QueryingEnumerable queryingEnumerable, CancellationTok _queryLogger = queryingEnumerable._queryLogger; _standAloneStateManager = queryingEnumerable._standAloneStateManager; _exceptionDetector = _cosmosQueryContext.ExceptionDetector; - _cancellationToken = cancellationToken; _concurrencyDetector = queryingEnumerable._threadSafetyChecksEnabled diff --git a/src/EFCore.Cosmos/Query/Internal/CosmosShapedQueryCompilingExpressionVisitor.ReadItemQueryingEnumerable.cs b/src/EFCore.Cosmos/Query/Internal/CosmosShapedQueryCompilingExpressionVisitor.ReadItemQueryingEnumerable.cs index d6df274a564..c1461cfa6fe 100644 --- a/src/EFCore.Cosmos/Query/Internal/CosmosShapedQueryCompilingExpressionVisitor.ReadItemQueryingEnumerable.cs +++ b/src/EFCore.Cosmos/Query/Internal/CosmosShapedQueryCompilingExpressionVisitor.ReadItemQueryingEnumerable.cs @@ -6,7 +6,6 @@ using System.Collections; using Microsoft.EntityFrameworkCore.Cosmos.Internal; using Microsoft.EntityFrameworkCore.Cosmos.Metadata.Internal; -using Microsoft.EntityFrameworkCore.Cosmos.Storage; using Newtonsoft.Json.Linq; namespace Microsoft.EntityFrameworkCore.Cosmos.Query.Internal; @@ -50,6 +49,7 @@ public ReadItemQueryingEnumerable( _queryLogger = _cosmosQueryContext.QueryLogger; _standAloneStateManager = standAloneStateManager; _threadSafetyChecksEnabled = threadSafetyChecksEnabled; + _cosmosContainer = rootEntityType.GetContainer() ?? throw new UnreachableException("Root entity type without a Cosmos container."); _cosmosPartitionKey = GeneratePartitionKey( From 96937a9927a7502b8bfed609fd2145b0d4371162 Mon Sep 17 00:00:00 2001 From: JoasE <32096708+JoasE@users.noreply.github.com> Date: Mon, 10 Nov 2025 19:43:03 +0100 Subject: [PATCH 37/37] Remove debug only test --- .../Storage/Internal/SessionTokenStorageTest.cs | 14 -------------- 1 file changed, 14 deletions(-) diff --git a/test/EFCore.Cosmos.Tests/Storage/Internal/SessionTokenStorageTest.cs b/test/EFCore.Cosmos.Tests/Storage/Internal/SessionTokenStorageTest.cs index fe9487af4a1..1783efeec50 100644 --- a/test/EFCore.Cosmos.Tests/Storage/Internal/SessionTokenStorageTest.cs +++ b/test/EFCore.Cosmos.Tests/Storage/Internal/SessionTokenStorageTest.cs @@ -942,20 +942,6 @@ public virtual void Constructor_AllContainersAreInitialized(SessionTokenManageme Assert.True(tokens.ContainsKey(_otherContainerName)); } - [ConditionalTheory] - [InlineData(SessionTokenManagementMode.SemiAutomatic)] - [InlineData(SessionTokenManagementMode.Manual)] - [InlineData(SessionTokenManagementMode.EnforcedManual)] - public virtual void Constructor_WhenDefaultContainerNotInContainerNames_ThrowsException(SessionTokenManagementMode mode) - { - var containers = new HashSet([_otherContainerName]); - Assert.True(!containers.Contains("default")); - Assert.ThrowsAny(() => - { - _ = new SessionTokenStorage("default", containers, mode); - }); - } - [ConditionalFact] public virtual void Constructor_WhenInitializing_AllContainersStartWithNullTokens() {