diff --git a/parser/ast.go b/parser/ast.go index ab170fe..7416eee 100644 --- a/parser/ast.go +++ b/parser/ast.go @@ -1480,6 +1480,7 @@ type AlterTableUpdate struct { UpdatePos Pos StatementEnd Pos Assignments []*UpdateAssignment + InPartition *PartitionClause WhereClause Expr } @@ -1504,6 +1505,10 @@ func (a *AlterTableUpdate) String() string { } builder.WriteString(assignment.String()) } + if a.InPartition != nil { + builder.WriteString(" IN ") + builder.WriteString(a.InPartition.String()) + } builder.WriteString(" WHERE ") builder.WriteString(a.WhereClause.String()) return builder.String() @@ -1517,6 +1522,11 @@ func (a *AlterTableUpdate) Accept(visitor ASTVisitor) error { return err } } + if a.InPartition != nil { + if err := a.InPartition.Accept(visitor); err != nil { + return err + } + } if err := a.WhereClause.Accept(visitor); err != nil { return err } diff --git a/parser/parser_alter.go b/parser/parser_alter.go index da9bd3e..abc77fd 100644 --- a/parser/parser_alter.go +++ b/parser/parser_alter.go @@ -860,7 +860,7 @@ func (p *Parser) parseAlterTableDelete(pos Pos) (AlterTableClause, error) { }, nil } -// Syntax: ALTER TABLE UPDATE column1 = expr1 [, column2 = expr2, ...] WHERE condition +// Syntax: ALTER TABLE UPDATE column1 = expr1 [, column2 = expr2, ...] [IN PARTITION partition_id] WHERE condition func (p *Parser) parseAlterTableUpdate(pos Pos) (AlterTableClause, error) { if err := p.expectKeyword(KeywordUpdate); err != nil { return nil, err @@ -883,6 +883,13 @@ func (p *Parser) parseAlterTableUpdate(pos Pos) (AlterTableClause, error) { assignments = append(assignments, assignment) } + var inPartition *PartitionClause + if p.tryConsumeKeywords(KeywordIn) { + inPartition, err = p.parsePartitionClause(p.Pos()) + if err != nil { + return nil, err + } + } if err := p.expectKeyword(KeywordWhere); err != nil { return nil, err } @@ -896,6 +903,7 @@ func (p *Parser) parseAlterTableUpdate(pos Pos) (AlterTableClause, error) { UpdatePos: pos, StatementEnd: whereExpr.End(), Assignments: assignments, + InPartition: inPartition, WhereClause: whereExpr, }, nil } @@ -911,7 +919,10 @@ func (p *Parser) parseUpdateAssignment(pos Pos) (*UpdateAssignment, error) { return nil, err } - expr, err := p.parseExpr(p.Pos()) + // Why don't we use parseExpr here? Because `ALTER TABLE UPDATE` syntax allows to + // use `IN PARTITION` keywords after assignments. So we need to limit the precedence + // to avoid parsing `IN PARTITION` as part of the expression. + expr, err := p.parseSubExpr(p.Pos(), precedenceIn) if err != nil { return nil, err } diff --git a/parser/testdata/ddl/alter_table_update_in_partition.sql b/parser/testdata/ddl/alter_table_update_in_partition.sql new file mode 100644 index 0000000..be33c47 --- /dev/null +++ b/parser/testdata/ddl/alter_table_update_in_partition.sql @@ -0,0 +1 @@ +ALTER TABLE test.users UPDATE status = 'inactive' IN PARTITION '2024-01-01' WHERE status = 'active'; diff --git a/parser/testdata/ddl/format/alter_table_update_in_partition.sql b/parser/testdata/ddl/format/alter_table_update_in_partition.sql new file mode 100644 index 0000000..0f37330 --- /dev/null +++ b/parser/testdata/ddl/format/alter_table_update_in_partition.sql @@ -0,0 +1,6 @@ +-- Origin SQL: +ALTER TABLE test.users UPDATE status = 'inactive' IN PARTITION '2024-01-01' WHERE status = 'active'; + + +-- Format SQL: +ALTER TABLE test.users UPDATE status = 'inactive' IN PARTITION '2024-01-01' WHERE status = 'active'; diff --git a/parser/testdata/ddl/output/alter_table_update.sql.golden.json b/parser/testdata/ddl/output/alter_table_update.sql.golden.json index b64d06d..efdcf49 100644 --- a/parser/testdata/ddl/output/alter_table_update.sql.golden.json +++ b/parser/testdata/ddl/output/alter_table_update.sql.golden.json @@ -71,6 +71,7 @@ } } ], + "InPartition": null, "WhereClause": { "LeftExpr": { "Name": "status", diff --git a/parser/testdata/ddl/output/alter_table_update_in_partition.sql.golden.json b/parser/testdata/ddl/output/alter_table_update_in_partition.sql.golden.json new file mode 100644 index 0000000..2871d91 --- /dev/null +++ b/parser/testdata/ddl/output/alter_table_update_in_partition.sql.golden.json @@ -0,0 +1,72 @@ +[ + { + "AlterPos": 0, + "StatementEnd": 98, + "TableIdentifier": { + "Database": { + "Name": "test", + "QuoteType": 1, + "NamePos": 12, + "NameEnd": 16 + }, + "Table": { + "Name": "users", + "QuoteType": 1, + "NamePos": 17, + "NameEnd": 22 + } + }, + "OnCluster": null, + "AlterExprs": [ + { + "UpdatePos": 23, + "StatementEnd": 98, + "Assignments": [ + { + "AssignmentPos": 30, + "Column": { + "Ident": { + "Name": "status", + "QuoteType": 1, + "NamePos": 30, + "NameEnd": 36 + }, + "DotIdent": null + }, + "Expr": { + "LiteralPos": 40, + "LiteralEnd": 48, + "Literal": "inactive" + } + } + ], + "InPartition": { + "PartitionPos": 53, + "Expr": { + "LiteralPos": 64, + "LiteralEnd": 74, + "Literal": "2024-01-01" + }, + "ID": null, + "All": false + }, + "WhereClause": { + "LeftExpr": { + "Name": "status", + "QuoteType": 1, + "NamePos": 82, + "NameEnd": 88 + }, + "Operation": "=", + "RightExpr": { + "LiteralPos": 92, + "LiteralEnd": 98, + "Literal": "active" + }, + "HasGlobal": false, + "HasNot": false + } + } + ] + } +] \ No newline at end of file diff --git a/parser/testdata/ddl/output/alter_table_update_with_cluster.sql.golden.json b/parser/testdata/ddl/output/alter_table_update_with_cluster.sql.golden.json index 19ee7ab..dfc20f9 100644 --- a/parser/testdata/ddl/output/alter_table_update_with_cluster.sql.golden.json +++ b/parser/testdata/ddl/output/alter_table_update_with_cluster.sql.golden.json @@ -60,6 +60,7 @@ } } ], + "InPartition": null, "WhereClause": { "LeftExpr": { "Name": "id", diff --git a/parser/walk.go b/parser/walk.go index d80a8d2..ec165ff 100644 --- a/parser/walk.go +++ b/parser/walk.go @@ -895,6 +895,9 @@ func Walk(node Expr, fn WalkFunc) bool { return false } } + if !Walk(n.InPartition, fn) { + return false + } if !Walk(n.WhereClause, fn) { return false }