Skip to content

Commit 3b4ee03

Browse files
JoasEJoasE
authored
Add Cosmos db manual session token management (#37027)
Implements #36504
1 parent d1731e0 commit 3b4ee03

25 files changed

+2998
-100
lines changed

‎src/EFCore.Cosmos/Extensions/CosmosDatabaseFacadeExtensions.cs‎

Lines changed: 73 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -25,6 +25,79 @@ public static class CosmosDatabaseFacadeExtensions
2525
public static CosmosClient GetCosmosClient(this DatabaseFacade databaseFacade)
2626
=> GetService<ISingletonCosmosClientWrapper>(databaseFacade).Client;
2727

28+
/// <summary>
29+
/// Gets the composite session token for the default container for this <see cref="DbContext" />.
30+
/// </summary>
31+
/// <remarks>See https://aka.ms/efcore-docs-cosmos-session for more information.</remarks>
32+
/// <param name="databaseFacade">The <see cref="DatabaseFacade" /> for the context.</param>
33+
/// <returns>The session token for the default container in the context, or <see langword="null"/> if none present.</returns>
34+
public static string? GetSessionToken(this DatabaseFacade databaseFacade)
35+
=> GetSessionTokenStorage(databaseFacade).GetDefaultContainerTrackedToken();
36+
37+
/// <summary>
38+
/// Gets a dictionary that contains the composite session token per container for this <see cref="DbContext" />.
39+
/// </summary>
40+
/// <remarks>See https://aka.ms/efcore-docs-cosmos-session for more information.</remarks>
41+
/// <param name="databaseFacade">The <see cref="DatabaseFacade" /> for the context.</param>
42+
/// <returns>The session token dictionary.</returns>
43+
public static IReadOnlyDictionary<string, string?> GetSessionTokens(this DatabaseFacade databaseFacade)
44+
=> GetSessionTokenStorage(databaseFacade).GetTrackedTokens();
45+
46+
/// <summary>
47+
/// Sets the composite session token for the default container for this <see cref="DbContext" />.
48+
/// </summary>
49+
/// <remarks>See https://aka.ms/efcore-docs-cosmos-session for more information.</remarks>
50+
/// <param name="databaseFacade">The <see cref="DatabaseFacade" /> for the context.</param>
51+
/// <param name="sessionToken">The session token to set.</param>
52+
public static void UseSessionToken(this DatabaseFacade databaseFacade, string sessionToken)
53+
=> GetSessionTokenStorage(databaseFacade).SetDefaultContainerSessionToken(sessionToken);
54+
55+
/// <summary>
56+
/// Sets the composite sessions token per container for this <see cref="DbContext" /> with the tokens specified in <paramref name="sessionTokens"/>.
57+
/// </summary>
58+
/// <remarks>See https://aka.ms/efcore-docs-cosmos-session for more information.</remarks>
59+
/// <param name="databaseFacade">The <see cref="DatabaseFacade" /> for the context.</param>
60+
/// <param name="sessionTokens">The session tokens to set per container.</param>
61+
public static void UseSessionTokens(this DatabaseFacade databaseFacade, IReadOnlyDictionary<string, string?> sessionTokens)
62+
{
63+
var sessionTokenStorage = GetSessionTokenStorage(databaseFacade);
64+
65+
sessionTokenStorage.SetSessionTokens(sessionTokens);
66+
}
67+
68+
/// <summary>
69+
/// Appends the composite session token for the default container for this <see cref="DbContext" />.
70+
/// </summary>
71+
/// <remarks>See https://aka.ms/efcore-docs-cosmos-session for more information.</remarks>
72+
/// <param name="databaseFacade">The <see cref="DatabaseFacade" /> for the context.</param>
73+
/// <param name="sessionToken">The session token to append.</param>
74+
public static void AppendSessionToken(this DatabaseFacade databaseFacade, string sessionToken)
75+
=> GetSessionTokenStorage(databaseFacade).AppendDefaultContainerSessionToken(sessionToken);
76+
77+
/// <summary>
78+
/// Appends the composite sessions token per container for this <see cref="DbContext" /> with the tokens specified in <paramref name="sessionTokens"/>.
79+
/// </summary>
80+
/// <remarks>See https://aka.ms/efcore-docs-cosmos-session for more information.</remarks>
81+
/// <param name="databaseFacade">The <see cref="DatabaseFacade" /> for the context.</param>
82+
/// <param name="sessionTokens">The session tokens to append per container.</param>
83+
public static void AppendSessionTokens(this DatabaseFacade databaseFacade, IReadOnlyDictionary<string, string> sessionTokens)
84+
{
85+
var sessionTokenStorage = GetSessionTokenStorage(databaseFacade);
86+
87+
sessionTokenStorage.AppendSessionTokens(sessionTokens);
88+
}
89+
90+
private static ISessionTokenStorage GetSessionTokenStorage(DatabaseFacade databaseFacade)
91+
{
92+
var db = GetService<IDatabase>(databaseFacade);
93+
if (db is not CosmosDatabaseWrapper dbWrapper)
94+
{
95+
throw new InvalidOperationException(CosmosStrings.CosmosNotInUse);
96+
}
97+
98+
return dbWrapper.SessionTokenStorage;
99+
}
100+
28101
private static TService GetService<TService>(IInfrastructure<IServiceProvider> databaseFacade)
29102
where TService : class
30103
{

‎src/EFCore.Cosmos/Extensions/CosmosServiceCollectionExtensions.cs‎

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -97,6 +97,7 @@ public static IServiceCollection AddEntityFrameworkCosmos(this IServiceCollectio
9797
.TryAdd<LoggingDefinitions, CosmosLoggingDefinitions>()
9898
.TryAdd<IDatabaseProvider, DatabaseProvider<CosmosOptionsExtension>>()
9999
.TryAdd<IDatabase, CosmosDatabaseWrapper>()
100+
.TryAdd<IResettableService, CosmosDatabaseWrapper>(sp => (CosmosDatabaseWrapper)sp.GetRequiredService<IDatabase>())
100101
.TryAdd<IExecutionStrategyFactory, CosmosExecutionStrategyFactory>()
101102
.TryAdd<IDbContextTransactionManager, CosmosTransactionManager>()
102103
.TryAdd<IModelValidator, CosmosModelValidator>()
@@ -121,7 +122,8 @@ public static IServiceCollection AddEntityFrameworkCosmos(this IServiceCollectio
121122
.TryAddScoped<ISqlExpressionFactory, SqlExpressionFactory>()
122123
.TryAddScoped<IMemberTranslatorProvider, CosmosMemberTranslatorProvider>()
123124
.TryAddScoped<IMethodCallTranslatorProvider, CosmosMethodCallTranslatorProvider>()
124-
.TryAddScoped<ICosmosClientWrapper, CosmosClientWrapper>());
125+
.TryAddScoped<ICosmosClientWrapper, CosmosClientWrapper>()
126+
.TryAddSingleton<ISessionTokenStorageFactory, SessionTokenStorageFactory>());
125127

126128
builder.TryAddCoreServices();
127129

‎src/EFCore.Cosmos/Infrastructure/CosmosDbContextOptionsBuilder.cs‎

Lines changed: 18 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -3,6 +3,7 @@
33

44
using System.ComponentModel;
55
using System.Net;
6+
using Microsoft.EntityFrameworkCore.Cosmos.Infrastructure;
67
using Microsoft.EntityFrameworkCore.Cosmos.Infrastructure.Internal;
78

89
namespace Microsoft.EntityFrameworkCore.Infrastructure;
@@ -211,6 +212,23 @@ public virtual CosmosDbContextOptionsBuilder MaxRequestsPerTcpConnection(int req
211212
public virtual CosmosDbContextOptionsBuilder ContentResponseOnWriteEnabled(bool enabled = true)
212213
=> WithOption(e => e.ContentResponseOnWriteEnabled(Check.NotNull(enabled)));
213214

215+
216+
/// <summary>
217+
/// Sets the <see cref="Cosmos.Infrastructure.SessionTokenManagementMode"/> to use.
218+
/// By default, <see cref="SessionTokenManagementMode.FullyAutomatic"/> will be used.
219+
/// Any other mode is only relevant when your application needs to manage session tokens manually.
220+
/// For example: If you're using a round-robin load balancer that doesn't maintain session affinity between requests.
221+
/// Manual session token management can break session consistency when not handled properly.
222+
/// See <see href="https://learn.microsoft.com/en-us/azure/cosmos-db/nosql/how-to-manage-consistency?tabs=portal%2Cdotnetv2%2Capi-async#utilize-session-tokens">Utilize session tokens</see> for more details.
223+
/// </summary>
224+
/// <remarks>
225+
/// See <see href="https://aka.ms/efcore-docs-dbcontext-options">Using DbContextOptions</see>, and
226+
/// <see href="https://aka.ms/efcore-docs-cosmos">Accessing Azure Cosmos DB with EF Core</see> for more information and examples.
227+
/// </remarks>
228+
/// <param name="mode">The <see cref="Cosmos.Infrastructure.SessionTokenManagementMode"/> to use.</param>
229+
public virtual CosmosDbContextOptionsBuilder SessionTokenManagementMode(SessionTokenManagementMode mode)
230+
=> WithOption(e => e.WithSessionTokenManagementMode(mode));
231+
214232
/// <summary>
215233
/// Sets an option by cloning the extension used to store the settings. This ensures the builder
216234
/// does not modify options that are already in use elsewhere.

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

Lines changed: 29 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -36,6 +36,7 @@ public class CosmosOptionsExtension : IDbContextOptionsExtension
3636
private bool? _enableContentResponseOnWrite;
3737
private DbContextOptionsExtensionInfo? _info;
3838
private Func<HttpClient>? _httpClientFactory;
39+
private SessionTokenManagementMode _sessionTokenManagementMode = SessionTokenManagementMode.FullyAutomatic;
3940

4041
/// <summary>
4142
/// This is an internal API that supports the Entity Framework Core infrastructure and not subject to
@@ -73,6 +74,7 @@ protected CosmosOptionsExtension(CosmosOptionsExtension copyFrom)
7374
_maxTcpConnectionsPerEndpoint = copyFrom._maxTcpConnectionsPerEndpoint;
7475
_maxRequestsPerTcpConnection = copyFrom._maxRequestsPerTcpConnection;
7576
_httpClientFactory = copyFrom._httpClientFactory;
77+
_sessionTokenManagementMode = copyFrom._sessionTokenManagementMode;
7678
}
7779

7880
/// <summary>
@@ -564,6 +566,30 @@ public virtual CosmosOptionsExtension WithHttpClientFactory(Func<HttpClient>? ht
564566
return clone;
565567
}
566568

569+
/// <summary>
570+
/// This is an internal API that supports the Entity Framework Core infrastructure and not subject to
571+
/// the same compatibility standards as public APIs. It may be changed or removed without notice in
572+
/// any release. You should only use it directly in your code with extreme caution and knowing that
573+
/// doing so can result in application failures when updating to a new Entity Framework Core release.
574+
/// </summary>
575+
public virtual SessionTokenManagementMode SessionTokenManagementMode
576+
=> _sessionTokenManagementMode;
577+
578+
/// <summary>
579+
/// This is an internal API that supports the Entity Framework Core infrastructure and not subject to
580+
/// the same compatibility standards as public APIs. It may be changed or removed without notice in
581+
/// any release. You should only use it directly in your code with extreme caution and knowing that
582+
/// doing so can result in application failures when updating to a new Entity Framework Core release.
583+
/// </summary>
584+
public virtual CosmosOptionsExtension WithSessionTokenManagementMode(SessionTokenManagementMode mode)
585+
{
586+
var clone = Clone();
587+
588+
clone._sessionTokenManagementMode = mode;
589+
590+
return clone;
591+
}
592+
567593
/// <summary>
568594
/// This is an internal API that supports the Entity Framework Core infrastructure and not subject to
569595
/// the same compatibility standards as public APIs. It may be changed or removed without notice in
@@ -632,6 +658,7 @@ public override int GetServiceProviderHashCode()
632658
hashCode.Add(Extension._maxTcpConnectionsPerEndpoint);
633659
hashCode.Add(Extension._maxRequestsPerTcpConnection);
634660
hashCode.Add(Extension._httpClientFactory);
661+
hashCode.Add(Extension._sessionTokenManagementMode);
635662

636663
_serviceProviderHash = hashCode.ToHashCode();
637664
}
@@ -656,7 +683,8 @@ public override bool ShouldUseSameServiceProvider(DbContextOptionsExtensionInfo
656683
&& Extension._gatewayModeMaxConnectionLimit == otherInfo.Extension._gatewayModeMaxConnectionLimit
657684
&& Extension._maxTcpConnectionsPerEndpoint == otherInfo.Extension._maxTcpConnectionsPerEndpoint
658685
&& Extension._maxRequestsPerTcpConnection == otherInfo.Extension._maxRequestsPerTcpConnection
659-
&& Extension._httpClientFactory == otherInfo.Extension._httpClientFactory;
686+
&& Extension._httpClientFactory == otherInfo.Extension._httpClientFactory
687+
&& Extension._sessionTokenManagementMode == otherInfo.Extension._sessionTokenManagementMode;
660688

661689
public override void PopulateDebugInfo(IDictionary<string, string> debugInfo)
662690
{

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

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -151,6 +151,14 @@ public class CosmosSingletonOptions : ICosmosSingletonOptions
151151
/// </summary>
152152
public virtual Func<HttpClient>? HttpClientFactory { get; private set; }
153153

154+
/// <summary>
155+
/// This is an internal API that supports the Entity Framework Core infrastructure and not subject to
156+
/// the same compatibility standards as public APIs. It may be changed or removed without notice in
157+
/// any release. You should only use it directly in your code with extreme caution and knowing that
158+
/// doing so can result in application failures when updating to a new Entity Framework Core release.
159+
/// </summary>
160+
public virtual SessionTokenManagementMode SessionTokenManagementMode { get; private set; }
161+
154162
/// <summary>
155163
/// This is an internal API that supports the Entity Framework Core infrastructure and not subject to
156164
/// the same compatibility standards as public APIs. It may be changed or removed without notice in

‎src/EFCore.Cosmos/Infrastructure/SessionTokenManagementMode.cs‎

Lines changed: 39 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,39 @@
1+
// Licensed to the .NET Foundation under one or more agreements.
2+
// The .NET Foundation licenses this file to you under the MIT license.
3+
4+
namespace Microsoft.EntityFrameworkCore.Cosmos.Infrastructure;
5+
6+
/// <summary>
7+
/// Defines the behavior of EF Core regarding the management of Cosmos DB session tokens.
8+
/// </summary>
9+
/// <remarks>
10+
/// See <see href="https://aka.ms/efcore-docs-cosmos-session">Cosmos session consistency</see> for more info.
11+
/// </remarks>
12+
public enum SessionTokenManagementMode
13+
{
14+
/// <summary>
15+
/// The default mode.
16+
/// Uses the underlying Cosmos DB SDK automatic session token management.
17+
/// EF will not track or parse session tokens returned from Cosmos DB. <see cref="CosmosDatabaseFacadeExtensions.UseSessionTokens(DatabaseFacade, IReadOnlyDictionary{string, string?})"/> and <see cref="CosmosDatabaseFacadeExtensions.GetSessionTokens(DatabaseFacade)"/> methods will throw when invoked.
18+
/// </summary>
19+
FullyAutomatic,
20+
21+
/// <summary>
22+
/// Allows the usage of <see cref="CosmosDatabaseFacadeExtensions.UseSessionTokens(DatabaseFacade, IReadOnlyDictionary{string, string?})"/> to overwrite the default Cosmos DB SDK automatic session token management.
23+
/// If 'UseSessionTokens' has not been invoked for a container, the default Cosmos DB SDK automatic session token management will be used.
24+
/// EF will track and parse session tokens returned from Cosmos DB, which can be retrieved via <see cref="CosmosDatabaseFacadeExtensions.GetSessionTokens(DatabaseFacade)"/>.
25+
/// </summary>
26+
SemiAutomatic,
27+
28+
/// <summary>
29+
/// Fully overwrites the Cosmos DB SDK automatic session token management, and only uses session tokens specified via <see cref="CosmosDatabaseFacadeExtensions.UseSessionTokens(DatabaseFacade, IReadOnlyDictionary{string, string?})"/>.
30+
/// If 'UseSessionTokens' has not been invoked for a container, no session token will be used.
31+
/// EF will track and parse session tokens returned from Cosmos DB, which can be retrieved via <see cref="CosmosDatabaseFacadeExtensions.GetSessionTokens(DatabaseFacade)"/>.
32+
/// </summary>
33+
Manual,
34+
35+
/// <summary>
36+
/// Same as <see cref="Manual"/>, but will throw an exception if <see cref="CosmosDatabaseFacadeExtensions.UseSessionTokens(DatabaseFacade, IReadOnlyDictionary{string, string?})"/> was not invoked before executing a read.
37+
/// </summary>
38+
EnforcedManual
39+
}

‎src/EFCore.Cosmos/Properties/CosmosStrings.Designer.cs‎

Lines changed: 22 additions & 0 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

‎src/EFCore.Cosmos/Properties/CosmosStrings.resx‎

Lines changed: 11 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -147,6 +147,10 @@
147147
<data name="ContainerContainingPropertyConflict" xml:space="preserve">
148148
<value>The entity type '{entityType}' is mapped to the container '{container}' but it is also configured as being contained in property '{property}'.</value>
149149
</data>
150+
<data name="ContainerNameDoesNotExist" xml:space="preserve">
151+
<value>The container with the name '{containerName}' does not exist.</value>
152+
<comment>string</comment>
153+
</data>
150154
<data name="ContainerNotOnRoot" xml:space="preserve">
151155
<value>An Azure Cosmos DB container name is defined on entity type '{entityType}', which inherits from '{baseEntityType}'. Container names must be defined on the root entity type of a hierarchy.</value>
152156
</data>
@@ -171,6 +175,9 @@
171175
<data name="ElementWithValueConverter" xml:space="preserve">
172176
<value>The property '{propertyType} {structuralType}.{property}' has element type '{elementType}', which requires a value converter. Elements types requiring value converters are not currently supported with the Azure Cosmos DB database provider.</value>
173177
</data>
178+
<data name="EnableManualSessionTokenManagement" xml:space="preserve">
179+
<value>Disable automatic session token management using 'options.SessionTokenManagementMode' to use this method.</value>
180+
</data>
174181
<data name="ETagNonStringStoreType" xml:space="preserve">
175182
<value>The type of the etag property '{property}' on '{entityType}' is '{propertyType}'. All etag properties must be strings or have a string value converter.</value>
176183
</data>
@@ -253,6 +260,10 @@
253260
<data name="MissingOrderingInSelectExpression" xml:space="preserve">
254261
<value>'Reverse' could not be translated to the server because there is no ordering on the server side.</value>
255262
</data>
263+
<data name="MissingSessionTokenEnforceManual" xml:space="preserve">
264+
<value>No session token has been set for container: '{container}'. While using 'EnforceManual' mode you must always set a session token for any container used on every context instance.</value>
265+
<comment>string</comment>
266+
</data>
256267
<data name="MultipleRootEntityTypesReferencedInQuery" xml:space="preserve">
257268
<value>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.</value>
258269
</data>

0 commit comments

Comments
 (0)