Implement THEN RETURN clause in DML statements (#214)

* Implement THEN RETURN

* Update testdata

* Fix Update.SQL()

* Fix Update.SQL()

* Fix needless sqlOpt

ast/ast.go

@@ -3165,20 +3165,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.
@@ -3231,31 +3259,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/sql.go

@@ -1385,20 +1385,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 {
@@ -1436,25 +1438,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 {

parser.go

@@ -4124,6 +4124,39 @@ func (p *Parser) parseDML() ast.DML {
 	panic(p.errorfAtToken(id, "expect pseudo keyword: INSERT, DELETE,  UPDATE but: %s", id.AsString))
 }
 
+func (p *Parser) tryParseWithAction() *ast.WithAction {
+	if p.Token.Kind != "WITH" {
+		return nil
+	}
+
+	with := p.expect("WITH").Pos
+	action := p.expectKeywordLike("ACTION").Pos
+	alias := p.tryParseAsAlias(withRequiredAs)
+
+	return &ast.WithAction{
+		With:  with,
+		Action: action,
+		Alias: alias,
+	}
+}
+
+func (p *Parser) tryParseThenReturn() *ast.ThenReturn {
+	if p.Token.Kind != "THEN" {
+		return nil
+	}
+
+	then := p.expect("THEN").Pos
+	p.expectKeywordLike("RETURN")
+	withAction := p.tryParseWithAction()
+	items := parseCommaSeparatedList(p, p.parseSelectItem)
+
+	return &ast.ThenReturn{
+		Then:       then,
+		WithAction: withAction,
+		Items:      items,
+	}
+}
+
 func (p *Parser) parseInsert(pos token.Pos) *ast.Insert {
 	var insertOrType ast.InsertOrType
 	if p.Token.Kind == "OR" {
@@ -4165,12 +4198,15 @@ func (p *Parser) parseInsert(pos token.Pos) *ast.Insert {
 		input = p.parseSubQueryInput()
 	}
 
+	thenReturn := p.tryParseThenReturn()
+
 	return &ast.Insert{
 		Insert:       pos,
 		InsertOrType: insertOrType,
 		TableName:    name,
 		Columns:      columns,
 		Input:        input,
+		ThenReturn:   thenReturn,
 	}
 }
 
@@ -4238,12 +4274,14 @@ func (p *Parser) parseDelete(pos token.Pos) *ast.Delete {
 	name := p.parseIdent()
 	as := p.tryParseAsAlias(withOptionalAs)
 	where := p.parseWhere()
+	thenReturn := p.tryParseThenReturn()
 
 	return &ast.Delete{
-		Delete:    pos,
-		TableName: name,
-		As:        as,
-		Where:     where,
+		Delete:     pos,
+		TableName:  name,
+		As:         as,
+		Where:      where,
+		ThenReturn: thenReturn,
 	}
 }
 
@@ -4256,13 +4294,15 @@ func (p *Parser) parseUpdate(pos token.Pos) *ast.Update {
 	items := parseCommaSeparatedList(p, p.parseUpdateItem)
 
 	where := p.parseWhere()
+	thenReturn := p.tryParseThenReturn()
 
 	return &ast.Update{
-		Update:    pos,
-		TableName: name,
-		As:        as,
-		Updates:   items,
-		Where:     where,
+		Update:     pos,
+		TableName:  name,
+		As:         as,
+		Updates:    items,
+		Where:      where,
+		ThenReturn: thenReturn,
 	}
 }
 

testdata/input/dml/delete_then_return.sql

@@ -0,0 +1 @@
+delete foo where foo = 1 and bar = 2 then return *

testdata/input/dml/insert_or_update_then_return_with_action_as.sql

@@ -0,0 +1,3 @@
+INSERT OR UPDATE INTO foo
+(foo, bar) VALUES (1, 2)
+THEN RETURN WITH ACTION AS act *

testdata/input/dml/update_then_return_with_action.sql

@@ -0,0 +1 @@
+update foo set foo = bar, bar = foo, baz = DEFAULT where foo = 1 then return with action *

testdata/result/dml/delete _from.sql.txt

@@ -43,6 +43,7 @@ delete from foo where foo = 1 and bar = 2
       },
     },
   },
+  ThenReturn: (*ast.ThenReturn)(nil),
 }
 
 --- SQL

testdata/result/dml/delete.sql.txt

@@ -43,6 +43,7 @@ delete foo where foo = 1 and bar = 2
       },
     },
   },
+  ThenReturn: (*ast.ThenReturn)(nil),
 }
 
 --- SQL

testdata/result/dml/delete_as.sql.txt

@@ -42,6 +42,7 @@ delete foo as F where F.foo = 1
       },
     },
   },
+  ThenReturn: (*ast.ThenReturn)(nil),
 }
 
 --- SQL