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>
uptrace · Nov 24, 2022 · 08876b4 · 08876b4
1 parent de858e2
commit 08876b4
Show file tree

Hide file tree

Showing 3 changed files with 21 additions and 17 deletions.
diff --git 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
@@ -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
@@ -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