speechAnalysis/models/db.go

package models

import (
	"fmt"
	"gorm.io/gorm/schema"
	"speechAnalysis/conf"
	"speechAnalysis/extend/util"
	"speechAnalysis/pkg/logx"
	"speechAnalysis/pkg/mysqlx"
	"speechAnalysis/pkg/snowflake"

	"gorm.io/gorm"
)

type BaseModelInt struct {
	ID        uint          `gorm:"comment:主键ID;primaryKey;" json:"id"`
	CreatedAt util.JSONTime `gorm:"index;comment:创建时间" json:"createdAt,omitempty" swaggerignore:"true"`
	UpdatedAt util.JSONTime `gorm:"comment:更新时间" json:"updatedAt,omitempty" swaggerignore:"true"`
}

type BaseModelString struct {
	ID        string        `gorm:"comment:主键ID;primaryKey;type:varchar(191);" json:"id"`
	CreatedAt util.JSONTime `gorm:"index;comment:创建时间" json:"createdAt,omitempty" swaggerignore:"true"`
	UpdatedAt util.JSONTime `gorm:"comment:更新时间" json:"updatedAt,omitempty" swaggerignore:"true"`
}

func (b *BaseModelString) BeforeCreate(tx *gorm.DB) (err error) {
	if b.ID == "" {
		b.ID = snowflake.GenerateIdStr()
	}
	return
}

func WithTransaction(fns ...func(*gorm.DB) error) error {
	var err error
	tx := mysqlx.GetDB().Begin()
	defer func() {
		if r := recover(); r != nil {
			tx.Rollback()
			return
		} else if err == nil {
			tx.Commit()
			return
		}
	}()

	for _, fn := range fns {
		err = fn(tx)
		if err != nil {
			tx.Rollback()
			return err
		}
	}

	return nil
}

func Init() error {
	fmt.Printf("dsn=%v\n", conf.DbConf.Dsn)
	if err := mysqlx.Init(conf.DbConf, logx.GetLogger()); err != nil {
		return err
	}

	if err := RegisterTables(); err != nil {
		return err
	}

	InsertDefaultData()

	return nil
}

func RegisterTables() error {
	db := mysqlx.GetDB()
	err := db.AutoMigrate(
		Audio{},
		AudioText{},
		Word{},
		TrainInfo{},
	)
	return err
}

type InitDefaultData interface {
	InitDefaultData() error
}

func InsertDefaultData() {
	models := []interface{}{}
	for _, model := range models {
		if id, ok := model.(InitDefaultData); ok {
			if err := id.InitDefaultData(); err != nil {
				logx.Errorf("InitDefaultData for table: %v, err :%v", model.(schema.Tabler).TableName(), err.Error())
			}
		}
	}
}