Cosmos db manual session token management

[JoasE]

What problem are you trying to solve?

When working with azure cosmos db in the context of a user, the session token is required for read your own write guarantees (or session consistency) to work.
https://learn.microsoft.com/en-us/azure/cosmos-db/consistency-levels#session-consistency
I could not find any documentation on setting or retrieving the session token in a dbcontext.

Describe the solution you'd like

Possibly you could define extension methods (or extension members when those are released) on the DbContext type to get and set the session token on the DbContext before any reads or writes.

myDbContext.SetSessionToken(sessionToken);
var result = await myDbContext.Documents.Where(x => x.Id == id).FirstOrDefaultAsync();

result.Property = "changes";
await myDbContext.SaveChangesAsync();
sessionToken = myDbContext.GetSessionToken();

[JoasE]

I would like to work on this if it doesn't turn out too compex.
Depending on the outcome of #17308 we might need a different api to set and retrieve the session token.
Session tokens are unique for each 'write action' in a collection and partition key. This write action could be an CreateItem/ReplaceItem or it could be a TransactionalBatch. The original suggested api would not make sense as the collection and partition key might differ per entity or entity type, and currently (without transactional batch) it will always be different per entity, as each change will be a unique CreateItem/ReplaceItem call.

We could add an extension method to IQueryable to set the session token and to EntityEntry to retrieve the session token.
But being able to set the session token for all queries executed on the context instance would also be nice.

var result = await myDbContext.Documents.WithSessionToken(sessionToken).Where(x => x.Id == id).FirstOrDefaultAsync();

result.Property = "changes";
await myDbContext.SaveChangesAsync();
sessionToken = myDbContext.Entry(result).GetSessionToken();

myDbContext.SetSessionTokenFor(x => x.Documents, sessionToken);
myDbContext.SetSessionTokenFor<Document>(sessionToken);
var result = await myDbContext.Documents.Where(x => x.Id == id).FirstOrDefaultAsync(); // WithSessionToken not needed

However, if the team decides to implement #17308 (TransactionalBatch) all entries in the same batch will have the same session token. An api that reflects that could also make sense

var result = await myDbContext.Documents.WithSessionToken(sessionToken).Where(x => x.Id == id).FirstOrDefaultAsync();

result.Property = "changes";
await myDbContext.SaveChangesAsync();
sessionToken = myDbContext.GetSessionTokenFor(x => x.Documents, "partitionKeyValue"?); // partition key value optional?
sessionTokens = myDbcontext.GetSessionTokens(); // Dictionary<object (partitionKey), string> ?

Input from the team is greatly appreciated :)

[AndriySvyryd]

Managing session tokens manually is a pretty advanced scenario and managing multiple session tokens per unit-of-work would be even more of a corner case, that would complicate the API beyond the value it would provide.
I think that a reasonable compromise is to add API similar to the proposed in the first comment:

var sessionToken = dbContext.Database.GetSessionToken();
dbContext.Database.SetSessionToken(sessionToken);

GetSessionToken() would throw if more than one session token has been returned since this context instance was created or leased from the pool or SetSessionToken(null) was called. This should work nicely with the atomicity exceptions enabled from #17308

[JoasE]

@AndriySvyryd I see what you’re saying, but in my experience managing session tokens manually is usually a common requirement rather than an “advanced” feature. To my understanding, horizontally scaled web applications that don’t use session affinity generally need some form of session token management to guarantee read-your-own-writes. Also, unless all documents in a given context are partitioned the same and stored in the same container, users might end up with multiple session tokens within a unit-of-work. Using AutoTransactionBehavior.Always ensures the former, though it’s not the default.

[AndriySvyryd]

Also, unless all documents in a given context are partitioned the same and stored in the same container, users might end up with multiple session tokens within a unit-of-work.

Even so, we would need to create something that's more user-friendly. For example, the object returned by dbContext.Database.GetSessionToken() could be something like a ConcurrentDictionary<Grouping, string>. That way, the user doesn't need to handle this for each collection.

[roji]

Just for my understanding, in which scenarios is it important for user code to be interacting with session tokens via the APIs proposed above? At least in an ideal world, shouldn't EF be managing session tokens internally, automatically providing read-your-writes behavior without the user having to manually read and set session tokens? In other words, I'd imagine a more high-level behavior where there's an EF enum setting managing the consistency level, with EF simply applying it under the hood.

Note that I'm not arguing against possibly also introducing a low-level API to manually manage session tokens (as proposed above); that might be necessary for really advanced/rare scenarios where users need to extract the session token in order to use it outside EF. But at least in the very common case, it doesn't feel like users should need to concern themselves with this and manually write such low-level code just to get read-your-writes behavior.

[JoasE]

Even so, we would need to create something that's more user-friendly. For example, the object returned by dbContext.Database.GetSessionToken() could be something like a ConcurrentDictionary<Grouping, string>. That way, the user doesn't need to handle this for each collection.

@AndriySvyryd That sounds like a better, simpler solution indeed! Would it be ok to expose the cosmos db sdk type PartitionKey in the grouping? I think in that case we would also need something like dbContext.Database.SetSessionToken(Grouping, string) to set the session token. And what are your thoughts on IQueryable.WithSessionToken(string) extension method? We could also implement both dbContext.Database.GetSessionToken**s**() and dbContext.Database.GetSessionToke**n**() in which the latter would throw if there are multiple session tokens available for the last SaveChanges.

Just for my understanding, in which scenarios is it important for user code to be interacting with session tokens via the APIs proposed above? At least in an ideal world, shouldn't EF be managing session tokens internally, automatically providing read-your-writes behavior without the user having to manually read and set session tokens?

@roji I understand! And actually, the azure cosmos db sdk does this for us in-memory already! However, since it is only in-memory, a distributed or horizontally scaled application that is not using session-affinity will have to implement some form of manual session token management to guarantee read-your-own-writes.
See: https://learn.microsoft.com/en-us/azure/cosmos-db/nosql/how-to-manage-consistency?tabs=portal%2Cdotnetv2%2Capi-async#utilize-session-tokens for more info

[roji]

However, since it is only in-memory, a distributed or horizontally scaled application that is not using session-affinity will have to implement some form of manual session token management to guarantee read-your-own-writes.

I don't follow (or you might have misunderstood my suggestion)... I understand that when interacting directly with the Cosmos SDK, one must manually manage session tokens. My question is why we need to surface that low-level management to EF Core users, rather than letting EF manage that for them. In other words, unless there's some specific technical blocker, I'd expect the EF provider would internally get the session tokens from the Cosmos SDK and set them back, providing users with a seamless read-what-you-wrote experience without them having to know anything about session tokens (of course, it might make sense to make this an opt-in behavior e.g. for performance reasons).

[JoasE]

My question is why we need to surface that low-level management to EF Core users, rather than letting EF manage that for them.

Sorry for the misunderstanding! @roji
In my experience, session token management is often implemented using cookies. This is also suggested in the [docs I linked earlier] (https://learn.microsoft.com/en-us/azure/cosmos-db/nosql/how-to-manage-consistency?tabs=portal%2Cdotnetv2%2Capi-async#utilize-session-tokens). This seems to me like a technical blocker, as we don't know which presentation framework/logic the user is using outside EF. Unless there is a EntityFrameworkCore.Aspnet integration package I am not aware of, we can not implement it for the user. And even if we would choose to publish such a package, it would need to leverage some of the suggested api's.
However, there are other ways a user might choose to implement the manual session token management (without session affinity). One option is using a distributed cache to store stateful sessions, but you would also need some cookie or token to be sent with the request to identify the session of the user, just like the manual cosmos session token management. I assumed (maybe incorrectly) that that would be beyond the scope of EF Core. We can not use cosmos as this distributed cache as we would need read-your-own-write guarantees in order for it to function.

[JoasE]

@roji In other words, to my understanding, in a distributed web application the session token retrieved from Cosmos DB must be managed (indirectly) by the user’s client code.