fix #83: Support RETURNING in InsertBuilder

Author: huandu

insert.go

@@ -14,6 +14,7 @@ const (
 	insertMarkerAfterCols
 	insertMarkerAfterValues
 	insertMarkerAfterSelect
+	insertMarkerAfterReturning
 )
 
 // NewInsertBuilder creates a new INSERT builder.
@@ -32,10 +33,11 @@ func newInsertBuilder() *InsertBuilder {
 
 // InsertBuilder is a builder to build INSERT.
 type InsertBuilder struct {
-	verb   string
-	table  string
-	cols   []string
-	values [][]string
+	verb      string
+	table     string
+	cols      []string
+	values    [][]string
+	returning []string
 
 	args *Args
 
@@ -112,6 +114,14 @@ func (ib *InsertBuilder) Values(value ...interface{}) *InsertBuilder {
 	return ib
 }
 
+// Returning sets returning columns.
+// For DBMS that doesn't support RETURNING, e.g. MySQL, it will be ignored.
+func (ib *InsertBuilder) Returning(col ...string) *InsertBuilder {
+	ib.returning = col
+	ib.marker = insertMarkerAfterReturning
+	return ib
+}
+
 // NumValue returns the number of values to insert.
 func (ib *InsertBuilder) NumValue() int {
 	return len(ib.values)
@@ -203,6 +213,15 @@ func (ib *InsertBuilder) BuildWithFlavor(flavor Flavor, initialArg ...interface{
 
 	ib.injection.WriteTo(buf, insertMarkerAfterValues)
 
+	if flavor == PostgreSQL || flavor == SQLite {
+		if len(ib.returning) > 0 {
+			buf.WriteLeadingString("RETURNING ")
+			buf.WriteStrings(ib.returning, ", ")
+		}
+
+		ib.injection.WriteTo(buf, insertMarkerAfterReturning)
+	}
+
 	return ib.args.CompileWithFlavor(buf.String(), flavor, initialArg...)
 }
 

insert_test.go

@@ -245,6 +245,54 @@ func ExampleInsertBuilder_NumValue() {
 	// 2
 }
 
+func ExampleInsertBuilder_Returning() {
+	sql, args := InsertInto("user").
+		Cols("name").Values("Huan Du").
+		Returning("id").
+		BuildWithFlavor(PostgreSQL)
+
+	fmt.Println(sql)
+	fmt.Println(args)
+
+	// Output:
+	// INSERT INTO user (name) VALUES ($1) RETURNING id
+	// [Huan Du]
+}
+
+func TestInsertBuilderReturning(test *testing.T) {
+	a := assert.New(test)
+	ib := InsertInto("user").
+		Cols("name").Values("Huan Du").
+		Returning("id")
+
+	sql, _ := ib.BuildWithFlavor(MySQL)
+	a.Equal("INSERT INTO user (name) VALUES (?)", sql)
+
+	sql, _ = ib.BuildWithFlavor(PostgreSQL)
+	a.Equal("INSERT INTO user (name) VALUES ($1) RETURNING id", sql)
+
+	sql, _ = ib.BuildWithFlavor(SQLite)
+	a.Equal("INSERT INTO user (name) VALUES (?) RETURNING id", sql)
+
+	sql, _ = ib.BuildWithFlavor(SQLServer)
+	a.Equal("INSERT INTO user (name) VALUES (@p1)", sql)
+
+	sql, _ = ib.BuildWithFlavor(CQL)
+	a.Equal("INSERT INTO user (name) VALUES (?)", sql)
+
+	sql, _ = ib.BuildWithFlavor(ClickHouse)
+	a.Equal("INSERT INTO user (name) VALUES (?)", sql)
+
+	sql, _ = ib.BuildWithFlavor(Presto)
+	a.Equal("INSERT INTO user (name) VALUES (?)", sql)
+
+	sql, _ = ib.BuildWithFlavor(Oracle)
+	a.Equal("INSERT INTO user (name) VALUES (:1)", sql)
+
+	sql, _ = ib.BuildWithFlavor(Informix)
+	a.Equal("INSERT INTO user (name) VALUES (?)", sql)
+}
+
 func TestInsertBuilderGetFlavor(t *testing.T) {
 	a := assert.New(t)
 	ib := newInsertBuilder()