Session: Fixes token being sent if request had a different consistency level on Gateway mode

Microsoft.Azure.Cosmos/src/GatewayStoreModel.cs

@@ -265,17 +265,18 @@ internal static async Task ApplySessionTokenAsync(
  }
 
  string requestConsistencyLevel = request.Headers[HttpConstants.HttpHeaders.ConsistencyLevel];
-
- bool sessionConsistency =
- defaultConsistencyLevel == ConsistencyLevel.Session ||
- (!string.IsNullOrEmpty(requestConsistencyLevel)
+ bool isConsideredAReadRequest = request.IsReadOnlyRequest || request.OperationType == OperationType.Batch;
+ bool requestHasConsistencySet = !string.IsNullOrEmpty(requestConsistencyLevel) && isConsideredAReadRequest; // Only read requests can have their consistency modified
+
+ bool sessionConsistencyApplies =
+ (!requestHasConsistencySet && defaultConsistencyLevel == ConsistencyLevel.Session) ||
+ (requestHasConsistencySet
  && string.Equals(requestConsistencyLevel, GatewayStoreModel.sessionConsistencyAsString, StringComparison.OrdinalIgnoreCase));
 
  bool isMultiMasterEnabledForRequest = globalEndpointManager.CanUseMultipleWriteLocations(request);
 
- if (!sessionConsistency 
- || (!request.IsReadOnlyRequest 
- && request.OperationType != OperationType.Batch
+ if (!sessionConsistencyApplies
+ || (!isConsideredAReadRequest
  && !isMultiMasterEnabledForRequest))
  {
  return; // Only apply the session token in case of session consistency and the request is read only or read/write on multimaster

Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/GatewayStoreModelTest.cs

@@ -414,6 +414,51 @@ await GatewayStoreModel.ApplySessionTokenAsync(
  });
  }
 
+ [DataTestMethod]
+ [DataRow(false, false, DisplayName = "Single master - Read")]
+ [DataRow(true, false, DisplayName = "Multi master - Read")]
+ [DataRow(false, true, DisplayName = "Single master - Write")]
+ [DataRow(true, true, DisplayName = "Multi master - Write")]
+ public async Task TestRequestOverloadRemovesSessionToken(bool multiMaster, bool isWriteRequest)
+ {
+ INameValueCollection headers = new StoreRequestNameValueCollection();
+ headers.Set(HttpConstants.HttpHeaders.ConsistencyLevel, ConsistencyLevel.Eventual.ToString());
+
+ DocumentServiceRequest dsrNoSessionToken = DocumentServiceRequest.CreateFromName(
+ isWriteRequest ? OperationType.Create : OperationType.Read,
+ "Test",
+ ResourceType.Document,
+ AuthorizationTokenType.PrimaryMasterKey,
+ headers);
+
+ string dsrSessionToken = Guid.NewGuid().ToString();
+ Mock<ISessionContainer> sMock = new Mock<ISessionContainer>();
+ sMock.Setup(x => x.ResolveGlobalSessionToken(dsrNoSessionToken)).Returns(dsrSessionToken);
+
+ Mock<IGlobalEndpointManager> globalEndpointManager = new Mock<IGlobalEndpointManager>();
+ globalEndpointManager.Setup(gem => gem.CanUseMultipleWriteLocations(It.Is<DocumentServiceRequest>(drs => drs == dsrNoSessionToken))).Returns(multiMaster);
+ await this.GetGatewayStoreModelForConsistencyTest(async (gatewayStoreModel) =>
+ {
+ await GatewayStoreModel.ApplySessionTokenAsync(
+ dsrNoSessionToken,
+ ConsistencyLevel.Session,
+ sMock.Object,
+ partitionKeyRangeCache: new Mock<PartitionKeyRangeCache>(null, null, null).Object,
+ clientCollectionCache: new Mock<ClientCollectionCache>(new SessionContainer("testhost"), gatewayStoreModel, null, null).Object,
+ globalEndpointManager: globalEndpointManager.Object);
+
+ if (isWriteRequest && multiMaster)
+ {
+ // Multi master write requests should not lower the consistency and remove the session token
+ Assert.AreEqual(dsrSessionToken, dsrNoSessionToken.Headers[HttpConstants.HttpHeaders.SessionToken]);
+ }
+ else
+ {
+ Assert.IsNull(dsrNoSessionToken.Headers[HttpConstants.HttpHeaders.SessionToken]);
+ }
+ });
+ }
+
  [TestMethod]
  public async Task TestErrorResponsesProvideBody()
  {