[Gateway V2][DO NOT MERGE]: Integrate Stored Procedure and Query Plan request routing to a Gateway V2 endpoint. #47759

jeet1995 · 2026-01-21T01:43:58Z

PR Description

Title

Add Query Plan and Execute Stored Procedure Support for Thin Client Mode in Azure Cosmos DB Java SDK

Summary

This PR adds support for routing QueryPlan and ExecuteJavaScript (stored procedure execution) operations through the thin client proxy endpoint, enabling these operations to use the same optimized pathway as other thin client operations. Additionally, it includes comprehensive end-to-end tests for thin client query and stored procedure functionality.

Changes

Core Changes

Stored Procedure Execution Support

RxDocumentServiceRequest.java:
- Added isExecuteStoredProcedureBasedRequest() method to identify stored procedure execution requests (ResourceType.StoredProcedure + OperationType.ExecuteJavaScript)
RxDocumentClientImpl.java:
- Updated useThinClientStoreModel() to allow stored procedure execution requests through thin client:
  - Added check for request.isExecuteStoredProcedureBasedRequest() in resource type validation
  - Added request.isExecuteStoredProcedureBasedRequest() to supported operations list

Query Plan Support

RxDocumentClientImpl.java:
- Added useThinClient(RxDocumentServiceRequest) method to IDocumentQueryClient interface implementation
- Modified getStoreModel() to route QueryPlan operations to gateway proxy when not using thin client
- Added OperationType.QueryPlan to supported thin client operations
ThinClientStoreModel.java:
- Added conditional handling for QueryPlan operations to allow null resolved partition key range (since query plans don't require partition resolution)
QueryPlanRetriever.java:
- Updated getQueryPlanThroughGatewayAsync() to accept DocumentCollection parameter
- Added logic to detect thin client mode via queryClient.useThinClient()
- Pass PartitionKeyDefinition to PartitionedQueryExecutionInfo for EPK conversion
PartitionedQueryExecutionInfo.java:
- Added useThinClientMode and partitionKeyDefinition fields
- Implemented parseQueryRangesForThinClient() method to manually parse query ranges from thin client proxy response format
- Added parsePartitionKeyInternal() to convert JSON representation of partition keys to PartitionKeyInternal objects
DocumentQueryExecutionContextFactory.java:
- Updated call to QueryPlanRetriever.getQueryPlanThroughGatewayAsync() to pass the collection parameter
IDocumentQueryClient.java:
- Added useThinClient(RxDocumentServiceRequest) method to the interface

Test Coverage

ThinClientE2ETest.java:

Added shared client/database/container setup with @BeforeClass/@AfterClass lifecycle methods
Container provisioned with 35,000 RU/s manual throughput

Basic Query Tests:

testThinClientQuerySelectAll() - Tests SELECT * FROM c query with full pagination draining
testThinClientQuerySelectById() - Tests SELECT * FROM c WHERE c.id = @id query
testThinClientQueryWithPartitionKeyOption() - Tests SELECT * FROM c with partition key in CosmosQueryRequestOptions
testThinClientQueryLegacy() - Legacy query test with parameterized query

Query Plan Feature Tests:

Test	Query	Expected Behavior
`testThinClientQueryPlanOrderBy()`	`SELECT * FROM c ORDER BY c.sortField`	hasOrderBy=true, results ordered
`testThinClientQueryPlanAggregate()`	`SELECT VALUE COUNT(1) FROM c`	hasAggregates=true, correct count
`testThinClientQueryPlanWithPartitionKeyFilterSingleRange()`	`SELECT * FROM c WHERE c.id = @id`	Single narrow range (not full range)
`testThinClientQueryPlanDistinct()`	`SELECT DISTINCT VALUE c.category FROM c`	hasDistinct=true, unique values
`testThinClientQueryPlanTop()`	`SELECT TOP 10 * FROM c`	hasTop=true, top=10
`testThinClientQueryPlanGroupBy()`	`SELECT c.category, COUNT(1) FROM c GROUP BY c.category`	hasGroupBy=true, hasAggregates=true
`testThinClientQueryPlanInvalidQuery()`	`SELEC * FORM c`	400 BadRequest error
`testThinClientQueryPlanOffsetLimit()`	`SELECT * FROM c ORDER BY c.idx OFFSET 5 LIMIT 5`	hasOffset=true, hasLimit=true

Stored Procedure Test:

testThinClientStoredProcedure() - Tests sproc creation and execution (creates a document via sproc)
All tests assert that only thin-client endpoints are used

Technical Details: QueryPlan Response Format Differences

The thin client proxy and standard gateway return query plan responses in different formats for the queryRanges field. This PR adds parsing logic to handle the thin client format.

Thin Client Proxy Response

Returns partition key values in raw PartitionKeyInternal JSON array format:

{
  "queryInfo": {
    "distinctType": "None",
    "groupByExpressions": [],
    "groupByAliases": [],
    "orderBy": [],
    "orderByExpressions": [],
    "aggregates": [],
    "hasSelectValue": 0,
    "rewrittenQuery": "",
    "groupByAliasToAggregateType": {},
    "hasNonStreamingOrderBy": 0
  },
  "queryRanges": [{
    "min": ["b3a69415-c069-478a-a553-1a3340b581f0"],
    "max": ["b3a69415-c069-478a-a553-1a3340b581f0"],
    "isMinInclusive": true,
    "isMaxInclusive": true
  }]
}

Key difference: min and max are JSON arrays representing PartitionKeyInternal values (e.g., ["b3a69415-c069-478a-a553-1a3340b581f0"]).

Standard Gateway Response

Returns pre-computed EPK (Effective Partition Key) hex strings:

{
  "partitionedQueryExecutionInfoVersion": 2,
  "queryInfo": {
    "distinctType": "None",
    "top": null,
    "offset": null,
    "limit": null,
    "orderBy": [],
    "orderByExpressions": [],
    "groupByExpressions": [],
    "groupByAliases": [],
    "aggregates": [],
    "groupByAliasToAggregateType": {},
    "rewrittenQuery": "",
    "hasSelectValue": false,
    "dCountInfo": null,
    "hasNonStreamingOrderBy": false
  },
  "queryRanges": [{
    "min": "24F3B4E0560B48E552D6CCDAA7837C94",
    "max": "24F3B4E0560B48E552D6CCDAA7837C94",
    "isMinInclusive": true,
    "isMaxInclusive": true
  }],
  "hybridSearchQueryInfo": null
}

Key difference: min and max are hex strings representing the computed EPK hash values.

Deserialization Logic

The PartitionedQueryExecutionInfo.getQueryRanges() method now branches based on thin client mode:

Mode	`queryRanges` Format	Deserialization Approach
Standard Gateway	`"min": "24F3B4E0..."` (hex string)	Direct deserialization via `getList()` to `List<Range<String>>`
Thin Client Proxy	`"min": ["value"]` (JSON array)	Manual parsing: deserialize to `PartitionKeyInternal`, then convert to EPK hex string via `PartitionKeyInternalHelper.getEffectivePartitionKeyString()`

The conversion uses the PartitionKeyDefinition from the DocumentCollection to properly compute the EPK hash based on the partition key type (Hash, MultiHash, etc.).

Operations Now Supported via Thin Client

Operation	Resource Type	Operation Type
Point Operations	Document	Read, Create, Replace, Delete, Upsert, Patch
Query	Document	Query
Batch	Document	Batch
Change Feed	Document	ReadFeed (non-AllVersionsAndDeletes)
Stored Procedure Execution (new)	StoredProcedure	ExecuteJavaScript
Query Plan (new)	Document	QueryPlan

Testing

Added 12 new end-to-end tests:
- 4 basic query tests (SELECT *, WHERE id=, partition key filter, legacy)
- 7 query plan feature tests (ORDER BY, COUNT, DISTINCT, TOP, GROUP BY, OFFSET LIMIT, invalid query)
- 1 stored procedure test
Tests verify both functional correctness and infrastructure correctness (thin client endpoint used)

Made changes.

All SDK Contribution checklist:

The pull request does not introduce [breaking changes]
CHANGELOG is updated for new features, bug fixes or other significant changes.
I have read the contribution guidelines.

General Guidelines and Best Practices

Title of the pull request is clear and informative.
There are a small number of commits, each of which have an informative message. This means that previously merged commits do not appear in the history of the PR. For more information on cleaning up the commits in your PR, see this page.

Testing Guidelines

Pull request includes test coverage for the included changes.

…ndpoint.

… GW Mode) cases.

jeet1995 · 2026-01-23T23:05:54Z

/azp run java - cosmos - tests

azure-pipelines · 2026-01-23T23:06:17Z

Azure Pipelines successfully started running 1 pipeline(s).

jeet1995 · 2026-01-29T22:24:16Z

/azp run java - cosmos - tests

azure-pipelines · 2026-01-29T22:24:39Z

Azure Pipelines successfully started running 1 pipeline(s).

jeet1995 · 2026-01-29T22:36:06Z

/azp run java - cosmos - tests

azure-pipelines · 2026-01-29T22:36:35Z

Azure Pipelines successfully started running 1 pipeline(s).

xinlian12 · 2026-01-30T17:31:28Z

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

+    @BeforeClass(groups = {"thinclient"})
+    public void beforeClass() {
+        // If running locally, uncomment these lines
+//        System.setProperty("COSMOS.THINCLIENT_ENABLED", "true");


do we still need these? probably can clean up now

At least COSMOS.THINCLIENT_ENABLED is always needed from an environment variable perspective (I'll clean the HTTP/2 enabled environment variable).

xinlian12 · 2026-01-30T17:33:15Z

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

+
+        if (allDocs != null && !allDocs.isEmpty()) {
+            for (ObjectNode doc : allDocs) {
+                String id = doc.get(ID_FIELD).asText();


seems can use the following deleteDocuments() method below. And also not sure how many docs need to clean each time, using bulk can be faster

xinlian12 · 2026-01-30T17:40:59Z

.../azure-cosmos/src/main/java/com/azure/cosmos/implementation/query/SingleGroupAggregator.java

            for (Map.Entry<String, AggregateOperator> aliasToAggregate : aggregateAliasToAggregateType.entrySet()) {
                String alias = aliasToAggregate.getKey();
                AggregateOperator aggregateOperator = null;
+                Object aliasAggregateOperator = aliasToAggregate.getValue();


defined but not used?

xinlian12 · 2026-01-30T17:44:23Z

...mos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/query/QueryPlanRetriever.java

                                                                                 requestHeaders);
-        queryPlanRequest.useGatewayMode = true;
+
+        // queryPlanRequest.useGatewayMode = true;


nit: remove

xinlian12 · 2026-01-30T17:48:33Z

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

+            }
+
+            // 3. Execute SELECT * FROM C WHERE c.id = @id query
+            String query = "SELECT * FROM c WHERE c.id = @id";


curious -> does full text search etc supported?does not see the tests here

Unsure - (likewise about vector search / hybrid search). I'll confirm with the proxy service team.

Should work because there is no specific logic in Gateway

We should have a pipeline running all the query tests in Gateway mode with thin client - let's sync-up at our JVM meeting today.

xinlian12

LGTM, thanks

mbhaskar · 2026-01-30T20:39:42Z

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

@@ -46,80 +54,360 @@
 public class ThinClientE2ETest {


Just a though: Why not use the existing tests and use the clientBuilders pattern instead and add a thinproxy clientbuilder. In this way, we can slowly add this provider to as many existing tests as we want and have them run on thinproxy?

@mbhaskar - sounds good. I'll make this change.

mbhaskar · 2026-01-30T22:24:06Z

...osmos/src/main/java/com/azure/cosmos/implementation/query/PartitionedQueryExecutionInfo.java

+        try {
+            // Use Jackson to deserialize using PartitionKeyInternal's custom deserializer
+            return Utils.getSimpleObjectMapper().treeToValue(node, PartitionKeyInternal.class);
+        } catch (Exception e) {


Catch specific exeption types? like JsonProcessingException

FabianMeiswinkel

LGTM - Thanks

jeet1995 added 2 commits January 20, 2026 18:20

Route Execute Stored Procedure requests to Thin Proxy endpoint.

4a83556

Route Execute Stored Procedure and QueryPlan requests to Thin Proxy e…

325ff0a

…ndpoint.

github-actions bot added the Cosmos label Jan 21, 2026

jeet1995 added 3 commits January 23, 2026 14:22

Ensure QueryPlan gets routed to Gateway Service Endpoint (in non-TC +…

eb9a83e

… GW Mode) cases.

Ensure QueryPlan gets routed to Gateway Service Endpoint (in non-TC +…

1492b36

… GW Mode) cases.

Ensure QueryPlan gets routed to Gateway Service Endpoint (in non-TC +…

2c8a5f3

… GW Mode) cases.

jeet1995 added 3 commits January 27, 2026 20:03

Obtain List<Range<String>> from List<PartitionKeyInternal>.

1e14a7b

Obtain List<Range<String>> from List<PartitionKeyInternal>.

91116c9

Adding query + thin-client tests.

688a5ac

jeet1995 changed the title ~~Az cosmos gateway v2 query plan support~~ [Gateway V2 / DO NOT MERGE]: Integrate Stored Procedure and Query Plan request routing to a Gateway V2 endpoint. Jan 29, 2026

Fixing tests.

8ae321f

Fixing tests.

e76c84c

jeet1995 changed the title ~~[Gateway V2 / DO NOT MERGE]: Integrate Stored Procedure and Query Plan request routing to a Gateway V2 endpoint.~~ [Gateway V2][DO NOT MERGE]: Integrate Stored Procedure and Query Plan request routing to a Gateway V2 endpoint. Jan 30, 2026

xinlian12 reviewed Jan 30, 2026

View reviewed changes

xinlian12 approved these changes Jan 30, 2026

View reviewed changes

mbhaskar reviewed Jan 30, 2026

View reviewed changes

Fixing tests.

7f3bad8

mbhaskar reviewed Jan 30, 2026

View reviewed changes

Addressing review comments.

7600598

FabianMeiswinkel approved these changes Feb 2, 2026

View reviewed changes

[Gateway V2][DO NOT MERGE]: Integrate Stored Procedure and Query Plan request routing to a Gateway V2 endpoint. #47759

Are you sure you want to change the base?

[Gateway V2][DO NOT MERGE]: Integrate Stored Procedure and Query Plan request routing to a Gateway V2 endpoint. #47759

Uh oh!

Conversation

jeet1995 commented Jan 21, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

PR Description

Title

Summary

Changes

Core Changes

Stored Procedure Execution Support

Query Plan Support

Test Coverage

Technical Details: QueryPlan Response Format Differences

Thin Client Proxy Response

Standard Gateway Response

Deserialization Logic

Operations Now Supported via Thin Client

Testing

All SDK Contribution checklist:

Uh oh!

jeet1995 commented Jan 23, 2026

Uh oh!

azure-pipelines bot commented Jan 23, 2026

Uh oh!

jeet1995 commented Jan 29, 2026

Uh oh!

azure-pipelines bot commented Jan 29, 2026

Uh oh!

jeet1995 commented Jan 29, 2026

Uh oh!

azure-pipelines bot commented Jan 29, 2026

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

xinlian12 Jan 30, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

xinlian12 left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

jeet1995 Jan 30, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

FabianMeiswinkel left a comment

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

jeet1995 commented Jan 21, 2026 •

edited

Loading

xinlian12 Jan 30, 2026 •

edited

Loading

jeet1995 Jan 30, 2026 •

edited

Loading