From 2f68074c90def96ab3aff14de1735c456b9efaf3 Mon Sep 17 00:00:00 2001 From: shuwenwei Date: Thu, 30 Jul 2026 11:29:06 +0800 Subject: [PATCH 1/4] Fix SHOW QUERIES privilege filtering --- .../iotdb/db/it/IoTDBShowQueriesIT.java | 116 ++++++++++++++++++ .../security/TreeAccessCheckVisitor.java | 2 +- 2 files changed, 117 insertions(+), 1 deletion(-) create mode 100644 integration-test/src/test/java/org/apache/iotdb/db/it/IoTDBShowQueriesIT.java diff --git a/integration-test/src/test/java/org/apache/iotdb/db/it/IoTDBShowQueriesIT.java b/integration-test/src/test/java/org/apache/iotdb/db/it/IoTDBShowQueriesIT.java new file mode 100644 index 000000000000..efd7ba7fc24f --- /dev/null +++ b/integration-test/src/test/java/org/apache/iotdb/db/it/IoTDBShowQueriesIT.java @@ -0,0 +1,116 @@ +/* + * 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.it; + +import org.apache.iotdb.commons.auth.entity.PrivilegeType; +import org.apache.iotdb.isession.ISession; +import org.apache.iotdb.isession.SessionConfig; +import org.apache.iotdb.isession.SessionDataSet; +import org.apache.iotdb.it.env.EnvFactory; +import org.apache.iotdb.it.framework.IoTDBTestRunner; +import org.apache.iotdb.itbase.category.ClusterIT; +import org.apache.iotdb.itbase.category.LocalStandaloneIT; + +import org.apache.tsfile.enums.TSDataType; +import org.apache.tsfile.write.record.Tablet; +import org.apache.tsfile.write.schema.MeasurementSchema; +import org.junit.AfterClass; +import org.junit.Assert; +import org.junit.BeforeClass; +import org.junit.Test; +import org.junit.experimental.categories.Category; +import org.junit.runner.RunWith; + +import java.util.List; + +import static org.apache.iotdb.commons.schema.column.ColumnHeaderConstant.STATEMENT; +import static org.apache.iotdb.db.it.utils.TestUtils.createUser; +import static org.apache.iotdb.db.it.utils.TestUtils.grantUserSeriesPrivilege; +import static org.apache.iotdb.db.it.utils.TestUtils.grantUserSystemPrivileges; + +@RunWith(IoTDBTestRunner.class) +@Category({LocalStandaloneIT.class, ClusterIT.class}) +public class IoTDBShowQueriesIT { + + private static final String DATABASE = "root.show_queries"; + private static final String DEVICE = DATABASE + ".d1"; + private static final String TIMESERIES = DEVICE + ".s1"; + private static final String QUERY_USER = "show_queries_query_user"; + private static final String SHOW_USER = "show_queries_show_user"; + private static final String PASSWORD = "password123456"; + private static final String PENDING_QUERY = "SELECT s1 FROM " + DEVICE; + // Leave one default 1,000-row TsBlock after the initial Session fetch. + private static final int STREAMING_QUERY_ROW_COUNT = SessionConfig.DEFAULT_FETCH_SIZE + 1_000; + + @BeforeClass + public static void setUp() throws Exception { + EnvFactory.getEnv().getConfig().getCommonConfig().setEnforceStrongPassword(false); + EnvFactory.getEnv().initClusterEnvironment(); + try (ISession session = EnvFactory.getEnv().getSessionConnection()) { + session.executeNonQueryStatement("CREATE DATABASE " + DATABASE); + session.executeNonQueryStatement("CREATE TIMESERIES " + TIMESERIES + " INT32"); + Tablet tablet = + new Tablet( + DEVICE, + List.of(new MeasurementSchema("s1", TSDataType.INT32)), + STREAMING_QUERY_ROW_COUNT); + for (int i = 0; i < STREAMING_QUERY_ROW_COUNT; i++) { + int rowIndex = tablet.getRowSize(); + tablet.addTimestamp(rowIndex, i); + tablet.addValue("s1", rowIndex, i); + } + session.insertTablet(tablet); + } + createUser(QUERY_USER, PASSWORD); + createUser(SHOW_USER, PASSWORD); + grantUserSeriesPrivilege(QUERY_USER, PrivilegeType.READ_DATA, DATABASE + ".**"); + } + + @AfterClass + public static void tearDown() throws Exception { + EnvFactory.getEnv().cleanClusterEnvironment(); + } + + @Test + public void showQueriesAuthTest() throws Exception { + try (ISession querySession = EnvFactory.getEnv().getSessionConnection(QUERY_USER, PASSWORD); + ISession showQueriesSession = + EnvFactory.getEnv().getSessionConnection(SHOW_USER, PASSWORD); + SessionDataSet pendingDataSet = querySession.executeQueryStatement(PENDING_QUERY)) { + // Keep QUERY_USER's query active by not consuming pendingDataSet. + Assert.assertFalse(containsQuery(showQueriesSession, PENDING_QUERY)); + + grantUserSystemPrivileges(SHOW_USER, PrivilegeType.SYSTEM); + Assert.assertTrue(containsQuery(showQueriesSession, PENDING_QUERY)); + } + } + + private static boolean containsQuery(ISession session, String query) throws Exception { + try (SessionDataSet dataSet = session.executeQueryStatement("SHOW QUERIES")) { + SessionDataSet.DataIterator iterator = dataSet.iterator(); + while (iterator.next()) { + if (query.equals(iterator.getString(STATEMENT))) { + return true; + } + } + return false; + } + } +} diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/relational/security/TreeAccessCheckVisitor.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/relational/security/TreeAccessCheckVisitor.java index 73d8518a7763..566898f7b8f8 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/relational/security/TreeAccessCheckVisitor.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/relational/security/TreeAccessCheckVisitor.java @@ -1963,7 +1963,7 @@ public TSStatus visitShowDataNodes( @Override public TSStatus visitShowQueries(ShowQueriesStatement statement, TreeAccessCheckContext context) { - if (checkHasGlobalAuth(context, PrivilegeType.MAINTAIN, () -> "")) { + if (!checkHasGlobalAuth(context, PrivilegeType.SYSTEM, () -> "")) { statement.setAllowedUsername(context.getUsername()); } return SUCCEED; From 7a7e1191a9fdcecb6e3bda8282f2088a88edf8d6 Mon Sep 17 00:00:00 2001 From: shuwenwei Date: Thu, 30 Jul 2026 11:45:12 +0800 Subject: [PATCH 2/4] Fix KILL QUERY privilege filtering --- .../apache/iotdb/db/it/IoTDBKillQueryIT.java | 127 ++++++++++++++++++ .../security/TreeAccessCheckVisitor.java | 4 +- 2 files changed, 129 insertions(+), 2 deletions(-) create mode 100644 integration-test/src/test/java/org/apache/iotdb/db/it/IoTDBKillQueryIT.java diff --git a/integration-test/src/test/java/org/apache/iotdb/db/it/IoTDBKillQueryIT.java b/integration-test/src/test/java/org/apache/iotdb/db/it/IoTDBKillQueryIT.java new file mode 100644 index 000000000000..27367c0a3f7a --- /dev/null +++ b/integration-test/src/test/java/org/apache/iotdb/db/it/IoTDBKillQueryIT.java @@ -0,0 +1,127 @@ +/* + * 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.it; + +import org.apache.iotdb.commons.auth.entity.PrivilegeType; +import org.apache.iotdb.isession.ISession; +import org.apache.iotdb.isession.SessionConfig; +import org.apache.iotdb.isession.SessionDataSet; +import org.apache.iotdb.it.env.EnvFactory; +import org.apache.iotdb.it.framework.IoTDBTestRunner; +import org.apache.iotdb.itbase.category.ClusterIT; +import org.apache.iotdb.itbase.category.LocalStandaloneIT; +import org.apache.iotdb.rpc.StatementExecutionException; + +import org.apache.tsfile.enums.TSDataType; +import org.apache.tsfile.write.record.Tablet; +import org.apache.tsfile.write.schema.MeasurementSchema; +import org.junit.AfterClass; +import org.junit.Assert; +import org.junit.BeforeClass; +import org.junit.Test; +import org.junit.experimental.categories.Category; +import org.junit.runner.RunWith; + +import java.util.List; + +import static org.apache.iotdb.commons.schema.column.ColumnHeaderConstant.QUERY_ID; +import static org.apache.iotdb.commons.schema.column.ColumnHeaderConstant.STATEMENT; +import static org.apache.iotdb.db.it.utils.TestUtils.createUser; +import static org.apache.iotdb.db.it.utils.TestUtils.grantUserSeriesPrivilege; +import static org.apache.iotdb.db.it.utils.TestUtils.grantUserSystemPrivileges; +import static org.apache.iotdb.rpc.TSStatusCode.NO_PERMISSION; + +@RunWith(IoTDBTestRunner.class) +@Category({LocalStandaloneIT.class, ClusterIT.class}) +public class IoTDBKillQueryIT { + + private static final String DATABASE = "root.kill_query"; + private static final String DEVICE = DATABASE + ".d1"; + private static final String TIMESERIES = DEVICE + ".s1"; + private static final String QUERY_USER = "kill_query_query_user"; + private static final String KILL_USER = "kill_query_kill_user"; + private static final String PASSWORD = "password123456"; + private static final String PENDING_QUERY = "SELECT s1 FROM " + DEVICE; + // Leave one default 1,000-row TsBlock after the initial Session fetch. + private static final int STREAMING_QUERY_ROW_COUNT = SessionConfig.DEFAULT_FETCH_SIZE + 1_000; + + @BeforeClass + public static void setUp() throws Exception { + EnvFactory.getEnv().getConfig().getCommonConfig().setEnforceStrongPassword(false); + EnvFactory.getEnv().initClusterEnvironment(); + try (ISession session = EnvFactory.getEnv().getSessionConnection()) { + session.executeNonQueryStatement("CREATE DATABASE " + DATABASE); + session.executeNonQueryStatement("CREATE TIMESERIES " + TIMESERIES + " INT32"); + Tablet tablet = + new Tablet( + DEVICE, + List.of(new MeasurementSchema("s1", TSDataType.INT32)), + STREAMING_QUERY_ROW_COUNT); + for (int i = 0; i < STREAMING_QUERY_ROW_COUNT; i++) { + int rowIndex = tablet.getRowSize(); + tablet.addTimestamp(rowIndex, i); + tablet.addValue("s1", rowIndex, i); + } + session.insertTablet(tablet); + } + createUser(QUERY_USER, PASSWORD); + createUser(KILL_USER, PASSWORD); + grantUserSeriesPrivilege(QUERY_USER, PrivilegeType.READ_DATA, DATABASE + ".**"); + } + + @AfterClass + public static void tearDown() throws Exception { + EnvFactory.getEnv().cleanClusterEnvironment(); + } + + @Test + public void killQueryAuthTest() throws Exception { + try (ISession querySession = EnvFactory.getEnv().getSessionConnection(QUERY_USER, PASSWORD); + ISession killSession = EnvFactory.getEnv().getSessionConnection(KILL_USER, PASSWORD); + ISession adminSession = EnvFactory.getEnv().getSessionConnection(); + SessionDataSet pendingDataSet = querySession.executeQueryStatement(PENDING_QUERY)) { + // Keep QUERY_USER's query active by not consuming pendingDataSet. + String queryId = getQueryId(adminSession, PENDING_QUERY); + + StatementExecutionException exception = + Assert.assertThrows( + StatementExecutionException.class, + () -> killSession.executeNonQueryStatement("KILL QUERY '" + queryId + "'")); + Assert.assertEquals(NO_PERMISSION.getStatusCode(), exception.getStatusCode()); + Assert.assertEquals(queryId, getQueryId(adminSession, PENDING_QUERY)); + + grantUserSystemPrivileges(KILL_USER, PrivilegeType.SYSTEM); + killSession.executeNonQueryStatement("KILL QUERY '" + queryId + "'"); + } + } + + private static String getQueryId(ISession session, String query) throws Exception { + try (SessionDataSet dataSet = session.executeQueryStatement("SHOW QUERIES")) { + SessionDataSet.DataIterator iterator = dataSet.iterator(); + while (iterator.next()) { + if (query.equals(iterator.getString(STATEMENT))) { + return iterator.getString(QUERY_ID); + } + } + } + Assert.fail("Query not found: " + query); + return null; + } +} diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/relational/security/TreeAccessCheckVisitor.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/relational/security/TreeAccessCheckVisitor.java index 566898f7b8f8..1132c7c990c7 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/relational/security/TreeAccessCheckVisitor.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/relational/security/TreeAccessCheckVisitor.java @@ -1783,9 +1783,9 @@ public TSStatus visitCountTimeSlotList( @Override public TSStatus visitKillQuery(KillQueryStatement statement, TreeAccessCheckContext context) { - if (checkHasGlobalAuth( + if (!checkHasGlobalAuth( context.setAuditLogOperation(AuditLogOperation.CONTROL), - PrivilegeType.MAINTAIN, + PrivilegeType.SYSTEM, () -> "")) { statement.setAllowedUsername(context.getUsername()); } From aebc2fc3883586f28789e9c51bdd8ed82d41c9ed Mon Sep 17 00:00:00 2001 From: shuwenwei Date: Thu, 30 Jul 2026 11:50:17 +0800 Subject: [PATCH 3/4] Restore MAINTAIN checks for query management --- .../plan/relational/security/TreeAccessCheckVisitor.java | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/relational/security/TreeAccessCheckVisitor.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/relational/security/TreeAccessCheckVisitor.java index 1132c7c990c7..48a60b11e885 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/relational/security/TreeAccessCheckVisitor.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/relational/security/TreeAccessCheckVisitor.java @@ -1785,7 +1785,7 @@ public TSStatus visitCountTimeSlotList( public TSStatus visitKillQuery(KillQueryStatement statement, TreeAccessCheckContext context) { if (!checkHasGlobalAuth( context.setAuditLogOperation(AuditLogOperation.CONTROL), - PrivilegeType.SYSTEM, + PrivilegeType.MAINTAIN, () -> "")) { statement.setAllowedUsername(context.getUsername()); } @@ -1963,7 +1963,7 @@ public TSStatus visitShowDataNodes( @Override public TSStatus visitShowQueries(ShowQueriesStatement statement, TreeAccessCheckContext context) { - if (!checkHasGlobalAuth(context, PrivilegeType.SYSTEM, () -> "")) { + if (!checkHasGlobalAuth(context, PrivilegeType.MAINTAIN, () -> "")) { statement.setAllowedUsername(context.getUsername()); } return SUCCEED; From b1ce5401bf69741045a046afb3819ceedf63c6d4 Mon Sep 17 00:00:00 2001 From: shuwenwei Date: Thu, 30 Jul 2026 11:53:41 +0800 Subject: [PATCH 4/4] Merge query management integration tests --- .../apache/iotdb/db/it/IoTDBKillQueryIT.java | 127 ------------------ .../iotdb/db/it/IoTDBShowQueriesIT.java | 39 ++++++ 2 files changed, 39 insertions(+), 127 deletions(-) delete mode 100644 integration-test/src/test/java/org/apache/iotdb/db/it/IoTDBKillQueryIT.java diff --git a/integration-test/src/test/java/org/apache/iotdb/db/it/IoTDBKillQueryIT.java b/integration-test/src/test/java/org/apache/iotdb/db/it/IoTDBKillQueryIT.java deleted file mode 100644 index 27367c0a3f7a..000000000000 --- a/integration-test/src/test/java/org/apache/iotdb/db/it/IoTDBKillQueryIT.java +++ /dev/null @@ -1,127 +0,0 @@ -/* - * 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.it; - -import org.apache.iotdb.commons.auth.entity.PrivilegeType; -import org.apache.iotdb.isession.ISession; -import org.apache.iotdb.isession.SessionConfig; -import org.apache.iotdb.isession.SessionDataSet; -import org.apache.iotdb.it.env.EnvFactory; -import org.apache.iotdb.it.framework.IoTDBTestRunner; -import org.apache.iotdb.itbase.category.ClusterIT; -import org.apache.iotdb.itbase.category.LocalStandaloneIT; -import org.apache.iotdb.rpc.StatementExecutionException; - -import org.apache.tsfile.enums.TSDataType; -import org.apache.tsfile.write.record.Tablet; -import org.apache.tsfile.write.schema.MeasurementSchema; -import org.junit.AfterClass; -import org.junit.Assert; -import org.junit.BeforeClass; -import org.junit.Test; -import org.junit.experimental.categories.Category; -import org.junit.runner.RunWith; - -import java.util.List; - -import static org.apache.iotdb.commons.schema.column.ColumnHeaderConstant.QUERY_ID; -import static org.apache.iotdb.commons.schema.column.ColumnHeaderConstant.STATEMENT; -import static org.apache.iotdb.db.it.utils.TestUtils.createUser; -import static org.apache.iotdb.db.it.utils.TestUtils.grantUserSeriesPrivilege; -import static org.apache.iotdb.db.it.utils.TestUtils.grantUserSystemPrivileges; -import static org.apache.iotdb.rpc.TSStatusCode.NO_PERMISSION; - -@RunWith(IoTDBTestRunner.class) -@Category({LocalStandaloneIT.class, ClusterIT.class}) -public class IoTDBKillQueryIT { - - private static final String DATABASE = "root.kill_query"; - private static final String DEVICE = DATABASE + ".d1"; - private static final String TIMESERIES = DEVICE + ".s1"; - private static final String QUERY_USER = "kill_query_query_user"; - private static final String KILL_USER = "kill_query_kill_user"; - private static final String PASSWORD = "password123456"; - private static final String PENDING_QUERY = "SELECT s1 FROM " + DEVICE; - // Leave one default 1,000-row TsBlock after the initial Session fetch. - private static final int STREAMING_QUERY_ROW_COUNT = SessionConfig.DEFAULT_FETCH_SIZE + 1_000; - - @BeforeClass - public static void setUp() throws Exception { - EnvFactory.getEnv().getConfig().getCommonConfig().setEnforceStrongPassword(false); - EnvFactory.getEnv().initClusterEnvironment(); - try (ISession session = EnvFactory.getEnv().getSessionConnection()) { - session.executeNonQueryStatement("CREATE DATABASE " + DATABASE); - session.executeNonQueryStatement("CREATE TIMESERIES " + TIMESERIES + " INT32"); - Tablet tablet = - new Tablet( - DEVICE, - List.of(new MeasurementSchema("s1", TSDataType.INT32)), - STREAMING_QUERY_ROW_COUNT); - for (int i = 0; i < STREAMING_QUERY_ROW_COUNT; i++) { - int rowIndex = tablet.getRowSize(); - tablet.addTimestamp(rowIndex, i); - tablet.addValue("s1", rowIndex, i); - } - session.insertTablet(tablet); - } - createUser(QUERY_USER, PASSWORD); - createUser(KILL_USER, PASSWORD); - grantUserSeriesPrivilege(QUERY_USER, PrivilegeType.READ_DATA, DATABASE + ".**"); - } - - @AfterClass - public static void tearDown() throws Exception { - EnvFactory.getEnv().cleanClusterEnvironment(); - } - - @Test - public void killQueryAuthTest() throws Exception { - try (ISession querySession = EnvFactory.getEnv().getSessionConnection(QUERY_USER, PASSWORD); - ISession killSession = EnvFactory.getEnv().getSessionConnection(KILL_USER, PASSWORD); - ISession adminSession = EnvFactory.getEnv().getSessionConnection(); - SessionDataSet pendingDataSet = querySession.executeQueryStatement(PENDING_QUERY)) { - // Keep QUERY_USER's query active by not consuming pendingDataSet. - String queryId = getQueryId(adminSession, PENDING_QUERY); - - StatementExecutionException exception = - Assert.assertThrows( - StatementExecutionException.class, - () -> killSession.executeNonQueryStatement("KILL QUERY '" + queryId + "'")); - Assert.assertEquals(NO_PERMISSION.getStatusCode(), exception.getStatusCode()); - Assert.assertEquals(queryId, getQueryId(adminSession, PENDING_QUERY)); - - grantUserSystemPrivileges(KILL_USER, PrivilegeType.SYSTEM); - killSession.executeNonQueryStatement("KILL QUERY '" + queryId + "'"); - } - } - - private static String getQueryId(ISession session, String query) throws Exception { - try (SessionDataSet dataSet = session.executeQueryStatement("SHOW QUERIES")) { - SessionDataSet.DataIterator iterator = dataSet.iterator(); - while (iterator.next()) { - if (query.equals(iterator.getString(STATEMENT))) { - return iterator.getString(QUERY_ID); - } - } - } - Assert.fail("Query not found: " + query); - return null; - } -} diff --git a/integration-test/src/test/java/org/apache/iotdb/db/it/IoTDBShowQueriesIT.java b/integration-test/src/test/java/org/apache/iotdb/db/it/IoTDBShowQueriesIT.java index efd7ba7fc24f..991cda2c7508 100644 --- a/integration-test/src/test/java/org/apache/iotdb/db/it/IoTDBShowQueriesIT.java +++ b/integration-test/src/test/java/org/apache/iotdb/db/it/IoTDBShowQueriesIT.java @@ -27,6 +27,7 @@ import org.apache.iotdb.it.framework.IoTDBTestRunner; import org.apache.iotdb.itbase.category.ClusterIT; import org.apache.iotdb.itbase.category.LocalStandaloneIT; +import org.apache.iotdb.rpc.StatementExecutionException; import org.apache.tsfile.enums.TSDataType; import org.apache.tsfile.write.record.Tablet; @@ -40,10 +41,12 @@ import java.util.List; +import static org.apache.iotdb.commons.schema.column.ColumnHeaderConstant.QUERY_ID; import static org.apache.iotdb.commons.schema.column.ColumnHeaderConstant.STATEMENT; import static org.apache.iotdb.db.it.utils.TestUtils.createUser; import static org.apache.iotdb.db.it.utils.TestUtils.grantUserSeriesPrivilege; import static org.apache.iotdb.db.it.utils.TestUtils.grantUserSystemPrivileges; +import static org.apache.iotdb.rpc.TSStatusCode.NO_PERMISSION; @RunWith(IoTDBTestRunner.class) @Category({LocalStandaloneIT.class, ClusterIT.class}) @@ -54,6 +57,7 @@ public class IoTDBShowQueriesIT { private static final String TIMESERIES = DEVICE + ".s1"; private static final String QUERY_USER = "show_queries_query_user"; private static final String SHOW_USER = "show_queries_show_user"; + private static final String KILL_USER = "show_queries_kill_user"; private static final String PASSWORD = "password123456"; private static final String PENDING_QUERY = "SELECT s1 FROM " + DEVICE; // Leave one default 1,000-row TsBlock after the initial Session fetch. @@ -80,6 +84,7 @@ public static void setUp() throws Exception { } createUser(QUERY_USER, PASSWORD); createUser(SHOW_USER, PASSWORD); + createUser(KILL_USER, PASSWORD); grantUserSeriesPrivilege(QUERY_USER, PrivilegeType.READ_DATA, DATABASE + ".**"); } @@ -102,6 +107,27 @@ public void showQueriesAuthTest() throws Exception { } } + @Test + public void killQueryAuthTest() throws Exception { + try (ISession querySession = EnvFactory.getEnv().getSessionConnection(QUERY_USER, PASSWORD); + ISession killSession = EnvFactory.getEnv().getSessionConnection(KILL_USER, PASSWORD); + ISession adminSession = EnvFactory.getEnv().getSessionConnection(); + SessionDataSet pendingDataSet = querySession.executeQueryStatement(PENDING_QUERY)) { + // Keep QUERY_USER's query active by not consuming pendingDataSet. + String queryId = getQueryId(adminSession, PENDING_QUERY); + + StatementExecutionException exception = + Assert.assertThrows( + StatementExecutionException.class, + () -> killSession.executeNonQueryStatement("KILL QUERY '" + queryId + "'")); + Assert.assertEquals(NO_PERMISSION.getStatusCode(), exception.getStatusCode()); + Assert.assertEquals(queryId, getQueryId(adminSession, PENDING_QUERY)); + + grantUserSystemPrivileges(KILL_USER, PrivilegeType.SYSTEM); + killSession.executeNonQueryStatement("KILL QUERY '" + queryId + "'"); + } + } + private static boolean containsQuery(ISession session, String query) throws Exception { try (SessionDataSet dataSet = session.executeQueryStatement("SHOW QUERIES")) { SessionDataSet.DataIterator iterator = dataSet.iterator(); @@ -113,4 +139,17 @@ private static boolean containsQuery(ISession session, String query) throws Exce return false; } } + + private static String getQueryId(ISession session, String query) throws Exception { + try (SessionDataSet dataSet = session.executeQueryStatement("SHOW QUERIES")) { + SessionDataSet.DataIterator iterator = dataSet.iterator(); + while (iterator.next()) { + if (query.equals(iterator.getString(STATEMENT))) { + return iterator.getString(QUERY_ID); + } + } + } + Assert.fail("Query not found: " + query); + return null; + } }