Add the support of the CREATE ROLE expression (#17)

AfterShip · Oct 20, 2023 · 3320b64 · 3320b64
1 parent fc9d12f
commit 3320b64
Show file tree

Hide file tree

Showing 7 changed files with 1,059 additions and 1 deletion.
diff --git a/parser/ast.go b/parser/ast.go
@@ -896,6 +896,143 @@ func (c *CreateFunction) String(level int) string {
 	return builder.String()
 }
 
+type RoleName struct {
+	Name      *Ident
+	OnCluster *OnClusterExpr
+}
+
+func (r *RoleName) Pos() Pos {
+	return r.Name.NamePos
+}
+
+func (r *RoleName) End() Pos {
+	if r.OnCluster != nil {
+		return r.OnCluster.End()
+	}
+	return r.Name.NameEnd
+}
+
+func (r *RoleName) String(level int) string {
+	var builder strings.Builder
+	builder.WriteString(r.Name.String(level))
+	if r.OnCluster != nil {
+		builder.WriteString(" ")
+		builder.WriteString(r.OnCluster.String(level))
+	}
+	return builder.String()
+}
+
+type SettingPair struct {
+	Name  *Ident
+	Value Expr
+}
+
+func (s *SettingPair) Pos() Pos {
+	return s.Name.NamePos
+}
+
+func (s *SettingPair) End() Pos {
+	return s.Value.End()
+}
+
+func (s *SettingPair) String(level int) string {
+	var builder strings.Builder
+	builder.WriteString(s.Name.String(level))
+	if s.Value != nil {
+		builder.WriteByte(' ')
+		builder.WriteString(s.Value.String(level))
+	}
+	return builder.String()
+}
+
+type RoleSetting struct {
+	SettingPairs []*SettingPair
+	Modifier     *Ident
+}
+
+func (r *RoleSetting) Pos() Pos {
+	if len(r.SettingPairs) > 0 {
+		return r.SettingPairs[0].Pos()
+	}
+	return r.Modifier.NamePos
+}
+
+func (r *RoleSetting) End() Pos {
+	if r.Modifier != nil {
+		return r.Modifier.NameEnd
+	}
+	return r.SettingPairs[len(r.SettingPairs)-1].End()
+}
+
+func (r *RoleSetting) String(level int) string {
+	var builder strings.Builder
+	for i, settingPair := range r.SettingPairs {
+		if i > 0 {
+			builder.WriteString(" ")
+		}
+		builder.WriteString(settingPair.String(level))
+	}
+	if r.Modifier != nil {
+		builder.WriteString(" ")
+		builder.WriteString(r.Modifier.String(level))
+	}
+	return builder.String()
+}
+
+type CreateRole struct {
+	CreatePos         Pos
+	StatementEnd      Pos
+	IfNotExists       bool
+	OrReplace         bool
+	RoleNames         []*RoleName
+	AccessStorageType *Ident
+	Settings          []*RoleSetting
+}
+
+func (c *CreateRole) Pos() Pos {
+	return c.CreatePos
+}
+
+func (c *CreateRole) End() Pos {
+	return c.StatementEnd
+}
+
+func (c *CreateRole) Type() string {
+	return "ROLE"
+}
+
+func (c *CreateRole) String(level int) string {
+	var builder strings.Builder
+	builder.WriteString("CREATE ROLE ")
+	if c.IfNotExists {
+		builder.WriteString("IF NOT EXISTS ")
+	}
+	if c.OrReplace {
+		builder.WriteString("OR REPLACE ")
+	}
+	for i, roleName := range c.RoleNames {
+		if i > 0 {
+			builder.WriteString(", ")
+		}
+		builder.WriteString(roleName.String(level))
+	}
+	if c.AccessStorageType != nil {
+		builder.WriteString(NewLine(level))
+		builder.WriteString("IN ")
+		builder.WriteString(c.AccessStorageType.String(level))
+	}
+	if len(c.Settings) > 0 {
+		builder.WriteString(" SETTINGS ")
+		for i, setting := range c.Settings {
+			if i > 0 {
+				builder.WriteString(", ")
+			}
+			builder.WriteString(setting.String(level))
+		}
+	}
+	return builder.String()
+}
+
 type DestinationExpr struct {
 	ToPos           Pos
 	TableIdentifier *TableIdentifier

diff --git a/parser/keyword.go b/parser/keyword.go
@@ -151,6 +151,7 @@ const (
 	KeywordReplica      = "REPLICA"
 	KeywordReplicated   = "REPLICATED"
 	KeywordRight        = "RIGHT"
+	KeywordRole         = "ROLE"
 	KeywordRollup       = "ROLLUP"
 	KeywordRow          = "ROW"
 	KeywordRows         = "ROWS"
@@ -355,6 +356,7 @@ var keywords = NewSet(
 	KeywordReplica,
 	KeywordReplicated,
 	KeywordRight,
+	KeywordRole,
 	KeywordRollup,
 	KeywordRow,
 	KeywordRows,

diff --git a/parser/parse_system.go b/parser/parse_system.go
@@ -340,3 +340,148 @@ func (p *Parser) parseCheckExpr(pos Pos) (*CheckExpr, error) {
 		Partition: partition,
 	}, nil
 }
+
+func (p *Parser) parseRoleName(_ Pos) (*RoleName, error) {
+	name, err := p.parseIdent()
+	if err != nil {
+		return nil, err
+	}
+	onCluster, err := p.tryParseOnCluster(p.Pos())
+	if err != nil {
+		return nil, err
+	}
+	return &RoleName{
+		Name:      name,
+		OnCluster: onCluster,
+	}, nil
+}
+
+func (p *Parser) tryParseRoleSettings(pos Pos) ([]*RoleSetting, error) {
+	if p.tryConsumeKeyword(KeywordSettings) == nil {
+		return nil, nil
+	}
+	return p.parseRoleSettings(pos)
+}
+
+func (p *Parser) parseRoleSetting(_ Pos) (*RoleSetting, error) {
+	pairs := make([]*SettingPair, 0)
+	for p.matchTokenKind(TokenIdent) {
+		name, err := p.parseIdent()
+		if err != nil {
+			return nil, err
+		}
+		switch name.Name {
+		case "NONE", "READABLE", "WRITABLE", "CONST", "CHANGEABLE_IN_READONLY":
+			return &RoleSetting{
+				Modifier:     name,
+				SettingPairs: pairs,
+			}, nil
+		}
+		switch {
+		case p.matchTokenKind("="),
+			p.matchTokenKind(TokenInt),
+			p.matchTokenKind(TokenFloat),
+			p.matchTokenKind(TokenString):
+			_ = p.tryConsumeTokenKind("=")
+			value, err := p.parseLiteral(p.Pos())
+			if err != nil {
+				return nil, err
+			}
+			pairs = append(pairs, &SettingPair{
+				Name:  name,
+				Value: value,
+			})
+		default:
+			pairs = append(pairs, &SettingPair{
+				Name: name,
+			})
+		}
+
+	}
+	return &RoleSetting{
+		SettingPairs: pairs,
+	}, nil
+}
+
+func (p *Parser) parseRoleSettings(_ Pos) ([]*RoleSetting, error) {
+	settings := make([]*RoleSetting, 0)
+	for {
+		setting, err := p.parseRoleSetting(p.Pos())
+		if err != nil {
+			return nil, err
+		}
+		settings = append(settings, setting)
+		if p.tryConsumeTokenKind(",") == nil {
+			break
+		}
+	}
+	return settings, nil
+}
+
+func (p *Parser) parseCreateRole(pos Pos) (*CreateRole, error) {
+	if err := p.consumeKeyword(KeywordRole); err != nil {
+		return nil, err
+	}
+
+	ifNotExists := false
+	orReplace := false
+	switch {
+	case p.matchKeyword(KeywordIf):
+		_ = p.lexer.consumeToken()
+		if err := p.consumeKeyword(KeywordNot); err != nil {
+			return nil, err
+		}
+		if err := p.consumeKeyword(KeywordExists); err != nil {
+			return nil, err
+		}
+		ifNotExists = true
+	case p.matchKeyword(KeywordOr):
+		_ = p.lexer.consumeToken()
+		if err := p.consumeKeyword(KeywordReplace); err != nil {
+			return nil, err
+		}
+		orReplace = true
+	}
+
+	roleNames := make([]*RoleName, 0)
+	roleName, err := p.parseRoleName(p.Pos())
+	if err != nil {
+		return nil, err
+	}
+	roleNames = append(roleNames, roleName)
+	for p.tryConsumeTokenKind(",") != nil {
+		roleName, err := p.parseRoleName(p.Pos())
+		if err != nil {
+			return nil, err
+		}
+		roleNames = append(roleNames, roleName)
+	}
+	statementEnd := roleNames[len(roleNames)-1].End()
+
+	var accessStorageType *Ident
+	if p.tryConsumeKeyword(KeywordIn) != nil {
+		accessStorageType, err = p.parseIdent()
+		if err != nil {
+			return nil, err
+		}
+		statementEnd = accessStorageType.NameEnd
+	}
+
+	settings, err := p.tryParseRoleSettings(p.Pos())
+	if err != nil {
+		return nil, err
+	}
+	if settings != nil {
+		statementEnd = settings[len(settings)-1].End()
+	}
+
+	return &CreateRole{
+		CreatePos:         pos,
+		StatementEnd:      statementEnd,
+		IfNotExists:       ifNotExists,
+		OrReplace:         orReplace,
+		RoleNames:         roleNames,
+		AccessStorageType: accessStorageType,
+		Settings:          settings,
+	}, nil
+}
diff --git a/parser/parser_table.go b/parser/parser_table.go
@@ -23,6 +23,8 @@ func (p *Parser) parseDDL(pos Pos) (DDL, error) {
 			return p.parseCreateLiveView(pos)
 		case p.matchKeyword(KeywordView):
 			return p.parseCreateView(pos)
+		case p.matchKeyword(KeywordRole):
+			return p.parseCreateRole(pos)
 		case p.matchKeyword(KeywordDictionary):
 		case p.matchKeyword(KeywordFunction):
 		case p.matchKeyword(KeywordRow):
@@ -726,7 +728,7 @@ func (p *Parser) parseSettingsExprList(pos Pos) (*SettingsExprList, error) {
 		return nil, err
 	}
 	items = append(items, expr)
-	for !p.lexer.isEOF() && p.tryConsumeTokenKind(",") != nil {
+	for p.tryConsumeTokenKind(",") != nil {
 		expr, err = p.parseSettingsExpr(p.Pos())
 		if err != nil {
 			return nil, err

diff --git a/parser/testdata/ddl/create_role.sql b/parser/testdata/ddl/create_role.sql
@@ -0,0 +1,18 @@
+-- Tags: no-parallel
+
+CREATE ROLE r1_01293;
+CREATE ROLE r1_01293 SETTINGS NONE;
+CREATE ROLE r2_01293 SETTINGS PROFILE 'default';
+CREATE ROLE r3_01293 SETTINGS max_memory_usage=5000000;
+CREATE ROLE r4_01293 SETTINGS max_memory_usage MIN=5000000;
+CREATE ROLE r5_01293 SETTINGS max_memory_usage MAX=5000000;
+CREATE ROLE r6_01293 SETTINGS max_memory_usage CONST;
+CREATE ROLE r7_01293 SETTINGS max_memory_usage WRITABLE;
+CREATE ROLE r8_01293 SETTINGS max_memory_usage=5000000 MIN 4000000 MAX 6000000 CONST;
+CREATE ROLE r9_01293 SETTINGS PROFILE 'default', max_memory_usage=5000000 WRITABLE;
+CREATE ROLE r1_01293, r2_01293;
+CREATE ROLE r1_01293 SETTINGS readonly=1;
+CREATE ROLE r2_01293 SETTINGS PROFILE 'default';
+CREATE ROLE r3_01293 SETTINGS max_memory_usage=5000000 MIN 4000000 MAX 6000000 WRITABLE;
+CREATE ROLE r4_01293 SETTINGS PROFILE 'default', max_memory_usage=5000000, readonly=1;
+CREATE ROLE r5_01293 SETTINGS NONE;
diff --git a/parser/testdata/ddl/format/create_role.sql b/parser/testdata/ddl/format/create_role.sql
@@ -0,0 +1,37 @@
+-- Origin SQL:
+-- Tags: no-parallel
+
+CREATE ROLE r1_01293;
+CREATE ROLE r1_01293 SETTINGS NONE;
+CREATE ROLE r2_01293 SETTINGS PROFILE 'default';
+CREATE ROLE r3_01293 SETTINGS max_memory_usage=5000000;
+CREATE ROLE r4_01293 SETTINGS max_memory_usage MIN=5000000;
+CREATE ROLE r5_01293 SETTINGS max_memory_usage MAX=5000000;
+CREATE ROLE r6_01293 SETTINGS max_memory_usage CONST;
+CREATE ROLE r7_01293 SETTINGS max_memory_usage WRITABLE;
+CREATE ROLE r8_01293 SETTINGS max_memory_usage=5000000 MIN 4000000 MAX 6000000 CONST;
+CREATE ROLE r9_01293 SETTINGS PROFILE 'default', max_memory_usage=5000000 WRITABLE;
+CREATE ROLE r1_01293, r2_01293;
+CREATE ROLE r1_01293 SETTINGS readonly=1;
+CREATE ROLE r2_01293 SETTINGS PROFILE 'default';
+CREATE ROLE r3_01293 SETTINGS max_memory_usage=5000000 MIN 4000000 MAX 6000000 WRITABLE;
+CREATE ROLE r4_01293 SETTINGS PROFILE 'default', max_memory_usage=5000000, readonly=1;
+CREATE ROLE r5_01293 SETTINGS NONE;
+
+-- Format SQL:
+CREATE ROLE r1_01293;
+CREATE ROLE r1_01293 SETTINGS  NONE;
+CREATE ROLE r2_01293 SETTINGS PROFILE 'default';
+CREATE ROLE r3_01293 SETTINGS max_memory_usage 5000000;
+CREATE ROLE r4_01293 SETTINGS max_memory_usage MIN 5000000;
+CREATE ROLE r5_01293 SETTINGS max_memory_usage MAX 5000000;
+CREATE ROLE r6_01293 SETTINGS max_memory_usage CONST;
+CREATE ROLE r7_01293 SETTINGS max_memory_usage WRITABLE;
+CREATE ROLE r8_01293 SETTINGS max_memory_usage 5000000 MIN 4000000 MAX 6000000 CONST;
+CREATE ROLE r9_01293 SETTINGS PROFILE 'default', max_memory_usage 5000000 WRITABLE;
+CREATE ROLE r1_01293, r2_01293;
+CREATE ROLE r1_01293 SETTINGS readonly 1;
+CREATE ROLE r2_01293 SETTINGS PROFILE 'default';
+CREATE ROLE r3_01293 SETTINGS max_memory_usage 5000000 MIN 4000000 MAX 6000000 WRITABLE;
+CREATE ROLE r4_01293 SETTINGS PROFILE 'default', max_memory_usage 5000000, readonly 1;
+CREATE ROLE r5_01293 SETTINGS  NONE;