Fix OnCluster expression should be parsed inside role statement (#26)

AfterShip · Nov 3, 2023 · 91f3ce2 · 91f3ce2
1 parent a6f475c
commit 91f3ce2
Show file tree

Hide file tree

Showing 9 changed files with 671 additions and 442 deletions.
diff --git a/parser/ast.go b/parser/ast.go
@@ -901,8 +901,9 @@ func (c *CreateFunction) String(level int) string {
 }
 
 type RoleName struct {
-	Name  Expr
-	Scope *StringLiteral
+	Name      Expr
+	Scope     *StringLiteral
+	OnCluster *OnClusterExpr
 }
 
 func (r *RoleName) Pos() Pos {
@@ -913,6 +914,9 @@ func (r *RoleName) End() Pos {
 	if r.Scope != nil {
 		return r.Scope.End()
 	}
+	if r.OnCluster != nil {
+		return r.OnCluster.End()
+	}
 	return r.Name.End()
 }
 
@@ -923,6 +927,10 @@ func (r *RoleName) String(level int) string {
 		builder.WriteString("@")
 		builder.WriteString(r.Scope.String(level))
 	}
+	if r.OnCluster != nil {
+		builder.WriteString(" ON ")
+		builder.WriteString(r.OnCluster.String(level))
+	}
 	return builder.String()
 }
 
@@ -989,7 +997,6 @@ type CreateRole struct {
 	IfNotExists       bool
 	OrReplace         bool
 	RoleNames         []*RoleName
-	OnCluster         *OnClusterExpr
 	AccessStorageType *Ident
 	Settings          []*RoleSetting
 }
@@ -3274,7 +3281,6 @@ type DropUserOrRole struct {
 	StatementEnd Pos
 	Names        []*RoleName
 	IfExists     bool
-	OnCluster    *OnClusterExpr
 	Modifier     string
 	From         *Ident
 }
@@ -3303,10 +3309,6 @@ func (d *DropUserOrRole) String(level int) string {
 		}
 		builder.WriteString(name.String(level))
 	}
-	if d.OnCluster != nil {
-		builder.WriteString(" ")
-		builder.WriteString(d.OnCluster.String(level))
-	}
 	if len(d.Modifier) != 0 {
 		builder.WriteString(" " + d.Modifier)
 	}

diff --git a/parser/parse_system.go b/parser/parse_system.go
@@ -355,17 +355,27 @@ func (p *Parser) parseRoleName(_ Pos) (*RoleName, error) {
 				return nil, err
 			}
 		}
+		onCluster, err := p.tryParseOnCluster(p.Pos())
+		if err != nil {
+			return nil, err
+		}
 		return &RoleName{
-			Name:  name,
-			Scope: scope,
+			Name:      name,
+			Scope:     scope,
+			OnCluster: onCluster,
 		}, nil
 	case p.matchTokenKind(TokenString):
 		name, err := p.parseString(p.Pos())
 		if err != nil {
 			return nil, err
 		}
+		onCluster, err := p.tryParseOnCluster(p.Pos())
+		if err != nil {
+			return nil, err
+		}
 		return &RoleName{
-			Name: name,
+			Name:      name,
+			OnCluster: onCluster,
 		}, nil
 	default:
 		return nil, fmt.Errorf("expected <ident> or <string>")
@@ -474,11 +484,6 @@ func (p *Parser) parseCreateRole(pos Pos) (*CreateRole, error) {
 	}
 	statementEnd := roleNames[len(roleNames)-1].End()
 
-	onCluster, err := p.tryParseOnCluster(p.Pos())
-	if err != nil {
-		return nil, err
-	}
-
 	var accessStorageType *Ident
 	if p.tryConsumeKeyword(KeywordIn) != nil {
 		accessStorageType, err = p.parseIdent()
@@ -502,7 +507,6 @@ func (p *Parser) parseCreateRole(pos Pos) (*CreateRole, error) {
 		IfNotExists:       ifNotExists,
 		OrReplace:         orReplace,
 		RoleNames:         roleNames,
-		OnCluster:         onCluster,
 		AccessStorageType: accessStorageType,
 		Settings:          settings,
 	}, nil

diff --git a/parser/testdata/ddl/alter_role.sql b/parser/testdata/ddl/alter_role.sql
@@ -1,6 +1,7 @@
 -- Tags: no-parallel
 
 ALTER ROLE r1_01293;
+ALTER ROLE r1_01293 ON CLUSTER cluster_1 RENAME TO r2_01293;
 ALTER ROLE r1_01293 RENAME TO r2_01293, r3_01293 RENAME TO r4_01293;
 ALTER ROLE r1_01293 SETTINGS NONE;
 ALTER ROLE r2_01293 SETTINGS PROFILE 'default';

diff --git a/parser/testdata/ddl/create_role.sql b/parser/testdata/ddl/create_role.sql
@@ -1,6 +1,10 @@
 -- Tags: no-parallel
 
 CREATE ROLE r1_01293;
+CREATE ROLE r1_01293 ON CLUSTER cluster_1;
+CREATE ROLE r1_01293, r2_01293;
+CREATE ROLE r1_01293 ON CLUSTER cluster_1, r2_01293;
+CREATE ROLE r1_01293 ON CLUSTER cluster_1, r2_01293 ON CLUSTER cluster_2;
 CREATE ROLE r1_01293 SETTINGS NONE;
 CREATE ROLE r2_01293 SETTINGS PROFILE 'default';
 CREATE ROLE r3_01293 SETTINGS max_memory_usage=5000000;

diff --git a/parser/testdata/ddl/format/alter_role.sql b/parser/testdata/ddl/format/alter_role.sql
@@ -2,6 +2,7 @@
 -- Tags: no-parallel
 
 ALTER ROLE r1_01293;
+ALTER ROLE r1_01293 ON CLUSTER cluster_1 RENAME TO r2_01293;
 ALTER ROLE r1_01293 RENAME TO r2_01293, r3_01293 RENAME TO r4_01293;
 ALTER ROLE r1_01293 SETTINGS NONE;
 ALTER ROLE r2_01293 SETTINGS PROFILE 'default';
@@ -23,6 +24,7 @@ ALTER ROLE r2_01293@'%.myhost.com';
 
 -- Format SQL:
 ALTER ROLE r1_01293;
+ALTER ROLE r1_01293 ON ON CLUSTER cluster_1 RENAME TO r2_01293;
 ALTER ROLE r1_01293 RENAME TO r2_01293, r3_01293 RENAME TO r4_01293;
 ALTER ROLE r1_01293 SETTINGS  NONE;
 ALTER ROLE r2_01293 SETTINGS PROFILE 'default';

diff --git a/parser/testdata/ddl/format/create_role.sql b/parser/testdata/ddl/format/create_role.sql
@@ -2,6 +2,10 @@
 -- Tags: no-parallel
 
 CREATE ROLE r1_01293;
+CREATE ROLE r1_01293 ON CLUSTER cluster_1;
+CREATE ROLE r1_01293, r2_01293;
+CREATE ROLE r1_01293 ON CLUSTER cluster_1, r2_01293;
+CREATE ROLE r1_01293 ON CLUSTER cluster_1, r2_01293 ON CLUSTER cluster_2;
 CREATE ROLE r1_01293 SETTINGS NONE;
 CREATE ROLE r2_01293 SETTINGS PROFILE 'default';
 CREATE ROLE r3_01293 SETTINGS max_memory_usage=5000000;
@@ -22,6 +26,10 @@ CREATE ROLE r2_01293@'%.myhost.com';
 
 -- Format SQL:
 CREATE ROLE r1_01293;
+CREATE ROLE r1_01293 ON ON CLUSTER cluster_1;
+CREATE ROLE r1_01293, r2_01293;
+CREATE ROLE r1_01293 ON ON CLUSTER cluster_1, r2_01293;
+CREATE ROLE r1_01293 ON ON CLUSTER cluster_1, r2_01293 ON ON CLUSTER cluster_2;
 CREATE ROLE r1_01293 SETTINGS  NONE;
 CREATE ROLE r2_01293 SETTINGS PROFILE 'default';
 CREATE ROLE r3_01293 SETTINGS max_memory_usage 5000000;