diff --git a/Microsoft.Azure.Cosmos/src/GatewayStoreModel.cs b/Microsoft.Azure.Cosmos/src/GatewayStoreModel.cs index daa85a46c2..d5623e9a8e 100644 --- a/Microsoft.Azure.Cosmos/src/GatewayStoreModel.cs +++ b/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 isReadOrBatchRequest = request.IsReadOnlyRequest || request.OperationType == OperationType.Batch; + bool requestHasConsistencySet = !string.IsNullOrEmpty(requestConsistencyLevel) && isReadOrBatchRequest; // 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 + || (!isReadOrBatchRequest && !isMultiMasterEnabledForRequest)) { return; // Only apply the session token in case of session consistency and the request is read only or read/write on multimaster diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/GatewayStoreModelTest.cs b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/GatewayStoreModelTest.cs index d70d9d2501..d832c70cbe 100644 --- a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/GatewayStoreModelTest.cs +++ b/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 sMock = new Mock(); + sMock.Setup(x => x.ResolveGlobalSessionToken(dsrNoSessionToken)).Returns(dsrSessionToken); + + Mock globalEndpointManager = new Mock(); + globalEndpointManager.Setup(gem => gem.CanUseMultipleWriteLocations(It.Is(drs => drs == dsrNoSessionToken))).Returns(multiMaster); + await this.GetGatewayStoreModelForConsistencyTest(async (gatewayStoreModel) => + { + await GatewayStoreModel.ApplySessionTokenAsync( + dsrNoSessionToken, + ConsistencyLevel.Session, + sMock.Object, + partitionKeyRangeCache: new Mock(null, null, null).Object, + clientCollectionCache: new Mock(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() {