util/db.go

package util

import (
	"database/sql"
	"errors"
	"fmt"
	"sync"
	"time"

	// PostgreSQL
	_ "github.com/lib/pq"
	// mysql
	_ "github.com/go-sql-driver/mysql"
	"github.com/jmoiron/sqlx"
)

const (
	// DBErrCodeOK ok
	DBErrCodeOK = 0
	// DBErrCodeException exception
	DBErrCodeException = 1
	// DBErrCodeExists exists
	DBErrCodeExists = 2
	// DBErrCodeNotFound not found
	DBErrCodeNotFound = 3
	// DBErrCodeAuthorized authorized
	DBErrCodeAuthorized = 4
	// DBErrCodeNotConnect connect error
	DBErrCodeNotConnect = 5
)

// DbConfig config
type DbConfig struct {
	Driver       string
	DNS          string
	MaxOpenConns int
	MaxIdle      int
	MaxLifetime  time.Duration
}

// DbReply db exec return insert/update/delete
type DbReply struct {
	OK           bool
	Err          error
	LastErr      error
	ErrCode      int
	LastID       int64
	RowsAffected int64
}

// DbReplyToReply db reply to respinse
func DbReplyToReply(reply DbReply) *ReplyData {
	status := ErrOk
	if !reply.OK {
		switch reply.ErrCode {
		case DBErrCodeException:
			status = ErrException
		case DBErrCodeExists:
			status = ErrDataExists
		case DBErrCodeNotFound:
			status = ErrDataNotFound
		case DBErrCodeAuthorized:
			status = ErrUnAuthorized
		case DBErrCodeNotConnect:
			status = ErrNotFound
		}
		return ErrReplyData(status, reply.LastErr.Error())
	}

	return OkReplyData()
}

var (
	config DbConfig
	db     *sqlx.DB
	err    error
	once   sync.Once
)

// DB define
type DB struct {
	Driver string
	conn   *sqlx.DB
	tx     *sqlx.Tx
}

// DNS
// postgres 'host=127.0.0.1 port=5432 user=postgres password=888888 dbname=park sslmode=disable'
// mysql user:password@tcp(host:port)/dbname?charset=utf8mb4

// SetDbConfig set
func SetDbConfig(cfg DbConfig) {
	config.Driver = cfg.Driver
	config.DNS = cfg.DNS
	config.MaxOpenConns = cfg.MaxOpenConns
	config.MaxIdle = cfg.MaxIdle
	config.MaxLifetime = cfg.MaxLifetime * time.Second
}

// ErrSQLNoRows check norows error
func ErrSQLNoRows(err error) bool {
	if err == sql.ErrNoRows {
		return true
	}
	return false
}

// DbReplyOk exec ok
func DbReplyOk(rowsAffected, lastID int64) DbReply {
	var reply DbReply

	reply.OK = true
	reply.ErrCode = 0
	reply.LastID = lastID
	reply.RowsAffected = rowsAffected

	return reply
}

// DbReplyFaild exec faild
func DbReplyFaild(errCode int, err, errText error) DbReply {
	var reply DbReply

	reply.OK = false
	reply.ErrCode = errCode
	reply.LastID = -1
	reply.RowsAffected = -1

	reply.Err = err
	reply.LastErr = errText

	return reply
}

// NewDB new DB object
func NewDB() *DB {
	return &DB{Driver: config.Driver}
}

// ReleaseDB free db connect
func ReleaseDB() {
	if db != nil {
		db.Close()
	}
}

// NewConfigDB new DB dynamic object
func NewConfigDB(config DbConfig) (dbx *DB, err error) {
	dbx = &DB{}

	dbx.Driver = config.Driver
	dbx.conn, err = sqlx.Connect(config.Driver, config.DNS)
	if err == nil {
		dbx.conn.SetMaxOpenConns(config.MaxOpenConns)
		dbx.conn.SetMaxIdleConns(config.MaxIdle)
		dbx.conn.SetConnMaxLifetime(config.MaxLifetime)
		dbx.conn.Ping()
	}

	return
}

// ReleaseConfigDB free db connect
func ReleaseConfigDB(dbx *DB) {
	if dbx.conn != nil {
		dbx.conn.Close()
	}
}

func connect() (dbx *sqlx.DB, err error) {
	once.Do(func() {
		db, err = sqlx.Connect(config.Driver, config.DNS)
		if err == nil {
			db.DB.SetMaxOpenConns(config.MaxOpenConns)
			db.DB.SetMaxIdleConns(config.MaxIdle)
			db.DB.SetConnMaxLifetime(config.MaxLifetime)
			db.Ping()
		}
	})
	dbx = db
	return
}

// Connect connect to database
func (d *DB) Connect() (err error) {
	if d.conn != nil {
		return
	}

	d.conn, err = connect()

	return
}

// Close close database connect
func (d *DB) Close() {
	//d.conn.Close()
}

// BeginTrans begin trans
func (d *DB) BeginTrans() (err error) {
	d.conn, err = connect()
	if err != nil {
		return
	}
	d.tx = d.conn.MustBegin()
	return
}

// Commit commit
func (d *DB) Commit() error {
	return d.tx.Commit()
}

// Rollback rollback
func (d *DB) Rollback() error {
	return d.tx.Rollback()
}

// TransExec trans execute
func (d *DB) TransExec(query string, args interface{}) (LastInsertId, RowsAffected int64, err error) {
	if rs, err := d.tx.NamedExec(query, args); err == nil {
		RowsAffected, _ = rs.RowsAffected()
		LastInsertId, _ = rs.LastInsertId()
	}
	return
}

// TransUpdate trans update
func (d *DB) TransUpdate(query string, args interface{}) (reply DbReply) {
	var (
		err error
		rs  sql.Result
	)

	if rs, err = d.tx.NamedExec(query, args); err == nil {
		a, _ := rs.RowsAffected()
		reply = DbReplyOk(a, 0)
	} else {
		reply = DbReplyFaild(DBErrCodeException, err, errors.New(`数据执行错误`))
	}
	return
}

// Rows get rows
func (d *DB) Rows(dest interface{}, query string, args interface{}) error {
	err := d.Connect()
	if err != nil {
		return err
	}
	defer d.Close()

	nstmt, err := d.conn.PrepareNamed(query)
	if err != nil {
		return err
	}
	defer nstmt.Close()

	err = nstmt.Select(dest, args)

	return err
}

// Row get row
func (d *DB) Row(dest interface{}, query string, args interface{}) error {
	err := d.Connect()
	if err != nil {
		return err
	}
	defer d.Close()

	nstmt, err := d.conn.PrepareNamed(query)
	if err != nil {
		return err
	}
	defer nstmt.Close()

	err = nstmt.Get(dest, args)

	return err
}

// InsertReply insert and return DbReply
func (d *DB) InsertReply(query string, args interface{}) (reply DbReply) {
	var (
		err error
		rs  sql.Result
	)
	err = d.Connect()
	if err != nil {
		reply = DbReplyFaild(DBErrCodeNotConnect, err, errors.New(`数据库连接错误`))
		return
	}
	defer d.Close()

	if rs, err = d.conn.NamedExec(query, args); err == nil {
		a, _ := rs.RowsAffected()
		n, _ := rs.LastInsertId()
		reply = DbReplyOk(a, n)
	} else {
		reply = DbReplyFaild(DBErrCodeException, err, errors.New(`数据执行错误`))
	}
	return
}

// UpdateReply update/delete and return DbReply
func (d *DB) UpdateReply(query string, args interface{}) (reply DbReply) {
	var (
		err error
		rs  sql.Result
	)
	err = d.Connect()
	if err != nil {
		reply = DbReplyFaild(DBErrCodeNotConnect, err, errors.New(`数据库连接错误`))
		return
	}
	defer d.Close()

	if rs, err = d.conn.NamedExec(query, args); err == nil {
		a, _ := rs.RowsAffected()
		reply = DbReplyOk(a, 0)
	} else {
		reply = DbReplyFaild(DBErrCodeException, err, errors.New(`数据执行错误`))
	}
	return
}

// Insert insert into
func (d *DB) Insert(query string, args interface{}) (LastInsertId, RowsAffected int64, err error) {
	err = d.Connect()
	if err != nil {
		return
	}
	defer d.Close()

	if rs, err := d.conn.NamedExec(query, args); err == nil {
		LastInsertId, _ = rs.LastInsertId()
		RowsAffected, _ = rs.RowsAffected()
	}
	return
}

// Update update/delete
func (d *DB) Update(query string, args interface{}) (RowsAffected int64, err error) {
	err = d.Connect()
	if err != nil {
		return
	}
	defer d.Close()

	if rs, err := d.conn.NamedExec(query, args); err == nil {
		RowsAffected, _ = rs.RowsAffected()
	}
	return
}

// Limit MySQL / PostgreSQL limit offset
func (d *DB) Limit(page, pagesize int) string {
	// PostgreSQL limit size offset 0
	// MySQL limit 0, size
	if d.Driver == `mysql` {
		return fmt.Sprintf(" limit %d, %d", (page-1)*pagesize, pagesize)
	}
	return fmt.Sprintf(" limit %d offset %d", pagesize, (page-1)*pagesize)
}