executor, plan: add revoke role support (#9771)

Author: Lingyu Song

executor/simple.go

@@ -85,6 +85,8 @@ func (e *SimpleExec) Next(ctx context.Context, req *chunk.RecordBatch) (err erro
 		err = e.executeDropStats(x)
 	case *ast.SetRoleStmt:
 		err = e.executeSetRole(x)
+	case *ast.RevokeRoleStmt:
+		err = e.executeRevokeRole(x)
 	}
 	e.done = true
 	return err
@@ -165,6 +167,52 @@ func (e *SimpleExec) executeBegin(ctx context.Context, s *ast.BeginStmt) error {
 	return nil
 }
 
+func (e *SimpleExec) executeRevokeRole(s *ast.RevokeRoleStmt) error {
+	for _, role := range s.Roles {
+		exists, err := userExists(e.ctx, role.Username, role.Hostname)
+		if err != nil {
+			return errors.Trace(err)
+		}
+		if !exists {
+			return ErrCannotUser.GenWithStackByArgs("REVOKE ROLE", role.String())
+		}
+	}
+
+	// begin a transaction to insert role graph edges.
+	if _, err := e.ctx.(sqlexec.SQLExecutor).Execute(context.Background(), "begin"); err != nil {
+		return errors.Trace(err)
+	}
+	for _, user := range s.Users {
+		exists, err := userExists(e.ctx, user.Username, user.Hostname)
+		if err != nil {
+			return errors.Trace(err)
+		}
+		if !exists {
+			if _, err := e.ctx.(sqlexec.SQLExecutor).Execute(context.Background(), "rollback"); err != nil {
+				return errors.Trace(err)
+			}
+			return ErrCannotUser.GenWithStackByArgs("REVOKE ROLE", user.String())
+		}
+		for _, role := range s.Roles {
+			if role.Hostname == "" {
+				role.Hostname = "%"
+			}
+			sql := fmt.Sprintf(`DELETE IGNORE FROM %s.%s WHERE FROM_HOST='%s' and FROM_USER='%s' and TO_HOST='%s' and TO_USER='%s'`, mysql.SystemDB, mysql.RoleEdgeTable, role.Hostname, role.Username, user.Hostname, user.Username)
+			if _, err := e.ctx.(sqlexec.SQLExecutor).Execute(context.Background(), sql); err != nil {
+				if _, err := e.ctx.(sqlexec.SQLExecutor).Execute(context.Background(), "rollback"); err != nil {
+					return errors.Trace(err)
+				}
+				return ErrCannotUser.GenWithStackByArgs("REVOKE ROLE", role.String())
+			}
+		}
+	}
+	if _, err := e.ctx.(sqlexec.SQLExecutor).Execute(context.Background(), "commit"); err != nil {
+		return err
+	}
+	err := domain.GetDomain(e.ctx).PrivilegeHandle().Update(e.ctx.(sessionctx.Context))
+	return errors.Trace(err)
+}
+
 func (e *SimpleExec) executeCommit(s *ast.CommitStmt) {
 	e.ctx.GetSessionVars().SetStatusFlag(mysql.ServerStatusInTrans, false)
 }

executor/simple_test.go

@@ -141,6 +141,23 @@ func (s *testSuite3) TestRole(c *C) {
 	tk.MustExec(dropRoleSQL)
 	dropRoleSQL = `DROP ROLE IF EXISTS 'r_3'@'localhost' ;`
 	tk.MustExec(dropRoleSQL)
+
+	// Test for revoke role
+	createRoleSQL = `CREATE ROLE 'test'@'localhost', r_1, r_2;`
+	tk.MustExec(createRoleSQL)
+	tk.MustExec("insert into mysql.role_edges (FROM_HOST,FROM_USER,TO_HOST,TO_USER) values ('localhost','test','%','root')")
+	tk.MustExec("insert into mysql.role_edges (FROM_HOST,FROM_USER,TO_HOST,TO_USER) values ('%','r_1','%','root')")
+	tk.MustExec("insert into mysql.role_edges (FROM_HOST,FROM_USER,TO_HOST,TO_USER) values ('%','r_2','%','root')")
+	_, err = tk.Exec("revoke test@localhost, r_1 from root;")
+	c.Check(err, IsNil)
+	_, err = tk.Exec("revoke `r_2`@`%` from root, u_2;")
+	c.Check(err, NotNil)
+	_, err = tk.Exec("revoke `r_2`@`%` from root;")
+	c.Check(err, IsNil)
+	result = tk.MustQuery(`SELECT * FROM mysql.default_roles WHERE DEFAULT_ROLE_USER="test" and DEFAULT_ROLE_HOST="localhost"`)
+	result.Check(nil)
+	dropRoleSQL = `DROP ROLE 'test'@'localhost', r_1, r_2;`
+	tk.MustExec(dropRoleSQL)
 }
 
 func (s *testSuite3) TestUser(c *C) {

planner/core/planbuilder.go

@@ -268,7 +268,7 @@ func (b *PlanBuilder) Build(node ast.Node) (Plan, error) {
 	case *ast.BinlogStmt, *ast.FlushStmt, *ast.UseStmt,
 		*ast.BeginStmt, *ast.CommitStmt, *ast.RollbackStmt, *ast.CreateUserStmt, *ast.SetPwdStmt,
 		*ast.GrantStmt, *ast.DropUserStmt, *ast.AlterUserStmt, *ast.RevokeStmt, *ast.KillStmt, *ast.DropStatsStmt,
-		*ast.GrantRoleStmt, *ast.SetRoleStmt:
+		*ast.GrantRoleStmt, *ast.RevokeRoleStmt, *ast.SetRoleStmt:
 		return b.buildSimple(node.(ast.StmtNode))
 	case ast.DDLNode:
 		return b.buildDDL(x)
@@ -1105,6 +1105,8 @@ func (b *PlanBuilder) buildSimple(node ast.StmtNode) (Plan, error) {
 		b.visitInfo = appendVisitInfo(b.visitInfo, mysql.GrantPriv, "", "", "", err)
 	case *ast.RevokeStmt:
 		b.visitInfo = appendVisitInfo(b.visitInfo, mysql.SuperPriv, "", "", "", nil)
+	case *ast.RevokeRoleStmt:
+		b.visitInfo = appendVisitInfo(b.visitInfo, mysql.SuperPriv, "", "", "", nil)
 	case *ast.KillStmt:
 		// If you have the SUPER privilege, you can kill all threads and statements.
 		// Otherwise, you can kill only your own threads and statements.