parser, execution: support kill connection_id() (#37650)

close #37261
pingcap · Oct 21, 2022 · 9d606cd · 9d606cd
1 parent 59886f5
commit 9d606cd
Show file tree

Hide file tree

Showing 7 changed files with 4,062 additions and 4,018 deletions.
diff --git a/executor/simple.go b/executor/simple.go
@@ -1554,6 +1554,14 @@ func (e *SimpleExec) executeSetPwd(ctx context.Context, s *ast.SetPwdStmt) error
 }
 
 func (e *SimpleExec) executeKillStmt(ctx context.Context, s *ast.KillStmt) error {
+	if x, ok := s.Expr.(*ast.FuncCallExpr); ok {
+		if x.FnName.L == ast.ConnectionID {
+			sm := e.ctx.GetSessionManager()
+			sm.Kill(e.ctx.GetSessionVars().ConnectionID, s.Query)
+			return nil
+		}
+		return errors.New("Invalid operation. Please use 'KILL TIDB [CONNECTION | QUERY] [connectionID | CONNECTION_ID()]' instead")
+	}
 	if !config.GetGlobalConfig().EnableGlobalKill {
 		conf := config.GetGlobalConfig()
 		if s.TiDBExtension || conf.CompatibleKillQuery {
@@ -1563,7 +1571,7 @@ func (e *SimpleExec) executeKillStmt(ctx context.Context, s *ast.KillStmt) error
 			}
 			sm.Kill(s.ConnectionID, s.Query)
 		} else {
-			err := errors.New("Invalid operation. Please use 'KILL TIDB [CONNECTION | QUERY] connectionID' instead")
+			err := errors.New("Invalid operation. Please use 'KILL TIDB [CONNECTION | QUERY] [connectionID | CONNECTION_ID()]' instead")
 			e.ctx.GetSessionVars().StmtCtx.AppendWarning(err)
 		}
 		return nil

diff --git a/executor/simple_test.go b/executor/simple_test.go
@@ -52,7 +52,7 @@ func TestKillStmt(t *testing.T) {
 	tk.MustExec("use test")
 	tk.MustExec(fmt.Sprintf("kill %d", connID))
 	result := tk.MustQuery("show warnings")
-	result.Check(testkit.Rows("Warning 1105 Invalid operation. Please use 'KILL TIDB [CONNECTION | QUERY] connectionID' instead"))
+	result.Check(testkit.Rows("Warning 1105 Invalid operation. Please use 'KILL TIDB [CONNECTION | QUERY] [connectionID | CONNECTION_ID()]' instead"))
 
 	newCfg2 := *originCfg
 	newCfg2.EnableGlobalKill = true
@@ -79,6 +79,7 @@ func TestKillStmt(t *testing.T) {
 	result = tk.MustQuery("show warnings")
 	result.Check(testkit.Rows())
 
+	tk.MustExecToErr("kill rand()", "Invalid operation. Please use 'KILL TIDB [CONNECTION | QUERY] [connectionID | CONNECTION_ID()]' instead")
 	// remote kill is tested in `tests/globalkilltest`
 }
 

diff --git a/parser/ast/misc.go b/parser/ast/misc.go
@@ -937,6 +937,8 @@ type KillStmt struct {
 	// So, "KILL TIDB" grammar is introduced, and it REQUIRES DIRECT client -> TiDB TOPOLOGY.
 	// TODO: The standard KILL grammar will be supported once we have global connectionID.
 	TiDBExtension bool
+
+	Expr ExprNode
 }
 
 // Restore implements Node interface.
@@ -948,7 +950,14 @@ func (n *KillStmt) Restore(ctx *format.RestoreCtx) error {
 	if n.Query {
 		ctx.WriteKeyWord(" QUERY")
 	}
-	ctx.WritePlainf(" %d", n.ConnectionID)
+	if n.Expr != nil {
+		ctx.WriteKeyWord(" ")
+		if err := n.Expr.Restore(ctx); err != nil {
+			return errors.Trace(err)
+		}
+	} else {
+		ctx.WritePlainf(" %d", n.ConnectionID)
+	}
 	return nil
 }