Merge pull request #212 from huandu/feature/returning-clause-issue-210

feat: implement Returning method for UpdateBuilder and DeleteBuilder

delete.go

@@ -10,6 +10,7 @@ const (
 	deleteMarkerAfterWhere
 	deleteMarkerAfterOrderBy
 	deleteMarkerAfterLimit
+	deleteMarkerAfterReturning
 )
 
 // NewDeleteBuilder creates a new DELETE builder.
@@ -47,6 +48,7 @@ type DeleteBuilder struct {
 	orderByCols []string
 	order       string
 	limitVar    string
+	returning   []string
 
 	args *Args
 
@@ -157,6 +159,14 @@ func (db *DeleteBuilder) Limit(limit int) *DeleteBuilder {
 	return db
 }
 
+// Returning sets returning columns.
+// For DBMS that doesn't support RETURNING, e.g. MySQL, it will be ignored.
+func (db *DeleteBuilder) Returning(col ...string) *DeleteBuilder {
+	db.returning = col
+	db.marker = deleteMarkerAfterReturning
+	return db
+}
+
 // String returns the compiled DELETE string.
 func (db *DeleteBuilder) String() string {
 	s, _ := db.Build()
@@ -218,6 +228,15 @@ func (db *DeleteBuilder) BuildWithFlavor(flavor Flavor, initialArg ...interface{
 		db.injection.WriteTo(buf, deleteMarkerAfterLimit)
 	}
 
+	if flavor == PostgreSQL || flavor == SQLite {
+		if len(db.returning) > 0 {
+			buf.WriteLeadingString("RETURNING ")
+			buf.WriteStrings(db.returning, ", ")
+		}
+
+		db.injection.WriteTo(buf, deleteMarkerAfterReturning)
+	}
+
 	return db.args.CompileWithFlavor(buf.String(), flavor, initialArg...)
 }
 

delete_test.go

@@ -96,3 +96,100 @@ func TestDeleteBuilderGetFlavor(t *testing.T) {
 	flavor = dbClick.Flavor()
 	a.Equal(ClickHouse, flavor)
 }
+
+func ExampleDeleteBuilder_Returning() {
+	db := NewDeleteBuilder()
+	db.DeleteFrom("user")
+	db.Where(db.Equal("id", 123))
+	db.Returning("id", "deleted_at")
+
+	sql, args := db.BuildWithFlavor(PostgreSQL)
+	fmt.Println(sql)
+	fmt.Println(args)
+
+	// Output:
+	// DELETE FROM user WHERE id = $1 RETURNING id, deleted_at
+	// [123]
+}
+
+func TestDeleteBuilderReturning(t *testing.T) {
+	a := assert.New(t)
+	db := NewDeleteBuilder()
+	db.DeleteFrom("user")
+	db.Where(db.Equal("id", 123))
+	db.Returning("id", "deleted_at")
+
+	sql, _ := db.BuildWithFlavor(MySQL)
+	a.Equal("DELETE FROM user WHERE id = ?", sql)
+
+	sql, _ = db.BuildWithFlavor(PostgreSQL)
+	a.Equal("DELETE FROM user WHERE id = $1 RETURNING id, deleted_at", sql)
+
+	sql, _ = db.BuildWithFlavor(SQLite)
+	a.Equal("DELETE FROM user WHERE id = ? RETURNING id, deleted_at", sql)
+
+	sql, _ = db.BuildWithFlavor(SQLServer)
+	a.Equal("DELETE FROM user WHERE id = @p1", sql)
+
+	sql, _ = db.BuildWithFlavor(CQL)
+	a.Equal("DELETE FROM user WHERE id = ?", sql)
+
+	sql, _ = db.BuildWithFlavor(ClickHouse)
+	a.Equal("DELETE FROM user WHERE id = ?", sql)
+
+	sql, _ = db.BuildWithFlavor(Presto)
+	a.Equal("DELETE FROM user WHERE id = ?", sql)
+
+	// Test with no returning columns
+	db2 := NewDeleteBuilder()
+	db2.DeleteFrom("user")
+	db2.Where(db2.Equal("id", 1))
+	db2.Returning() // Empty returning
+
+	sql, _ = db2.BuildWithFlavor(PostgreSQL)
+	a.Equal("DELETE FROM user WHERE id = $1", sql)
+
+	// Test with single column
+	db3 := NewDeleteBuilder()
+	db3.DeleteFrom("user")
+	db3.Where(db3.Equal("id", 1))
+	db3.Returning("id")
+
+	sql, _ = db3.BuildWithFlavor(PostgreSQL)
+	a.Equal("DELETE FROM user WHERE id = $1 RETURNING id", sql)
+
+	// Test with ORDER BY and LIMIT
+	db4 := NewDeleteBuilder()
+	db4.DeleteFrom("user")
+	db4.Where(db4.Equal("status", 1))
+	db4.OrderBy("id").Asc()
+	db4.Limit(5)
+	db4.Returning("id", "name")
+
+	sql, _ = db4.BuildWithFlavor(PostgreSQL)
+	a.Equal("DELETE FROM user WHERE status = $1 ORDER BY id ASC LIMIT $2 RETURNING id, name", sql)
+
+	// Test chaining
+	db5 := NewDeleteBuilder().DeleteFrom("user").Where("status = 0").Returning("id").Returning("name", "deleted_at")
+	sql, _ = db5.BuildWithFlavor(PostgreSQL)
+	a.Equal("DELETE FROM user WHERE status = 0 RETURNING name, deleted_at", sql) // Last Returning call overwrites
+
+	// Test SQL injection after RETURNING
+	db6 := NewDeleteBuilder()
+	db6.DeleteFrom("user")
+	db6.Where(db6.Equal("id", 1))
+	db6.Returning("id", "name")
+	db6.SQL("/* comment after returning */")
+
+	sql, _ = db6.BuildWithFlavor(PostgreSQL)
+	a.Equal("DELETE FROM user WHERE id = $1 RETURNING id, name /* comment after returning */", sql)
+
+	// Test with CTE (WITH clause)
+	cte := With(CTETable("temp_user").As(Select("id").From("inactive_users")))
+	db7 := cte.DeleteFrom("user")
+	db7.Where("user.id IN (SELECT id FROM temp_user)")
+	db7.Returning("id", "deleted_at")
+
+	sql, _ = db7.BuildWithFlavor(PostgreSQL)
+	a.Equal("WITH temp_user AS (SELECT id FROM inactive_users) DELETE FROM user, temp_user WHERE user.id IN (SELECT id FROM temp_user) RETURNING id, deleted_at", sql)
+}

update.go

@@ -15,6 +15,7 @@ const (
 	updateMarkerAfterWhere
 	updateMarkerAfterOrderBy
 	updateMarkerAfterLimit
+	updateMarkerAfterReturning
 )
 
 // NewUpdateBuilder creates a new UPDATE builder.
@@ -53,6 +54,7 @@ type UpdateBuilder struct {
 	orderByCols []string
 	order       string
 	limitVar    string
+	returning   []string
 
 	args *Args
 
@@ -216,6 +218,14 @@ func (ub *UpdateBuilder) Limit(limit int) *UpdateBuilder {
 	return ub
 }
 
+// Returning sets returning columns.
+// For DBMS that doesn't support RETURNING, e.g. MySQL, it will be ignored.
+func (ub *UpdateBuilder) Returning(col ...string) *UpdateBuilder {
+	ub.returning = col
+	ub.marker = updateMarkerAfterReturning
+	return ub
+}
+
 // NumAssignment returns the number of assignments to update.
 func (ub *UpdateBuilder) NumAssignment() int {
 	return len(ub.assignments)
@@ -310,6 +320,15 @@ func (ub *UpdateBuilder) BuildWithFlavor(flavor Flavor, initialArg ...interface{
 		ub.injection.WriteTo(buf, updateMarkerAfterLimit)
 	}
 
+	if flavor == PostgreSQL || flavor == SQLite {
+		if len(ub.returning) > 0 {
+			buf.WriteLeadingString("RETURNING ")
+			buf.WriteStrings(ub.returning, ", ")
+		}
+
+		ub.injection.WriteTo(buf, updateMarkerAfterReturning)
+	}
+
 	return ub.args.CompileWithFlavor(buf.String(), flavor, initialArg...)
 }
 

update_test.go

@@ -162,3 +162,107 @@ func TestUpdateBuilderGetFlavor(t *testing.T) {
 	flavor = ubClick.Flavor()
 	a.Equal(ClickHouse, flavor)
 }
+
+func ExampleUpdateBuilder_Returning() {
+	ub := NewUpdateBuilder()
+	ub.Update("user")
+	ub.Set(ub.Assign("name", "Huan Du"))
+	ub.Where(ub.Equal("id", 123))
+	ub.Returning("id", "updated_at")
+
+	sql, args := ub.BuildWithFlavor(PostgreSQL)
+	fmt.Println(sql)
+	fmt.Println(args)
+
+	// Output:
+	// UPDATE user SET name = $1 WHERE id = $2 RETURNING id, updated_at
+	// [Huan Du 123]
+}
+
+func TestUpdateBuilderReturning(t *testing.T) {
+	a := assert.New(t)
+	ub := NewUpdateBuilder()
+	ub.Update("user")
+	ub.Set(ub.Assign("name", "Huan Du"))
+	ub.Where(ub.Equal("id", 123))
+	ub.Returning("id", "updated_at")
+
+	sql, _ := ub.BuildWithFlavor(MySQL)
+	a.Equal("UPDATE user SET name = ? WHERE id = ?", sql)
+
+	sql, _ = ub.BuildWithFlavor(PostgreSQL)
+	a.Equal("UPDATE user SET name = $1 WHERE id = $2 RETURNING id, updated_at", sql)
+
+	sql, _ = ub.BuildWithFlavor(SQLite)
+	a.Equal("UPDATE user SET name = ? WHERE id = ? RETURNING id, updated_at", sql)
+
+	sql, _ = ub.BuildWithFlavor(SQLServer)
+	a.Equal("UPDATE user SET name = @p1 WHERE id = @p2", sql)
+
+	sql, _ = ub.BuildWithFlavor(CQL)
+	a.Equal("UPDATE user SET name = ? WHERE id = ?", sql)
+
+	sql, _ = ub.BuildWithFlavor(ClickHouse)
+	a.Equal("UPDATE user SET name = ? WHERE id = ?", sql)
+
+	sql, _ = ub.BuildWithFlavor(Presto)
+	a.Equal("UPDATE user SET name = ? WHERE id = ?", sql)
+
+	// Test with no returning columns
+	ub2 := NewUpdateBuilder()
+	ub2.Update("user")
+	ub2.Set(ub2.Assign("name", "Test"))
+	ub2.Where(ub2.Equal("id", 1))
+	ub2.Returning() // Empty returning
+
+	sql, _ = ub2.BuildWithFlavor(PostgreSQL)
+	a.Equal("UPDATE user SET name = $1 WHERE id = $2", sql)
+
+	// Test with single column
+	ub3 := NewUpdateBuilder()
+	ub3.Update("user")
+	ub3.Set(ub3.Assign("name", "Test"))
+	ub3.Where(ub3.Equal("id", 1))
+	ub3.Returning("id")
+
+	sql, _ = ub3.BuildWithFlavor(PostgreSQL)
+	a.Equal("UPDATE user SET name = $1 WHERE id = $2 RETURNING id", sql)
+
+	// Test with ORDER BY and LIMIT
+	ub4 := NewUpdateBuilder()
+	ub4.Update("user")
+	ub4.Set(ub4.Assign("name", "Test"))
+	ub4.Where(ub4.Equal("status", 1))
+	ub4.OrderBy("id").Asc()
+	ub4.Limit(5)
+	ub4.Returning("id", "name")
+
+	sql, _ = ub4.BuildWithFlavor(PostgreSQL)
+	a.Equal("UPDATE user SET name = $1 WHERE status = $2 ORDER BY id ASC LIMIT $3 RETURNING id, name", sql)
+
+	// Test chaining
+	ub5 := NewUpdateBuilder().Update("user").Set("status = 1").Returning("id").Returning("name", "updated_at")
+	sql, _ = ub5.BuildWithFlavor(PostgreSQL)
+	a.Equal("UPDATE user SET status = 1 RETURNING name, updated_at", sql) // Last Returning call overwrites
+
+	// Test SQL injection after RETURNING
+	ub6 := NewUpdateBuilder()
+	ub6.Update("user")
+	ub6.Set(ub6.Assign("name", "Test"))
+	ub6.Where(ub6.Equal("id", 1))
+	ub6.Returning("id", "name")
+	ub6.SQL("/* comment after returning */")
+
+	sql, _ = ub6.BuildWithFlavor(PostgreSQL)
+	a.Equal("UPDATE user SET name = $1 WHERE id = $2 RETURNING id, name /* comment after returning */", sql)
+
+	// Test with CTE (WITH clause)
+	cte := With(CTETable("temp_user").As(Select("id").From("active_users")))
+	ub7 := cte.Update("user")
+	ub7.Set(ub7.Assign("status", "active"))
+	ub7.Where("user.id IN (SELECT id FROM temp_user)")
+	ub7.Returning("id", "status")
+
+	sql, _ = ub7.BuildWithFlavor(PostgreSQL)
+	a.Equal("WITH temp_user AS (SELECT id FROM active_users) UPDATE user SET status = $1 FROM temp_user WHERE user.id IN (SELECT id FROM temp_user) RETURNING id, status", sql)
+}