Merge remote-tracking branch 'origin/main' into feature/ml-predict

ast/ast.go

@@ -173,6 +173,7 @@ func (CallExpr) isExpr()              {}
 func (CountStarExpr) isExpr()         {}
 func (CastExpr) isExpr()              {}
 func (ExtractExpr) isExpr()           {}
+func (ReplaceFieldsExpr) isExpr()     {}
 func (CaseExpr) isExpr()              {}
 func (IfExpr) isExpr()                {}
 func (ParenExpr) isExpr()             {}
@@ -218,6 +219,7 @@ type Arg interface {
 func (ExprArg) isArg()     {}
 func (IntervalArg) isArg() {}
 func (SequenceArg) isArg() {}
+func (LambdaArg) isArg()   {}
 
 type TVFArg interface {
 	Node
@@ -1224,6 +1226,21 @@ type SequenceArg struct {
 	Expr Expr
 }
 
+// LambdaArg is lambda expression argument of the generic function call.
+//
+//	{{if .Lparen.Invalid}}{{.Args | sqlJoin ", "}}{{else}}({{.Args | sqlJoin ", "}}) -> {{.Expr | sql}}
+//
+// Note: Args won't be empty. If Lparen is not appeared, Args have exactly one element.
+type LambdaArg struct {
+	// pos = Lparen || Args[0].pos
+	// end = Expr.end
+
+	Lparen token.Pos // optional
+
+	Args []*Ident // if Lparen.Invalid() then len(Args) = 1 else len(Args) > 0
+	Expr Expr
+}
+
 // ModelArg is argument of model function call.
 //
 //	MODEL {{.Name | sql}}
@@ -1329,6 +1346,31 @@ type ExtractExpr struct {
 	AtTimeZone *AtTimeZone // optional
 }
 
+// ReplaceFieldsArg is value AS field_path node in ReplaceFieldsExpr.
+//
+//	{{.Expr | sql}} AS {{.Field | sql}}
+type ReplaceFieldsArg struct {
+	// pos = Expr.pos
+	// end = Field.end
+
+	Expr  Expr
+	Field *Path
+}
+
+// ReplaceFieldsExpr is REPLACE_FIELDS call expression node.
+//
+//	REPLACE_FIELDS({{.Expr.| sql}}, {{.Fields | sqlJoin ", "}})
+type ReplaceFieldsExpr struct {
+	// pos = ReplaceFields
+	// end = Rparen + 1
+
+	ReplaceFields token.Pos // position of "REPLACE_FIELDS" keyword
+	Rparen        token.Pos // position of ")"
+
+	Expr   Expr
+	Fields []*ReplaceFieldsArg
+}
+
 // AtTimeZone is AT TIME ZONE clause in EXTRACT call.
 //
 //	AT TIME ZONE {{.Expr | sql}}
@@ -3179,20 +3221,48 @@ type AlterSearchIndex struct {
 //
 // ================================================================================
 
+// WithAction is WITH ACTION clause in ThenReturn.
+//
+//	WITH ACTION {{.Alias | sqlOpt}}
+type WithAction struct {
+	// pos = With
+	// end = Alias.end || Action + 6
+
+	With   token.Pos // position of "WITH" keyword
+	Action token.Pos // position of "ACTION" keyword
+
+	Alias *AsAlias // optional
+}
+
+// ThenReturn is THEN RETURN clause in DML.
+//
+//	THEN RETURN {{.WithAction | sqlOpt}} {{.Items | sqlJoin ", "}}
+type ThenReturn struct {
+	// pos = Then
+	// end = Items[$].end
+
+	Then token.Pos // position of "THEN" keyword
+
+	WithAction *WithAction // optional
+	Items      []SelectItem
+}
+
 // Insert is INSERT statement node.
 //
 //	INSERT {{if .InsertOrType}}OR .InsertOrType{{end}}INTO {{.TableName | sql}} ({{.Columns | sqlJoin ","}}) {{.Input | sql}}
+//	{{.ThenReturn | sqlOpt}}
 type Insert struct {
 	// pos = Insert
-	// end = Input.end
+	// end = (ThenReturn ?? Input).end
 
 	Insert token.Pos // position of "INSERT" keyword
 
 	InsertOrType InsertOrType
 
-	TableName *Ident
-	Columns   []*Ident
-	Input     InsertInput
+	TableName  *Ident
+	Columns    []*Ident
+	Input      InsertInput
+	ThenReturn *ThenReturn // optional
 }
 
 // ValuesInput is VALUES clause in INSERT.
@@ -3245,31 +3315,35 @@ type SubQueryInput struct {
 // Delete is DELETE statement.
 //
 //	DELETE FROM {{.TableName | sql}} {{.As | sqlOpt}} {{.Where | sql}}
+//	{{.ThenReturn | sqlOpt}}
 type Delete struct {
 	// pos = Delete
-	// end = Where.end
+	// end = (ThenReturn ?? Where).end
 
 	Delete token.Pos // position of "DELETE" keyword
 
-	TableName *Ident
-	As        *AsAlias // optional
-	Where     *Where
+	TableName  *Ident
+	As         *AsAlias // optional
+	Where      *Where
+	ThenReturn *ThenReturn // optional
 }
 
 // Update is UPDATE statement.
 //
 //	UPDATE {{.TableName | sql}} {{.As | sqlOpt}}
 //	SET {{.Updates | sqlJoin ","}} {{.Where | sql}}
+//	{{.ThenReturn | sqlOpt}}
 type Update struct {
 	// pos = Update
-	// end = Where.end
+	// end = (ThenReturn ?? Where).end
 
 	Update token.Pos // position of "UPDATE" keyword
 
-	TableName *Ident
-	As        *AsAlias      // optional
-	Updates   []*UpdateItem // len(Updates) > 0
-	Where     *Where
+	TableName  *Ident
+	As         *AsAlias      // optional
+	Updates    []*UpdateItem // len(Updates) > 0
+	Where      *Where
+	ThenReturn *ThenReturn // optional
 }
 
 // UpdateItem is SET clause items in UPDATE.

ast/const.go

@@ -14,12 +14,12 @@ const (
 	JoinTypeJoinHint       JoinHintKey = "JOIN_TYPE"
 )
 
+// JoinMethod is used for prefix of JOIN, not for hint.
 type JoinMethod string
 
 const (
-	HashJoinMethod  JoinMethod = "HASH"
-	ApplyJoinMethod JoinMethod = "APPLY"
-	LoopJoinMethod  JoinMethod = "LOOP" // Undocumented, but the Spanner accept this value at least.
+	HashJoinMethod   JoinMethod = "HASH"
+	LookupJoinMethod JoinMethod = "LOOKUP" // Undocumented, but the GoogleSQL can parse this value.
 )
 
 type PositionKeyword string

ast/sql.go

@@ -506,6 +506,15 @@ func (c *CallExpr) SQL() string {
 		")"
 }
 
+func (l *LambdaArg) SQL() string {
+	// This implementation is not exactly matched with the doc comment for simplicity.
+	return strOpt(!l.Lparen.Invalid(), "(") +
+		sqlJoin(l.Args, ", ") +
+		strOpt(!l.Lparen.Invalid(), ")") +
+		" -> " +
+		l.Expr.SQL()
+}
+
 func (c *TVFCallExpr) SQL() string {
 	return c.Name.SQL() + "(" +
 		sqlJoin(c.Args, ", ") +
@@ -566,6 +575,12 @@ func (a *AtTimeZone) SQL() string {
 	return "AT TIME ZONE " + a.Expr.SQL()
 }
 
+func (r *ReplaceFieldsArg) SQL() string { return r.Expr.SQL() + " AS " + r.Field.SQL() }
+
+func (r *ReplaceFieldsExpr) SQL() string {
+	return "REPLACE_FIELDS(" + r.Expr.SQL() + ", " + sqlJoin(r.Fields, ", ") + ")"
+}
+
 func (c *CastExpr) SQL() string {
 	return strOpt(c.Safe, "SAFE_") + "CAST(" + c.Expr.SQL() + " AS " + c.Type.SQL() + ")"
 }
@@ -1387,20 +1402,22 @@ func (a *AlterSearchIndex) SQL() string {
 //
 // ================================================================================
 
+func (w *WithAction) SQL() string {
+	return "WITH ACTION" + sqlOpt(" ", w.Alias, "")
+}
+
+func (t *ThenReturn) SQL() string {
+	return "THEN RETURN " + sqlOpt("", t.WithAction, " ") + sqlJoin(t.Items, ", ")
+}
+
 func (i *Insert) SQL() string {
-	sql := "INSERT "
-	if i.InsertOrType != "" {
-		sql += "OR " + string(i.InsertOrType) + " "
-	}
-	sql += "INTO " + i.TableName.SQL() + " ("
-	for i, c := range i.Columns {
-		if i != 0 {
-			sql += ", "
-		}
-		sql += c.SQL()
-	}
-	sql += ") " + i.Input.SQL()
-	return sql
+	return "INSERT " +
+		strOpt(i.InsertOrType != "", "OR "+string(i.InsertOrType)+" ") +
+		"INTO " + i.TableName.SQL() + " (" +
+		sqlJoin(i.Columns, ", ") +
+		") " +
+		i.Input.SQL() +
+		sqlOpt(" ", i.ThenReturn, "")
 }
 
 func (v *ValuesInput) SQL() string {
@@ -1438,25 +1455,20 @@ func (s *SubQueryInput) SQL() string {
 }
 
 func (d *Delete) SQL() string {
-	sql := "DELETE FROM " + d.TableName.SQL()
-	if d.As != nil {
-		sql += " " + d.As.SQL()
-	}
-	sql += " " + d.Where.SQL()
-	return sql
+	return "DELETE FROM " +
+		d.TableName.SQL() + " " +
+		sqlOpt("", d.As, " ") +
+		d.Where.SQL() +
+		sqlOpt(" ", d.ThenReturn, "")
 }
 
 func (u *Update) SQL() string {
-	sql := "UPDATE " + u.TableName.SQL()
-	if u.As != nil {
-		sql += " " + u.As.SQL()
-	}
-	sql += " SET " + u.Updates[0].SQL()
-	for _, item := range u.Updates[1:] {
-		sql += ", " + item.SQL()
-	}
-	sql += " " + u.Where.SQL()
-	return sql
+	return "UPDATE " + u.TableName.SQL() + " " +
+		sqlOpt("", u.As, " ") +
+		"SET " +
+		sqlJoin(u.Updates, ", ") +
+		" " + u.Where.SQL() +
+		sqlOpt(" ", u.ThenReturn, "")
 }
 
 func (u *UpdateItem) SQL() string {