gtsocial-umbx

query_merge.go (7316B)

---

 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 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
 }
 
 func (q *MergeQuery) Err(err error) *MergeQuery {
 	q.setErr(err)
 	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) 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
 }
 
 // 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
 }
 
 // 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
 }
 
 // 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) 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) Scan(ctx context.Context, dest ...interface{}) error {
 	_, err := q.scanOrExec(ctx, dest, true)
 	return err
 }
 
 func (q *MergeQuery) Exec(ctx context.Context, dest ...interface{}) (sql.Result, error) {
 	return q.scanOrExec(ctx, dest, len(dest) > 0)
 }
 
 func (q *MergeQuery) scanOrExec(
 	ctx context.Context, dest []interface{}, hasDest bool,
 ) (sql.Result, error) {
 	if q.err != nil {
 		return nil, q.err
 	}
 
 	// Run append model hooks before generating the query.
 	if err := q.beforeAppendModel(ctx, q); err != nil {
 		return nil, err
 	}
 
 	// Generate the query before checking hasReturning.
 	queryBytes, err := q.AppendQuery(q.db.fmter, q.db.makeQueryBytes())
 	if err != nil {
 		return nil, err
 	}
 
 	useScan := hasDest || (q.hasReturning() && q.hasFeature(feature.InsertReturning|feature.Output))
 	var model Model
 
 	if useScan {
 		var err error
 		model, err = q.getModel(dest)
 		if err != nil {
 			return nil, err
 		}
 	}
 
 	query := internal.String(queryBytes)
 	var res sql.Result
 
 	if useScan {
 		res, err = q.scan(ctx, q, query, model, true)
 		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)
 }
 
 //------------------------------------------------------------------------------
 
 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
 }