Allow to use the optional IN PARTITION expression in ALTER TABLE (#206)

Author: git-hulk

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
 	}

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
 	}

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';

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';

parser/testdata/ddl/output/alter_table_update.sql.golden.json

@@ -71,6 +71,7 @@
             }
           }
         ],
+        "InPartition": null,
         "WhereClause": {
           "LeftExpr": {
             "Name": "status",

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
+        }
+      }
+    ]
+  }
+]

parser/testdata/ddl/output/alter_table_update_with_cluster.sql.golden.json

@@ -60,6 +60,7 @@
             }
           }
         ],
+        "InPartition": null,
         "WhereClause": {
           "LeftExpr": {
             "Name": "id",

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
 		}