diff --git a/delete.go b/delete.go index a96295b..1858e21 100644 --- a/delete.go +++ b/delete.go @@ -223,6 +223,15 @@ func (db *DeleteBuilder) BuildWithFlavor(flavor Flavor, initialArg ...interface{ db.injection.WriteTo(buf, deleteMarkerAfterDeleteFrom) + if flavor == SQLServer { + if len(db.returning) > 0 { + buf.WriteLeadingString("OUTPUT ") + buf.WriteStringsPrefixed("DELETED.", db.returning, ", ") + } + + db.injection.WriteTo(buf, insertMarkerAfterReturning) + } + if db.WhereClause != nil { db.whereClauseProxy.WhereClause = db.WhereClause defer func() { diff --git a/delete_test.go b/delete_test.go index 400bd8f..15ff183 100644 --- a/delete_test.go +++ b/delete_test.go @@ -129,7 +129,7 @@ func TestDeleteBuilderReturning(t *testing.T) { 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) + a.Equal("DELETE FROM user OUTPUT DELETED.id, DELETED.deleted_at WHERE id = @p1", sql) sql, _ = db.BuildWithFlavor(CQL) a.Equal("DELETE FROM user WHERE id = ?", sql) diff --git a/insert.go b/insert.go index 7698139..cd4a860 100644 --- a/insert.go +++ b/insert.go @@ -200,6 +200,15 @@ func (ib *InsertBuilder) BuildWithFlavor(flavor Flavor, initialArg ...interface{ ib.injection.WriteTo(buf, insertMarkerAfterCols) } + if flavor == SQLServer { + if len(ib.returning) > 0 { + buf.WriteLeadingString("OUTPUT ") + buf.WriteStringsPrefixed("INSERTED.", ib.returning, ", ") + } + + ib.injection.WriteTo(buf, insertMarkerAfterReturning) + } + if ib.sbHolder != "" { buf.WriteString(" ") buf.WriteString(ib.sbHolder) diff --git a/insert_test.go b/insert_test.go index 6c11c40..527df27 100644 --- a/insert_test.go +++ b/insert_test.go @@ -275,7 +275,7 @@ func TestInsertBuilderReturning(test *testing.T) { a.Equal("INSERT INTO user (name) VALUES (?) RETURNING id", sql) sql, _ = ib.BuildWithFlavor(SQLServer) - a.Equal("INSERT INTO user (name) VALUES (@p1)", sql) + a.Equal("INSERT INTO user (name) OUTPUT INSERTED.id VALUES (@p1)", sql) sql, _ = ib.BuildWithFlavor(CQL) a.Equal("INSERT INTO user (name) VALUES (?)", sql) diff --git a/stringbuilder.go b/stringbuilder.go index 6fd37df..02f9189 100644 --- a/stringbuilder.go +++ b/stringbuilder.go @@ -56,6 +56,16 @@ func (sb *stringBuilder) WriteStrings(ss []string, sep string) { } } +func (sb *stringBuilder) WriteStringsPrefixed(prefix string, ss []string, sep string) { + prefixedSs := []string{} + + for _, s := range ss { + prefixedSs = append(prefixedSs, prefix+s) + } + + sb.WriteStrings(prefixedSs, sep) +} + func (sb *stringBuilder) WriteRune(r rune) { sb.builder.WriteRune(r) } diff --git a/update.go b/update.go index 56e9ed7..735e168 100644 --- a/update.go +++ b/update.go @@ -301,6 +301,15 @@ func (ub *UpdateBuilder) BuildWithFlavor(flavor Flavor, initialArg ...interface{ ub.injection.WriteTo(buf, updateMarkerAfterSet) + if flavor == SQLServer { + if len(ub.returning) > 0 { + buf.WriteLeadingString("OUTPUT ") + buf.WriteStringsPrefixed("INSERTED.", ub.returning, ", ") + } + + ub.injection.WriteTo(buf, insertMarkerAfterReturning) + } + if flavor != MySQL { // For ISO SQL, CTE table names should be written after FROM keyword. if ub.cteBuilder != nil { diff --git a/update_test.go b/update_test.go index acce757..b1778cc 100644 --- a/update_test.go +++ b/update_test.go @@ -197,7 +197,7 @@ func TestUpdateBuilderReturning(t *testing.T) { 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) + a.Equal("UPDATE user SET name = @p1 OUTPUT INSERTED.id, INSERTED.updated_at WHERE id = @p2", sql) sql, _ = ub.BuildWithFlavor(CQL) a.Equal("UPDATE user SET name = ? WHERE id = ?", sql)