parser: add cast to array and member of support (#39820)

close #39821
pingcap · Dec 12, 2022 · 3183582 · 3183582
1 parent fa792b3
commit 3183582
Show file tree

Hide file tree

Showing 8 changed files with 10,235 additions and 10,117 deletions.
diff --git a/executor/seqtest/seq_executor_test.go b/executor/seqtest/seq_executor_test.go
@@ -940,7 +940,7 @@ func TestBatchInsertDelete(t *testing.T) {
 		atomic.StoreUint64(&kv.TxnTotalSizeLimit, originLimit)
 	}()
 	// Set the limitation to a small value, make it easier to reach the limitation.
-	atomic.StoreUint64(&kv.TxnTotalSizeLimit, 5800)
+	atomic.StoreUint64(&kv.TxnTotalSizeLimit, 5900)
 
 	tk := testkit.NewTestKit(t, store)
 	tk.MustExec("use test")

diff --git a/parser/ast/functions.go b/parser/ast/functions.go
@@ -332,6 +332,7 @@ const (
 	JSONReplace       = "json_replace"
 	JSONRemove        = "json_remove"
 	JSONContains      = "json_contains"
+	JSONMemberOf      = "json_memberof"
 	JSONContainsPath  = "json_contains_path"
 	JSONValid         = "json_valid"
 	JSONArrayAppend   = "json_array_append"
@@ -380,21 +381,9 @@ type FuncCallExpr struct {
 
 // Restore implements Node interface.
 func (n *FuncCallExpr) Restore(ctx *format.RestoreCtx) error {
-	var specialLiteral string
-	switch n.FnName.L {
-	case DateLiteral:
-		specialLiteral = "DATE "
-	case TimeLiteral:
-		specialLiteral = "TIME "
-	case TimestampLiteral:
-		specialLiteral = "TIMESTAMP "
-	}
-	if specialLiteral != "" {
-		ctx.WritePlain(specialLiteral)
-		if err := n.Args[0].Restore(ctx); err != nil {
-			return errors.Annotatef(err, "An error occurred while restore FuncCastExpr.Expr")
-		}
-		return nil
+	done, err := n.customRestore(ctx)
+	if done {
+		return err
 	}
 
 	if len(n.Schema.String()) != 0 {
@@ -495,29 +484,70 @@ func (n *FuncCallExpr) Restore(ctx *format.RestoreCtx) error {
 	return nil
 }
 
+func (n *FuncCallExpr) customRestore(ctx *format.RestoreCtx) (bool, error) {
+	var specialLiteral string
+	switch n.FnName.L {
+	case DateLiteral:
+		specialLiteral = "DATE "
+	case TimeLiteral:
+		specialLiteral = "TIME "
+	case TimestampLiteral:
+		specialLiteral = "TIMESTAMP "
+	}
+	if specialLiteral != "" {
+		ctx.WritePlain(specialLiteral)
+		if err := n.Args[0].Restore(ctx); err != nil {
+			return true, errors.Annotatef(err, "An error occurred while restore FuncCallExpr.Expr")
+		}
+		return true, nil
+	}
+	if n.FnName.L == JSONMemberOf {
+		if err := n.Args[0].Restore(ctx); err != nil {
+			return true, errors.Annotatef(err, "An error occurred while restore FuncCallExpr.(MEMBER OF).Args[0]")
+		}
+		ctx.WriteKeyWord(" MEMBER OF ")
+		ctx.WritePlain("(")
+		if err := n.Args[1].Restore(ctx); err != nil {
+			return true, errors.Annotatef(err, "An error occurred while restore FuncCallExpr.(MEMBER OF).Args[1]")
+		}
+		ctx.WritePlain(")")
+		return true, nil
+	}
+	return false, nil
+}
+
 // Format the ExprNode into a Writer.
 func (n *FuncCallExpr) Format(w io.Writer) {
-	fmt.Fprintf(w, "%s(", n.FnName.L)
 	if !n.specialFormatArgs(w) {
+		fmt.Fprintf(w, "%s(", n.FnName.L)
 		for i, arg := range n.Args {
 			arg.Format(w)
 			if i != len(n.Args)-1 {
 				fmt.Fprint(w, ", ")
 			}
 		}
+		fmt.Fprint(w, ")")
 	}
-	fmt.Fprint(w, ")")
 }
 
 // specialFormatArgs formats argument list for some special functions.
 func (n *FuncCallExpr) specialFormatArgs(w io.Writer) bool {
 	switch n.FnName.L {
 	case DateAdd, DateSub, AddDate, SubDate:
+		fmt.Fprintf(w, "%s(", n.FnName.L)
 		n.Args[0].Format(w)
 		fmt.Fprint(w, ", INTERVAL ")
 		n.Args[1].Format(w)
 		fmt.Fprint(w, " ")
 		n.Args[2].Format(w)
+		fmt.Fprint(w, ")")
+		return true
+	case JSONMemberOf:
+		n.Args[0].Format(w)
+		fmt.Fprint(w, " MEMBER OF ")
+		fmt.Fprint(w, " (")
+		n.Args[1].Format(w)
+		fmt.Fprint(w, ")")
 		return true
 	}
 	return false

diff --git a/parser/lexer.go b/parser/lexer.go
@@ -252,13 +252,17 @@ func (s *Scanner) Lex(v *yySymType) int {
 	if tok == not && s.sqlMode.HasHighNotPrecedenceMode() {
 		return not2
 	}
-	if tok == as && s.getNextToken() == of {
+	if (tok == as || tok == member) && s.getNextToken() == of {
 		_, pos, lit = s.scan()
 		v.ident = fmt.Sprintf("%s %s", v.ident, lit)
-		s.lastKeyword = asof
 		s.lastScanOffset = pos.Offset
 		v.offset = pos.Offset
-		return asof
+		if tok == as {
+			s.lastKeyword = asof
+			return asof
+		}
+		s.lastKeyword = memberof
+		return memberof
 	}
 	if tok == to {
 		tok1, tok2 := s.getNextTwoTokens()

diff --git a/parser/misc.go b/parser/misc.go
@@ -153,6 +153,7 @@ var tokenMap = map[string]int{
 	"ANY":                      any,
 	"APPROX_COUNT_DISTINCT":    approxCountDistinct,
 	"APPROX_PERCENTILE":        approxPercentile,
+	"ARRAY":                    array,
 	"AS":                       as,
 	"ASC":                      asc,
 	"ASCII":                    ascii,
@@ -474,6 +475,7 @@ var tokenMap = map[string]int{
 	"MEDIUMINT":                mediumIntType,
 	"MEDIUMTEXT":               mediumtextType,
 	"MEMORY":                   memory,
+	"MEMBER":                   member,
 	"MERGE":                    merge,
 	"MICROSECOND":              microsecond,
 	"MIN_ROWS":                 minRows,