Added debug check moved query plan check to IsMasterOperation

Azure · Aug 11, 2020 · 6c09061 · 6c09061
1 parent 3d85fd1
commit 6c09061
Show file tree

Hide file tree

Showing 2 changed files with 23 additions and 10 deletions.
diff --git a/Microsoft.Azure.Cosmos/src/GatewayStoreModel.cs b/Microsoft.Azure.Cosmos/src/GatewayStoreModel.cs
@@ -6,6 +6,7 @@ namespace Microsoft.Azure.Cosmos
 {
     using System;
     using System.Collections.Generic;
+    using System.Diagnostics;
     using System.Linq;
     using System.Net;
     using System.Net.Http;
@@ -197,6 +198,7 @@ internal static void ApplySessionToken(
         {
             if (request.Headers == null)
             {
+                Debug.Fail("DocumentServiceRequest does not have headers.");
                 return;
             }
 
@@ -216,15 +218,6 @@ internal static void ApplySessionToken(
                 return; // User is explicitly controlling the session.
             }
 
-            if (request.OperationType == OperationType.QueryPlan)
-            {
-                string isPlanOnlyString = request.Headers[HttpConstants.HttpHeaders.IsQueryPlanRequest];
-                if (bool.TryParse(isPlanOnlyString, out bool isPlanOnly) && isPlanOnly)
-                {
-                    return; // for query plan session token is not needed
-                }
-            }
-
             string requestConsistencyLevel = request.Headers[HttpConstants.HttpHeaders.ConsistencyLevel];
 
             bool sessionConsistency =
@@ -256,7 +249,8 @@ internal static bool IsMasterOperation(
             return ReplicatedResourceClient.IsMasterResource(resourceType) ||
                    GatewayStoreModel.IsStoredProcedureCrudOperation(resourceType, operationType) ||
                    resourceType == ResourceType.Trigger ||
-                   resourceType == ResourceType.UserDefinedFunction;
+                   resourceType == ResourceType.UserDefinedFunction ||
+                   operationType == OperationType.QueryPlan;
         }
 
         internal static bool IsStoredProcedureCrudOperation(

diff --git a/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/GatewayStoreModelTest.cs b/Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/GatewayStoreModelTest.cs
@@ -197,6 +197,25 @@ public void TestApplySessionForMasterOperation()
                     Assert.IsNull(dsr.Headers[HttpConstants.HttpHeaders.SessionToken]);
                 }
             }
+
+            Assert.IsTrue(GatewayStoreModel.IsMasterOperation(
+                    ResourceType.Document,
+                    OperationType.QueryPlan));
+
+            DocumentServiceRequest dsrQueryPlan = DocumentServiceRequest.CreateFromName(
+                OperationType.QueryPlan,
+                "Test",
+                ResourceType.Document,
+                AuthorizationTokenType.PrimaryMasterKey);
+
+            dsrQueryPlan.Headers.Add(HttpConstants.HttpHeaders.SessionToken, Guid.NewGuid().ToString());
+
+            GatewayStoreModel.ApplySessionToken(
+                dsrQueryPlan,
+                ConsistencyLevel.Session,
+                new Mock<ISessionContainer>().Object);
+
+            Assert.IsNull(dsrQueryPlan.Headers[HttpConstants.HttpHeaders.SessionToken]);
         }
 
         [TestMethod]