chunhao_receive/utils/db.go

package utils

import (
	"database/sql"
	"fmt"

	"gorm.io/gorm"
)

type WtDB struct {
	*gorm.DB
}

func (db *WtDB) Table(name string, args ...interface{}) (tx *WtDB) {
	return &WtDB{db.DB.Table(name, args...)}
}

func (db *WtDB) Distinct(args ...interface{}) (tx *WtDB) {
	return &WtDB{db.DB.Distinct(args...)}
}

func (db *WtDB) Select(query interface{}, args ...interface{}) (tx *WtDB) {
	return &WtDB{db.DB.Select(query, args...)}
}

func (db *WtDB) WhereRaw(query interface{}, args ...interface{}) (tx *WtDB) {
	return &WtDB{db.DB.Where(query, args...)}
}

func (db *WtDB) Where(field string, args ...interface{}) (tx *WtDB) {
	n := len(args)
	if n == 0 {
		return &WtDB{db.DB.Where(field)}
	}

	if n == 1 {
		return &WtDB{db.DB.Where(field+" = ?", args...)}
	}

	if n == 2 {
		var opList = []string{"=", ">", "<", ">=", "<=", "!=", "like"}
		if v, ok := args[0].(string); ok {
			if InArray(v, opList) {
				return &WtDB{db.DB.Where(fmt.Sprintf("%s %s ?", field, v), args[1:]...)}
			}
		}
	}

	return &WtDB{db.DB.Where(field, args...)}
}

func (db *WtDB) WhereIn(field string, args ...interface{}) (tx *WtDB) {
	return &WtDB{db.DB.Where(fmt.Sprintf("%s in ?", field), args...)}
}

func (db *WtDB) WhereNotIn(field string, args ...interface{}) (tx *WtDB) {
	return &WtDB{db.DB.Where(fmt.Sprintf("%s not in ?", field), args...)}
}

func (db *WtDB) Not(query interface{}, args ...interface{}) (tx *WtDB) {
	return &WtDB{db.DB.Not(query, args...)}
}

func (db *WtDB) Or(query interface{}, args ...interface{}) (tx *WtDB) {
	return &WtDB{db.DB.Or(query, args...)}
}

func (db *WtDB) Join(table string, query string, args ...interface{}) (tx *WtDB) {
	return &WtDB{db.DB.Joins(fmt.Sprintf("inner join %s on %s", table, query), args...)}
}

func (db *WtDB) LeftJoin(table string, query string, args ...interface{}) (tx *WtDB) {
	return &WtDB{db.DB.Joins(fmt.Sprintf("left join %s on %s", table, query), args...)}
}

func (db *WtDB) RightJoin(table string, query string, args ...interface{}) (tx *WtDB) {
	return &WtDB{db.DB.Joins(fmt.Sprintf("right join %s on %s", table, query), args...)}
}

func (db *WtDB) JoinRaw(query string, args ...interface{}) (tx *WtDB) {
	return &WtDB{db.DB.Joins(query, args...)}
}

func (db *WtDB) Group(name string) (tx *WtDB) {
	return &WtDB{db.DB.Group(name)}
}

func (db *WtDB) Having(query interface{}, args ...interface{}) (tx *WtDB) {
	return &WtDB{db.DB.Having(query, args...)}
}

func (db *WtDB) Order(value interface{}) (tx *WtDB) {
	return &WtDB{db.DB.Order(value)}
}

func (db *WtDB) Limit(limit int) (tx *WtDB) {
	return &WtDB{db.DB.Limit(limit)}
}

func (db *WtDB) Offset(offset int) (tx *WtDB) {
	return &WtDB{db.DB.Offset(offset)}
}

func (db *WtDB) Raw(sql string, values ...interface{}) (tx *WtDB) {
	return &WtDB{db.DB.Raw(sql, values...)}
}

func (db *WtDB) SubQuery() *gorm.DB {
	return db.DB
}

func (db *WtDB) Begin(opts ...*sql.TxOptions) (tx *WtDB) {
	return &WtDB{DB: db.DB.Begin(opts...)}
}

func (db *WtDB) Commit() (tx *WtDB) {
	return &WtDB{DB: db.DB.Commit()}
}

func (db *WtDB) Rollback() (tx *WtDB) {
	return &WtDB{DB: db.DB.Rollback()}
}

func (db *WtDB) Transaction(fc func(tx *WtDB) error) (err error) {
	return db.DB.Transaction(func(t *gorm.DB) error {
		return fc(&WtDB{DB: t})
	})
}