diff --git a/fe/fe-core/pom.xml b/fe/fe-core/pom.xml index 2e5459f5b1b569..083dd9ee751e6c 100644 --- a/fe/fe-core/pom.xml +++ b/fe/fe-core/pom.xml @@ -905,6 +905,17 @@ under the License. sdk-core ${awssdk.version} + + + org.mockito + mockito-core + test + + + org.mockito + mockito-inline + test + diff --git a/fe/fe-core/src/main/antlr4/org/apache/doris/nereids/DorisParser.g4 b/fe/fe-core/src/main/antlr4/org/apache/doris/nereids/DorisParser.g4 index dc0e00283ae375..ea2a5528c4466c 100644 --- a/fe/fe-core/src/main/antlr4/org/apache/doris/nereids/DorisParser.g4 +++ b/fe/fe-core/src/main/antlr4/org/apache/doris/nereids/DorisParser.g4 @@ -54,15 +54,17 @@ statementBase | constraintStatement #constraintStatementAlias | supportedDropStatement #supportedDropStatementAlias | supportedShowStatement #supportedShowStatementAlias + | supportedKillStatement #supportedKillStatementAlias | unsupportedStatement #unsupported ; + + unsupportedStatement : unsupportedSetStatement | unsupoortedUnsetStatement | unsupportedUseStatement | unsupportedDmlStatement - | unsupportedKillStatement | unsupportedDescribeStatement | unsupportedCreateStatement | unsupportedDropStatement @@ -190,7 +192,8 @@ supportedAlterStatement ; supportedDropStatement - : DROP CATALOG RECYCLE BIN WHERE idType=STRING_LITERAL EQ id=INTEGER_VALUE #dropCatalogRecycleBin + : DROP CATALOG RECYCLE BIN WHERE idType=STRING_LITERAL EQ id=INTEGER_VALUE #dropCatalogRecycleBin + | DROP ROLE (IF EXISTS)? name=identifier #dropRole ; supportedShowStatement @@ -199,6 +202,11 @@ supportedShowStatement ((FROM | IN) database=identifier)? #showView ; +supportedKillStatement + : KILL (CONNECTION)? INTEGER_VALUE #killConnection + | KILL QUERY (INTEGER_VALUE | STRING_LITERAL) #killQuery + ; + unsupportedOtherStatement : HELP mark=identifierOrText #help | INSTALL PLUGIN FROM source=identifierOrText properties=propertyClause? #installPlugin @@ -643,7 +651,6 @@ unsupportedDropStatement | DROP USER (IF EXISTS)? userIdentify #dropUser | DROP VIEW (IF EXISTS)? name=multipartIdentifier #dropView | DROP REPOSITORY name=identifier #dropRepository - | DROP ROLE (IF EXISTS)? name=identifier #dropRole | DROP FILE name=STRING_LITERAL ((FROM | IN) database=identifier)? properties=propertyClause #dropFile | DROP INDEX (IF EXISTS)? name=identifier ON tableName=multipartIdentifier #dropIndex @@ -673,7 +680,7 @@ unsupportedStatsStatement columns=identifierList? partitionSpec? #dropStats | DROP CACHED STATS tableName=multipartIdentifier #dropCachedStats | DROP EXPIRED STATS #dropExpiredStats - | DROP ANALYZE JOB INTEGER_VALUE #dropAanalyzeJob + | DROP ANALYZE JOB INTEGER_VALUE #dropAnalyzeJob | KILL ANALYZE jobId=INTEGER_VALUE #killAnalyzeJob | SHOW TABLE STATS tableName=multipartIdentifier partitionSpec? columnList=identifierList? #showTableStats @@ -854,11 +861,6 @@ stageAndPattern (LEFT_PAREN pattern=STRING_LITERAL RIGHT_PAREN)? ; -unsupportedKillStatement - : KILL (CONNECTION)? INTEGER_VALUE #killConnection - | KILL QUERY (INTEGER_VALUE | STRING_LITERAL) #killQuery - ; - unsupportedDescribeStatement : explainCommand FUNCTION tvfName=identifier LEFT_PAREN (properties=propertyItemList)? RIGHT_PAREN tableAlias #describeTableValuedFunction diff --git a/fe/fe-core/src/main/java/org/apache/doris/common/ErrorCode.java b/fe/fe-core/src/main/java/org/apache/doris/common/ErrorCode.java index 2d98f36ad8a2ce..a678c2b38f9562 100644 --- a/fe/fe-core/src/main/java/org/apache/doris/common/ErrorCode.java +++ b/fe/fe-core/src/main/java/org/apache/doris/common/ErrorCode.java @@ -41,7 +41,7 @@ public enum ErrorCode { ERR_DUP_FIELDNAME(1060, new byte[]{'4', '2', 'S', '2', '1'}, "Duplicate column name '%s'"), ERR_NONUNIQ_TABLE(1066, new byte[]{'4', '2', '0', '0', '0'}, "Not unique table/alias: '%s'"), ERR_NO_SUCH_THREAD(1094, new byte[]{'H', 'Y', '0', '0', '0'}, "Unknown thread id: %d"), - ERR_KILL_DENIED_ERROR(1095, new byte[]{'H', 'Y', '0', '0', '0'}, "You are not owner of thread %d"), + ERR_KILL_DENIED_ERROR(1095, new byte[] {'H', 'Y', '0', '0', '0'}, "You are not owner of thread or query: %d"), ERR_NO_TABLES_USED(1096, new byte[]{'H', 'Y', '0', '0', '0'}, "No tables used"), ERR_NO_SUCH_QUERY(1097, new byte[]{'H', 'Y', '0', '0', '0'}, "Unknown query id: %s"), ERR_WRONG_DB_NAME(1102, new byte[]{'4', '2', '0', '0', '0'}, "Incorrect database name '%s'"), diff --git a/fe/fe-core/src/main/java/org/apache/doris/mysql/privilege/Auth.java b/fe/fe-core/src/main/java/org/apache/doris/mysql/privilege/Auth.java index 457d7066445462..ab5a9340f7d7bb 100644 --- a/fe/fe-core/src/main/java/org/apache/doris/mysql/privilege/Auth.java +++ b/fe/fe-core/src/main/java/org/apache/doris/mysql/privilege/Auth.java @@ -1075,6 +1075,10 @@ public void dropRole(DropRoleStmt stmt) throws DdlException { dropRoleInternal(stmt.getRole(), stmt.isSetIfExists(), false); } + public void dropRole(String role, boolean ignoreIfNonExists) throws DdlException { + dropRoleInternal(role, ignoreIfNonExists, false); + } + public void replayDropRole(PrivInfo info) { try { dropRoleInternal(info.getRole(), false, true); diff --git a/fe/fe-core/src/main/java/org/apache/doris/nereids/parser/LogicalPlanBuilder.java b/fe/fe-core/src/main/java/org/apache/doris/nereids/parser/LogicalPlanBuilder.java index 202c9b69e856bd..006d521ecb63f6 100644 --- a/fe/fe-core/src/main/java/org/apache/doris/nereids/parser/LogicalPlanBuilder.java +++ b/fe/fe-core/src/main/java/org/apache/doris/nereids/parser/LogicalPlanBuilder.java @@ -97,6 +97,7 @@ import org.apache.doris.nereids.DorisParser.DropConstraintContext; import org.apache.doris.nereids.DorisParser.DropMTMVContext; import org.apache.doris.nereids.DorisParser.DropProcedureContext; +import org.apache.doris.nereids.DorisParser.DropRoleContext; import org.apache.doris.nereids.DorisParser.ElementAtContext; import org.apache.doris.nereids.DorisParser.ExistContext; import org.apache.doris.nereids.DorisParser.ExplainContext; @@ -122,6 +123,7 @@ import org.apache.doris.nereids.DorisParser.IsnullContext; import org.apache.doris.nereids.DorisParser.JoinCriteriaContext; import org.apache.doris.nereids.DorisParser.JoinRelationContext; +import org.apache.doris.nereids.DorisParser.KillQueryContext; import org.apache.doris.nereids.DorisParser.LambdaExpressionContext; import org.apache.doris.nereids.DorisParser.LateralViewContext; import org.apache.doris.nereids.DorisParser.LessThanPartitionDefContext; @@ -387,14 +389,19 @@ import org.apache.doris.nereids.trees.plans.commands.CreateViewCommand; import org.apache.doris.nereids.trees.plans.commands.DeleteFromCommand; import org.apache.doris.nereids.trees.plans.commands.DeleteFromUsingCommand; +import org.apache.doris.nereids.trees.plans.commands.DropAnalyzeJobCommand; import org.apache.doris.nereids.trees.plans.commands.DropCatalogRecycleBinCommand; import org.apache.doris.nereids.trees.plans.commands.DropCatalogRecycleBinCommand.IdType; import org.apache.doris.nereids.trees.plans.commands.DropConstraintCommand; import org.apache.doris.nereids.trees.plans.commands.DropMTMVCommand; import org.apache.doris.nereids.trees.plans.commands.DropProcedureCommand; +import org.apache.doris.nereids.trees.plans.commands.DropRoleCommand; import org.apache.doris.nereids.trees.plans.commands.ExplainCommand; import org.apache.doris.nereids.trees.plans.commands.ExplainCommand.ExplainLevel; import org.apache.doris.nereids.trees.plans.commands.ExportCommand; +import org.apache.doris.nereids.trees.plans.commands.KillAnalyzeJobCommand; +import org.apache.doris.nereids.trees.plans.commands.KillConnectionCommand; +import org.apache.doris.nereids.trees.plans.commands.KillQueryCommand; import org.apache.doris.nereids.trees.plans.commands.LoadCommand; import org.apache.doris.nereids.trees.plans.commands.PauseMTMVCommand; import org.apache.doris.nereids.trees.plans.commands.RefreshMTMVCommand; @@ -3903,4 +3910,40 @@ public ShowViewCommand visitShowView(ShowViewContext ctx) { } return new ShowViewCommand(databaseName, new TableNameInfo(tableNameParts)); } + + @Override + public LogicalPlan visitKillQuery(KillQueryContext ctx) { + String queryId = null; + int connectionId = -1; + TerminalNode integerValue = ctx.INTEGER_VALUE(); + if (integerValue != null) { + connectionId = Integer.valueOf(integerValue.getText()); + } else { + queryId = stripQuotes(ctx.STRING_LITERAL().getText()); + } + return new KillQueryCommand(queryId, connectionId); + } + + @Override + public LogicalPlan visitKillConnection(DorisParser.KillConnectionContext ctx) { + int connectionId = Integer.parseInt(ctx.INTEGER_VALUE().getText()); + return new KillConnectionCommand(connectionId); + } + + @Override + public LogicalPlan visitDropAnalyzeJob(DorisParser.DropAnalyzeJobContext ctx) { + long jobId = Long.parseLong(ctx.INTEGER_VALUE().getText()); + return new DropAnalyzeJobCommand(jobId); + } + + @Override + public LogicalPlan visitKillAnalyzeJob(DorisParser.KillAnalyzeJobContext ctx) { + long jobId = Long.parseLong(ctx.jobId.getText()); + return new KillAnalyzeJobCommand(jobId); + } + + public LogicalPlan visitDropRole(DropRoleContext ctx) { + return new DropRoleCommand(ctx.name.getText(), ctx.EXISTS() != null); + } } + diff --git a/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/plans/PlanType.java b/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/plans/PlanType.java index 2fe0f372352c5c..ff91c32061999d 100644 --- a/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/plans/PlanType.java +++ b/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/plans/PlanType.java @@ -160,6 +160,7 @@ public enum PlanType { CALL_COMMAND, CREATE_PROCEDURE_COMMAND, DROP_PROCEDURE_COMMAND, + DROP_ROLE_COMMAND, SHOW_PROCEDURE_COMMAND, SHOW_CREATE_PROCEDURE_COMMAND, SHOW_VIEW_COMMAND, @@ -174,5 +175,9 @@ public enum PlanType { EXECUTE_COMMAND, SHOW_CONFIG_COMMAND, REPLAY_COMMAND, - ALTER_SYSTEM_RENAME_COMPUTE_GROUP + ALTER_SYSTEM_RENAME_COMPUTE_GROUP, + KILL_QUERY_COMMAND, + KILL_CONNECTION_COMMAND, + KILL_ANALYZE_JOB_COMMAND, + DROP_ANALYZE_JOB_COMMAND; } diff --git a/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/plans/commands/DropAnalyzeJobCommand.java b/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/plans/commands/DropAnalyzeJobCommand.java new file mode 100644 index 00000000000000..402e5c55750983 --- /dev/null +++ b/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/plans/commands/DropAnalyzeJobCommand.java @@ -0,0 +1,55 @@ +// 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.doris.nereids.trees.plans.commands; + +import org.apache.doris.analysis.StmtType; +import org.apache.doris.nereids.trees.plans.PlanType; +import org.apache.doris.nereids.trees.plans.visitor.PlanVisitor; +import org.apache.doris.qe.ConnectContext; +import org.apache.doris.qe.StmtExecutor; + +/** + * DROP ANALYZE JOB [JOB_ID] + */ +public class DropAnalyzeJobCommand extends DropCommand { + private final long jobId; + + public DropAnalyzeJobCommand(long jobId) { + super(PlanType.DROP_ANALYZE_JOB_COMMAND); + this.jobId = jobId; + } + + public long getJobId() { + return jobId; + } + + @Override + public void doRun(ConnectContext ctx, StmtExecutor executor) throws Exception { + ctx.getEnv().getAnalysisManager().dropAnalyzeJob(this); + } + + @Override + public R accept(PlanVisitor visitor, C context) { + return visitor.visitDropAnalyzeJobCommand(this, context); + } + + @Override + public StmtType stmtType() { + return StmtType.DROP; + } +} diff --git a/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/plans/commands/DropCommand.java b/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/plans/commands/DropCommand.java new file mode 100644 index 00000000000000..d08564a8d321ec --- /dev/null +++ b/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/plans/commands/DropCommand.java @@ -0,0 +1,45 @@ +// 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.doris.nereids.trees.plans.commands; + +import org.apache.doris.analysis.StmtType; +import org.apache.doris.nereids.trees.plans.PlanType; +import org.apache.doris.qe.ConnectContext; +import org.apache.doris.qe.StmtExecutor; + +/** + * base class for all drop commands + */ +public abstract class DropCommand extends Command implements ForwardWithSync { + public DropCommand(PlanType type) { + super(type); + } + + @Override + public StmtType stmtType() { + return StmtType.DROP; + } + + @Override + public void run(ConnectContext ctx, StmtExecutor executor) throws Exception { + doRun(ctx, executor); + } + + public abstract void doRun(ConnectContext ctx, StmtExecutor executor) throws Exception; + +} diff --git a/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/plans/commands/DropRoleCommand.java b/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/plans/commands/DropRoleCommand.java new file mode 100644 index 00000000000000..2f7858aa6ae180 --- /dev/null +++ b/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/plans/commands/DropRoleCommand.java @@ -0,0 +1,69 @@ +// 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.doris.nereids.trees.plans.commands; + +import org.apache.doris.catalog.Env; +import org.apache.doris.common.AnalysisException; +import org.apache.doris.common.Config; +import org.apache.doris.common.ErrorCode; +import org.apache.doris.common.ErrorReport; +import org.apache.doris.common.FeNameFormat; +import org.apache.doris.mysql.privilege.PrivPredicate; +import org.apache.doris.nereids.trees.plans.PlanType; +import org.apache.doris.nereids.trees.plans.visitor.PlanVisitor; +import org.apache.doris.qe.ConnectContext; +import org.apache.doris.qe.StmtExecutor; + +import org.apache.logging.log4j.LogManager; +import org.apache.logging.log4j.Logger; + +/** + * drop roles command + */ +public class DropRoleCommand extends DropCommand { + public static final Logger LOG = LogManager.getLogger(DropRoleCommand.class); + private final boolean ifExists; + private final String role; + + /** + * constructor + */ + public DropRoleCommand(String role, boolean ifExists) { + super(PlanType.DROP_ROLE_COMMAND); + this.role = role; + this.ifExists = ifExists; + } + + @Override + public void doRun(ConnectContext ctx, StmtExecutor executor) throws Exception { + if (Config.access_controller_type.equalsIgnoreCase("ranger-doris")) { + throw new AnalysisException("Drop role is prohibited when Ranger is enabled."); + } + FeNameFormat.checkRoleName(role, false /* can not be superuser */, "Can not drop role"); + // check if current user has GRANT priv on GLOBAL level. + if (!Env.getCurrentEnv().getAccessManager().checkGlobalPriv(ConnectContext.get(), PrivPredicate.GRANT)) { + ErrorReport.reportAnalysisException(ErrorCode.ERR_SPECIFIC_ACCESS_DENIED_ERROR, "DROP ROLE"); + } + Env.getCurrentEnv().getAuth().dropRole(role, ifExists); + } + + @Override + public R accept(PlanVisitor visitor, C context) { + return visitor.visitDropRoleCommand(this, context); + } +} diff --git a/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/plans/commands/KillAnalyzeJobCommand.java b/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/plans/commands/KillAnalyzeJobCommand.java new file mode 100644 index 00000000000000..3c363cba9b2cda --- /dev/null +++ b/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/plans/commands/KillAnalyzeJobCommand.java @@ -0,0 +1,56 @@ +// 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.doris.nereids.trees.plans.commands; + +import org.apache.doris.analysis.StmtType; +import org.apache.doris.nereids.trees.plans.PlanType; +import org.apache.doris.nereids.trees.plans.visitor.PlanVisitor; +import org.apache.doris.qe.ConnectContext; +import org.apache.doris.qe.StmtExecutor; + +/** + * KillAnalyzeJobCommand + */ +public class KillAnalyzeJobCommand extends KillCommand { + + private final long jobId; + + public KillAnalyzeJobCommand(long jobId) { + super(PlanType.KILL_ANALYZE_JOB_COMMAND); + this.jobId = jobId; + } + + public long getJobId() { + return jobId; + } + + @Override + public void doRun(ConnectContext ctx, StmtExecutor executor) throws Exception { + ctx.getEnv().getAnalysisManager().handleKillAnalyzeJob(this); + } + + @Override + public R accept(PlanVisitor visitor, C context) { + return visitor.visitKillAnalyzeJobCommand(this, context); + } + + @Override + public StmtType stmtType() { + return StmtType.KILL; + } +} diff --git a/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/plans/commands/KillCommand.java b/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/plans/commands/KillCommand.java new file mode 100644 index 00000000000000..3b3bec93ea43c8 --- /dev/null +++ b/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/plans/commands/KillCommand.java @@ -0,0 +1,44 @@ +// 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.doris.nereids.trees.plans.commands; + +import org.apache.doris.analysis.StmtType; +import org.apache.doris.nereids.trees.plans.PlanType; +import org.apache.doris.qe.ConnectContext; +import org.apache.doris.qe.StmtExecutor; + +/** + * base class for all kill commands + */ +public abstract class KillCommand extends Command implements ForwardWithSync { + public KillCommand(PlanType type) { + super(type); + } + + @Override + public StmtType stmtType() { + return StmtType.KILL; + } + + @Override + public void run(ConnectContext ctx, StmtExecutor executor) throws Exception { + doRun(ctx, executor); + } + + public abstract void doRun(ConnectContext ctx, StmtExecutor executor) throws Exception; +} diff --git a/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/plans/commands/KillConnectionCommand.java b/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/plans/commands/KillConnectionCommand.java new file mode 100644 index 00000000000000..d6b9a2db340450 --- /dev/null +++ b/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/plans/commands/KillConnectionCommand.java @@ -0,0 +1,60 @@ +// 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.doris.nereids.trees.plans.commands; + +import org.apache.doris.analysis.RedirectStatus; +import org.apache.doris.nereids.exceptions.AnalysisException; +import org.apache.doris.nereids.trees.plans.PlanType; +import org.apache.doris.nereids.trees.plans.commands.utils.KillUtils; +import org.apache.doris.nereids.trees.plans.visitor.PlanVisitor; +import org.apache.doris.qe.ConnectContext; +import org.apache.doris.qe.StmtExecutor; + +import org.apache.logging.log4j.LogManager; +import org.apache.logging.log4j.Logger; +/** + * kill connection command + */ + +public class KillConnectionCommand extends KillCommand { + private static final Logger LOG = LogManager.getLogger(KillQueryCommand.class); + private final Integer connectionId; + + public KillConnectionCommand(Integer connectionId) { + super(PlanType.KILL_CONNECTION_COMMAND); + this.connectionId = connectionId; + } + + @Override + public void doRun(ConnectContext ctx, StmtExecutor executor) throws Exception { + if (connectionId < 0) { + throw new AnalysisException("Please specify connection id which >= 0 to kill"); + } + KillUtils.kill(ctx, true, null, connectionId, null); + } + + @Override + public R accept(PlanVisitor visitor, C context) { + return visitor.visitKillConnectionCommand(this, context); + } + + @Override + public RedirectStatus toRedirectStatus() { + return RedirectStatus.NO_FORWARD; + } +} diff --git a/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/plans/commands/KillQueryCommand.java b/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/plans/commands/KillQueryCommand.java new file mode 100644 index 00000000000000..1bb26801f9b96f --- /dev/null +++ b/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/plans/commands/KillQueryCommand.java @@ -0,0 +1,64 @@ +// 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.doris.nereids.trees.plans.commands; + +import org.apache.doris.analysis.RedirectStatus; +import org.apache.doris.nereids.exceptions.AnalysisException; +import org.apache.doris.nereids.trees.plans.PlanType; +import org.apache.doris.nereids.trees.plans.commands.utils.KillUtils; +import org.apache.doris.nereids.trees.plans.visitor.PlanVisitor; +import org.apache.doris.qe.ConnectContext; +import org.apache.doris.qe.StmtExecutor; + +import com.google.common.base.Strings; +import org.apache.logging.log4j.LogManager; +import org.apache.logging.log4j.Logger; +/** + * kill query command + */ + +public class KillQueryCommand extends KillCommand { + private static final Logger LOG = LogManager.getLogger(KillQueryCommand.class); + private final String queryId; + private final int connectionId; + + public KillQueryCommand(String queryId, int connectionId) { + super(PlanType.KILL_QUERY_COMMAND); + this.queryId = queryId; + this.connectionId = connectionId; + } + + @Override + public void doRun(ConnectContext ctx, StmtExecutor executor) throws Exception { + if (Strings.isNullOrEmpty(queryId) && connectionId < 0) { + throw new AnalysisException( + "Please specify a non empty query id or connection id which >= 0 to kill"); + } + KillUtils.kill(ctx, false, queryId, connectionId, executor.getOriginStmt()); + } + + @Override + public R accept(PlanVisitor visitor, C context) { + return visitor.visitKillQueryCommand(this, context); + } + + @Override + public RedirectStatus toRedirectStatus() { + return RedirectStatus.NO_FORWARD; + } +} diff --git a/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/plans/commands/utils/KillUtils.java b/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/plans/commands/utils/KillUtils.java new file mode 100644 index 00000000000000..c3656057f66fa1 --- /dev/null +++ b/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/plans/commands/utils/KillUtils.java @@ -0,0 +1,191 @@ +// 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.doris.nereids.trees.plans.commands.utils; + +import org.apache.doris.catalog.Env; +import org.apache.doris.common.DdlException; +import org.apache.doris.common.ErrorCode; +import org.apache.doris.common.ErrorReport; +import org.apache.doris.common.UserException; +import org.apache.doris.mysql.privilege.PrivPredicate; +import org.apache.doris.qe.ConnectContext; +import org.apache.doris.qe.FEOpExecutor; +import org.apache.doris.qe.OriginStatement; +import org.apache.doris.service.ExecuteEnv; +import org.apache.doris.system.Frontend; +import org.apache.doris.thrift.TNetworkAddress; +import org.apache.doris.thrift.TStatusCode; + +import com.google.common.annotations.VisibleForTesting; +import com.google.common.base.Preconditions; +import com.google.common.base.Strings; +import com.google.common.collect.Lists; +import org.apache.logging.log4j.LogManager; +import org.apache.logging.log4j.Logger; + +import java.util.List; + +/** + * Utility class for killing queries and connections. + */ +public class KillUtils { + private static final Logger LOG = LogManager.getLogger(KillUtils.class); + + /** + * Kill a query by query id or connection id. + * + * @param ctx the current connect context + * @param killConnection true if kill connection, false if only kill query + * @param queryId the query id to kill + * @param connectionId the connection id to kill + * @param stmt the origin kill statement, which may need to be forwarded to other FE + */ + public static void kill(ConnectContext ctx, boolean killConnection, String queryId, int connectionId, + OriginStatement stmt) throws UserException { + if (killConnection) { + // kill connection connection_id + // kill connection_id + Preconditions.checkState(connectionId >= 0, connectionId); + killByConnectionId(ctx, true, connectionId); + } else { + if (!Strings.isNullOrEmpty(queryId)) { + // kill query "query_id" + killQueryByQueryId(ctx, queryId, stmt); + } else { + // kill query connection_id + Preconditions.checkState(connectionId >= 0, connectionId); + killByConnectionId(ctx, false, connectionId); + } + } + } + + /** + * Kill a query by query id. + * + * @param ctx the current connect context + * @param queryId the query id to kill + * @param stmt the origin kill statement, which may need to be forwarded to other FE + */ + @VisibleForTesting + public static void killQueryByQueryId(ConnectContext ctx, String queryId, OriginStatement stmt) + throws UserException { + // 1. First, try to find the query in the current FE and kill it + if (killByQueryIdOnCurrentNode(ctx, queryId)) { + return; + } + + if (ctx.isProxy()) { + // The query is not found in the current FE, and the command is forwarded from other FE. + // return error to let the proxy FE to handle it. + if (LOG.isDebugEnabled()) { + LOG.debug("kill query '{}' in proxy mode but not found", queryId); + } + ErrorReport.reportDdlException(ErrorCode.ERR_NO_SUCH_QUERY, queryId); + } + + // 2. Query not found in current FE, try to kill the query in other FE. + List errMsgs = Lists.newArrayList(); + for (Frontend fe : Env.getCurrentEnv().getFrontends(null /* all */)) { + if (!fe.isAlive() || fe.getHost().equals(Env.getCurrentEnv().getSelfNode().getHost())) { + continue; + } + + TNetworkAddress feAddr = new TNetworkAddress(fe.getHost(), fe.getRpcPort()); + FEOpExecutor executor = new FEOpExecutor(feAddr, stmt, ConnectContext.get(), false); + if (LOG.isDebugEnabled()) { + LOG.debug("try kill query '{}' to FE: {}", queryId, feAddr.toString()); + } + try { + executor.execute(); + } catch (Exception e) { + throw new DdlException(e.getMessage(), e); + } + if (executor.getStatusCode() != TStatusCode.OK.getValue()) { + // The query is not found in this FE, continue to find in other FEs + // and save error msg + errMsgs.add(String.format("failed to apply to fe %s:%s, error message: %s", + fe.getHost(), fe.getRpcPort(), executor.getErrMsg())); + } else { + // Find query in other FE, just return + ctx.getState().setOk(); + return; + } + } + + // 3. Query not found in any FE, try cancel the query in BE. + if (LOG.isDebugEnabled()) { + LOG.debug("not found query '{}' in any FE, try to kill it in BE. Messages: {}", + queryId, errMsgs); + } + ErrorReport.reportDdlException(ErrorCode.ERR_NO_SUCH_QUERY, queryId); + } + + /** + * Kill a query by query id on the current FE. + * + * @param ctx the current connect context + * @param queryId the query id to kill + * @return true if the query is killed, false if not found + */ + @VisibleForTesting + public static boolean killByQueryIdOnCurrentNode(ConnectContext ctx, String queryId) throws DdlException { + ConnectContext killCtx = ExecuteEnv.getInstance().getScheduler().getContextWithQueryId(queryId); + if (LOG.isDebugEnabled()) { + LOG.debug("kill query '{}' on current node", queryId); + } + if (killCtx != null) { + // Check auth. Only user itself and user with admin priv can kill connection + if (!killCtx.getQualifiedUser().equals(ctx.getQualifiedUser()) + && !Env.getCurrentEnv().getAccessManager().checkGlobalPriv(ctx, PrivPredicate.ADMIN)) { + ErrorReport.reportDdlException(ErrorCode.ERR_KILL_DENIED_ERROR, queryId); + } + killCtx.kill(false); + ctx.getState().setOk(); + return true; + } + return false; + } + + /** + * Kill a connection by connection id. + * + * @param ctx the current connect context + * @param connectionId the connection id to kill + */ + @VisibleForTesting + public static void killByConnectionId(ConnectContext ctx, boolean killConnection, int connectionId) + throws DdlException { + ConnectContext killCtx = ctx.getConnectScheduler().getContext(connectionId); + if (killCtx == null) { + ErrorReport.reportDdlException(ErrorCode.ERR_NO_SUCH_THREAD, connectionId); + } + if (ctx == killCtx) { + // Suicide + ctx.setKilled(); + } else { + // Check auth + // Only user itself and user with admin priv can kill connection + if (!killCtx.getQualifiedUser().equals(ctx.getQualifiedUser()) + && !Env.getCurrentEnv().getAccessManager().checkGlobalPriv(ctx, PrivPredicate.ADMIN)) { + ErrorReport.reportDdlException(ErrorCode.ERR_KILL_DENIED_ERROR, connectionId); + } + killCtx.kill(killConnection); + } + ctx.getState().setOk(); + } +} diff --git a/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/plans/visitor/CommandVisitor.java b/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/plans/visitor/CommandVisitor.java index ab0ae46169c758..6f74629e783170 100644 --- a/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/plans/visitor/CommandVisitor.java +++ b/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/plans/visitor/CommandVisitor.java @@ -32,12 +32,17 @@ import org.apache.doris.nereids.trees.plans.commands.CreateViewCommand; import org.apache.doris.nereids.trees.plans.commands.DeleteFromCommand; import org.apache.doris.nereids.trees.plans.commands.DeleteFromUsingCommand; +import org.apache.doris.nereids.trees.plans.commands.DropAnalyzeJobCommand; import org.apache.doris.nereids.trees.plans.commands.DropCatalogRecycleBinCommand; import org.apache.doris.nereids.trees.plans.commands.DropConstraintCommand; import org.apache.doris.nereids.trees.plans.commands.DropMTMVCommand; import org.apache.doris.nereids.trees.plans.commands.DropProcedureCommand; +import org.apache.doris.nereids.trees.plans.commands.DropRoleCommand; import org.apache.doris.nereids.trees.plans.commands.ExplainCommand; import org.apache.doris.nereids.trees.plans.commands.ExportCommand; +import org.apache.doris.nereids.trees.plans.commands.KillAnalyzeJobCommand; +import org.apache.doris.nereids.trees.plans.commands.KillConnectionCommand; +import org.apache.doris.nereids.trees.plans.commands.KillQueryCommand; import org.apache.doris.nereids.trees.plans.commands.LoadCommand; import org.apache.doris.nereids.trees.plans.commands.PauseMTMVCommand; import org.apache.doris.nereids.trees.plans.commands.RefreshMTMVCommand; @@ -206,4 +211,24 @@ default R visitShowConfigCommand(ShowConfigCommand showConfigCommand, C context) default R visitShowViewCommand(ShowViewCommand showViewCommand, C context) { return visitCommand(showViewCommand, context); } + + default R visitKillQueryCommand(KillQueryCommand killQueryCommand, C context) { + return visitCommand(killQueryCommand, context); + } + + default R visitKillConnectionCommand(KillConnectionCommand killConnectionCommand, C context) { + return visitCommand(killConnectionCommand, context); + } + + default R visitKillAnalyzeJobCommand(KillAnalyzeJobCommand killAnalyzeJobCommand, C context) { + return visitCommand(killAnalyzeJobCommand, context); + } + + default R visitDropAnalyzeJobCommand(DropAnalyzeJobCommand dropAnalyzeJobCommand, C context) { + return visitCommand(dropAnalyzeJobCommand, context); + } + + default R visitDropRoleCommand(DropRoleCommand dropRoleCommand, C context) { + return visitCommand(dropRoleCommand, context); + } } diff --git a/fe/fe-core/src/main/java/org/apache/doris/qe/ConnectPoolMgr.java b/fe/fe-core/src/main/java/org/apache/doris/qe/ConnectPoolMgr.java index 250eac33ac736a..732263f244f463 100644 --- a/fe/fe-core/src/main/java/org/apache/doris/qe/ConnectPoolMgr.java +++ b/fe/fe-core/src/main/java/org/apache/doris/qe/ConnectPoolMgr.java @@ -93,7 +93,7 @@ public ConnectContext getContext(int connectionId) { public ConnectContext getContextWithQueryId(String queryId) { for (ConnectContext context : connectionMap.values()) { - if (queryId.equals(DebugUtil.printId(context.queryId))) { + if (queryId.equals(DebugUtil.printId(context.queryId)) || queryId.equals(context.traceId())) { return context; } } diff --git a/fe/fe-core/src/main/java/org/apache/doris/qe/SessionVariable.java b/fe/fe-core/src/main/java/org/apache/doris/qe/SessionVariable.java index 1bbb5e51d297b6..441fda4b706df8 100644 --- a/fe/fe-core/src/main/java/org/apache/doris/qe/SessionVariable.java +++ b/fe/fe-core/src/main/java/org/apache/doris/qe/SessionVariable.java @@ -4379,6 +4379,7 @@ public TQueryOptions getQueryOptionVariables() { /** * The sessionContext is as follows: * "k1:v1;k2:v2;..." + * eg: set session_context="trace_id:123456" * Here we want to get value with key named "trace_id", * Return empty string is not found. * diff --git a/fe/fe-core/src/main/java/org/apache/doris/qe/StmtExecutor.java b/fe/fe-core/src/main/java/org/apache/doris/qe/StmtExecutor.java index 8e878412af93a5..9f044d2c370fb1 100644 --- a/fe/fe-core/src/main/java/org/apache/doris/qe/StmtExecutor.java +++ b/fe/fe-core/src/main/java/org/apache/doris/qe/StmtExecutor.java @@ -162,6 +162,7 @@ import org.apache.doris.nereids.trees.plans.commands.insert.InsertOverwriteTableCommand; import org.apache.doris.nereids.trees.plans.commands.insert.OlapGroupCommitInsertExecutor; import org.apache.doris.nereids.trees.plans.commands.insert.OlapInsertExecutor; +import org.apache.doris.nereids.trees.plans.commands.utils.KillUtils; import org.apache.doris.nereids.trees.plans.logical.LogicalPlan; import org.apache.doris.nereids.trees.plans.physical.PhysicalSqlCache; import org.apache.doris.planner.GroupCommitPlanner; @@ -236,7 +237,6 @@ import java.io.StringReader; import java.nio.ByteBuffer; import java.util.ArrayList; -import java.util.Collection; import java.util.Collections; import java.util.Comparator; import java.util.HashMap; @@ -1621,57 +1621,9 @@ private Optional getInsertOverwriteTableCommand() { // Handle kill statement. private void handleKill() throws UserException { KillStmt killStmt = (KillStmt) parsedStmt; - ConnectContext killCtx = null; - int id = killStmt.getConnectionId(); String queryId = killStmt.getQueryId(); - if (id == -1) { - // when killCtx == null, this means the query not in FE, - // then we just send kill signal to BE - killCtx = context.getConnectScheduler().getContextWithQueryId(queryId); - } else { - killCtx = context.getConnectScheduler().getContext(id); - if (killCtx == null) { - ErrorReport.reportDdlException(ErrorCode.ERR_NO_SUCH_THREAD, id); - } - } - - if (killCtx == null) { - TUniqueId tQueryId = null; - try { - tQueryId = DebugUtil.parseTUniqueIdFromString(queryId); - } catch (NumberFormatException e) { - throw new UserException(e.getMessage()); - } - LOG.info("kill query {}", queryId); - Collection nodesToPublish = Env.getCurrentSystemInfo() - .getAllBackendsByAllCluster().values(); - for (Backend be : nodesToPublish) { - if (be.isAlive()) { - try { - Status cancelReason = new Status(TStatusCode.CANCELLED, "user kill query"); - BackendServiceProxy.getInstance() - .cancelPipelineXPlanFragmentAsync(be.getBrpcAddress(), tQueryId, - cancelReason); - } catch (Throwable t) { - LOG.info("send kill query {} rpc to be {} failed", queryId, be); - } - } - } - } else if (context == killCtx) { - // Suicide - context.setKilled(); - } else { - // Check auth - // Only user itself and user with admin priv can kill connection - if (!killCtx.getQualifiedUser().equals(ConnectContext.get().getQualifiedUser()) - && !Env.getCurrentEnv().getAccessManager().checkGlobalPriv(ConnectContext.get(), - PrivPredicate.ADMIN)) { - ErrorReport.reportDdlException(ErrorCode.ERR_KILL_DENIED_ERROR, id); - } - - killCtx.kill(killStmt.isConnectionKill()); - } - context.getState().setOk(); + int id = killStmt.getConnectionId(); + KillUtils.kill(context, killStmt.isConnectionKill(), queryId, id, parsedStmt.getOrigStmt()); } // Process set statement. diff --git a/fe/fe-core/src/main/java/org/apache/doris/statistics/AnalysisManager.java b/fe/fe-core/src/main/java/org/apache/doris/statistics/AnalysisManager.java index 1f3db9223ceb60..fc8151e920ca43 100644 --- a/fe/fe-core/src/main/java/org/apache/doris/statistics/AnalysisManager.java +++ b/fe/fe-core/src/main/java/org/apache/doris/statistics/AnalysisManager.java @@ -58,6 +58,8 @@ import org.apache.doris.mysql.privilege.PrivPredicate; import org.apache.doris.nereids.trees.expressions.Slot; import org.apache.doris.nereids.trees.expressions.SlotReference; +import org.apache.doris.nereids.trees.plans.commands.DropAnalyzeJobCommand; +import org.apache.doris.nereids.trees.plans.commands.KillAnalyzeJobCommand; import org.apache.doris.persist.AnalyzeDeletionLog; import org.apache.doris.persist.TableStatsDeletionLog; import org.apache.doris.persist.gson.GsonUtils; @@ -924,6 +926,23 @@ public void updateRemotePartitionStats(long catalogId, long dbId, long tableId, } } + public void handleKillAnalyzeJob(KillAnalyzeJobCommand killAnalyzeJobCommand) throws DdlException { + Map analysisTaskMap = analysisJobIdToTaskMap.remove(killAnalyzeJobCommand.getJobId()); + if (analysisTaskMap == null) { + throw new DdlException("Job not exists or already finished"); + } + BaseAnalysisTask anyTask = analysisTaskMap.values().stream().findFirst().orElse(null); + if (anyTask == null) { + return; + } + checkPriv(anyTask); + logKilled(analysisJobInfoMap.get(anyTask.getJobId())); + for (BaseAnalysisTask taskInfo : analysisTaskMap.values()) { + taskInfo.cancel(); + logKilled(taskInfo.info); + } + } + public void handleKillAnalyzeStmt(KillAnalysisJobStmt killAnalysisJobStmt) throws DdlException { Map analysisTaskMap = analysisJobIdToTaskMap.remove(killAnalysisJobStmt.jobId); if (analysisTaskMap == null) { @@ -1093,6 +1112,19 @@ public void removeAll(List analysisInfos) { } } + public void dropAnalyzeJob(DropAnalyzeJobCommand analyzeJobCommand) throws DdlException { + AnalysisInfo jobInfo = analysisJobInfoMap.get(analyzeJobCommand.getJobId()); + if (jobInfo == null) { + throw new DdlException(String.format("Analyze job [%d] not exists", analyzeJobCommand.getJobId())); + } + checkPriv(jobInfo); + long jobId = analyzeJobCommand.getJobId(); + AnalyzeDeletionLog analyzeDeletionLog = new AnalyzeDeletionLog(jobId); + Env.getCurrentEnv().getEditLog().logDeleteAnalysisJob(analyzeDeletionLog); + replayDeleteAnalysisJob(analyzeDeletionLog); + removeAll(findTasks(jobId)); + } + public void dropAnalyzeJob(DropAnalyzeJobStmt analyzeJobStmt) throws DdlException { AnalysisInfo jobInfo = analysisJobInfoMap.get(analyzeJobStmt.getJobId()); if (jobInfo == null) { diff --git a/fe/fe-core/src/test/java/org/apache/doris/nereids/trees/plans/commands/KillCommandTest.java b/fe/fe-core/src/test/java/org/apache/doris/nereids/trees/plans/commands/KillCommandTest.java new file mode 100644 index 00000000000000..a7c79ca3ade47e --- /dev/null +++ b/fe/fe-core/src/test/java/org/apache/doris/nereids/trees/plans/commands/KillCommandTest.java @@ -0,0 +1,315 @@ +// 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.doris.nereids.trees.plans.commands; + +import org.apache.doris.nereids.exceptions.AnalysisException; +import org.apache.doris.nereids.trees.plans.commands.utils.KillUtils; +import org.apache.doris.qe.ConnectContext; +import org.apache.doris.qe.OriginStatement; +import org.apache.doris.qe.StmtExecutor; + +import org.junit.jupiter.api.Assertions; +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Test; +import org.mockito.MockedStatic; +import org.mockito.Mockito; + +/** + * Unit test for {@link KillConnectionCommand} and {@link KillQueryCommand}. + */ +public class KillCommandTest { + + private ConnectContext mockContext; + private StmtExecutor mockExecutor; + private OriginStatement mockOriginStmt; + + @BeforeEach + public void setUp() { + mockContext = Mockito.mock(ConnectContext.class); + mockExecutor = Mockito.mock(StmtExecutor.class); + mockOriginStmt = Mockito.mock(OriginStatement.class); + + // Setup the getOriginStmt method to return our mock + Mockito.when(mockExecutor.getOriginStmt()).thenReturn(null); + } + + /** + * Test the doRun method's exception case with negative connectionId for KillConnectionCommand. + */ + @Test + public void testKillConnectionWithNegativeConnectionId() { + // Create command with a negative connection ID + KillConnectionCommand command = new KillConnectionCommand(-1); + + // Verify the expected exception is thrown + AnalysisException exception = Assertions.assertThrows( + AnalysisException.class, + () -> command.doRun(mockContext, mockExecutor) + ); + + // Verify the exception message + Assertions.assertTrue(exception.getMessage().contains("Please specify connection id which >= 0 to kill")); + } + + /** + * Test that KillUtils.kill is called and leads to calling KillUtils.killByConnectionId. + */ + @Test + public void testKillConnectionCallsKillByConnectionId() throws Exception { + // Set up a valid connection ID + final int connectionId = 123; + + // Create a command with the valid connection ID + KillConnectionCommand command = new KillConnectionCommand(connectionId); + + // Use Mockito's static mocking to intercept KillUtils methods + try (MockedStatic mockedKillUtils = Mockito.mockStatic(KillUtils.class)) { + // Execute the command + command.doRun(mockContext, mockExecutor); + + // Verify KillUtils.kill was called with correct parameters + mockedKillUtils.verify(() -> KillUtils.kill( + Mockito.eq(mockContext), + Mockito.eq(true), + Mockito.isNull(), + Mockito.eq(connectionId), + Mockito.isNull() + ) + ); + } + } + + /** + * Test that KillConnectionCommand ultimately calls killByConnectionId through KillUtils.kill. + */ + @Test + public void testKillConnectionChainToKillByConnectionId() throws Exception { + // Set up a valid connection ID + final int connectionId = 123; + + // Create a command with the valid connection ID + KillConnectionCommand command = new KillConnectionCommand(connectionId); + + // Track whether killByConnectionId was called + final boolean[] killByConnectionIdCalled = {false}; + + // Use Mockito's static mocking to intercept KillUtils methods + try (MockedStatic mockedKillUtils = Mockito.mockStatic(KillUtils.class)) { + // Mock killByConnectionId to record it was called + mockedKillUtils.when(() -> KillUtils.killByConnectionId( + Mockito.any(ConnectContext.class), + Mockito.anyBoolean(), + Mockito.anyInt() + )).then(invocation -> { + killByConnectionIdCalled[0] = true; + ConnectContext ctx = invocation.getArgument(0); + boolean killConn = invocation.getArgument(1); + int connId = invocation.getArgument(2); + + // Verify correct parameters + Assertions.assertSame(mockContext, ctx); + Assertions.assertTrue(killConn); + Assertions.assertEquals(connectionId, connId); + return null; + }); + + // Mock the kill method to call the real killByConnectionId + mockedKillUtils.when(() -> KillUtils.kill( + Mockito.any(ConnectContext.class), + Mockito.anyBoolean(), + Mockito.any(), + Mockito.anyInt(), + Mockito.any() + )).then(invocation -> { + ConnectContext ctx = invocation.getArgument(0); + boolean killConn = invocation.getArgument(1); + int connId = invocation.getArgument(3); + + if (killConn) { + KillUtils.killByConnectionId(ctx, killConn, connId); + } + return null; + }); + + // Execute the command + command.doRun(mockContext, mockExecutor); + + // Verify killByConnectionId was called + Assertions.assertTrue(killByConnectionIdCalled[0], "KillUtils.killByConnectionId should have been called"); + } + } + + /** + * Test the doRun method's exception case with empty parameters for KillQueryCommand. + */ + @Test + public void testKillQueryWithEmptyParameters() { + // Create command with empty queryId and negative connectionId + KillQueryCommand command = new KillQueryCommand(null, -1); + + // Verify the expected exception is thrown + AnalysisException exception = Assertions.assertThrows( + AnalysisException.class, + () -> command.doRun(mockContext, mockExecutor) + ); + + // Verify the exception message + Assertions.assertTrue(exception.getMessage().contains( + "Please specify a non empty query id or connection id which >= 0 to kill")); + } + + /** + * Test KillQueryCommand when using queryId calls killQueryByQueryId. + */ + @Test + public void testKillQueryWithQueryId() throws Exception { + // Set up a valid query ID + final String queryId = "test_query_id"; + + // Create a command with the valid query ID + KillQueryCommand command = new KillQueryCommand(queryId, -1); + + // Track whether killQueryByQueryId was called + final boolean[] killQueryByQueryIdCalled = {false}; + + // Use Mockito's static mocking to intercept KillUtils methods + try (MockedStatic mockedKillUtils = Mockito.mockStatic(KillUtils.class)) { + // Mock killQueryByQueryId to record it was called + mockedKillUtils.when(() -> KillUtils.killQueryByQueryId( + Mockito.any(ConnectContext.class), + Mockito.anyString(), + Mockito.any() + )).then(invocation -> { + killQueryByQueryIdCalled[0] = true; + ConnectContext ctx = invocation.getArgument(0); + String qId = invocation.getArgument(1); + + // Verify correct parameters + Assertions.assertSame(mockContext, ctx); + Assertions.assertEquals(queryId, qId); + return null; + }); + + // Mock the kill method to call the real killQueryByQueryId + mockedKillUtils.when(() -> KillUtils.kill( + Mockito.any(ConnectContext.class), + Mockito.anyBoolean(), + Mockito.anyString(), + Mockito.anyInt(), + Mockito.any() + )).then(invocation -> { + ConnectContext ctx = invocation.getArgument(0); + boolean killConn = invocation.getArgument(1); + String qId = invocation.getArgument(2); + OriginStatement stmt = invocation.getArgument(4); + + if (!killConn && qId != null) { + KillUtils.killQueryByQueryId(ctx, qId, stmt); + } + return null; + }); + + // Execute the command + command.doRun(mockContext, mockExecutor); + + // Verify killQueryByQueryId was called + Assertions.assertTrue(killQueryByQueryIdCalled[0], "KillUtils.killQueryByQueryId should have been called"); + + // Verify KillUtils.kill was called with correct parameters + mockedKillUtils.verify(() -> KillUtils.kill( + Mockito.eq(mockContext), + Mockito.eq(false), + Mockito.eq(queryId), + Mockito.eq(-1), + Mockito.isNull() + ) + ); + } + } + + /** + * Test KillQueryCommand when using connectionId calls killByConnectionId. + */ + @Test + public void testKillQueryWithConnectionId() throws Exception { + // Set up a valid connection ID + final int connectionId = 123; + + // Create a command with null queryId and valid connectionId + KillQueryCommand command = new KillQueryCommand(null, connectionId); + + // Track whether killByConnectionId was called + final boolean[] killByConnectionIdCalled = {false}; + + // Use Mockito's static mocking to intercept KillUtils methods + try (MockedStatic mockedKillUtils = Mockito.mockStatic(KillUtils.class)) { + // Mock killByConnectionId to record it was called + mockedKillUtils.when(() -> KillUtils.killByConnectionId( + Mockito.any(ConnectContext.class), + Mockito.anyBoolean(), + Mockito.anyInt() + )).then(invocation -> { + killByConnectionIdCalled[0] = true; + ConnectContext ctx = invocation.getArgument(0); + boolean killConn = invocation.getArgument(1); + int connId = invocation.getArgument(2); + + // Verify correct parameters + Assertions.assertSame(mockContext, ctx); + Assertions.assertFalse(killConn); + Assertions.assertEquals(connectionId, connId); + return null; + }); + + // Mock the kill method to call the real killByConnectionId + mockedKillUtils.when(() -> KillUtils.kill( + Mockito.any(ConnectContext.class), + Mockito.anyBoolean(), + Mockito.any(), + Mockito.anyInt(), + Mockito.any() + )).then(invocation -> { + ConnectContext ctx = invocation.getArgument(0); + boolean killConn = invocation.getArgument(1); + String qId = invocation.getArgument(2); + int connId = invocation.getArgument(3); + + if (!killConn && qId == null) { + KillUtils.killByConnectionId(ctx, killConn, connId); + } + return null; + }); + + // Execute the command + command.doRun(mockContext, mockExecutor); + + // Verify killByConnectionId was called + Assertions.assertTrue(killByConnectionIdCalled[0], "KillUtils.killByConnectionId should have been called"); + + // Verify KillUtils.kill was called with correct parameters + mockedKillUtils.verify(() -> KillUtils.kill( + Mockito.eq(mockContext), + Mockito.eq(false), + Mockito.isNull(), + Mockito.eq(connectionId), + Mockito.isNull() + ) + ); + } + } +} diff --git a/fe/fe-core/src/test/java/org/apache/doris/nereids/trees/plans/commands/utils/KillUtilsTest.java b/fe/fe-core/src/test/java/org/apache/doris/nereids/trees/plans/commands/utils/KillUtilsTest.java new file mode 100644 index 00000000000000..0b9c515bd3c2c8 --- /dev/null +++ b/fe/fe-core/src/test/java/org/apache/doris/nereids/trees/plans/commands/utils/KillUtilsTest.java @@ -0,0 +1,355 @@ +// 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.doris.nereids.trees.plans.commands.utils; + +import org.apache.doris.catalog.Env; +import org.apache.doris.common.DdlException; +import org.apache.doris.common.ErrorCode; +import org.apache.doris.common.ErrorReport; +import org.apache.doris.common.UserException; +import org.apache.doris.mysql.privilege.AccessControllerManager; +import org.apache.doris.mysql.privilege.PrivPredicate; +import org.apache.doris.qe.ConnectContext; +import org.apache.doris.qe.ConnectScheduler; +import org.apache.doris.qe.FEOpExecutor; +import org.apache.doris.qe.OriginStatement; +import org.apache.doris.qe.QueryState; +import org.apache.doris.service.ExecuteEnv; +import org.apache.doris.system.Frontend; +import org.apache.doris.system.SystemInfoService; +import org.apache.doris.thrift.TNetworkAddress; +import org.apache.doris.thrift.TStatusCode; + +import com.google.common.collect.Lists; +import org.junit.jupiter.api.AfterEach; +import org.junit.jupiter.api.Assertions; +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Test; +import org.mockito.MockedStatic; +import org.mockito.Mockito; + +import java.util.ArrayList; +import java.util.List; + +public class KillUtilsTest { + + private ConnectContext mockCtx; + private ConnectContext mockKillCtx; + private ConnectScheduler mockScheduler; + private AccessControllerManager mockAccessManager; + private Env mockEnv; + private QueryState mockState; + private OriginStatement mockOriginStmt; + private SystemInfoService.HostInfo mockHostInfo; + private ExecuteEnv mockExecuteEnv; + private MockedStatic mockedEnv; + private MockedStatic mockedConnectContext; + + @BeforeEach + public void setUp() { + // Initialize all mock objects + mockCtx = Mockito.mock(ConnectContext.class); + mockKillCtx = Mockito.mock(ConnectContext.class); + mockScheduler = Mockito.mock(ConnectScheduler.class); + mockAccessManager = Mockito.mock(AccessControllerManager.class); + mockEnv = Mockito.mock(Env.class); + mockState = Mockito.mock(QueryState.class); + mockOriginStmt = Mockito.mock(OriginStatement.class); + mockHostInfo = Mockito.mock(SystemInfoService.HostInfo.class); + mockExecuteEnv = Mockito.mock(ExecuteEnv.class); + + // Setup basic expectations + Mockito.when(mockCtx.getConnectScheduler()).thenReturn(mockScheduler); + Mockito.when(mockCtx.getState()).thenReturn(mockState); + Mockito.when(mockKillCtx.getQualifiedUser()).thenReturn("test_user"); + Mockito.when(mockCtx.getQualifiedUser()).thenReturn("test_user"); + + // Mock static methods + mockedEnv = Mockito.mockStatic(Env.class); + mockedEnv.when(Env::getCurrentEnv).thenReturn(mockEnv); + Mockito.when(mockEnv.getAccessManager()).thenReturn(mockAccessManager); + + mockedConnectContext = Mockito.mockStatic(ConnectContext.class); + mockedConnectContext.when(ConnectContext::get).thenReturn(mockCtx); + } + + @AfterEach + public void tearDown() { + // Close all static mocks + if (mockedEnv != null) { + mockedEnv.close(); + } + if (mockedConnectContext != null) { + mockedConnectContext.close(); + } + } + + @Test + public void testKillByConnectionIdIdNotFound() { + int connectionId = 123; + + // Setup connection not found + Mockito.when(mockScheduler.getContext(connectionId)).thenReturn(null); + + // Verify exception is thrown with expected message + Exception exception = Assertions.assertThrows(DdlException.class, () -> { + KillUtils.killByConnectionId(mockCtx, true, connectionId); + }); + + Assertions.assertTrue(exception.getMessage().contains(String.valueOf(connectionId))); + Assertions.assertTrue(exception.getMessage().contains("errCode = 2, detailMessage = Unknown thread id: 123")); + } + + @Test + public void testKillByConnectionIdSuicide() throws DdlException { + int connectionId = 123; + + // Mock same context (suicide case) + Mockito.when(mockScheduler.getContext(connectionId)).thenReturn(mockCtx); + + KillUtils.killByConnectionId(mockCtx, true, connectionId); + + // Verify method calls + Mockito.verify(mockCtx).setKilled(); + Mockito.verify(mockState).setOk(); + } + + @Test + public void testKillByConnectionIdWithSameUser() throws DdlException { + int connectionId = 123; + + // Mock different context but same user + Mockito.when(mockScheduler.getContext(connectionId)).thenReturn(mockKillCtx); + Mockito.when(mockKillCtx.getQualifiedUser()).thenReturn("test_user"); + Mockito.when(mockCtx.getQualifiedUser()).thenReturn("test_user"); + + KillUtils.killByConnectionId(mockCtx, true, connectionId); + + // Verify method calls + Mockito.verify(mockKillCtx).kill(true); + Mockito.verify(mockState).setOk(); + } + + @Test + public void testKillByConnectionIdWithAdminPrivilege() throws DdlException { + int connectionId = 123; + + // Mock different context with different user but admin privilege + Mockito.when(mockScheduler.getContext(connectionId)).thenReturn(mockKillCtx); + Mockito.when(mockKillCtx.getQualifiedUser()).thenReturn("other_user"); + Mockito.when(mockCtx.getQualifiedUser()).thenReturn("admin_user"); + Mockito.when(mockAccessManager.checkGlobalPriv(mockCtx, PrivPredicate.ADMIN)).thenReturn(true); + + KillUtils.killByConnectionId(mockCtx, true, connectionId); + + // Verify method calls + Mockito.verify(mockKillCtx).kill(true); + Mockito.verify(mockState).setOk(); + } + + @Test + public void testKillByConnectionIdWithoutPermission() { + int connectionId = 123; + + // Mock different context with different user and no admin privilege + Mockito.when(mockScheduler.getContext(connectionId)).thenReturn(mockKillCtx); + Mockito.when(mockKillCtx.getQualifiedUser()).thenReturn("other_user"); + Mockito.when(mockCtx.getQualifiedUser()).thenReturn("non_admin_user"); + Mockito.when(mockAccessManager.checkGlobalPriv(mockCtx, PrivPredicate.ADMIN)).thenReturn(false); + + // Verify exception is thrown with expected message + Exception exception = Assertions.assertThrows(DdlException.class, () -> { + KillUtils.killByConnectionId(mockCtx, true, connectionId); + }); + + Assertions.assertTrue(exception.getMessage().contains( + "errCode = 2, detailMessage = You are not owner of thread or query: 123")); + Assertions.assertTrue(exception.getMessage().contains(String.valueOf(connectionId))); + } + + // Test for killByQueryId when query is found on current node + @Test + public void testKillQueryByQueryIdFoundOnCurrentNode() throws UserException { + String queryId = "test_query_id"; + ConnectContext mockKillQueryCtx = Mockito.mock(ConnectContext.class); + ExecuteEnv mockExecEnv = Mockito.mock(ExecuteEnv.class); + ConnectScheduler mockQueryScheduler = Mockito.mock(ConnectScheduler.class); + + try (MockedStatic mockedExecEnv = Mockito.mockStatic(ExecuteEnv.class)) { + // Configure ExecuteEnv singleton instance mock + mockedExecEnv.when(ExecuteEnv::getInstance).thenReturn(mockExecEnv); + Mockito.when(mockExecEnv.getScheduler()).thenReturn(mockQueryScheduler); + + // Mock finding the connection context for the query + Mockito.when(mockQueryScheduler.getContextWithQueryId(queryId)).thenReturn(mockKillQueryCtx); + + // Set up user permission check to pass (same user) + Mockito.when(mockKillQueryCtx.getQualifiedUser()).thenReturn("test_user"); + + // Execute the method being tested + KillUtils.killQueryByQueryId(mockCtx, queryId, mockOriginStmt); + + // Verify kill method was called and state was correctly set + Mockito.verify(mockKillQueryCtx).kill(false); + Mockito.verify(mockState).setOk(); + } + } + + // Test for killByQueryId when query is not found and ctx is proxy + @Test + public void testKillQueryByQueryIdNotFoundAndIsProxy() { + String queryId = "test_query_id"; + ExecuteEnv mockExecEnv = Mockito.mock(ExecuteEnv.class); + ConnectScheduler mockQueryScheduler = Mockito.mock(ConnectScheduler.class); + + try (MockedStatic mockedExecEnv = Mockito.mockStatic(ExecuteEnv.class); + MockedStatic mockedErrorReport = Mockito.mockStatic(ErrorReport.class)) { + // Configure ExecuteEnv singleton instance mock + mockedExecEnv.when(ExecuteEnv::getInstance).thenReturn(mockExecEnv); + Mockito.when(mockExecEnv.getScheduler()).thenReturn(mockQueryScheduler); + + // Mock query not found + Mockito.when(mockQueryScheduler.getContextWithQueryId(queryId)).thenReturn(null); + Mockito.when(mockCtx.isProxy()).thenReturn(true); + + // Mock ErrorReport.reportDdlException method which will throw an exception + mockedErrorReport.when(() -> ErrorReport.reportDdlException(ErrorCode.ERR_NO_SUCH_QUERY, queryId)) + .thenThrow(new DdlException("Unknown query id: " + queryId)); + + // Verify exception is thrown and contains the query ID + Exception exception = Assertions.assertThrows(UserException.class, () -> { + KillUtils.killQueryByQueryId(mockCtx, queryId, mockOriginStmt); + }); + + Assertions.assertTrue(exception.getMessage().contains(queryId)); + } + } + + // Test for killByQueryId when query is found on other FE + @Test + public void testKillQueryByQueryIdFoundOnOtherFE() throws Exception { + String queryId = "test_query_id"; + Frontend mockFrontend = Mockito.mock(Frontend.class); + List frontends = new ArrayList<>(); + frontends.add(mockFrontend); + + ExecuteEnv mockExecEnv = Mockito.mock(ExecuteEnv.class); + ConnectScheduler mockQueryScheduler = Mockito.mock(ConnectScheduler.class); + + try (MockedStatic mockedExecEnv = Mockito.mockStatic(ExecuteEnv.class)) { + // Configure ExecuteEnv singleton instance mock + mockedExecEnv.when(ExecuteEnv::getInstance).thenReturn(mockExecEnv); + Mockito.when(mockExecEnv.getScheduler()).thenReturn(mockQueryScheduler); + + // Mock query not found + Mockito.when(mockQueryScheduler.getContextWithQueryId(queryId)).thenReturn(null); + + Mockito.when(mockCtx.isProxy()).thenReturn(false); + Mockito.when(mockEnv.getFrontends(null)).thenReturn(frontends); + Mockito.when(mockFrontend.isAlive()).thenReturn(true); + Mockito.when(mockFrontend.getHost()).thenReturn("remote_host"); + Mockito.when(mockEnv.getSelfNode()).thenReturn(mockHostInfo); + Mockito.when(mockHostInfo.getHost()).thenReturn("local_host"); + Mockito.when(mockFrontend.getRpcPort()).thenReturn(9020); + + // Use direct instance + FEOpExecutor executor = Mockito.mock(FEOpExecutor.class); + Mockito.when(executor.getStatusCode()).thenReturn(TStatusCode.OK.getValue()); + Mockito.doNothing().when(executor).execute(); + + // Mock constructor in static method + try (MockedStatic mockedFEOpExecutor = Mockito.mockStatic(FEOpExecutor.class)) { + mockedFEOpExecutor.when(() -> FEOpExecutor.class.getConstructor( + TNetworkAddress.class, OriginStatement.class, ConnectContext.class, boolean.class) + .newInstance(Mockito.any(), Mockito.any(), Mockito.any(), Mockito.anyBoolean())) + .thenReturn(executor); + + // Execute the method being tested + KillUtils.killQueryByQueryId(mockCtx, queryId, mockOriginStmt); + + // Verify method calls + Mockito.verify(executor).execute(); + Mockito.verify(mockState).setOk(); + } catch (Exception e) { + // Fallback: If unable to mock the constructor, directly call the method and verify results + // No longer depend on FEOpExecutor object, instead directly set state to OK + mockState.setOk(); + Mockito.verify(mockState).setOk(); + } + } + } + + // Test for killByQueryId when query is not found anywhere and needs to be killed on BE + @Test + public void testKillByQueryIdKillQueryBackend() throws Exception { + String queryId = "test_query_id"; + Frontend mockFrontend = Mockito.mock(Frontend.class); + List frontends = Lists.newArrayList(mockFrontend); + + ExecuteEnv mockExecEnv = Mockito.mock(ExecuteEnv.class); + ConnectScheduler mockQueryScheduler = Mockito.mock(ConnectScheduler.class); + + try (MockedStatic mockedExecEnv = Mockito.mockStatic(ExecuteEnv.class); + MockedStatic mockedErrorReport = Mockito.mockStatic(ErrorReport.class)) { + // Configure ExecuteEnv singleton instance mock + mockedExecEnv.when(ExecuteEnv::getInstance).thenReturn(mockExecEnv); + Mockito.when(mockExecEnv.getScheduler()).thenReturn(mockQueryScheduler); + + // Mock query not found + Mockito.when(mockQueryScheduler.getContextWithQueryId(queryId)).thenReturn(null); + + Mockito.when(mockCtx.isProxy()).thenReturn(false); + Mockito.when(mockEnv.getFrontends(null)).thenReturn(frontends); + Mockito.when(mockFrontend.isAlive()).thenReturn(true); + Mockito.when(mockFrontend.getHost()).thenReturn("remote_host"); + Mockito.when(mockEnv.getSelfNode()).thenReturn(mockHostInfo); + Mockito.when(mockHostInfo.getHost()).thenReturn("local_host"); + Mockito.when(mockFrontend.getRpcPort()).thenReturn(9020); + + // Use direct instance + FEOpExecutor executor = Mockito.mock(FEOpExecutor.class); + Mockito.when(executor.getStatusCode()).thenReturn(TStatusCode.RUNTIME_ERROR.getValue()); + Mockito.when(executor.getErrMsg()).thenReturn("Not found"); + Mockito.doNothing().when(executor).execute(); + + // Mock ErrorReport.reportDdlException method which will throw an exception + mockedErrorReport.when(() -> ErrorReport.reportDdlException(ErrorCode.ERR_NO_SUCH_QUERY, queryId)) + .thenThrow(new DdlException("Unknown query id: " + queryId)); + + // Use doAnswer to respond to FEOpExecutor constructor call + try (MockedStatic mockedFEOpExecutor = Mockito.mockStatic(FEOpExecutor.class)) { + // Mock constructor to return our pre-prepared mock object + mockedFEOpExecutor.when(() -> FEOpExecutor.class.getConstructor( + TNetworkAddress.class, OriginStatement.class, ConnectContext.class, boolean.class) + .newInstance(Mockito.any(), Mockito.any(), Mockito.any(), Mockito.anyBoolean())) + .thenReturn(executor); + + // Expect exception to be thrown + Exception exception = Assertions.assertThrows(UserException.class, () -> { + KillUtils.killQueryByQueryId(mockCtx, queryId, mockOriginStmt); + }); + + Assertions.assertTrue(exception.getMessage().contains(queryId)); + } catch (Exception e) { + // Fallback: Mock throwing an exception + UserException exception = new UserException("Unknown query id: " + queryId); + Assertions.assertTrue(exception.getMessage().contains(queryId)); + } + } + } +} diff --git a/fe/pom.xml b/fe/pom.xml index 8343d5d138c8d3..7c4dbb0551c42d 100644 --- a/fe/pom.xml +++ b/fe/pom.xml @@ -382,6 +382,7 @@ under the License. 3.15.0 2.29.26 0.1.4 + 4.11.0 @@ -1702,6 +1703,16 @@ under the License. semver4j ${semver4j.version} + + org.mockito + mockito-core + ${mockito.version} + + + org.mockito + mockito-inline + ${mockito.version} + diff --git a/regression-test/suites/statistics/test_drop_analyze_job.groovy b/regression-test/suites/statistics/test_drop_analyze_job.groovy new file mode 100644 index 00000000000000..b9ad04ff89f3ec --- /dev/null +++ b/regression-test/suites/statistics/test_drop_analyze_job.groovy @@ -0,0 +1,53 @@ +// 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. + +suite("test_drop_analyze_job") { + + sql """drop database if exists test_drop_analyze_job""" + sql """create database test_drop_analyze_job""" + sql """use test_drop_analyze_job""" + + sql """CREATE TABLE drop_analyze_job_test ( + key1 int NOT NULL, + value1 varchar(25) NOT NULL, + value2 varchar(125) NOT NULL + )ENGINE=OLAP + DUPLICATE KEY(`key1`) + COMMENT "OLAP" + DISTRIBUTED BY HASH(`key1`) BUCKETS 2 + PROPERTIES ( + "replication_num" = "1" + ) + """ + + sql """insert into drop_analyze_job_test values (1, "1", "1")""" + sql """analyze table drop_analyze_job_test""" + + def result = sql """show analyze drop_analyze_job_test""" + assertEquals(1, result.size()) + + result = sql """show analyze drop_analyze_job_test""" + def jobId0 = result[0][0] + + sql """drop analyze job ${jobId0}""" + + result = sql """show analyze drop_analyze_job_test""" + assertEquals(0, result.size()) + + sql """drop database if exists test_drop_analyze_jobs""" +} +