feat: mssql and pg merge query (#723)

* chore: support MergeQuery for MSSQL and PostgreSQL

Co-authored-by: gfk <gfk@bb.io>

db.go

@@ -82,6 +82,10 @@ func (db *DB) NewValues(model interface{}) *ValuesQuery {
 	return NewValuesQuery(db, model)
 }
 
+func (db *DB) NewMerge() *MergeQuery {
+	return NewMergeQuery(db)
+}
+
 func (db *DB) NewSelect() *SelectQuery {
 	return NewSelectQuery(db)
 }
@@ -330,6 +334,10 @@ func (c Conn) NewValues(model interface{}) *ValuesQuery {
 	return NewValuesQuery(c.db, model).Conn(c)
 }
 
+func (c Conn) NewMerge() *MergeQuery {
+	return NewMergeQuery(c.db).Conn(c)
+}
+
 func (c Conn) NewSelect() *SelectQuery {
 	return NewSelectQuery(c.db).Conn(c)
 }
@@ -640,6 +648,10 @@ func (tx Tx) NewValues(model interface{}) *ValuesQuery {
 	return NewValuesQuery(tx.db, model).Conn(tx)
 }
 
+func (tx Tx) NewMerge() *MergeQuery {
+	return NewMergeQuery(tx.db).Conn(tx)
+}
+
 func (tx Tx) NewSelect() *SelectQuery {
 	return NewSelectQuery(tx.db).Conn(tx)
 }

internal/dbtest/mssql_test.go

@@ -0,0 +1,53 @@
+package dbtest_test
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/require"
+)
+
+func TestMssqlMerge(t *testing.T) {
+	db := mssql2019(t)
+	defer db.Close()
+
+	type Model struct {
+		ID int64 `bun:",pk,autoincrement"`
+
+		Name  string
+		Value string
+	}
+
+	err := db.ResetModel(ctx, (*Model)(nil))
+	require.NoError(t, err)
+
+	_, err = db.NewInsert().Model(&Model{Name: "A", Value: "hello"}).Exec(ctx)
+	require.NoError(t, err)
+
+	newModels := []*Model{
+		{
+			Name:  "A",
+			Value: "world",
+		},
+		{
+			Name:  "B",
+			Value: "test",
+		},
+	}
+
+	changes := []string{}
+	_, err = db.NewMerge().
+		Model(&Model{}).
+		With("_data", db.NewValues(&newModels)).
+		Using("_data").
+		On("?TableAlias.name = _data.name").
+		When("MATCHED THEN UPDATE SET ?TableAlias.value = _data.value").
+		When("NOT MATCHED THEN INSERT (name, value) VALUES (_data.name, _data.value)").
+		Returning("$action").
+		Exec(ctx, &changes)
+	require.NoError(t, err)
+
+	require.Len(t, changes, 2)
+	require.Equal(t, "UPDATE", changes[0])
+	require.Equal(t, "INSERT", changes[1])
+
+}

internal/dbtest/query_test.go

@@ -923,6 +923,52 @@ func TestQuery(t *testing.T) {
 			}
 			return db.NewSelect().Model(new(Model)).Relation("SoftDelete")
 		},
+		func(db *bun.DB) schema.QueryAppender {
+			type Model struct {
+				ID    int64 `bun:",pk,autoincrement"`
+				Name  string
+				Value string
+			}
+
+			newModels := []*Model{
+				{Name: "A", Value: "world"},
+				{Name: "B", Value: "test"},
+			}
+
+			return db.NewMerge().
+				Model(new(Model)).
+				With("_data", db.NewValues(&newModels)).
+				Using("_data").
+				On("?TableAlias.name = _data.name").
+				WhenUpdate("MATCHED", func(q *bun.UpdateQuery) *bun.UpdateQuery {
+					return q.Set("value = _data.value")
+				}).
+				WhenInsert("NOT MATCHED", func(q *bun.InsertQuery) *bun.InsertQuery {
+					return q.Value("name", "_data.name").Value("value", "_data.value")
+				}).
+				Returning("$action")
+		},
+		func(db *bun.DB) schema.QueryAppender {
+			type Model struct {
+				ID    int64 `bun:",pk,autoincrement"`
+				Name  string
+				Value string
+			}
+
+			newModels := []*Model{
+				{Name: "A", Value: "world"},
+				{Name: "B", Value: "test"},
+			}
+
+			return db.NewMerge().
+				Model(new(Model)).
+				With("_data", db.NewValues(&newModels)).
+				Using("_data").
+				On("?TableAlias.name = _data.name").
+				WhenDelete("MATCHED").
+				When("NOT MATCHED THEN INSERT (name, value) VALUES (_data.name, _data.value)").
+				Returning("$action")
+		},
 	}
 
 	timeRE := regexp.MustCompile(`'2\d{3}-\d{2}-\d{2} \d{2}:\d{2}:\d{2}(\.\d+)?(\+\d{2}:\d{2})?'`)

internal/dbtest/testdata/snapshots/TestQuery-mariadb-154

@@ -0,0 +1 @@
+bun: merge not supported for current dialect

internal/dbtest/testdata/snapshots/TestQuery-mariadb-155

@@ -0,0 +1 @@
+bun: merge not supported for current dialect

internal/dbtest/testdata/snapshots/TestQuery-mssql2019-154

@@ -0,0 +1 @@
+WITH "_data" AS (SELECT * FROM (VALUES (NULL, 'A', 'world'), (NULL, 'B', 'test')) AS t ("id", "name", "value")) MERGE "models" AS "model" USING _data ON "model".name = _data.name WHEN MATCHED THEN UPDATE SET value = _data.value WHEN NOT MATCHED THEN INSERT ("name", "value") VALUES (_data.name, _data.value) OUTPUT $action;

internal/dbtest/testdata/snapshots/TestQuery-mssql2019-155

@@ -0,0 +1 @@
+WITH "_data" AS (SELECT * FROM (VALUES (NULL, 'A', 'world'), (NULL, 'B', 'test')) AS t ("id", "name", "value")) MERGE "models" AS "model" USING _data ON "model".name = _data.name WHEN MATCHED THEN DELETE WHEN NOT MATCHED THEN INSERT (name, value) VALUES (_data.name, _data.value) OUTPUT $action;

internal/dbtest/testdata/snapshots/TestQuery-mysql5-154

@@ -0,0 +1 @@
+bun: merge not supported for current dialect

internal/dbtest/testdata/snapshots/TestQuery-mysql5-155

@@ -0,0 +1 @@
+bun: merge not supported for current dialect

internal/dbtest/testdata/snapshots/TestQuery-mysql8-154

@@ -0,0 +1 @@
+bun: merge not supported for current dialect

internal/dbtest/testdata/snapshots/TestQuery-mysql8-155

@@ -0,0 +1 @@
+bun: merge not supported for current dialect

internal/dbtest/testdata/snapshots/TestQuery-pg-154

@@ -0,0 +1 @@
+WITH "_data" ("id", "name", "value") AS (VALUES (NULL::BIGINT, 'A'::VARCHAR, 'world'::VARCHAR), (NULL::BIGINT, 'B'::VARCHAR, 'test'::VARCHAR)) MERGE INTO "models" AS "model" USING _data ON "model".name = _data.name WHEN MATCHED THEN UPDATE SET value = _data.value WHEN NOT MATCHED THEN INSERT ("id", "name", "value") VALUES (DEFAULT, _data.name, _data.value);

internal/dbtest/testdata/snapshots/TestQuery-pg-155

@@ -0,0 +1 @@
+WITH "_data" ("id", "name", "value") AS (VALUES (NULL::BIGINT, 'A'::VARCHAR, 'world'::VARCHAR), (NULL::BIGINT, 'B'::VARCHAR, 'test'::VARCHAR)) MERGE INTO "models" AS "model" USING _data ON "model".name = _data.name WHEN MATCHED THEN DELETE WHEN NOT MATCHED THEN INSERT (name, value) VALUES (_data.name, _data.value);

internal/dbtest/testdata/snapshots/TestQuery-pgx-154

@@ -0,0 +1 @@
+WITH "_data" ("id", "name", "value") AS (VALUES (NULL::BIGINT, 'A'::VARCHAR, 'world'::VARCHAR), (NULL::BIGINT, 'B'::VARCHAR, 'test'::VARCHAR)) MERGE INTO "models" AS "model" USING _data ON "model".name = _data.name WHEN MATCHED THEN UPDATE SET value = _data.value WHEN NOT MATCHED THEN INSERT ("id", "name", "value") VALUES (DEFAULT, _data.name, _data.value);

internal/dbtest/testdata/snapshots/TestQuery-pgx-155

@@ -0,0 +1 @@
+WITH "_data" ("id", "name", "value") AS (VALUES (NULL::BIGINT, 'A'::VARCHAR, 'world'::VARCHAR), (NULL::BIGINT, 'B'::VARCHAR, 'test'::VARCHAR)) MERGE INTO "models" AS "model" USING _data ON "model".name = _data.name WHEN MATCHED THEN DELETE WHEN NOT MATCHED THEN INSERT (name, value) VALUES (_data.name, _data.value);

internal/dbtest/testdata/snapshots/TestQuery-sqlite-154

@@ -0,0 +1 @@
+bun: merge not supported for current dialect

internal/dbtest/testdata/snapshots/TestQuery-sqlite-155

@@ -0,0 +1 @@
+bun: merge not supported for current dialect

query_insert.go

@@ -192,7 +192,7 @@ func (q *InsertQuery) AppendQuery(fmter schema.Formatter, b []byte) (_ []byte, e
 		return nil, err
 	}
 
-	b, err = q.appendColumnsValues(fmter, b)
+	b, err = q.appendColumnsValues(fmter, b, false)
 	if err != nil {
 		return nil, err
 	}
@@ -214,7 +214,7 @@ func (q *InsertQuery) AppendQuery(fmter schema.Formatter, b []byte) (_ []byte, e
 }
 
 func (q *InsertQuery) appendColumnsValues(
-	fmter schema.Formatter, b []byte,
+	fmter schema.Formatter, b []byte, skipOutput bool,
 ) (_ []byte, err error) {
 	if q.hasMultiTables() {
 		if q.columns != nil {
@@ -275,7 +275,7 @@ func (q *InsertQuery) appendColumnsValues(
 	b = q.appendFields(fmter, b, fields)
 	b = append(b, ")"...)
 
-	if q.hasFeature(feature.Output) && q.hasReturning() {
+	if q.hasFeature(feature.Output) && q.hasReturning() && !skipOutput {
 		b = append(b, " OUTPUT "...)
 		b, err = q.appendOutput(fmter, b)
 		if err != nil {