From 08876b4d420e761cbfa658aa6bb89b3f7c62c240 Mon Sep 17 00:00:00 2001
From: bgdnxt <95686883+bgdnxt@users.noreply.github.com>
Date: Thu, 24 Nov 2022 22:00:41 +0800
Subject: [PATCH] feat:  mssql output support for update or delete query (#718)

feat: query update & deleted add feature.Output supported

Co-authored-by: gfk <gfk@bb.io>
---
 .../testdata/snapshots/TestQuery-mssql2019-130 |  2 +-
 query_delete.go                                | 18 ++++++++++--------
 query_update.go                                | 18 ++++++++++--------
 3 files changed, 21 insertions(+), 17 deletions(-)

diff --git a/internal/dbtest/testdata/snapshots/TestQuery-mssql2019-130 b/internal/dbtest/testdata/snapshots/TestQuery-mssql2019-130
index 32f7d0914..1f315dcb6 100644
--- a/internal/dbtest/testdata/snapshots/TestQuery-mssql2019-130
+++ b/internal/dbtest/testdata/snapshots/TestQuery-mssql2019-130
@@ -1 +1 @@
-UPDATE "models" SET "updated_at" = NOW() WHERE ("id" = 0)
+UPDATE "models" SET "updated_at" = NOW() OUTPUT * WHERE ("id" = 0)
diff --git a/query_delete.go b/query_delete.go
index 1572ac328..daceadfb3 100644
--- a/query_delete.go
+++ b/query_delete.go
@@ -127,13 +127,6 @@ func (q *DeleteQuery) Returning(query string, args ...interface{}) *DeleteQuery
 	return q
 }
 
-func (q *DeleteQuery) hasReturning() bool {
-	if !q.db.features.Has(feature.Returning) {
-		return false
-	}
-	return q.returningQuery.hasReturning()
-}
-
 //------------------------------------------------------------------------------
 
 func (q *DeleteQuery) Operation() string {
@@ -189,6 +182,14 @@ func (q *DeleteQuery) AppendQuery(fmter schema.Formatter, b []byte) (_ []byte, e
 		}
 	}
 
+	if q.hasFeature(feature.Output) && q.hasReturning() {
+		b = append(b, " OUTPUT "...)
+		b, err = q.appendOutput(fmter, b)
+		if err != nil {
+			return nil, err
+		}
+	}
+
 	b, err = q.mustAppendWhere(fmter, b, withAlias)
 	if err != nil {
 		return nil, err
@@ -247,7 +248,8 @@ func (q *DeleteQuery) Exec(ctx context.Context, dest ...interface{}) (sql.Result
 
 	var res sql.Result
 
-	if hasDest := len(dest) > 0; hasDest || q.hasReturning() {
+	if hasDest := len(dest) > 0; hasDest ||
+		(q.hasReturning() && q.hasFeature(feature.Returning|feature.Output)) {
 		model, err := q.getModel(dest)
 		if err != nil {
 			return nil, err
diff --git a/query_update.go b/query_update.go
index 1662245a8..3b578efe1 100644
--- a/query_update.go
+++ b/query_update.go
@@ -174,13 +174,6 @@ func (q *UpdateQuery) Returning(query string, args ...interface{}) *UpdateQuery
 	return q
 }
 
-func (q *UpdateQuery) hasReturning() bool {
-	if !q.db.features.Has(feature.Returning) {
-		return false
-	}
-	return q.returningQuery.hasReturning()
-}
-
 //------------------------------------------------------------------------------
 
 func (q *UpdateQuery) Operation() string {
@@ -229,6 +222,14 @@ func (q *UpdateQuery) AppendQuery(fmter schema.Formatter, b []byte) (_ []byte, e
 		}
 	}
 
+	if q.hasFeature(feature.Output) && q.hasReturning() {
+		b = append(b, " OUTPUT "...)
+		b, err = q.appendOutput(fmter, b)
+		if err != nil {
+			return nil, err
+		}
+	}
+
 	b, err = q.mustAppendWhere(fmter, b, q.hasTableAlias(fmter))
 	if err != nil {
 		return nil, err
@@ -450,7 +451,8 @@ func (q *UpdateQuery) Exec(ctx context.Context, dest ...interface{}) (sql.Result
 
 	var res sql.Result
 
-	if hasDest := len(dest) > 0; hasDest || q.hasReturning() {
+	if hasDest := len(dest) > 0; hasDest ||
+		(q.hasReturning() && q.hasFeature(feature.Returning|feature.Output)) {
 		model, err := q.getModel(dest)
 		if err != nil {
 			return nil, err