[parser] introduce a new Restore Flag RestoreBracketAroundBinaryOpera…

…tion (pingcap#1332)
xhebox · Oct 8, 2021 · d0699d1 · d0699d1
1 parent 9bff5e6
commit d0699d1
Show file tree

Hide file tree

Showing 3 changed files with 53 additions and 1 deletion.
diff --git a/parser/ast/expressions.go b/parser/ast/expressions.go
@@ -173,6 +173,9 @@ func restoreBinaryOpWithSpacesAround(ctx *format.RestoreCtx, op opcode.Op) error
 
 // Restore implements Node interface.
 func (n *BinaryOperationExpr) Restore(ctx *format.RestoreCtx) error {
+	if ctx.Flags.HasRestoreBracketAroundBinaryOperation() {
+		ctx.WritePlain("(")
+	}
 	if err := n.L.Restore(ctx); err != nil {
 		return errors.Annotate(err, "An error occurred when restore BinaryOperationExpr.L")
 	}
@@ -182,7 +185,9 @@ func (n *BinaryOperationExpr) Restore(ctx *format.RestoreCtx) error {
 	if err := n.R.Restore(ctx); err != nil {
 		return errors.Annotate(err, "An error occurred when restore BinaryOperationExpr.R")
 	}
-
+	if ctx.Flags.HasRestoreBracketAroundBinaryOperation() {
+		ctx.WritePlain(")")
+	}
 	return nil
 }
 

diff --git a/parser/format/format.go b/parser/format/format.go
@@ -218,6 +218,7 @@ const (
 	RestoreNameBackQuotes
 
 	RestoreSpacesAroundBinaryOperation
+	RestoreBracketAroundBinaryOperation
 
 	RestoreStringWithoutCharset
 	RestoreStringWithoutDefaultCharset
@@ -283,6 +284,10 @@ func (rf RestoreFlags) HasSpacesAroundBinaryOperationFlag() bool {
 	return rf.has(RestoreSpacesAroundBinaryOperation)
 }
 
+func (rf RestoreFlags) HasRestoreBracketAroundBinaryOperation() bool {
+	return rf.has(RestoreBracketAroundBinaryOperation)
+}
+
 func (rf RestoreFlags) HasStringWithoutDefaultCharset() bool {
 	return rf.has(RestoreStringWithoutDefaultCharset)
 }

diff --git a/parser/parser_test.go b/parser/parser_test.go
@@ -6186,6 +6186,48 @@ func (s *testParserSuite) TestHelp(c *C) {
 	s.RunTest(c, table)
 }
 
+func (s *testParserSuite) TestRestoreBinOpWithBrackets(c *C) {
+	cases := []testCase{
+		{"select mod(a+b, 4)+1", true, "SELECT (((`a` + `b`) % 4) + 1)"},
+		{"select mod( year(a) - abs(weekday(a) + dayofweek(a)), 4) + 1", true, "SELECT (((year(`a`) - abs((weekday(`a`) + dayofweek(`a`)))) % 4) + 1)"},
+	}
+
+	parser := parser.New()
+	parser.EnableWindowFunc(s.enableWindowFunc)
+	for _, t := range cases {
+		_, _, err := parser.Parse(t.src, "", "")
+		comment := Commentf("source %v", t.src)
+		if !t.ok {
+			c.Assert(err, NotNil, comment)
+			continue
+		}
+		c.Assert(err, IsNil, comment)
+		// restore correctness test
+		if t.ok {
+			var sb strings.Builder
+			comment := Commentf("source %v", t.src)
+			stmts, _, err := parser.Parse(t.src, "", "")
+			c.Assert(err, IsNil, comment)
+			restoreSQLs := ""
+			for _, stmt := range stmts {
+				sb.Reset()
+				ctx := NewRestoreCtx(RestoreStringSingleQuotes|RestoreSpacesAroundBinaryOperation|RestoreBracketAroundBinaryOperation|RestoreStringWithoutCharset|RestoreNameBackQuotes, &sb)
+				ctx.DefaultDB = "test"
+				err = stmt.Restore(ctx)
+				c.Assert(err, IsNil, comment)
+				restoreSQL := sb.String()
+				comment = Commentf("source %v; restore %v", t.src, restoreSQL)
+				if restoreSQLs != "" {
+					restoreSQLs += "; "
+				}
+				restoreSQLs += restoreSQL
+			}
+			comment = Commentf("restore %v; expect %v", restoreSQLs, t.restore)
+			c.Assert(restoreSQLs, Equals, t.restore, comment)
+		}
+	}
+}
+
 // For CTE bindings.
 func (s *testParserSuite) TestCTEBindings(c *C) {
 	table := []testCase{