pm/pkg/db/db.go

package db

import (
	"database/sql"
	"errors"
	"fmt"
	"strconv"
	"time"
	"zhiyuan/pkg/config"
	"zhiyuan/pkg/logger"
	"zhiyuan/pkg/utils"

	"github.com/didi/gendry/builder"
	"github.com/didi/gendry/scanner"
	_ "github.com/go-sql-driver/mysql"
)

var db *sql.DB
var GeoDB *sql.DB

func Setup() {
	var err error
	db, err = sql.Open("mysql", fmt.Sprintf("%s:%s@tcp(%s)/%s?charset=%s&parseTime=True&loc=Local",
		config.Cfg.Database.User,
		config.Cfg.Database.Password,
		config.Cfg.Database.Host,
		config.Cfg.Database.Name,
		config.Cfg.Database.Charset))
	if err != nil {
		logger.Sugar.Error(err)
	}
	GeoDB, err = sql.Open("mysql", fmt.Sprintf("%s:%s@tcp(%s)/%s?charset=%s&parseTime=True&loc=Local",
		config.Cfg.Database.User,
		config.Cfg.Database.Password,
		config.Cfg.Database.Host,
		"geo",
		config.Cfg.Database.Charset))
	if err != nil {
		logger.Sugar.Error(err)
	}
	scanner.SetTagName("json")
	return
}

func Query(query string, params map[string]interface{}, retVal interface{}, dbs *sql.DB) error {
	if dbs == nil {
		dbs = db
	}
	cond, values, err := builder.NamedQuery(query, params)
	if err != nil {
		return err
	}
	logger.Sugar.Infof("query: %v; params: %v", cond, values)
	rows, err := dbs.Query(cond, values...)
	// defer rows.Close()
	if err != nil {
		logger.Sugar.Error(err)
		return err
	}
	return scanner.ScanClose(rows, retVal)
}

func resolveDataFromRows(rows *sql.Rows) ([]map[string]interface{}, error) {
	if nil == rows {
		return nil, errors.New("[scanner]: rows can't be nil")
	}
	columns, err := rows.Columns()
	if nil != err {
		return nil, err
	}
	length := len(columns)
	var result []map[string]interface{}
	//unnecessary to put below into rows.Next loop,reduce allocating
	values := make([]interface{}, length)
	for i := 0; i < length; i++ {
		values[i] = new(interface{})
	}
	for rows.Next() {
		err = rows.Scan(values...)
		if nil != err {
			return nil, err
		}
		mp := make(map[string]interface{})
		for idx, name := range columns {
			//mp[name] = reflect.ValueOf(values[idx]).Elem().Interface()
			mp[name] = *(values[idx].(*interface{}))
		}
		result = append(result, mp)
	}
	return result, nil
}

type CloseErr struct {
	err error
}

func (e CloseErr) Error() string {
	return e.err.Error()
}

func newCloseErr(err error) error {
	if err == nil {
		return nil
	}
	return CloseErr{err}
}

func QueryMap(query string, params map[string]interface{}, dbs *sql.DB) ([]map[string]interface{}, error) {
	if dbs == nil {
		dbs = db
	}
	cond, values, err := builder.NamedQuery(query, params)
	if err != nil {
		return nil, err
	}
	//logger.Sugar.Infof("query: %v; params: %v", cond, values)
	rows, err := dbs.Query(cond, values...)
	// defer rows.Close()
	if err != nil {
		logger.Sugar.Error(err)
		return nil, err
	}
	//result, err := scanner.ScanMapDecodeClose(rows)
	results, err := resolveDataFromRows(rows)
	if nil != err {
		return nil, err
	}
	for i := 0; i < len(results); i++ {
		for k, v := range results[i] {
			rv, ok := v.([]uint8)
			if !ok {
				continue
			}
			s := string(rv)
			results[i][k] = s
			if len(s) > 8 {
				continue
			}
			//convert to int
			intVal, err := strconv.Atoi(s)
			if err == nil {
				results[i][k] = intVal
				continue
			}
			//convert to float64
			floatVal, err := strconv.ParseFloat(s, 64)
			if err == nil {
				results[i][k] = floatVal
				continue
			}
		}
	}
	if nil != rows {
		errClose := rows.Close()
		if err == nil {
			err = newCloseErr(errClose)
		}
	}
	return results, err
}

func BuildInsert(table string, data []map[string]interface{}, dbs *sql.DB) (int64, error) {
	if dbs == nil {
		dbs = db
	}

	cond, values, err := builder.BuildInsert(table, data)
	if nil != err {
		logger.Sugar.Error(err)
		return 0, err
	}
	//logger.Sugar.Infof("query: %v; params: %v", cond, values)
	result, err := dbs.Exec(cond, values...)
	if nil != err || nil == result {
		logger.Sugar.Error(err)
		return 0, err
	}
	return result.LastInsertId()
}

func BuildUpdate(table string, where, data map[string]interface{}, dbs *sql.DB) (int64, error) {
	if dbs == nil {
		dbs = db
	}
	cond, values, err := builder.BuildUpdate(table, where, data)
	if nil != err {
		return 0, err
	}
	//logger.Sugar.Infof("query: %v; params: %v", cond, values)
	result, err := dbs.Exec(cond, values...)
	if nil != err || nil == result {
		logger.Sugar.Error(err)
		return 0, err
	}
	return result.RowsAffected()
}

func GetMultiMap(table string, where map[string]interface{}, fields []string) ([]map[string]interface{}, error) {
	rows, err := BuildSelectRows(table, where, fields)
	if err != nil {
		return nil, err
	}
	result, mapErr := scanner.ScanMapDecodeClose(rows)
	return result, mapErr
}

func GetMulti(table string, where map[string]interface{}, fields []string, retVal interface{}) error {
	rows, err := BuildSelectRows(table, where, fields)
	if err != nil {
		return err
	}
	return ScanClose(rows, retVal)
}

func GetOneMap(table string, where map[string]interface{}, fields []string) (map[string]interface{}, error) {
	wheres := make(map[string]interface{})
	for key, value := range where {
		wheres[key] = value
	}
	wheres["_limit"] = []uint{1}
	result, err := GetMultiMap(table, wheres, fields)
	if err != nil {
		return nil, err
	}
	if len(result) == 0 {
		return nil, nil
	}
	return result[0], nil
}

func GetOne(table string, where map[string]interface{}, fields []string, retVal interface{}) error {
	wheres := make(map[string]interface{})
	for key, value := range where {
		wheres[key] = value
	}
	wheres["_limit"] = []uint{1}
	return GetMulti(table, where, fields, retVal)
}

func ScanClose(rows *sql.Rows, target interface{}) error {
	err := scanner.ScanClose(rows, target)
	if err == scanner.ErrEmptyResult {
		return nil
	}
	return err
}

func Count(table string, where map[string]interface{}) (int64, error) {
	res, err := GetOneMap(table, where, []string{"count(*) as count"})
	if err != nil {
		return 0, err
	}
	return res["count"].(int64), err
}

func BuildSelectRows(table string, where map[string]interface{}, fields []string) (*sql.Rows, error) {
	cond, values, err := BuildSelect(table, where, fields)
	if err != nil {
		return nil, err
	}
	return GetRows(cond, values)
}

func BuildSelect(table string, where map[string]interface{}, fields []string) (string, []interface{}, error) {
	cond, values, err := builder.BuildSelect(table, where, fields)
	if err != nil {
		logger.Sugar.Error(err)
	}
	return cond, values, err
}

func GetRows(cond string, values []interface{}) (*sql.Rows, error) {
	logger.Sugar.Infof("query: %v; params: %v", cond, values)
	rows, err := db.Query(cond, values...)
	// defer rows.Close()
	if err != nil {
		logger.Sugar.Error(err)
		return nil, errors.New("未知错误")
	}
	return rows, nil
}

func CountRaw(query string, params map[string]interface{}) (int64, error) {
	res, err := GetOneMapRaw("SELECT count(*) as count FROM "+query, params)
	if err != nil {
		return 0, err
	}
	return utils.ToInt64(res["count"]), err
}

func GetMultiRaw(query string, where map[string]string, params map[string]interface{}, retVal interface{}) error {
	cond, values, err := builder.NamedQuery(JoinQuery(query, where), params)
	if err != nil {
		return err
	}
	rows, err := GetRows(cond, values)
	if err != nil {
		return err
	}
	return scanner.ScanClose(rows, retVal)
}

func JoinQuery(query string, where map[string]string) string {
	if where["where"] != "" {
		query = query + " WHERE " + where["where"]
	}
	if where["_group_by"] != "" {
		query = query + " GROUP BY " + where["_group_by"]
	}
	if where["_having"] != "" {
		query = query + " HAVING " + where["_having"]
	}
	if where["_order_by"] != "" {
		query = query + " ORDER BY " + where["_order_by"]
	}
	if where["_page_size"] != "" && where["_page_num"] != "" {
		pageNum := utils.ToInt(where["_page_num"])
		pageSize := utils.ToInt(where["_page_size"])
		query = query + " LIMIT " + utils.ToStr((pageNum-1)*pageSize) + "," + utils.ToStr(pageSize)
	}
	return query
}

func GetMultiMapRaw(query string, params map[string]interface{}) ([]map[string]interface{}, error) {
	cond, values, err := builder.NamedQuery(query, params)
	if err != nil {
		return nil, err
	}
	rows, err := GetRows(cond, values)
	if err != nil {
		return nil, err
	}
	result, mapErr := scanner.ScanMapDecodeClose(rows)
	return result, mapErr
}

func GetOneMapRaw(query string, params map[string]interface{}) (map[string]interface{}, error) {
	result, err := GetMultiMapRaw(query, params)
	if err != nil {
		return nil, err
	}
	if len(result) == 0 {
		return nil, errors.New("empty record1")
	}
	return result[0], err
}

func UpdateRaw(query string, params map[string]interface{}) (int64, error) {
	cond, values, err := builder.NamedQuery(query, params)
	if err != nil {
		return 0, err
	}
	return GetUpdateResult(cond, values)
}

func UpdateRawSimple(table string, querySet, queryWhere string, params map[string]interface{}) (int64, error) {
	querySet = querySet + ",updated_at=" + utils.ToStr(time.Now().Unix())
	query := fmt.Sprintf("UPDATE %s SET %s WHERE %s", table, querySet, queryWhere)
	return UpdateRaw(query, params)
}

func InsertRaw(query string, params map[string]interface{}) (int64, error) {
	cond, values, err := builder.NamedQuery(query, params)
	if err != nil {
		return 0, err
	}
	return GetInsertResult(cond, values)
}

func Update(table string, where, data map[string]interface{}) (int64, error) {
	data["updated_at"] = time.Now().Unix()
	cond, values, err := builder.BuildUpdate(table, where, data)
	if nil != err {
		return 0, err
	}
	return GetUpdateResult(cond, values)
}

func Delete(table string, where map[string]interface{}) (int64, error) {
	cond, values, err := builder.BuildDelete(table, where)
	if nil != err {
		return 0, err
	}
	return GetUpdateResult(cond, values)
}

func DeleteSoft(table string, where map[string]interface{}) (int64, error) {
	cond, values, err := builder.BuildUpdate(table, where, map[string]interface{}{"deleted_at": time.Now().Unix()})
	if nil != err {
		return 0, err
	}
	return GetUpdateResult(cond, values)
}

func Insert(table string, data []map[string]interface{}) (int64, error) {
	for k, v := range data {
		v["created_at"] = time.Now().Unix()
		v["updated_at"] = time.Now().Unix()
		data[k] = v
	}
	cond, values, err := builder.BuildInsert(table, data)
	if nil != err {
		logger.Sugar.Error(err)
		return 0, err
	}
	return GetInsertResult(cond, values)
}

func InsertOne(table string, data map[string]interface{}) (int64, error) {
	return Insert(table, []map[string]interface{}{data})
}

func GetUpdateResult(cond string, values []interface{}) (int64, error) {
	result, err := Exec(cond, values)
	if err != nil {
		return 0, err
	}
	return result.RowsAffected()
}

func GetInsertResult(cond string, values []interface{}) (int64, error) {
	result, err := Exec(cond, values)
	if err != nil {
		return 0, err
	}
	return result.LastInsertId()
}

func Exec(cond string, values []interface{}) (sql.Result, error) {
	//logger.Sugar.Infof("query: %v; params: %v", cond, values)
	result, err := db.Exec(cond, values...)
	if nil != err || nil == result {
		logger.Sugar.Error(err)
		return nil, errors.New("未知错误")
	}
	return result, err
}

func GetOffset(pageNum uint, pageSize uint) []uint {
	return []uint{(pageNum - 1) * pageSize, pageSize}
}