support "returning" on insert

Author: deoxxa

insert.go

@@ -3,14 +3,16 @@ package sqlbuilder
 type InsertColumns map[*BasicColumn]AsExpr
 
 type InsertStatement struct {
-	table   *Table
-	columns InsertColumns
+	table     *Table
+	columns   InsertColumns
+	returning []AsExpr
 }
 
 func (s *InsertStatement) clone() *InsertStatement {
 	return &InsertStatement{
-		table:   s.table,
-		columns: s.columns,
+		table:     s.table,
+		columns:   s.columns,
+		returning: s.returning,
 	}
 }
 
@@ -30,6 +32,12 @@ func (s *InsertStatement) Columns(columns InsertColumns) *InsertStatement {
 	return c
 }
 
+func (s *InsertStatement) Returning(returning ...AsExpr) *InsertStatement {
+	c := s.clone()
+	c.returning = returning
+	return c
+}
+
 func (q *InsertStatement) AsStatement(s *Serializer) {
 	s.D("INSERT INTO ").F(q.table.AsNamed).D(" ")
 
@@ -50,4 +58,18 @@ func (q *InsertStatement) AsStatement(s *Serializer) {
 		s.F(k.AsExpr).DC(", ", i < len(keys)-1)
 	}
 	s.D(")")
+
+	if len(q.returning) > 0 {
+		s.D(" RETURNING ")
+
+		for i, c := range q.returning {
+			if a, ok := c.(AsResultColumn); ok {
+				s.F(a.AsResultColumn)
+			} else {
+				s.F(c.AsExpr)
+			}
+
+			s.DC(", ", i < len(q.returning)-1)
+		}
+	}
 }

sqlbuilder_test.go

@@ -7,6 +7,40 @@ import (
 	"github.com/stretchr/testify/assert"
 )
 
+func TestInsert(t *testing.T) {
+	a := assert.New(t)
+
+	tbl := NewTable("users", "id", "name", "email")
+
+	q := Insert().Table(tbl).Columns(InsertColumns{
+		tbl.C("name"):  Bind("jim"),
+		tbl.C("email"): Bind("[email protected]"),
+	})
+
+	qs, qv, err := NewSerializer(DialectPostgres{}).F(q.AsStatement).ToSQL()
+
+	a.NoError(err)
+	a.Equal("INSERT INTO \"users\" (\"name\", \"email\") VALUES ($1, $2)", qs)
+	a.Equal([]interface{}{"jim", "[email protected]"}, qv)
+}
+
+func TestInsertReturning(t *testing.T) {
+	a := assert.New(t)
+
+	tbl := NewTable("users", "id", "name", "email")
+
+	q := Insert().Table(tbl).Columns(InsertColumns{
+		tbl.C("name"):  Bind("jim"),
+		tbl.C("email"): Bind("[email protected]"),
+	}).Returning(tbl.C("id"))
+
+	qs, qv, err := NewSerializer(DialectPostgres{}).F(q.AsStatement).ToSQL()
+
+	a.NoError(err)
+	a.Equal("INSERT INTO \"users\" (\"name\", \"email\") VALUES ($1, $2) RETURNING \"users\".\"id\"", qs)
+	a.Equal([]interface{}{"jim", "[email protected]"}, qv)
+}
+
 func TestUpdate(t *testing.T) {
 	a := assert.New(t)