lxutils/lxDb/db.go

package lxDb

import (
	"context"
	"fmt"
	"git.listensoft.net/tool/lxutils/lxzap"
	"github.com/gin-gonic/gin"
	"go.uber.org/zap"
	"gorm.io/driver/mysql"
	"gorm.io/gorm"
	"gorm.io/gorm/logger"
	"gorm.io/gorm/schema"
	"os"
)

// DB2.0 数据库连接 gorm2.0
var DBS map[string]*gorm.DB //gs 数据库
type DbConfig struct {
	Name     string
	Host     string
	Port     string
	User     string
	Password string
	Database string
	Charset  string
}

func GetDB(c *gin.Context, dbName ...string) *gorm.DB {
	v, _ := c.Get("X-Span-ID")
	spanId := fmt.Sprintf("%v", v)
	ctx := context.WithValue(context.Background(), "X-Span-ID", spanId)
	d := "one"
	if len(dbName) != 0 {
		d = dbName[0]
	}
	return DBS[d].WithContext(ctx)
}

func GetDB2(dbName ...string) *gorm.DB {
	d := "one"
	if len(dbName) != 0 {
		d = dbName[0]
	}
	return DBS[d]
}

func InitDBS(env string, conf []DbConfig) {
	DBS = map[string]*gorm.DB{}
	if len(conf) == 1 {
		conf[0].Name = "one"
	}
	for _, config := range conf {
		InitDB(env, config)
	}
}

func InitDB(env string, conf DbConfig) {
	if conf.Name == "" {
		DBS = nil
		fmt.Println("db name 错误")
		return
	}
	if conf.Host == "" {
		DBS = nil
		fmt.Println("未配置Db连接")
		return
	}
	if conf.Charset == "" {
		conf.Charset = "utf8mb4"
	}
	dsn := fmt.Sprintf("%s:%s@tcp(%s:%s)/%s?charset=%s&parseTime=True&loc=Local",
		conf.User, conf.Password, conf.Host, conf.Port, conf.Database, conf.Charset)
	if env == "dev" {
		db, err := gorm.Open(mysql.Open(dsn), &gorm.Config{
			Logger: logger.Default.LogMode(logger.Info),
			NamingStrategy: schema.NamingStrategy{
				SingularTable: true, // 使用单数表名，启用该选项，此时，`User` 的表名应该是 `user`
			},
		})
		if err != nil {
			fmt.Println(err.Error())
			os.Exit(-1)
		}
		DBS[conf.Name] = db
	} else {
		logger2 := lxzap.NewGormZap(zap.L())
		db, err := gorm.Open(mysql.Open(dsn), &gorm.Config{
			//Logger: logger.Default.LogMode(logger.Info),
			Logger: logger2,
			NamingStrategy: schema.NamingStrategy{
				SingularTable: true, // 使用单数表名，启用该选项，此时，`User` 的表名应该是 `user`
			},
		})
		if err != nil {
			fmt.Println(err.Error())
			os.Exit(-1)
		}
		DBS[conf.Name] = db
	}
	//db.AutoMigrate(TaskData{}, Task{}, Version{})
}