Skip to content

Commit

Permalink
add set default role
Browse files Browse the repository at this point in the history
  • Loading branch information
imtbkcat committed Apr 9, 2019
1 parent cdceeb2 commit 03c66e4
Show file tree
Hide file tree
Showing 4 changed files with 69 additions and 1 deletion.
53 changes: 53 additions & 0 deletions ast/misc.go
Original file line number Diff line number Diff line change
Expand Up @@ -42,6 +42,7 @@ var (
_ StmtNode = &RollbackStmt{}
_ StmtNode = &SetPwdStmt{}
_ StmtNode = &SetRoleStmt{}
_ StmtNode = &SetDefaultRoleStmt{}
_ StmtNode = &SetStmt{}
_ StmtNode = &UseStmt{}
_ StmtNode = &FlushStmt{}
Expand Down Expand Up @@ -847,6 +848,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
Expand Down
9 changes: 9 additions & 0 deletions parser.go

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

6 changes: 6 additions & 0 deletions parser.y
Original file line number Diff line number Diff line change
Expand Up @@ -5680,6 +5680,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:
Expand Down
2 changes: 1 addition & 1 deletion parser_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -771,7 +771,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'"},
Expand Down

0 comments on commit 03c66e4

Please sign in to comment.