Datanode: validate that a query id belongs to the calling session by ColinLeeo · Pull Request #18643 · apache/iotdb

ColinLeeo · 2026-09-16T01:41:45Z

Summary

fetchResultsV2, fetchResults and closeOperation take a queryId from the client and resolve it in the coordinator wide map of running queries. The request was only checked for login, not for whether that queryId was issued to the calling session, so a request could stream the result of a query, advance its cursor, or release it, on behalf of another session.

This PR adds a session level validation of the client supplied queryId. When the query is not owned by the calling session, the request returns NO_PERMISSION and the running query is left untouched. A queryId that is no longer running keeps the previous behaviour, so a closeOperation sent after a result set was fully consumed still succeeds as usual.

Changes

IClientSession#containsQueryId(Long statementId, long queryId) added. ClientSession and InternalClientSession implement it over the statementId -> queryId bookkeeping the sessions already keep; MqttClientSession and RestClientSession cannot submit queries and return false. A null statement id scans all of the session's query ids, which is what V1 fetchResults sends (the JDBC data set does not set a statement id on fetch).
ClientRPCServiceImpl: the check is applied in fetchResultsV2, fetchResults and closeOperation. The check runs after the existing lookup, so an unknown or killed queryId still returns QUERY_WAS_KILLED as before, and a rejected fetch neither records operation latency nor cleans the query up.
closeOperation only rejects when the query is still running; releasing an already released queryId remains a no-op.
Added the en/zh message and QueryOwnershipTest, covering the session bookkeeping, both fetch APIs, and the close path.

Verification

mvn test-compile -DskipTests and mvn test-compile -DskipTests -P with-zh-locale
mvn -pl iotdb-core/datanode -am -Dtest=QueryOwnershipTest -DfailIfNoTests=false test -- 7 tests, all pass
mvn spotless:apply -pl iotdb-core/datanode

fetchResultsV2, fetchResults and closeOperation resolve the queryId sent by the client in the coordinator wide map of running queries, without checking which session submitted that query. Add a per session validation of the queryId and return NO_PERMISSION when it was not issued to the calling session, so that the running query is neither read nor released by another session. - add IClientSession.containsQueryId, implemented on ClientSession and InternalClientSession over the statementId -> queryId bookkeeping that already exists, and on MqttClientSession/RestClientSession which cannot submit queries - check the queryId in fetchResultsV2, fetchResults and closeOperation; a queryId that is no longer running keeps the previous behaviour - a rejected fetch does not record latency or clean up the query - add the en/zh message and a unit test covering the two fetch APIs, the close path, the session level bookkeeping, and a fetch that leaves the result set unconsumed

Copilot

Copilot review overview

🟡 Changes recommended

The ownership tests need corrected setup and explicit coverage of the V1 null-statement-id path.

Get a fresh assessment by requesting another Copilot review.

Review effort: Lite
Findings: 1

Open (1)

V1 fetch test incorrectly sets statement ID · New

What changed in this PR

This PR adds session-level ownership validation for DataNode query fetching and closing, preventing cross-session access to running queries.

Changes:

Validates query ownership across fetch and close RPCs.
Extends session bookkeeping and localized messages.
Adds ownership and lifecycle tests.

File	Summary
`iotdb-core/datanode/src/test/java/org/apache/iotdb/db/protocol/session/QueryOwnershipTest.java`	Tests ownership behavior. Moderate findings: the V1 test must omit `statementId` (2 votes); fetch and close tests must register the query with a distinct owner session (1 vote each).
`iotdb-core/datanode/src/main/java/org/apache/iotdb/db/protocol/thrift/impl/ClientRPCServiceImpl.java`	Enforces ownership checks for fetch and close operations.
`iotdb-core/datanode/src/main/java/org/apache/iotdb/db/protocol/session/RestClientSession.java`	Returns no query ownership for unsupported submissions.
`iotdb-core/datanode/src/main/java/org/apache/iotdb/db/protocol/session/MqttClientSession.java`	Returns no query ownership for unsupported submissions.
`iotdb-core/datanode/src/main/java/org/apache/iotdb/db/protocol/session/InternalClientSession.java`	Implements query ownership lookup.
`iotdb-core/datanode/src/main/java/org/apache/iotdb/db/protocol/session/IClientSession.java`	Adds the query ownership API.
`iotdb-core/datanode/src/main/java/org/apache/iotdb/db/protocol/session/ClientSession.java`	Implements session query ownership lookup.
`iotdb-core/datanode/src/main/i18n/zh/org/apache/iotdb/db/i18n/DataNodeMiscMessages.java`	Adds the localized permission message.
`iotdb-core/datanode/src/main/i18n/en/org/apache/iotdb/db/i18n/DataNodeMiscMessages.java`	Adds the localized permission message.

💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.

+            TSFetchResultsReq request = createFetchResultsReq(anotherSession);
+            Assert.assertEquals(
+                TSStatusCode.NO_PERMISSION.getStatusCode(),
+                service.fetchResults(request).getStatus().getCode());
+            Assert.assertEquals(
+                TSStatusCode.NO_PERMISSION.getStatusCode(),
+                service.fetchResultsV2(request).getStatus().getCode());


JackieTien97

Requesting changes for the two inline findings: the existence-check/cleanup race in closeOperation and avoidable per-statement allocation in the fetch path when statementId is omitted.

Validation against 5153854: the reactor clean test build and all 7 existing QueryOwnershipTest cases passed. Of 9 additional handler-level regression tests, 8 passed and the deterministic concurrent-close regression failed as described below. Allocation figures are local JMH measurements on OpenJDK 21.0.2 / Apple M2 Pro, not end-to-end database throughput measurements.

JackieTien97 · 2026-09-22T11:18:38Z

+        && COORDINATOR.getQueryExecution(req.queryId) != null
+        && !clientSession.containsQueryId(
+            req.isSetStatementId() ? req.getStatementId() : null, req.queryId)) {


[P2] Do not pass an unowned queryId to global cleanup

When this lookup returns null, the ownership check is skipped, but the request still reaches SessionManager.closeDataset, which unconditionally invokes Coordinator.cleanupQueryExecution(queryId). A query belonging to another session can be registered between this lookup and the subsequent map.remove: query IDs are allocated before the execution is published in queryExecutionMap. The foreign close then removes and stops that query and returns SUCCESS.

I reproduced this with a deterministic two-thread test using the actual handler, SessionManager and Coordinator: pause closeOperation after the null lookup, register the other session's execution, then resume. The result was status=200, victimExecutionRemaining=false, while the owner's session binding remained present. This leaves a concurrency gap in the new close protection.

Please ensure that a queryId not owned by the calling session never reaches global cleanup. If it is absent, return the compatible no-op result directly rather than falling through to cleanup, while preserving existing malformed-request handling. Please also add a regression test for registration between the existence check and cleanup.

JackieTien97 · 2026-09-22T11:18:38Z

+      for (Set<Long> queryIds : statementIdToQueryId.values()) {
+        if (queryIds != null && queryIds.contains(queryId)) {


[P2] Box queryId once before scanning the statement sets

queryId is a primitive long, so contains(queryId) boxes it on every loop iteration. For IDs outside the Long cache, local JMH measurements against this implementation show roughly 24 bytes allocated per visited statement set, including empty sets whose statements remain open. This path is used by the legacy Java V1 dataset and the current C++ V2 fetch implementation, both of which omit statementId.

With the matching query in the last set, 1,000 statements added approximately 8.6 us and 24 KB per fetch; 10,000 added approximately 73 us and 240 KB. These are local lookup/ownership microbenchmarks, not full RPC latency. The allocations recur on every fetched page and can create avoidable GC pressure on connections with many statements.

Please create a single Long boxedQueryId before the loop and reuse it for contains. A separate comparison reduced the 1,000-statement scan from approximately 24,056 B/op to 80 B/op without adding an index or changing ownership semantics. The scan remains linear, but its temporary allocation no longer needs to grow with the number of statements.

closeOperation consulted the coordinator map before the session binding and fell through to the global cleanup when the lookup returned null. Query ids are allocated before their execution is published, so a query of another session can be registered between the lookup and the cleanup, and the foreign request released it. Check the session binding first and return the compatible no-op result directly when the query id is not running. Also box the primitive query id once in ClientSession#containsQueryId instead of once per visited statement set, and align the tests with the V1 and V2 fetch contracts: the V1 request carries no statement id, and a query registered while a foreign close is served must survive.

ColinLeeo force-pushed the colin/query-id-session-ownership branch from b1fdd6e to b1807ca Compare September 22, 2026 07:53

ColinLeeo force-pushed the colin/query-id-session-ownership branch from b1807ca to 5153854 Compare September 22, 2026 08:44

ColinLeeo requested review from HTHou and JackieTien97 and a lite review from Copilot September 22, 2026 08:53

Copilot started reviewing on behalf of ColinLeeo September 22, 2026 09:13 View session

Copilot AI reviewed Sep 22, 2026

View reviewed changes

JackieTien97 requested changes Sep 22, 2026

View reviewed changes

Navigation Menu

Datanode: validate that a query id belongs to the calling session - #18643

Datanode: validate that a query id belongs to the calling session#18643
ColinLeeo wants to merge 2 commits into
apache:masterfrom
ColinLeeo:colin/query-id-session-ownership

ColinLeeo commented Sep 16, 2026 •

edited

Loading

Uh oh!

Copilot AI left a comment

Uh oh!

JackieTien97 left a comment

Uh oh!

JackieTien97 Sep 22, 2026

Uh oh!

JackieTien97 Sep 22, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

		for (Set<Long> queryIds : statementIdToQueryId.values()) {
		if (queryIds != null && queryIds.contains(queryId)) {

Conversation

ColinLeeo commented Sep 16, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Summary

Changes

Verification

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Copilot review overview

🟡 Changes recommended

Uh oh!

JackieTien97 left a comment

Choose a reason for hiding this comment

Uh oh!

JackieTien97 Sep 22, 2026

Choose a reason for hiding this comment

Uh oh!

JackieTien97 Sep 22, 2026

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

ColinLeeo commented Sep 16, 2026 •

edited

Loading