parser, ast: add SET ROLE support

ast/functions.go

@@ -233,6 +233,7 @@ const (
 	Collation      = "collation"
 	ConnectionID   = "connection_id"
 	CurrentUser    = "current_user"
+	CurrentRole    = "current_role"
 	Database       = "database"
 	FoundRows      = "found_rows"
 	LastInsertId   = "last_insert_id"

ast/misc.go

@@ -40,6 +40,7 @@ var (
 	_ StmtNode = &PrepareStmt{}
 	_ StmtNode = &RollbackStmt{}
 	_ StmtNode = &SetPwdStmt{}
+	_ StmtNode = &SetRoleStmt{}
 	_ StmtNode = &SetStmt{}
 	_ StmtNode = &UseStmt{}
 	_ StmtNode = &FlushStmt{}
@@ -758,6 +759,60 @@ func (n *ChangeStmt) Accept(v Visitor) (Node, bool) {
 	return v.Leave(n)
 }
 
+// SetRoleStmtType is the type for FLUSH statement.
+type SetRoleStmtType int
+
+// SetRole statement types.
+const (
+	SetRoleDefault SetRoleStmtType = iota
+	SetRoleNone
+	SetRoleAll
+	SetRoleAllExcept
+	SetRoleRegular
+)
+
+type SetRoleStmt struct {
+	stmtNode
+
+	SetRoleOpt SetRoleStmtType
+	RoleList   []*auth.RoleIdentity
+}
+
+func (n *SetRoleStmt) Restore(ctx *RestoreCtx) error {
+	ctx.WriteKeyWord("SET ROLE")
+	switch n.SetRoleOpt {
+	case SetRoleDefault:
+		ctx.WriteKeyWord(" DEFAULT")
+	case SetRoleNone:
+		ctx.WriteKeyWord(" NONE")
+	case SetRoleAll:
+		ctx.WriteKeyWord(" ALL")
+	case SetRoleAllExcept:
+		ctx.WriteKeyWord(" ALL EXCEPT")
+	}
+	for i, role := range n.RoleList {
+		ctx.WritePlain(" ")
+		err := role.Restore(ctx)
+		if err != nil {
+			return errors.Annotate(err, "An error occurred while restore SetRoleStmt.RoleList")
+		}
+		if i != len(n.RoleList)-1 {
+			ctx.WritePlain(",")
+		}
+	}
+	return nil
+}
+
+// Accept implements Node Accept interface.
+func (n *SetRoleStmt) Accept(v Visitor) (Node, bool) {
+	newNode, skipChildren := v.Enter(n)
+	if skipChildren {
+		return v.Leave(newNode)
+	}
+	n = newNode.(*SetRoleStmt)
+	return v.Leave(n)
+}
+
 // UserSpec is used for parsing create user statement.
 type UserSpec struct {
 	User    *auth.UserIdentity

auth/auth.go

@@ -33,11 +33,6 @@ type UserIdentity struct {
 	AuthHostname string // Match in privs system (i.e. could be a wildcard)
 }
 
-type RoleIdentity struct {
-	Username string
-	Hostname string
-}
-
 // Restore implements Node interface.
 func (user *UserIdentity) Restore(ctx *RestoreCtx) error {
 	if user.CurrentUser {
@@ -67,6 +62,26 @@ func (user *UserIdentity) AuthIdentityString() string {
 	return fmt.Sprintf("%s@%s", user.AuthUsername, user.AuthHostname)
 }
 
+type RoleIdentity struct {
+	Username string
+	Hostname string
+}
+
+func (role *RoleIdentity) Restore(ctx *RestoreCtx) error {
+	ctx.WriteName(role.Username)
+	if role.Hostname != "" {
+		ctx.WritePlain("@")
+		ctx.WriteName(role.Hostname)
+	}
+	return nil
+}
+
+// String converts UserIdentity to the format user@host.
+func (role *RoleIdentity) String() string {
+	// TODO: Escape username and hostname.
+	return fmt.Sprintf("`%s`@`%s`", role.Username, role.Hostname)
+}
+
 // CheckScrambledPassword check scrambled password received from client.
 // The new authentication is performed in following manner:
 //   SERVER:  public_seed=create_random_string()

go.sum1

@@ -94,6 +94,8 @@ github.com/remyoudompheng/bigfft v0.0.0-20170806203942-52369c62f446 h1:/NRJ5vAYo
 github.com/remyoudompheng/bigfft v0.0.0-20170806203942-52369c62f446/go.mod h1:uYEyJGbgTkfkS4+E/PavXkNJcbFIpEtjt2B0KDQ5+9M=
 github.com/shirou/gopsutil v2.18.10+incompatible h1:cy84jW6EVRPa5g9HAHrlbxMSIjBhDSX0OFYyMYminYs=
 github.com/shirou/gopsutil v2.18.10+incompatible/go.mod h1:5b4v6he4MtMOwMlS0TUMTu2PcXUg8+E1lC7eC3UO/RA=
+github.com/shirou/gopsutil v2.18.12+incompatible h1:1eaJvGomDnH74/5cF4CTmTbLHAriGFsTZppLXDX93OM=
+github.com/shirou/gopsutil v2.18.12+incompatible/go.mod h1:5b4v6he4MtMOwMlS0TUMTu2PcXUg8+E1lC7eC3UO/RA=
 github.com/shurcooL/httpfs v0.0.0-20171119174359-809beceb2371/go.mod h1:ZY1cvUeJuFPAdZ/B6v7RHavJWZn2YPVFQ1OSXhCGOkg=
 github.com/shurcooL/vfsgen v0.0.0-20181020040650-a97a25d856ca/go.mod h1:TrYk7fJVaAttu97ZZKrO9UbRa8izdowaMIZcxYMbVaw=
 github.com/sirupsen/logrus v1.2.0/go.mod h1:LxeOpSwHxABJmUn/MG1IvRgCAasNZTLOkJPxbbu5VWo=

misc.go

@@ -204,6 +204,7 @@ var tokenMap = map[string]int{
 	"CURRENT_TIME":             currentTime,
 	"CURRENT_TIMESTAMP":        currentTs,
 	"CURRENT_USER":             currentUser,
+	"CURRENT_ROLE":             currentRole,
 	"CURTIME":                  curTime,
 	"DATA":                     data,
 	"DATABASE":                 database,

mysql/errcode.go

@@ -894,6 +894,7 @@ const (
 	ErrInvalidJSONPathWildcard                                      = 3149
 	ErrInvalidJSONContainsPathType                                  = 3150
 	ErrJSONUsedAsKey                                                = 3152
+	ErrRoleNotGranted                                               = 3530
 	ErrWindowNoSuchWindow                                           = 3579
 	ErrWindowCircularityInWindowGraph                               = 3580
 	ErrWindowNoChildPartitioning                                    = 3581

mysql/errname.go

@@ -912,6 +912,7 @@ var MySQLErrName = map[uint16]string{
 	ErrWindowNoGroupOrderUnused:                              "ASC or DESC with GROUP BY isn't allowed with window functions; put ASC or DESC in ORDER BY",
 	ErrWindowExplainJson:                                     "To get information about window functions use EXPLAIN FORMAT=JSON",
 	ErrWindowFunctionIgnoresFrame:                            "Window function '%s' ignores the frame clause of window '%s' and aggregates over the whole partition",
+	ErrRoleNotGranted:                                        "%s is is not granted to %s",
 
 	// TiDB errors.
 	ErrMemExceedThreshold: "%s holds %dB memory, exceeds threshold %dB.%s",