Integrate Stored Procedure and Query Plan request routing to Gateway V2 endpoint

sdk/cosmos/azure-cosmos-tests/pom.xml

@@ -930,6 +930,8 @@ Licensed under the MIT License.
               </suiteXmlFiles>
               <systemPropertyVariables>
                 <COSMOS.CLIENT_LEAK_DETECTION_ENABLED>true</COSMOS.CLIENT_LEAK_DETECTION_ENABLED>
+                <COSMOS.THINCLIENT_ENABLED>true</COSMOS.THINCLIENT_ENABLED>
+                <COSMOS.HTTP2_ENABLED>true</COSMOS.HTTP2_ENABLED>
                 <io.netty.leakDetection.samplingInterval>1</io.netty.leakDetection.samplingInterval>
                 <io.netty.leakDetection.targetRecords>256</io.netty.leakDetection.targetRecords>
                 <io.netty.leakDetection.level>paranoid</io.netty.leakDetection.level>

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

@@ -0,0 +1,162 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+package com.azure.cosmos.implementation;
+
+import com.azure.cosmos.implementation.query.DocumentQueryExecutionContextFactory;
+import com.azure.cosmos.implementation.query.IDocumentQueryClient;
+import com.azure.cosmos.models.CosmosQueryRequestOptions;
+import com.azure.cosmos.models.PartitionKeyDefinition;
+import com.azure.cosmos.models.SqlQuerySpec;
+import org.mockito.ArgumentCaptor;
+import org.mockito.Mockito;
+import org.testng.annotations.Test;
+import reactor.core.publisher.Mono;
+
+import java.util.Collections;
+
+import static org.assertj.core.api.Assertions.assertThat;
+
+/**
+ * Unit tests that pin the bifurcation of the validation-only QueryPlan request issued by
+ * {@code readManyByPartitionKeys} when a custom query is supplied.
+ *
+ * <p>The fork lives in {@code QueryPlanRetriever.getQueryPlanThroughGatewayAsync}:
+ * <pre>queryPlanRequest.useGatewayMode = !(queryClient.useThinClient(queryPlanRequest)
+ *     &amp;&amp; Configs.isThinClientQueryPlanEnabled());</pre>
+ *
+ * <p>Routing is driven by the thin-client opt-in ({@link IDocumentQueryClient#useThinClient}),
+ * not by collection metadata. These tests assert both halves of that contract:
+ * <ul>
+ *   <li>thin-client not eligible ⇒ request pinned to Gateway V1 (Compute Gateway);</li>
+ *   <li>thin-client eligible ⇒ request is thin-client eligible (Gateway V2 / proxy) — the
+ *       {@code useGatewayMode=false} flag is the prerequisite the store-model layer reads
+ *       to bifurcate.</li>
+ * </ul>
+ */
+public class ReadManyByPartitionKeyQueryPlanRoutingTest {
+
+    @Test(groups = { "unit" })
+    public void validationQueryPlanRoutesToGatewayV1WhenThinClientNotEligible() {
+        ArgumentCaptor<RxDocumentServiceRequest> requestCaptor =
+            ArgumentCaptor.forClass(RxDocumentServiceRequest.class);
+        IDocumentQueryClient queryClient = mockQueryClient(requestCaptor, /* useThinClient */ false);
+
+        DocumentQueryExecutionContextFactory
+            .fetchQueryPlanForValidation(
+                Mockito.mock(DiagnosticsClientContext.class),
+                queryClient,
+                new SqlQuerySpec("SELECT * FROM c"),
+                "dbs/db/colls/col",
+                new CosmosQueryRequestOptions(),
+                collectionWithPartitionKey(),
+                /* queryPlanCachingEnabled */ false,
+                Collections.emptyMap())
+            .block();
+
+        assertThat(requestCaptor.getAllValues())
+            .as("a single validation query-plan request must be issued")
+            .hasSize(1);
+        assertThat(requestCaptor.getValue().useGatewayMode)
+            .as("validation query-plan must pin to Gateway V1 when thin client is not eligible")
+            .isTrue();
+    }
+
+    @Test(groups = { "unit" })
+    public void validationQueryPlanIsThinClientEligibleWhenOptedIn() {
+        ArgumentCaptor<RxDocumentServiceRequest> requestCaptor =
+            ArgumentCaptor.forClass(RxDocumentServiceRequest.class);
+        IDocumentQueryClient queryClient = mockQueryClient(requestCaptor, /* useThinClient */ true);
+
+        DocumentQueryExecutionContextFactory
+            .fetchQueryPlanForValidation(
+                Mockito.mock(DiagnosticsClientContext.class),
+                queryClient,
+                new SqlQuerySpec("SELECT * FROM c"),
+                "dbs/db/colls/col",
+                new CosmosQueryRequestOptions(),
+                collectionWithPartitionKey(),
+                /* queryPlanCachingEnabled */ false,
+                Collections.emptyMap())
+            .block();
+
+        assertThat(requestCaptor.getAllValues())
+            .as("a single validation query-plan request must be issued")
+            .hasSize(1);
+        assertThat(requestCaptor.getValue().useGatewayMode)
+            .as("validation query-plan must be thin-client eligible (useGatewayMode=false) "
+                + "when the thin-client opt-in is satisfied")
+            .isFalse();
+    }
+
+    @Test(groups = { "unit" })
+    public void validationQueryPlanForcedToGatewayV1WhenKillSwitchDisabled() {
+        // The kill-switch (Configs.isThinClientQueryPlanEnabled(), backed by the
+        // COSMOS.THINCLIENT_QUERY_PLAN_ENABLED system property / COSMOS_THINCLIENT_QUERY_PLAN_ENABLED
+        // env var) must dominate the thin-client opt-in: even when the request is thin-client eligible
+        // (useThinClient=true), disabling the switch pins the validation query-plan back onto Gateway V1
+        // (useGatewayMode=true). This is the inverse of validationQueryPlanIsThinClientEligibleWhenOptedIn.
+        String previous = System.getProperty("COSMOS.THINCLIENT_QUERY_PLAN_ENABLED");
+        System.setProperty("COSMOS.THINCLIENT_QUERY_PLAN_ENABLED", "false");
+        try {
+            ArgumentCaptor<RxDocumentServiceRequest> requestCaptor =
+                ArgumentCaptor.forClass(RxDocumentServiceRequest.class);
+            IDocumentQueryClient queryClient = mockQueryClient(requestCaptor, /* useThinClient */ true);
+
+            DocumentQueryExecutionContextFactory
+                .fetchQueryPlanForValidation(
+                    Mockito.mock(DiagnosticsClientContext.class),
+                    queryClient,
+                    new SqlQuerySpec("SELECT * FROM c"),
+                    "dbs/db/colls/col",
+                    new CosmosQueryRequestOptions(),
+                    collectionWithPartitionKey(),
+                    /* queryPlanCachingEnabled */ false,
+                    Collections.emptyMap())
+                .block();
+
+            assertThat(requestCaptor.getAllValues())
+                .as("a single validation query-plan request must be issued")
+                .hasSize(1);
+            assertThat(requestCaptor.getValue().useGatewayMode)
+                .as("validation query-plan must be forced onto Gateway V1 (useGatewayMode=true) when "
+                    + "the thin-client query-plan kill-switch is disabled, even though the request is "
+                    + "thin-client eligible")
+                .isTrue();
+        } finally {
+            if (previous == null) {
+                System.clearProperty("COSMOS.THINCLIENT_QUERY_PLAN_ENABLED");
+            } else {
+                System.setProperty("COSMOS.THINCLIENT_QUERY_PLAN_ENABLED", previous);
+            }
+        }
+    }
+
+    private static DocumentCollection collectionWithPartitionKey() {
+        PartitionKeyDefinition partitionKeyDefinition = new PartitionKeyDefinition();
+        partitionKeyDefinition.setPaths(Collections.singletonList("/pk"));
+        DocumentCollection collection = new DocumentCollection();
+        collection.setPartitionKey(partitionKeyDefinition);
+        return collection;
+    }
+
+    private static IDocumentQueryClient mockQueryClient(
+        ArgumentCaptor<RxDocumentServiceRequest> requestCaptor,
+        boolean useThinClient) {
+        IDocumentQueryClient queryClient = Mockito.mock(IDocumentQueryClient.class);
+        Mockito.when(queryClient.useThinClient(Mockito.any())).thenReturn(useThinClient);
+        // executeFeedOperationWithAvailabilityStrategy is a generic <T> method; use doAnswer so the
+        // returned Mono.empty() is supplied at invocation time without forcing a generic cast on the
+        // stubbing call site. We only care about the captured request, not the downstream payload.
+        Mockito
+            .doAnswer(invocation -> Mono.empty())
+            .when(queryClient)
+            .executeFeedOperationWithAvailabilityStrategy(
+                Mockito.any(),
+                Mockito.any(),
+                Mockito.any(),
+                requestCaptor.capture(),
+                Mockito.any(),
+                Mockito.anyString());
+        return queryClient;
+    }
+}

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

@@ -15,7 +15,9 @@
 import com.azure.cosmos.implementation.routing.RegionalRoutingContext;
 import io.netty.buffer.ByteBuf;
 import io.netty.buffer.PooledByteBufAllocator;
+import io.netty.buffer.Unpooled;
 import io.netty.channel.ConnectTimeoutException;
+import io.netty.handler.codec.http.HttpResponseStatus;
 import io.netty.handler.timeout.ReadTimeoutException;
 import org.mockito.ArgumentCaptor;
 import org.mockito.Mockito;
@@ -28,6 +30,7 @@
 import java.lang.reflect.Method;
 import java.net.SocketException;
 import java.net.URI;
+import java.nio.charset.StandardCharsets;
 import java.time.Duration;
 import java.util.HashMap;
 import java.util.Map;
@@ -1041,6 +1044,94 @@ private static int getGatewayRetryWithTimeoutInSeconds(RxGatewayStoreModel store
         return (int) getGatewayRetryWithTimeoutInSeconds.invoke(storeModel);
     }
 
+    /**
+     * Error bodies that {@code validateOrThrow} must turn into a CosmosException carrying the
+     * real HTTP status (400 here) rather than leaking the parse failure as statusCode 0.
+     *
+     * Columns: description, rawBody, expectMessageContains (null = don't assert message content).
+     * The Gateway V2 / thin-client proxy can return a non-JSON body (optionally NUL-padded) or a
+     * valid-but-non-object JSON value (scalar / array) for query-plan failures; both must fall back
+     * cleanly. A valid CosmosError JSON body must still parse and preserve its message.
+     */
+    @DataProvider(name = "validateOrThrowErrorBodyProvider")
+    public Object[][] validateOrThrowErrorBodyProvider() {
+        return new Object[][]{
+            {"NUL-padded plain text", "Query plan generation failed\u0000\u0000\u0000\u0000", "Query plan generation failed"},
+            {"JSON scalar string", "\"BadRequest\"", null},
+            {"JSON array", "[1,2,3]", null},
+            {"JSON number", "42", null},
+            {"valid CosmosError JSON", "{\"code\":\"BadRequest\",\"message\":\"bad query syntax\"}", "bad query syntax"},
+            {"empty body", "", null},
+        };
+    }
+
+    /**
+     * Locks the contract of the {@code validateOrThrow} fallback: regardless of the shape of the
+     * error body the proxy returns, the thrown {@link CosmosException} must carry the real HTTP
+     * status code (400) and never the leaked statusCode 0. The non-object-JSON rows (scalar/array/
+     * number) specifically guard the ClassCastException path that previously escaped the catch.
+     */
+    @Test(groups = "unit", dataProvider = "validateOrThrowErrorBodyProvider")
+    public void validateOrThrowPreservesStatusCodeForMalformedBodies(
+        String description,
+        String rawBody,
+        String expectMessageContains) throws Exception {
+
+        RxGatewayStoreModel storeModel = new RxGatewayStoreModel(
+            mockDiagnosticsClientContext(),
+            Mockito.mock(ISessionContainer.class),
+            ConsistencyLevel.SESSION,
+            QueryCompatibilityMode.Default,
+            new UserAgentContainer(),
+            Mockito.mock(GlobalEndpointManager.class),
+            Mockito.mock(HttpClient.class),
+            null,
+            null);
+
+        RxDocumentServiceRequest dsr = RxDocumentServiceRequest.createFromName(
+            mockDiagnosticsClientContext(),
+            OperationType.QueryPlan,
+            "/dbs/db/colls/col",
+            ResourceType.Document);
+        dsr.requestContext = new DocumentServiceRequestContext();
+
+        ByteBuf bodyBuf = Unpooled.wrappedBuffer(rawBody.getBytes(StandardCharsets.UTF_8));
+
+        try {
+            invokeValidateOrThrow(storeModel, dsr, HttpResponseStatus.BAD_REQUEST, new HttpHeaders(), bodyBuf);
+            fail("validateOrThrow should have thrown a CosmosException for: " + description);
+        } catch (java.lang.reflect.InvocationTargetException invocationTargetException) {
+            Throwable cause = invocationTargetException.getCause();
+            assertThat(cause)
+                .as("validateOrThrow should throw CosmosException for: " + description)
+                .isInstanceOf(CosmosException.class);
+            CosmosException cosmosException = (CosmosException) cause;
+            assertThat(cosmosException.getStatusCode())
+                .as("status code must be preserved (never statusCode 0) for: " + description)
+                .isEqualTo(HttpConstants.StatusCodes.BADREQUEST);
+            if (expectMessageContains != null) {
+                assertThat(cosmosException.getMessage())
+                    .as("error message should be preserved for: " + description)
+                    .contains(expectMessageContains);
+            }
+        }
+    }
+
+    private static void invokeValidateOrThrow(RxGatewayStoreModel storeModel,
+                                              RxDocumentServiceRequest request,
+                                              HttpResponseStatus status,
+                                              HttpHeaders headers,
+                                              ByteBuf retainedBodyAsByteBuf) throws Exception {
+        Method validateOrThrow = RxGatewayStoreModel.class.getDeclaredMethod(
+            "validateOrThrow",
+            RxDocumentServiceRequest.class,
+            HttpResponseStatus.class,
+            HttpHeaders.class,
+            ByteBuf.class);
+        validateOrThrow.setAccessible(true);
+        validateOrThrow.invoke(storeModel, request, status, headers, retainedBodyAsByteBuf);
+    }
+
     enum SessionTokenType {
         NONE, // no session token applied
         USER, // userControlled session token