From 478576d4bf034b15945454804af195c075b15fae Mon Sep 17 00:00:00 2001
From: imtbkcat <songlingyu@pingcap.com>
Date: Tue, 26 Mar 2019 21:31:23 +0800
Subject: [PATCH] add set default role

---
 ast/misc.go    | 53 ++++++++++++++++++++++++++++++++++++++++++++++++++
 parser.go      |  9 +++++++++
 parser.y       |  6 ++++++
 parser_test.go |  2 +-
 4 files changed, 69 insertions(+), 1 deletion(-)

diff --git a/ast/misc.go b/ast/misc.go
index 4056017d9..97ca3374f 100755
--- a/ast/misc.go
+++ b/ast/misc.go
@@ -41,6 +41,7 @@ var (
 	_ StmtNode = &RollbackStmt{}
 	_ StmtNode = &SetPwdStmt{}
 	_ StmtNode = &SetRoleStmt{}
+	_ StmtNode = &SetDefaultRoleStmt{}
 	_ StmtNode = &SetStmt{}
 	_ StmtNode = &UseStmt{}
 	_ StmtNode = &FlushStmt{}
@@ -813,6 +814,58 @@ func (n *SetRoleStmt) Accept(v Visitor) (Node, bool) {
 	return v.Leave(n)
 }
 
+type SetDefaultRoleStmt struct {
+	stmtNode
+
+	SetRoleOpt SetRoleStmtType
+	RoleList   []*auth.RoleIdentity
+	UserList   []*auth.UserIdentity
+}
+
+func (n *SetDefaultRoleStmt) Restore(ctx *RestoreCtx) error {
+	ctx.WriteKeyWord("SET DEFAULT ROLE")
+	switch n.SetRoleOpt {
+	case SetRoleNone:
+		ctx.WriteKeyWord(" NONE")
+	case SetRoleAll:
+		ctx.WriteKeyWord(" ALL")
+	default:
+		ctx.WriteKeyWord("")
+	}
+	for i, role := range n.RoleList {
+		ctx.WritePlain(" ")
+		err := role.Restore(ctx)
+		if err != nil {
+			return errors.Annotate(err, "An error occurred while restore SetDefaultRoleStmt.RoleList")
+		}
+		if i != len(n.RoleList)-1 {
+			ctx.WritePlain(",")
+		}
+	}
+	ctx.WritePlain(" TO")
+	for i, user := range n.UserList {
+		ctx.WritePlain(" ")
+		err := user.Restore(ctx)
+		if err != nil {
+			return errors.Annotate(err, "An error occurred while restore SetDefaultRoleStmt.UserList")
+		}
+		if i != len(n.UserList)-1 {
+			ctx.WritePlain(",")
+		}
+	}
+	return nil
+}
+
+// Accept implements Node Accept interface.
+func (n *SetDefaultRoleStmt) Accept(v Visitor) (Node, bool) {
+	newNode, skipChildren := v.Enter(n)
+	if skipChildren {
+		return v.Leave(newNode)
+	}
+	n = newNode.(*SetDefaultRoleStmt)
+	return v.Leave(n)
+}
+
 // UserSpec is used for parsing create user statement.
 type UserSpec struct {
 	User    *auth.UserIdentity
diff --git a/parser.go b/parser.go
index 5142c4f61..723e5f971 100644
--- a/parser.go
+++ b/parser.go
@@ -11169,6 +11169,15 @@ yynewstate:
 		{
 			parser.yyVAL.statement = yyS[yypt-0].item.(*ast.SetRoleStmt)
 		}
+	case 1046:
+		{
+			tmp := yyS[yypt-2].item.(*ast.SetRoleStmt)
+			parser.yyVAL.statement = &ast.SetDefaultRoleStmt{
+				SetRoleOpt: tmp.SetRoleOpt,
+				RoleList:   tmp.RoleList,
+				UserList:   yyS[yypt-0].item.([]*auth.UserIdentity),
+			}
+		}
 	case 1047:
 		{
 			parser.yyVAL.item = &ast.SetRoleStmt{SetRoleOpt: ast.SetRoleNone, RoleList: nil}
diff --git a/parser.y b/parser.y
index 5d94b1d64..77b7ee4b6 100644
--- a/parser.y
+++ b/parser.y
@@ -5613,6 +5613,12 @@ SetRoleStmt:
 SetDefaultRoleStmt:
 	"SET" "DEFAULT" "ROLE" SetDefaultRoleOpt "TO" UsernameList
 	{
+		tmp := $4.(*ast.SetRoleStmt)
+		$$ = &ast.SetDefaultRoleStmt{
+			SetRoleOpt: tmp.SetRoleOpt,
+			RoleList: tmp.RoleList,
+			UserList: $6.([]*auth.UserIdentity),
+		}
 	}
 
 SetDefaultRoleOpt:
diff --git a/parser_test.go b/parser_test.go
index 87698a3dc..f671fd3f4 100755
--- a/parser_test.go
+++ b/parser_test.go
@@ -779,7 +779,7 @@ func (s *testParserSuite) TestDBAStmt(c *C) {
 		{"SET ROLE DEFAULT", true, "SET ROLE DEFAULT"},
 		{"SET ROLE ALL", true, "SET ROLE ALL"},
 		{"SET ROLE ALL EXCEPT `role1`, `role2`", true, "SET ROLE ALL EXCEPT `role1`@`%`, `role2`@`%`"},
-		{"SET DEFAULT ROLE administrator, developer TO 'joe'@'10.0.0.1'", true, ""},
+		{"SET DEFAULT ROLE administrator, developer TO `joe`@`10.0.0.1`", true, "SET DEFAULT ROLE `administrator`@`%`, `developer`@`%` TO `joe`@`10.0.0.1`"},
 		// for set names and set vars
 		{"set names utf8, @@session.sql_mode=1;", true, "SET NAMES 'utf8', @@SESSION.`sql_mode`=1"},
 		{"set @@session.sql_mode=1, names utf8, charset utf8;", true, "SET @@SESSION.`sql_mode`=1, NAMES 'utf8', NAMES 'utf8'"},