Fixed issue by not passing session token to query plan request (#374)

Azure · Jan 21, 2021 · 97eac92 · 97eac92
1 parent 50bf6a7
commit 97eac92
Show file tree

Hide file tree

Showing 2 changed files with 31 additions and 6 deletions.
diff --git a/gateway/src/main/java/com/microsoft/azure/cosmosdb/rx/internal/RxGatewayStoreModel.java b/gateway/src/main/java/com/microsoft/azure/cosmosdb/rx/internal/RxGatewayStoreModel.java
@@ -537,11 +537,18 @@ private void captureSessionToken(RxDocumentServiceRequest request, Map<String, S
     private void applySessionToken(RxDocumentServiceRequest request) {
         Map<String, String> headers = request.getHeaders();
 
-        if (headers != null &&
-                !Strings.isNullOrEmpty(request.getHeaders().get(HttpConstants.HttpHeaders.SESSION_TOKEN))) {
-            if (ReplicatedResourceClientUtils.isMasterResource(request.getResourceType())) {
+        // Master resource operations don't require session token.
+        if (isMasterOperation(request.getResourceType(), request.getOperationType())) {
+            if (isSessionTokenPresent(request)) {
                 request.getHeaders().remove(HttpConstants.HttpHeaders.SESSION_TOKEN);
             }
+            return;
+        }
+
+        //  NOTE: Words of wisdom: do not combine these two checks
+        //  We want to make sure if session token is not present then we don't add it in case of master operation.
+        //  Which the above check takes care of.
+        if (isSessionTokenPresent(request)) {
             return; //User is explicitly controlling the session.
         }
 
@@ -552,7 +559,7 @@ private void applySessionToken(RxDocumentServiceRequest request) {
                         (!Strings.isNullOrEmpty(requestConsistencyLevel)
                                 && Strings.areEqual(requestConsistencyLevel, ConsistencyLevel.Session.name()));
 
-        if (!sessionConsistency || ReplicatedResourceClientUtils.isMasterResource(request.getResourceType())) {
+        if (!sessionConsistency) {
             return; // Only apply the session token in case of session consistency and when resource is not a master resource
         }
 
@@ -562,5 +569,21 @@ private void applySessionToken(RxDocumentServiceRequest request) {
         if (!Strings.isNullOrEmpty(sessionToken)) {
             headers.put(HttpConstants.HttpHeaders.SESSION_TOKEN, sessionToken);
         }
-    }    
+    }
+
+    private boolean isSessionTokenPresent(RxDocumentServiceRequest request) {
+        return request.getHeaders() != null && !Strings.isNullOrEmpty(request.getHeaders().get(HttpConstants.HttpHeaders.SESSION_TOKEN));
+    }
+
+    private static boolean isMasterOperation(ResourceType resourceType, OperationType operationType) {
+        // Stored procedures CRUD operations are done on master so they do not require the session token.
+        // Stored procedures execute is not a master operation
+        return ReplicatedResourceClientUtils.isMasterResource(resourceType) ||
+            isStoredProcedureMasterOperation(resourceType, operationType) ||
+            operationType == OperationType.QueryPlan;
+    }
+
+    private static boolean isStoredProcedureMasterOperation(ResourceType resourceType, OperationType operationType) {
+        return resourceType == ResourceType.StoredProcedure && operationType != OperationType.ExecuteJavaScript;
+    }
 }
diff --git a/sdk/src/main/java/com/microsoft/azure/cosmosdb/rx/internal/RxDocumentClientImpl.java b/sdk/src/main/java/com/microsoft/azure/cosmosdb/rx/internal/RxDocumentClientImpl.java
@@ -3163,7 +3163,9 @@ private RxStoreModel getStoreProxy(RxDocumentServiceRequest request) {
         } else {
             if ((request.getOperationType() == OperationType.Query || request.getOperationType() == OperationType.SqlQuery) &&
                     Utils.isCollectionChild(request.getResourceType())) {
-                if (request.getPartitionKeyRangeIdentity() == null) {
+                // Go to gateway only when partition key range and partition key are not set. This should be very rare
+                if (request.getPartitionKeyRangeIdentity() == null &&
+                    request.getHeaders().get(HttpConstants.HttpHeaders.PARTITION_KEY) == null) {
                     return this.gatewayProxy;
                 }
             }