[Internal / Gateway V 2_0 / DO NOT MERGE]: Route Execute Stored Procedure requests to Gateway V2 endpoint.

sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/ThinClientE2ETest.java

@@ -27,6 +27,9 @@
 import com.azure.cosmos.models.CosmosItemResponse;
 import com.azure.cosmos.models.CosmosItemRequestOptions;
 import com.azure.cosmos.models.CosmosPatchOperations;
+import com.azure.cosmos.models.CosmosStoredProcedureProperties;
+import com.azure.cosmos.models.CosmosStoredProcedureRequestOptions;
+import com.azure.cosmos.models.CosmosStoredProcedureResponse;
 import com.fasterxml.jackson.databind.ObjectMapper;
 import com.fasterxml.jackson.databind.node.ObjectNode;
 import org.slf4j.Logger;
@@ -375,4 +378,95 @@ public void testThinClientDocumentPointOperations() {
             }
         }
     }
+
+    @Test(groups = {"thinclient"}, retryAnalyzer = FlakyTestRetryAnalyzer.class)
+    public void testThinClientStoredProcedure() {
+        CosmosAsyncClient client = null;
+        try {
+            // If running locally, uncomment these lines
+             System.setProperty("COSMOS.THINCLIENT_ENABLED", "true");
+             System.setProperty("COSMOS.HTTP2_ENABLED", "true");
+
+            client = new CosmosClientBuilder()
+                .endpoint(TestConfigurations.HOST)
+                .key(TestConfigurations.MASTER_KEY)
+                .gatewayMode()
+                .consistencyLevel(ConsistencyLevel.SESSION)
+                .buildAsyncClient();
+
+            String idName = "id";
+            String partitionKeyName = "partitionKey";
+
+            client.createDatabaseIfNotExists("db1").block();
+
+            CosmosContainerProperties containerDef =
+                new CosmosContainerProperties("c2", "/" + partitionKeyName);
+            ThroughputProperties ruCfg = ThroughputProperties.createManualThroughput(35_000);
+
+            client.getDatabase("db1").createContainerIfNotExists(containerDef, ruCfg).block();
+
+            CosmosAsyncContainer container = client.getDatabase("db1").getContainer("c2");
+
+            // Create a stored procedure that creates a document
+            String sprocId = "createDocSproc_" + UUID.randomUUID().toString();
+            String pkValue = UUID.randomUUID().toString();
+            CosmosStoredProcedureProperties storedProcedureDef = new CosmosStoredProcedureProperties(
+                sprocId,
+                "function createDocument(docToCreate) {" +
+                    "    var context = getContext();" +
+                    "    var container = context.getCollection();" +
+                    "    var response = context.getResponse();" +
+                    "    var accepted = container.createDocument(" +
+                    "        container.getSelfLink()," +
+                    "        docToCreate," +
+                    "        function(err, docCreated) {" +
+                    "            if (err) throw new Error('Error creating document: ' + err.message);" +
+                    "            response.setBody(docCreated);" +
+                    "        }" +
+                    "    );" +
+                    "    if (!accepted) throw new Error('Document creation was not accepted');" +
+                    "}"
+            );
+
+            // Create stored procedure
+            CosmosStoredProcedureResponse createResponse = container.getScripts()
+                .createStoredProcedure(storedProcedureDef)
+                .block();
+            assertThat(createResponse).isNotNull();
+            assertThat(createResponse.getStatusCode()).isEqualTo(201);
+
+            // Execute stored procedure with a specific partition key to create a document
+            CosmosStoredProcedureRequestOptions options = new CosmosStoredProcedureRequestOptions();
+            options.setPartitionKey(new PartitionKey(pkValue));
+
+            String docId = UUID.randomUUID().toString();
+            String docToCreate = String.format("{\"%s\": \"%s\", \"%s\": \"%s\"}", idName, docId, partitionKeyName, pkValue);
+
+            CosmosStoredProcedureResponse executeResponse = container.getScripts()
+                .getStoredProcedure(sprocId)
+                .execute(Arrays.asList(docToCreate), options)
+                .block();
+
+            assertThat(executeResponse).isNotNull();
+            assertThat(executeResponse.getStatusCode()).isEqualTo(200);
+            assertThat(executeResponse.getRequestCharge()).isGreaterThan(0.0);
+            assertThinClientEndpointUsed(executeResponse.getDiagnostics());
+
+            // Verify the document was created by reading it
+            CosmosItemResponse<ObjectNode> readResponse = container.readItem(docId, new PartitionKey(pkValue), ObjectNode.class).block();
+            assertThat(readResponse).isNotNull();
+            assertThat(readResponse.getStatusCode()).isEqualTo(200);
+            assertThat(readResponse.getItem().get(idName).asText()).isEqualTo(docId);
+            assertThat(readResponse.getItem().get(partitionKeyName).asText()).isEqualTo(pkValue);
+
+            // Clean up - delete the created document and stored procedure
+            container.deleteItem(docId, new PartitionKey(pkValue)).block();
+            container.getScripts().getStoredProcedure(sprocId).delete().block();
+
+        } finally {
+            if (client != null) {
+                client.close();
+            }
+        }
+    }
 }

sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/RxDocumentClientImpl.java

@@ -8009,7 +8009,7 @@ public boolean useThinClient() {
     private boolean useThinClientStoreModel(RxDocumentServiceRequest request) {
         if (!useThinClient
             || !this.globalEndpointManager.hasThinClientReadLocations()
-            || request.getResourceType() != ResourceType.Document) {
+            || request.getResourceType() != ResourceType.Document && !request.isExecuteStoredProcedureBasedRequest()) {
 
             return false;
         }
@@ -8019,7 +8019,8 @@ private boolean useThinClientStoreModel(RxDocumentServiceRequest request) {
         return operationType.isPointOperation()
                     || operationType == OperationType.Query
                     || operationType == OperationType.Batch
-                    || request.isChangeFeedRequest() && !request.isAllVersionsAndDeletesChangeFeedMode();
+                    || request.isChangeFeedRequest() && !request.isAllVersionsAndDeletesChangeFeedMode()
+                    || request.isExecuteStoredProcedureBasedRequest();
     }
 
     private DocumentClientRetryPolicy getRetryPolicyForPointOperation(

sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/RxDocumentServiceRequest.java

@@ -955,6 +955,10 @@ public boolean isChangeFeedRequest() {
         return this.headers.containsKey(HttpConstants.HttpHeaders.A_IM);
     }
 
+    public boolean isExecuteStoredProcedureBasedRequest() {
+        return this.resourceType == ResourceType.StoredProcedure && this.operationType == OperationType.ExecuteJavaScript;
+    }
+
     public boolean isAllVersionsAndDeletesChangeFeedMode() {
         String aImHeader = this.headers.get(HttpConstants.HttpHeaders.A_IM);
         return this.headers.containsKey(HttpConstants.HttpHeaders.A_IM) && HttpConstants.A_IMHeaderValues.FULL_FIDELITY_FEED.equals(aImHeader);