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 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

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()
         {