diff --git a/db.go b/db.go index 9b8e391dc..106dfe905 100644 --- a/db.go +++ b/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) } diff --git a/internal/dbtest/mssql_test.go b/internal/dbtest/mssql_test.go new file mode 100644 index 000000000..579cee42b --- /dev/null +++ b/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]) + +} diff --git a/internal/dbtest/query_test.go b/internal/dbtest/query_test.go index c0d785ccb..0c8448568 100644 --- a/internal/dbtest/query_test.go +++ b/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})?'`) diff --git a/internal/dbtest/testdata/snapshots/TestQuery-mariadb-154 b/internal/dbtest/testdata/snapshots/TestQuery-mariadb-154 new file mode 100644 index 000000000..1f669bf4f --- /dev/null +++ b/internal/dbtest/testdata/snapshots/TestQuery-mariadb-154 @@ -0,0 +1 @@ +bun: merge not supported for current dialect diff --git a/internal/dbtest/testdata/snapshots/TestQuery-mariadb-155 b/internal/dbtest/testdata/snapshots/TestQuery-mariadb-155 new file mode 100644 index 000000000..1f669bf4f --- /dev/null +++ b/internal/dbtest/testdata/snapshots/TestQuery-mariadb-155 @@ -0,0 +1 @@ +bun: merge not supported for current dialect diff --git a/internal/dbtest/testdata/snapshots/TestQuery-mssql2019-154 b/internal/dbtest/testdata/snapshots/TestQuery-mssql2019-154 new file mode 100644 index 000000000..94fbda0ad --- /dev/null +++ b/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; diff --git a/internal/dbtest/testdata/snapshots/TestQuery-mssql2019-155 b/internal/dbtest/testdata/snapshots/TestQuery-mssql2019-155 new file mode 100644 index 000000000..8b9a1d2aa --- /dev/null +++ b/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; diff --git a/internal/dbtest/testdata/snapshots/TestQuery-mysql5-154 b/internal/dbtest/testdata/snapshots/TestQuery-mysql5-154 new file mode 100644 index 000000000..1f669bf4f --- /dev/null +++ b/internal/dbtest/testdata/snapshots/TestQuery-mysql5-154 @@ -0,0 +1 @@ +bun: merge not supported for current dialect diff --git a/internal/dbtest/testdata/snapshots/TestQuery-mysql5-155 b/internal/dbtest/testdata/snapshots/TestQuery-mysql5-155 new file mode 100644 index 000000000..1f669bf4f --- /dev/null +++ b/internal/dbtest/testdata/snapshots/TestQuery-mysql5-155 @@ -0,0 +1 @@ +bun: merge not supported for current dialect diff --git a/internal/dbtest/testdata/snapshots/TestQuery-mysql8-154 b/internal/dbtest/testdata/snapshots/TestQuery-mysql8-154 new file mode 100644 index 000000000..1f669bf4f --- /dev/null +++ b/internal/dbtest/testdata/snapshots/TestQuery-mysql8-154 @@ -0,0 +1 @@ +bun: merge not supported for current dialect diff --git a/internal/dbtest/testdata/snapshots/TestQuery-mysql8-155 b/internal/dbtest/testdata/snapshots/TestQuery-mysql8-155 new file mode 100644 index 000000000..1f669bf4f --- /dev/null +++ b/internal/dbtest/testdata/snapshots/TestQuery-mysql8-155 @@ -0,0 +1 @@ +bun: merge not supported for current dialect diff --git a/internal/dbtest/testdata/snapshots/TestQuery-pg-154 b/internal/dbtest/testdata/snapshots/TestQuery-pg-154 new file mode 100644 index 000000000..9d8581d01 --- /dev/null +++ b/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); diff --git a/internal/dbtest/testdata/snapshots/TestQuery-pg-155 b/internal/dbtest/testdata/snapshots/TestQuery-pg-155 new file mode 100644 index 000000000..d8318bb63 --- /dev/null +++ b/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); diff --git a/internal/dbtest/testdata/snapshots/TestQuery-pgx-154 b/internal/dbtest/testdata/snapshots/TestQuery-pgx-154 new file mode 100644 index 000000000..9d8581d01 --- /dev/null +++ b/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); diff --git a/internal/dbtest/testdata/snapshots/TestQuery-pgx-155 b/internal/dbtest/testdata/snapshots/TestQuery-pgx-155 new file mode 100644 index 000000000..d8318bb63 --- /dev/null +++ b/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); diff --git a/internal/dbtest/testdata/snapshots/TestQuery-sqlite-154 b/internal/dbtest/testdata/snapshots/TestQuery-sqlite-154 new file mode 100644 index 000000000..1f669bf4f --- /dev/null +++ b/internal/dbtest/testdata/snapshots/TestQuery-sqlite-154 @@ -0,0 +1 @@ +bun: merge not supported for current dialect diff --git a/internal/dbtest/testdata/snapshots/TestQuery-sqlite-155 b/internal/dbtest/testdata/snapshots/TestQuery-sqlite-155 new file mode 100644 index 000000000..1f669bf4f --- /dev/null +++ b/internal/dbtest/testdata/snapshots/TestQuery-sqlite-155 @@ -0,0 +1 @@ +bun: merge not supported for current dialect diff --git a/query_insert.go b/query_insert.go index 5eac05579..32c97070f 100644 --- a/query_insert.go +++ b/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 { diff --git a/query_merge.go b/query_merge.go new file mode 100644 index 000000000..08bfec7fc --- /dev/null +++ b/query_merge.go @@ -0,0 +1,296 @@ +package bun + +import ( + "context" + "database/sql" + "errors" + + "github.com/uptrace/bun/dialect" + "github.com/uptrace/bun/dialect/feature" + "github.com/uptrace/bun/internal" + "github.com/uptrace/bun/schema" +) + +type whenInsert struct { + expr string + query *InsertQuery +} + +func (w whenInsert) AppendQuery(fmter schema.Formatter, b []byte) (_ []byte, err error) { + b = append(b, w.expr...) + if w.query != nil { + b = append(b, " THEN INSERT"...) + b, err = w.query.appendColumnsValues(fmter, b, true) + if err != nil { + return nil, err + } + } + return b, nil +} + +type whenUpdate struct { + expr string + query *UpdateQuery +} + +func (w whenUpdate) AppendQuery(fmter schema.Formatter, b []byte) (_ []byte, err error) { + b = append(b, w.expr...) + if w.query != nil { + b = append(b, " THEN UPDATE SET "...) + b, err = w.query.appendSet(fmter, b) + if err != nil { + return nil, err + } + } + return b, nil +} + +type whenDelete struct { + expr string +} + +func (w whenDelete) AppendQuery(fmter schema.Formatter, b []byte) (_ []byte, err error) { + b = append(b, w.expr...) + b = append(b, " THEN DELETE"...) + return b, nil +} + +type MergeQuery struct { + baseQuery + returningQuery + + using schema.QueryWithArgs + on schema.QueryWithArgs + when []schema.QueryAppender +} + +var _ Query = (*MergeQuery)(nil) + +func NewMergeQuery(db *DB) *MergeQuery { + q := &MergeQuery{ + baseQuery: baseQuery{ + db: db, + conn: db.DB, + }, + } + if !(q.db.dialect.Name() == dialect.MSSQL || q.db.dialect.Name() == dialect.PG) { + q.err = errors.New("bun: merge not supported for current dialect") + } + return q +} + +func (q *MergeQuery) Conn(db IConn) *MergeQuery { + q.setConn(db) + return q +} + +func (q *MergeQuery) Model(model interface{}) *MergeQuery { + q.setModel(model) + return q +} + +// Apply calls the fn passing the MergeQuery as an argument. +func (q *MergeQuery) Apply(fn func(*MergeQuery) *MergeQuery) *MergeQuery { + if fn != nil { + return fn(q) + } + return q +} + +func (q *MergeQuery) With(name string, query schema.QueryAppender) *MergeQuery { + q.addWith(name, query, false) + return q +} + +func (q *MergeQuery) WithRecursive(name string, query schema.QueryAppender) *MergeQuery { + q.addWith(name, query, true) + return q +} + +//------------------------------------------------------------------------------ + +func (q *MergeQuery) Table(tables ...string) *MergeQuery { + for _, table := range tables { + q.addTable(schema.UnsafeIdent(table)) + } + return q +} + +func (q *MergeQuery) TableExpr(query string, args ...interface{}) *MergeQuery { + q.addTable(schema.SafeQuery(query, args)) + return q +} + +func (q *MergeQuery) ModelTableExpr(query string, args ...interface{}) *MergeQuery { + q.modelTableName = schema.SafeQuery(query, args) + return q +} + +//------------------------------------------------------------------------------ + +// Returning adds a RETURNING clause to the query. +// +// To suppress the auto-generated RETURNING clause, use `Returning("NULL")`. +// Only for mssql output, postgres not supported returning in merge query +func (q *MergeQuery) Returning(query string, args ...interface{}) *MergeQuery { + q.addReturning(schema.SafeQuery(query, args)) + return q +} + +//------------------------------------------------------------------------------ + +func (q *MergeQuery) Operation() string { + return "MERGE" +} + +func (q *MergeQuery) AppendQuery(fmter schema.Formatter, b []byte) (_ []byte, err error) { + if q.err != nil { + return nil, q.err + } + + fmter = formatterWithModel(fmter, q) + + b, err = q.appendWith(fmter, b) + if err != nil { + return nil, err + } + + b = append(b, "MERGE "...) + if q.db.dialect.Name() == dialect.PG { + b = append(b, "INTO "...) + } + + b, err = q.appendFirstTableWithAlias(fmter, b) + if err != nil { + return nil, err + } + + b = append(b, " USING "...) + b, err = q.using.AppendQuery(fmter, b) + if err != nil { + return nil, err + } + + b = append(b, " ON "...) + b, err = q.on.AppendQuery(fmter, b) + if err != nil { + return nil, err + } + + for _, w := range q.when { + b = append(b, " WHEN "...) + b, err = w.AppendQuery(fmter, b) + if err != nil { + return nil, err + } + } + + if q.hasFeature(feature.Output) && q.hasReturning() { + b = append(b, " OUTPUT "...) + b, err = q.appendOutput(fmter, b) + if err != nil { + return nil, err + } + } + + //A MERGE statement must be terminated by a semi-colon (;). + b = append(b, ";"...) + + return b, nil +} + +//------------------------------------------------------------------------------ + +func (q *MergeQuery) Using(s string, args ...interface{}) *MergeQuery { + q.using = schema.SafeQuery(s, args) + return q +} + +func (q *MergeQuery) On(s string, args ...interface{}) *MergeQuery { + q.on = schema.SafeQuery(s, args) + return q +} + +// WhenUpdate for when update clause +func (q *MergeQuery) WhenUpdate(expr string, fn func(q *UpdateQuery) *UpdateQuery) *MergeQuery { + sq := NewUpdateQuery(q.db) + //apply the model as default into sub query + if q.model != nil { + sq = sq.Model(q.model) + } + sq = sq.Apply(fn) + q.when = append(q.when, whenUpdate{expr: expr, query: sq}) + return q +} + +// WhenInsert for when insert clause +func (q *MergeQuery) WhenInsert(expr string, fn func(q *InsertQuery) *InsertQuery) *MergeQuery { + sq := NewInsertQuery(q.db) + //apply the model as default into sub query, since appendColumnsValues required + if q.model != nil { + sq = sq.Model(q.model) + } + sq = sq.Apply(fn) + q.when = append(q.when, whenInsert{expr: expr, query: sq}) + return q +} + +// WhenDelete for when delete clause +func (q *MergeQuery) WhenDelete(expr string) *MergeQuery { + q.when = append(q.when, whenDelete{expr: expr}) + return q +} + +// When for raw expression clause +func (q *MergeQuery) When(expr string, args ...interface{}) *MergeQuery { + q.when = append(q.when, schema.SafeQuery(expr, args)) + return q +} + +//------------------------------------------------------------------------------ + +func (q *MergeQuery) Exec(ctx context.Context, dest ...interface{}) (sql.Result, error) { + if q.err != nil { + return nil, q.err + } + if err := q.beforeAppendModel(ctx, q); err != nil { + return nil, err + } + + queryBytes, err := q.AppendQuery(q.db.fmter, q.db.makeQueryBytes()) + if err != nil { + return nil, err + } + + query := internal.String(queryBytes) + var res sql.Result + + if hasDest := len(dest) > 0; hasDest || + (q.hasReturning() && q.hasFeature(feature.InsertReturning|feature.Output)) { + model, err := q.getModel(dest) + if err != nil { + return nil, err + } + + res, err = q.scan(ctx, q, query, model, hasDest) + if err != nil { + return nil, err + } + } else { + res, err = q.exec(ctx, q, query) + if err != nil { + return nil, err + } + } + + return res, nil +} + +func (q *MergeQuery) String() string { + buf, err := q.AppendQuery(q.db.Formatter(), nil) + if err != nil { + panic(err) + } + + return string(buf) +}