From 595699250798cb5ef4a9ef8ef324c32233729d4d Mon Sep 17 00:00:00 2001
From: deoxxa <deoxxa@fknsrs.biz>
Date: Fri, 11 Dec 2015 11:24:17 +1100
Subject: [PATCH] add RETURNING option to insert query for postgres

---
 insert.go      | 24 +++++++++++++++++++++---
 insert_test.go |  7 +++++++
 2 files changed, 28 insertions(+), 3 deletions(-)

diff --git a/insert.go b/insert.go
index a8576c9..80b7761 100644
--- a/insert.go
+++ b/insert.go
@@ -2,9 +2,10 @@ package sqlbuilder
 
 // InsertStatement represents a INSERT statement.
 type InsertStatement struct {
-	columns ColumnList
-	values  []literal
-	into    Table
+	columns   ColumnList
+	values    []literal
+	into      Table
+	returning Column
 
 	err error
 }
@@ -72,6 +73,18 @@ func (b *InsertStatement) Set(column Column, value interface{}) *InsertStatement
 	return b
 }
 
+func (b *InsertStatement) Returning(column Column) *InsertStatement {
+	if b.err != nil {
+		return b
+	}
+	if !b.into.hasColumn(column) {
+		b.err = newError("column not found in FROM.")
+		return b
+	}
+	b.returning = column
+	return b
+}
+
 // ToSql generates query string, placeholder arguments, and returns err on errors.
 func (b *InsertStatement) ToSql() (query string, args []interface{}, err error) {
 	bldr := newBuilder()
@@ -119,5 +132,10 @@ func (b *InsertStatement) ToSql() (query string, args []interface{}, err error)
 	bldr.AppendItems(values, ", ")
 	bldr.Append(" )")
 
+	if b.returning != nil {
+		bldr.Append(" RETURNING ")
+		bldr.AppendItem(b.returning)
+	}
+
 	return
 }
diff --git a/insert_test.go b/insert_test.go
index 65606d8..3bd5d44 100644
--- a/insert_test.go
+++ b/insert_test.go
@@ -46,6 +46,13 @@ func TestInsert(t *testing.T) {
 		query:  `INSERT INTO "TABLE_A" ( "str", "bool", "float", "date", "bytes" ) VALUES ( ?, ?, ?, ?, ? );`,
 		args:   []interface{}{"hoge", true, 0.1, time.Unix(0, 0).UTC(), []byte{0x01}},
 		errmsg: "",
+	}, {
+		stmt: Insert(table1).
+			Set(table1.C("str"), "x").
+			Returning(table1.C("id")),
+		query:  `INSERT INTO "TABLE_A" ( "str" ) VALUES ( ? ) RETURNING "TABLE_A"."id";`,
+		args:   []interface{}{"x"},
+		errmsg: "",
 	}, {
 		stmt:   Insert(table1).Values(1, "hoge", true, 0.1, time.Unix(0, 0).UTC(), []byte{0x01}),
 		query:  `INSERT INTO "TABLE_A" ( "id", "str", "bool", "float", "date", "bytes" ) VALUES ( ?, ?, ?, ?, ?, ? );`,