diff --git a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/ThinClientE2ETest.java b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/ThinClientE2ETest.java index 5589d147783b..c72909265398 100644 --- a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/ThinClientE2ETest.java +++ b/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 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(); + } + } + } } diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/RxDocumentClientImpl.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/RxDocumentClientImpl.java index 025e6c96e657..81809450ab3a 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/RxDocumentClientImpl.java +++ b/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( diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/RxDocumentServiceRequest.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/RxDocumentServiceRequest.java index 4c71d49fd5e6..92e6f563a3aa 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/RxDocumentServiceRequest.java +++ b/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);