This is an automated email from the ASF dual-hosted git repository.
JackieTien97 pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/iotdb.git
The following commit(s) were added to refs/heads/master by this push:
new e7ec99f5d17 Datanode: validate that a query id belongs to the calling
session (#18643)
e7ec99f5d17 is described below
commit e7ec99f5d1717dbee749474ff149b97b87f4675c
Author: Colin Lee <[email protected]>
AuthorDate: Wed Sep 23 10:59:51 2026 +0800
Datanode: validate that a query id belongs to the calling session (#18643)
---
.../apache/iotdb/db/i18n/DataNodeMiscMessages.java | 2 +
.../apache/iotdb/db/i18n/DataNodeMiscMessages.java | 2 +
.../iotdb/db/protocol/session/ClientSession.java | 23 ++
.../iotdb/db/protocol/session/IClientSession.java | 3 +
.../db/protocol/session/InternalClientSession.java | 5 +
.../db/protocol/session/MqttClientSession.java | 5 +
.../db/protocol/session/RestClientSession.java | 5 +
.../protocol/thrift/impl/ClientRPCServiceImpl.java | 97 ++++--
.../db/protocol/session/QueryOwnershipTest.java | 343 +++++++++++++++++++++
9 files changed, 459 insertions(+), 26 deletions(-)
diff --git
a/iotdb-core/datanode/src/main/i18n/en/org/apache/iotdb/db/i18n/DataNodeMiscMessages.java
b/iotdb-core/datanode/src/main/i18n/en/org/apache/iotdb/db/i18n/DataNodeMiscMessages.java
index c961031b38e..405a5139e34 100644
---
a/iotdb-core/datanode/src/main/i18n/en/org/apache/iotdb/db/i18n/DataNodeMiscMessages.java
+++
b/iotdb-core/datanode/src/main/i18n/en/org/apache/iotdb/db/i18n/DataNodeMiscMessages.java
@@ -21,6 +21,8 @@ package org.apache.iotdb.db.i18n;
/** Compile-time i18n constants for DataNode misc subsystems (English). */
public final class DataNodeMiscMessages {
+ public static final String
MESSAGE_QUERY_DOES_NOT_BELONG_TO_CURRENT_SESSION_A1198237 =
+ "The requested query does not belong to the current session.";
public static final String
MESSAGE_MISSING_LOAD_TSFILE_SLICE_METADATA_ARG_DE4333DA =
"Missing Load TsFile slice metadata: %s";
diff --git
a/iotdb-core/datanode/src/main/i18n/zh/org/apache/iotdb/db/i18n/DataNodeMiscMessages.java
b/iotdb-core/datanode/src/main/i18n/zh/org/apache/iotdb/db/i18n/DataNodeMiscMessages.java
index 11bc2fc579f..a87a49eaaf5 100644
---
a/iotdb-core/datanode/src/main/i18n/zh/org/apache/iotdb/db/i18n/DataNodeMiscMessages.java
+++
b/iotdb-core/datanode/src/main/i18n/zh/org/apache/iotdb/db/i18n/DataNodeMiscMessages.java
@@ -21,6 +21,8 @@ package org.apache.iotdb.db.i18n;
/** 编译时国际化常量 - DataNode 杂项子系统（中文）。 */
public final class DataNodeMiscMessages {
+ public static final String
MESSAGE_QUERY_DOES_NOT_BELONG_TO_CURRENT_SESSION_A1198237 =
+ "请求的查询不属于当前会话。";
public static final String
MESSAGE_MISSING_LOAD_TSFILE_SLICE_METADATA_ARG_DE4333DA =
"缺少 Load TsFile 分片元数据：%s";
diff --git
a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/protocol/session/ClientSession.java
b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/protocol/session/ClientSession.java
index bad4ddd6c7d..2f08920b8d2 100644
---
a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/protocol/session/ClientSession.java
+++
b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/protocol/session/ClientSession.java
@@ -90,6 +90,29 @@ public class ClientSession extends IClientSession {
queryIds.add(queryId);
}
+ @Override
+ public boolean containsQueryId(Long statementId, long queryId) {
+ return containsQueryId(statementIdToQueryId, statementId, queryId);
+ }
+
+ public static boolean containsQueryId(
+ Map<Long, Set<Long>> statementIdToQueryId, Long statementId, long
queryId) {
+ // Set#contains takes an Object, so box the primitive queryId once: a
client that does not send
+ // a statement id makes this method visit every statement set of the
session, and a box per
+ // visited set would allocate once per statement on every fetched page.
+ Long boxedQueryId = queryId;
+ if (statementId == null) {
+ for (Set<Long> queryIds : statementIdToQueryId.values()) {
+ if (queryIds != null && queryIds.contains(boxedQueryId)) {
+ return true;
+ }
+ }
+ return false;
+ }
+ Set<Long> queryIds = statementIdToQueryId.get(statementId);
+ return queryIds != null && queryIds.contains(boxedQueryId);
+ }
+
@Override
public void removeQueryId(Long statementId, Long queryId) {
removeQueryId(statementIdToQueryId, statementId, queryId);
diff --git
a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/protocol/session/IClientSession.java
b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/protocol/session/IClientSession.java
index bac4b15e342..5114ecb0d78 100644
---
a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/protocol/session/IClientSession.java
+++
b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/protocol/session/IClientSession.java
@@ -164,6 +164,9 @@ public abstract class IClientSession {
public abstract void addQueryId(Long statementId, long queryId);
+ // statementId could be null
+ public abstract boolean containsQueryId(Long statementId, long queryId);
+
// statementId could be null
public abstract void removeQueryId(Long statementId, Long queryId);
diff --git
a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/protocol/session/InternalClientSession.java
b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/protocol/session/InternalClientSession.java
index 460ec9319f2..8bbb248e989 100644
---
a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/protocol/session/InternalClientSession.java
+++
b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/protocol/session/InternalClientSession.java
@@ -90,6 +90,11 @@ public class InternalClientSession extends IClientSession {
queryIds.add(queryId);
}
+ @Override
+ public boolean containsQueryId(Long statementId, long queryId) {
+ return ClientSession.containsQueryId(statementIdToQueryId, statementId,
queryId);
+ }
+
@Override
public void removeQueryId(Long statementId, Long queryId) {
ClientSession.removeQueryId(statementIdToQueryId, statementId, queryId);
diff --git
a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/protocol/session/MqttClientSession.java
b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/protocol/session/MqttClientSession.java
index 65e2c9a5b5d..0fb830bcca3 100644
---
a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/protocol/session/MqttClientSession.java
+++
b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/protocol/session/MqttClientSession.java
@@ -77,6 +77,11 @@ public class MqttClientSession extends IClientSession {
throw new UnsupportedOperationException();
}
+ @Override
+ public boolean containsQueryId(Long statementId, long queryId) {
+ return false;
+ }
+
@Override
public void removeQueryId(Long statementId, Long queryId) {
throw new UnsupportedOperationException();
diff --git
a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/protocol/session/RestClientSession.java
b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/protocol/session/RestClientSession.java
index 58bae05fffe..7fe2ceece08 100644
---
a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/protocol/session/RestClientSession.java
+++
b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/protocol/session/RestClientSession.java
@@ -79,6 +79,11 @@ public class RestClientSession extends IClientSession {
throw new UnsupportedOperationException();
}
+ @Override
+ public boolean containsQueryId(Long statementId, long queryId) {
+ return false;
+ }
+
@Override
public void removeQueryId(Long statementId, Long queryId) {
throw new UnsupportedOperationException();
diff --git
a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/protocol/thrift/impl/ClientRPCServiceImpl.java
b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/protocol/thrift/impl/ClientRPCServiceImpl.java
index ce0c5d67450..692b5bedd2d 100644
---
a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/protocol/thrift/impl/ClientRPCServiceImpl.java
+++
b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/protocol/thrift/impl/ClientRPCServiceImpl.java
@@ -1533,6 +1533,7 @@ public class ClientRPCServiceImpl implements
IClientRPCServiceWithHandler {
String statementType = null;
Throwable t = null;
IQueryExecution queryExecution = null;
+ boolean queryOwnedBySession = false;
IClientSession clientSession =
SESSION_MANAGER.getCurrSessionAndUpdateIdleTime();
Long statementId = req.isSetStatementId() ? req.getStatementId() : null;
try {
@@ -1542,13 +1543,22 @@ public class ClientRPCServiceImpl implements
IClientRPCServiceWithHandler {
}
queryExecution = COORDINATOR.getQueryExecution(req.queryId);
-
if (queryExecution == null) {
TSStatus noQueryExecutionStatus = new
TSStatus(QUERY_WAS_KILLED.getStatusCode());
noQueryExecutionStatus.setMessage(NO_QUERY_EXECUTION_ERR_MSG);
return RpcUtils.getTSFetchResultsResp(noQueryExecutionStatus);
}
+ if (!clientSession.containsQueryId(statementId, req.queryId)) {
+ // The query is still running, but it was submitted by another
session: do not stream its
+ // result and do not release it, so that the query which owns it is
left untouched.
+ return RpcUtils.getTSFetchResultsResp(
+ RpcUtils.getStatus(
+ TSStatusCode.NO_PERMISSION,
+
DataNodeMiscMessages.MESSAGE_QUERY_DOES_NOT_BELONG_TO_CURRENT_SESSION_A1198237));
+ }
+ queryOwnedBySession = true;
+
TSFetchResultsResp resp =
RpcUtils.getTSFetchResultsResp(TSStatusCode.SUCCESS_STATUS);
queryExecution.updateCurrentRpcStartTime(startTime);
@@ -1577,19 +1587,21 @@ public class ClientRPCServiceImpl implements
IClientRPCServiceWithHandler {
throw error;
} finally {
- long currentOperationCost = System.nanoTime() - startTime;
- COORDINATOR.recordExecutionTime(req.queryId, currentOperationCost);
-
- // record each operation time cost
- CommonUtils.addStatementExecutionLatency(
- OperationType.FETCH_RESULTS, statementType, currentOperationCost);
+ if (queryOwnedBySession) {
+ long currentOperationCost = System.nanoTime() - startTime;
+ COORDINATOR.recordExecutionTime(req.queryId, currentOperationCost);
- if (finished) {
- // record total time cost for one query
- long executionTime = COORDINATOR.getTotalExecutionTime(req.queryId);
- CommonUtils.addQueryLatency(
- StatementType.QUERY, executionTime > 0 ? executionTime :
currentOperationCost);
- clearUp(clientSession, statementId, req.queryId, req, t);
+ // record each operation time cost
+ CommonUtils.addStatementExecutionLatency(
+ OperationType.FETCH_RESULTS, statementType, currentOperationCost);
+
+ if (finished) {
+ // record total time cost for one query
+ long executionTime = COORDINATOR.getTotalExecutionTime(req.queryId);
+ CommonUtils.addQueryLatency(
+ StatementType.QUERY, executionTime > 0 ? executionTime :
currentOperationCost);
+ clearUp(clientSession, statementId, req.queryId, req, t);
+ }
}
SESSION_MANAGER.updateIdleTime();
@@ -1683,8 +1695,27 @@ public class ClientRPCServiceImpl implements
IClientRPCServiceWithHandler {
@Override
public TSStatus closeOperation(TSCloseOperationReq req) {
+ IClientSession clientSession = SESSION_MANAGER.getCurrSession();
+ if (req.isSetQueryId()
+ && req.isSetStatementId()
+ && clientSession != null
+ && clientSession.isLogin()
+ && !clientSession.containsQueryId(req.getStatementId(), req.queryId)) {
+ // The queryId indexes the process-wide map of running queries, so only
the session that
+ // submitted the query may release it.
+ if (COORDINATOR.getQueryExecution(req.queryId) != null) {
+ return RpcUtils.getStatus(
+ TSStatusCode.NO_PERMISSION,
+
DataNodeMiscMessages.MESSAGE_QUERY_DOES_NOT_BELONG_TO_CURRENT_SESSION_A1198237);
+ }
+ // A queryId that is no longer running keeps the previous behaviour:
releasing it stays a
+ // no-op. It must not fall through to the global cleanup below: query
ids are allocated
+ // before their execution is published, so the session that owns this
queryId can register
+ // it between the lookup above and the cleanup.
+ return RpcUtils.getStatus(TSStatusCode.SUCCESS_STATUS);
+ }
return SESSION_MANAGER.closeOperation(
- SESSION_MANAGER.getCurrSession(),
+ clientSession,
req.queryId,
req.statementId,
req.isSetStatementId(),
@@ -2291,6 +2322,7 @@ public class ClientRPCServiceImpl implements
IClientRPCServiceWithHandler {
String statementType = null;
Throwable t = null;
IQueryExecution queryExecution = null;
+ boolean queryOwnedBySession = false;
IClientSession clientSession =
SESSION_MANAGER.getCurrSessionAndUpdateIdleTime();
Long statementId = req.isSetStatementId() ? req.getStatementId() : null;
try {
@@ -2305,6 +2337,17 @@ public class ClientRPCServiceImpl implements
IClientRPCServiceWithHandler {
noQueryExecutionStatus.setMessage(NO_QUERY_EXECUTION_ERR_MSG);
return RpcUtils.getTSFetchResultsResp(noQueryExecutionStatus);
}
+
+ if (!clientSession.containsQueryId(statementId, req.queryId)) {
+ // The query is still running, but it was submitted by another
session: do not stream its
+ // result and do not release it, so that the query which owns it is
left untouched.
+ return RpcUtils.getTSFetchResultsResp(
+ RpcUtils.getStatus(
+ TSStatusCode.NO_PERMISSION,
+
DataNodeMiscMessages.MESSAGE_QUERY_DOES_NOT_BELONG_TO_CURRENT_SESSION_A1198237));
+ }
+ queryOwnedBySession = true;
+
queryExecution.updateCurrentRpcStartTime(startTime);
statementType = queryExecution.getStatementType();
@@ -2332,19 +2375,21 @@ public class ClientRPCServiceImpl implements
IClientRPCServiceWithHandler {
throw error;
} finally {
- long currentOperationCost = System.nanoTime() - startTime;
- COORDINATOR.recordExecutionTime(req.queryId, currentOperationCost);
+ if (queryOwnedBySession) {
+ long currentOperationCost = System.nanoTime() - startTime;
+ COORDINATOR.recordExecutionTime(req.queryId, currentOperationCost);
- // record each operation time cost
- CommonUtils.addStatementExecutionLatency(
- OperationType.FETCH_RESULTS, statementType, currentOperationCost);
-
- if (finished) {
- // record total time cost for one query
- long executionTime = COORDINATOR.getTotalExecutionTime(req.queryId);
- CommonUtils.addQueryLatency(
- StatementType.QUERY, executionTime > 0 ? executionTime :
currentOperationCost);
- clearUp(clientSession, statementId, req.queryId, req, t);
+ // record each operation time cost
+ CommonUtils.addStatementExecutionLatency(
+ OperationType.FETCH_RESULTS, statementType, currentOperationCost);
+
+ if (finished) {
+ // record total time cost for one query
+ long executionTime = COORDINATOR.getTotalExecutionTime(req.queryId);
+ CommonUtils.addQueryLatency(
+ StatementType.QUERY, executionTime > 0 ? executionTime :
currentOperationCost);
+ clearUp(clientSession, statementId, req.queryId, req, t);
+ }
}
SESSION_MANAGER.updateIdleTime();
diff --git
a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/protocol/session/QueryOwnershipTest.java
b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/protocol/session/QueryOwnershipTest.java
new file mode 100644
index 00000000000..b8c53352c09
--- /dev/null
+++
b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/protocol/session/QueryOwnershipTest.java
@@ -0,0 +1,343 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.iotdb.db.protocol.session;
+
+import org.apache.iotdb.db.conf.IoTDBDescriptor;
+import org.apache.iotdb.db.protocol.thrift.impl.ClientRPCServiceImpl;
+import org.apache.iotdb.db.queryengine.plan.Coordinator;
+import org.apache.iotdb.db.queryengine.plan.execution.IQueryExecution;
+import org.apache.iotdb.rpc.TSStatusCode;
+import org.apache.iotdb.service.rpc.thrift.TSCloseOperationReq;
+import org.apache.iotdb.service.rpc.thrift.TSFetchResultsReq;
+import org.apache.iotdb.service.rpc.thrift.TSFetchResultsResp;
+
+import org.junit.AfterClass;
+import org.junit.Assert;
+import org.junit.BeforeClass;
+import org.junit.Test;
+import org.mockito.Mockito;
+
+import java.lang.reflect.Field;
+import java.net.Socket;
+import java.util.Map;
+import java.util.Optional;
+
+public class QueryOwnershipTest {
+
+ private static final long STATEMENT_ID = 1L;
+ private static final long QUERY_ID = 2L;
+
+ private static int previousDataNodeId;
+
+ @BeforeClass
+ public static void setUp() {
+ // the coordinator builds its query id generator from the data node id of
this node
+ previousDataNodeId =
IoTDBDescriptor.getInstance().getConfig().getDataNodeId();
+ IoTDBDescriptor.getInstance().getConfig().setDataNodeId(0);
+ }
+
+ @AfterClass
+ public static void tearDown() {
+
IoTDBDescriptor.getInstance().getConfig().setDataNodeId(previousDataNodeId);
+ }
+
+ @Test
+ public void testQueryIdsAreBoundToTheSessionThatSubmittedTheQuery() {
+ ClientSession owner = createSession("user");
+ owner.addStatementId(STATEMENT_ID);
+ owner.addQueryId(STATEMENT_ID, QUERY_ID);
+
+ ClientSession anotherSession = createSession("user");
+ anotherSession.addStatementId(STATEMENT_ID);
+
+ Assert.assertTrue(owner.containsQueryId(STATEMENT_ID, QUERY_ID));
+ // clients that do not send a statement id together with the query id are
still served
+ Assert.assertTrue(owner.containsQueryId(null, QUERY_ID));
+ Assert.assertFalse(anotherSession.containsQueryId(STATEMENT_ID, QUERY_ID));
+ Assert.assertFalse(anotherSession.containsQueryId(null, QUERY_ID));
+ Assert.assertFalse(owner.containsQueryId(STATEMENT_ID + 1, QUERY_ID));
+ }
+
+ @Test
+ public void testFetchResultsRejectsQueryOfAnotherSession() throws Exception {
+ ClientSession anotherSession = createSession("user");
+ anotherSession.addStatementId(STATEMENT_ID);
+ anotherSession.setLogin(true);
+
+ Map<Long, IQueryExecution> queryExecutionMap = getQueryExecutionMap();
+ queryExecutionMap.put(QUERY_ID, mockQueryExecution());
+ try {
+ withCurrentSession(
+ anotherSession,
+ () -> {
+ ClientRPCServiceImpl service = new ClientRPCServiceImpl();
+ Assert.assertEquals(
+ TSStatusCode.NO_PERMISSION.getStatusCode(),
+
service.fetchResults(createFetchResultsReq(anotherSession)).getStatus().getCode());
+ Assert.assertEquals(
+ TSStatusCode.NO_PERMISSION.getStatusCode(),
+ service
+
.fetchResultsV2(createFetchResultsReqWithStatementId(anotherSession))
+ .getStatus()
+ .getCode());
+ });
+ // the rejected requests must not release the query of the session that
submitted it
+ Assert.assertTrue(queryExecutionMap.containsKey(QUERY_ID));
+ } finally {
+ queryExecutionMap.remove(QUERY_ID);
+ }
+ }
+
+ @Test
+ public void testFetchResultsOfOwnQueryIsStillServed() throws Exception {
+ ClientSession owner = createSession("user");
+ owner.addStatementId(STATEMENT_ID);
+ owner.addQueryId(STATEMENT_ID, QUERY_ID);
+ owner.setLogin(true);
+
+ IQueryExecution queryExecution = mockQueryExecution();
+ Map<Long, IQueryExecution> queryExecutionMap = getQueryExecutionMap();
+ queryExecutionMap.put(QUERY_ID, queryExecution);
+ try {
+ withCurrentSession(
+ owner,
+ () ->
+ Assert.assertEquals(
+ TSStatusCode.SUCCESS_STATUS.getStatusCode(),
+ new ClientRPCServiceImpl()
+
.fetchResultsV2(createFetchResultsReqWithStatementId(owner))
+ .getStatus()
+ .getCode()));
+ // a fully consumed query is released, and the client closes it
afterwards
+ Assert.assertFalse(queryExecutionMap.containsKey(QUERY_ID));
+ Assert.assertFalse(owner.containsQueryId(STATEMENT_ID, QUERY_ID));
+ } finally {
+ queryExecutionMap.remove(QUERY_ID);
+ }
+ }
+
+ @Test
+ public void testFetchResultsOfOwnQueryWithoutStatementIdIsStillServed()
throws Exception {
+ ClientSession owner = createSession("user");
+ owner.addStatementId(STATEMENT_ID);
+ owner.addQueryId(STATEMENT_ID, QUERY_ID);
+ owner.setLogin(true);
+
+ Map<Long, IQueryExecution> queryExecutionMap = getQueryExecutionMap();
+ queryExecutionMap.put(QUERY_ID, mockQueryExecution());
+ try {
+ // the query id is bound to a statement of this session, the request
does not mention it
+ withCurrentSession(
+ owner,
+ () ->
+ Assert.assertEquals(
+ TSStatusCode.SUCCESS_STATUS.getStatusCode(),
+ new ClientRPCServiceImpl()
+ .fetchResults(createFetchResultsReq(owner))
+ .getStatus()
+ .getCode()));
+ // a fully consumed query is released for the session that submitted it
+ Assert.assertFalse(queryExecutionMap.containsKey(QUERY_ID));
+ Assert.assertFalse(owner.containsQueryId(STATEMENT_ID, QUERY_ID));
+ } finally {
+ queryExecutionMap.remove(QUERY_ID);
+ }
+ }
+
+ @Test
+ public void testFetchResultsOfQueryWithMoreDataKeepsQueryAndSessionBinding()
throws Exception {
+ ClientSession owner = createSession("user");
+ owner.addStatementId(STATEMENT_ID);
+ owner.addQueryId(STATEMENT_ID, QUERY_ID);
+ owner.setLogin(true);
+
+ Map<Long, IQueryExecution> queryExecutionMap = getQueryExecutionMap();
+ queryExecutionMap.put(QUERY_ID, mockQueryExecution(true));
+ try {
+ withCurrentSession(
+ owner,
+ () -> {
+ TSFetchResultsResp response =
+ new ClientRPCServiceImpl()
+
.fetchResultsV2(createFetchResultsReqWithStatementId(owner));
+ Assert.assertEquals(
+ TSStatusCode.SUCCESS_STATUS.getStatusCode(),
response.getStatus().getCode());
+ Assert.assertTrue(response.isMoreData());
+ });
+ // the result set is not consumed yet, so the query has to stay
fetchable by its owner
+ Assert.assertTrue(queryExecutionMap.containsKey(QUERY_ID));
+ Assert.assertTrue(owner.containsQueryId(STATEMENT_ID, QUERY_ID));
+ } finally {
+ queryExecutionMap.remove(QUERY_ID);
+ }
+ }
+
+ @Test
+ public void testCloseOperationRejectsQueryOfAnotherSession() throws
Exception {
+ ClientSession anotherSession = createSession("user");
+ anotherSession.addStatementId(STATEMENT_ID);
+ anotherSession.setLogin(true);
+
+ Map<Long, IQueryExecution> queryExecutionMap = getQueryExecutionMap();
+ queryExecutionMap.put(QUERY_ID, mockQueryExecution());
+ try {
+ withCurrentSession(
+ anotherSession,
+ () ->
+ Assert.assertEquals(
+ TSStatusCode.NO_PERMISSION.getStatusCode(),
+ new
ClientRPCServiceImpl().closeOperation(createCloseOperationReq()).getCode()));
+ // the rejected request must not release the query of the session that
submitted it
+ Assert.assertTrue(queryExecutionMap.containsKey(QUERY_ID));
+ } finally {
+ queryExecutionMap.remove(QUERY_ID);
+ }
+ }
+
+ @Test
+ public void
testCloseOperationDoesNotReleaseQueryRegisteredWhileTheRequestIsServed()
+ throws Exception {
+ Map<Long, IQueryExecution> queryExecutionMap = getQueryExecutionMap();
+ IQueryExecution queryOfAnotherSession = mockQueryExecution();
+ // Query ids are allocated before their execution is published, so the
session that owns this
+ // queryId can register its execution between the ownership check of this
request and the point
+ // where the request would release the query. Simulate that publication
from inside the
+ // ownership check, which is where the request decides who may release the
queryId.
+ ClientSession anotherSession =
+ new ClientSession(Mockito.mock(Socket.class)) {
+ @Override
+ public boolean containsQueryId(Long statementId, long queryId) {
+ queryExecutionMap.putIfAbsent(queryId, queryOfAnotherSession);
+ return super.containsQueryId(statementId, queryId);
+ }
+ };
+ anotherSession.setUsername("user");
+ anotherSession.addStatementId(STATEMENT_ID);
+ anotherSession.setLogin(true);
+
+ try {
+ withCurrentSession(
+ anotherSession,
+ () ->
+ Assert.assertEquals(
+ TSStatusCode.NO_PERMISSION.getStatusCode(),
+ new
ClientRPCServiceImpl().closeOperation(createCloseOperationReq()).getCode()));
+ // the request must never release the query that was just registered by
its owner
+ Assert.assertTrue(queryExecutionMap.containsKey(QUERY_ID));
+ } finally {
+ queryExecutionMap.remove(QUERY_ID);
+ }
+ }
+
+ @Test
+ public void testCloseOperationReleasesQueryOfOwnSession() throws Exception {
+ ClientSession owner = createSession("user");
+ owner.addStatementId(STATEMENT_ID);
+ owner.addQueryId(STATEMENT_ID, QUERY_ID);
+ owner.setLogin(true);
+
+ Map<Long, IQueryExecution> queryExecutionMap = getQueryExecutionMap();
+ queryExecutionMap.put(QUERY_ID, mockQueryExecution());
+ try {
+ withCurrentSession(
+ owner,
+ () ->
+ Assert.assertEquals(
+ TSStatusCode.SUCCESS_STATUS.getStatusCode(),
+ new
ClientRPCServiceImpl().closeOperation(createCloseOperationReq()).getCode()));
+ Assert.assertFalse(queryExecutionMap.containsKey(QUERY_ID));
+ Assert.assertFalse(owner.containsQueryId(STATEMENT_ID, QUERY_ID));
+ } finally {
+ queryExecutionMap.remove(QUERY_ID);
+ }
+ }
+
+ @Test
+ public void testCloseOperationOfQueryThatIsNoLongerRunningStaysANoOp() {
+ // a client that consumed a result set completely sends closeOperation
after the query has
+ // already been released, and that request has to succeed as it always did
+ ClientSession session = createSession("user");
+ session.addStatementId(STATEMENT_ID);
+ session.setLogin(true);
+
+ withCurrentSession(
+ session,
+ () ->
+ Assert.assertEquals(
+ TSStatusCode.SUCCESS_STATUS.getStatusCode(),
+ new
ClientRPCServiceImpl().closeOperation(createCloseOperationReq()).getCode()));
+ }
+
+ private IQueryExecution mockQueryExecution() throws Exception {
+ return mockQueryExecution(false);
+ }
+
+ /**
+ * A mocked execution. {@code hasNextResult} says whether the mocked query
still holds data after
+ * the batch that is about to be read, which decides if a single fetch
consumes the whole result
+ * set and therefore releases the query.
+ */
+ private IQueryExecution mockQueryExecution(boolean hasNextResult) throws
Exception {
+ IQueryExecution queryExecution = Mockito.mock(IQueryExecution.class);
+ Mockito.when(queryExecution.getQueryId()).thenReturn("query");
+ // the mock carries no buffered batch, only the "is there anything left"
flag
+
Mockito.when(queryExecution.getByteBufferBatchResult()).thenReturn(Optional.empty());
+ Mockito.when(queryExecution.hasNextResult()).thenReturn(hasNextResult);
+ return queryExecution;
+ }
+
+ /** The V1 fetch request of the legacy JDBC data set, which does not send a
statement id. */
+ private TSFetchResultsReq createFetchResultsReq(ClientSession session) {
+ return new TSFetchResultsReq(session.getId(), "select 1", 1024, QUERY_ID,
true);
+ }
+
+ /** The V2 fetch request, which always carries the statement id. */
+ private TSFetchResultsReq createFetchResultsReqWithStatementId(ClientSession
session) {
+ return createFetchResultsReq(session).setStatementId(STATEMENT_ID);
+ }
+
+ private TSCloseOperationReq createCloseOperationReq() {
+ return new
TSCloseOperationReq().setStatementId(STATEMENT_ID).setQueryId(QUERY_ID);
+ }
+
+ private void withCurrentSession(ClientSession session, Runnable body) {
+ SessionManager sessionManager = SessionManager.getInstance();
+ IClientSession previousSession = sessionManager.getCurrSession();
+ sessionManager.setCurrSession(session);
+ try {
+ body.run();
+ } finally {
+ sessionManager.restoreSession(previousSession, session);
+ }
+ }
+
+ private ClientSession createSession(String username) {
+ ClientSession session = new ClientSession(Mockito.mock(Socket.class));
+ session.setUsername(username);
+ return session;
+ }
+
+ @SuppressWarnings("unchecked")
+ private Map<Long, IQueryExecution> getQueryExecutionMap() throws Exception {
+ Field field = Coordinator.class.getDeclaredField("queryExecutionMap");
+ field.setAccessible(true);
+ return (Map<Long, IQueryExecution>) field.get(Coordinator.getInstance());
+ }
+}
