yx-dataset-server/boot/gorm.go

package boot

import (
	"fmt"
	"time"
	"yx-dataset-server/app/model"
	"yx-dataset-server/app/model/entity"
	iModel "yx-dataset-server/app/model/impl/model"
	"yx-dataset-server/library/logger"
	"yx-dataset-server/library/utils"

	"go.uber.org/dig"
	"gorm.io/driver/mysql"
	"gorm.io/gorm"
	"gorm.io/gorm/schema"
)

// 初始化gorm
func initGorm() (*gorm.DB, error) {
	dsn := fmt.Sprintf("%s:%s@tcp(%s:%d)/%s?%s",
		utils.GetConfig("mysql.user").String(),
		utils.GetConfig("mysql.password").String(),
		utils.GetConfig("mysql.host").String(),
		utils.GetConfig("mysql.port").Int(),
		utils.GetConfig("mysql.db_name").String(),
		utils.GetConfig("mysql.parameters").String())

	db, err := gorm.Open(mysql.Open(dsn), &gorm.Config{
		SkipDefaultTransaction: true, // 跳过默认事务
		NamingStrategy: schema.NamingStrategy{
			TablePrefix:   utils.GetConfig("gorm.table_prefix").String(),
			SingularTable: true, // 使用单数表名
		},
	})
	if err != nil {
		return nil, err
	}
	db.Logger = logger.NewEntry(logger.GetLogger())
	if utils.GetConfig("common.run_mode").String() == "debug" {
		db.Debug()
	}

	sqlDb, err := db.DB()
	if err != nil {
		return nil, err
	}
	sqlDb.SetMaxIdleConns(utils.GetConfig("gorm.max_idle_conns").Int())
	sqlDb.SetMaxOpenConns(utils.GetConfig("gorm.max_open_conns").Int())
	sqlDb.SetConnMaxLifetime(time.Duration(utils.GetConfig("gorm.max_open_conns").Int()) * time.Second)
	return db, nil
}

// 自动创建数据表映射
func autoMigrate(db *gorm.DB) error {
	return db.AutoMigrate(
		new(entity.Dataset),
		new(entity.DatasetFile),
		new(entity.LoginHistory),
		new(entity.Menu),

		new(entity.Organization),

		new(entity.Role),
		new(entity.RoleMenu),

		new(entity.User),
		new(entity.UserDataset),
		new(entity.UserRole),

		new(entity.ChatAssistant),
		new(entity.ChatSession),
		new(entity.ChatMessage),
		new(entity.ChatDataset),

		new(entity.RobotConfig),
		new(entity.RobotDataset),

	)
}

func InjectModel(container *dig.Container) error {
	_ = container.Provide(iModel.NewTrans)
	_ = container.Provide(func(m *iModel.Trans) model.ITrans { return m })
	_ = container.Provide(iModel.NewDataset)
	_ = container.Provide(func(m *iModel.Dataset) model.IDataset { return m })
	_ = container.Provide(iModel.NewDatasetFile)
	_ = container.Provide(func(m *iModel.DatasetFile) model.IDatasetFile { return m })
	_ = container.Provide(iModel.NewLoginHistory)
	_ = container.Provide(func(m *iModel.LoginHistory) model.ILoginHistory { return m })

	_ = container.Provide(iModel.NewOrganization)
	_ = container.Provide(func(m *iModel.Organization) model.IOrganization { return m })

	_ = container.Provide(iModel.NewUser)
	_ = container.Provide(func(m *iModel.User) model.IUser { return m })
	_ = container.Provide(iModel.NewUserDataset)
	_ = container.Provide(func(m *iModel.UserDataset) model.IUserDataset { return m })
	_ = container.Provide(iModel.NewUserRole)
	_ = container.Provide(func(m *iModel.UserRole) model.IUserRole { return m })

	_ = container.Provide(iModel.NewRole)
	_ = container.Provide(func(m *iModel.Role) model.IRole { return m })
	_ = container.Provide(iModel.NewMenu)
	_ = container.Provide(func(m *iModel.Menu) model.IMenu { return m })
	_ = container.Provide(iModel.NewRoleMenu)
	_ = container.Provide(func(m *iModel.RoleMenu) model.IRoleMenu { return m })

	_ = container.Provide(iModel.NewChatAssistant)
	_ = container.Provide(func(m *iModel.ChatAssistant) model.IChatAssistant { return m })
	_ = container.Provide(iModel.NewChatSession)
	_ = container.Provide(func(m *iModel.ChatSession) model.IChatSession { return m })
	_ = container.Provide(iModel.NewChatMessage)
	_ = container.Provide(func(m *iModel.ChatMessage) model.IChatMessage { return m })
	_ = container.Provide(iModel.NewChatDataset)
	_ = container.Provide(func(m *iModel.ChatDataset) model.IChatDataset { return m })

	_ = container.Provide(iModel.NewRobotConfig)
	_ = container.Provide(func(m *iModel.RobotConfig) model.IRobotConfig { return m })
	_ = container.Provide(iModel.NewRobotDataset)
	_ = container.Provide(func(m *iModel.RobotDataset) model.IRobotDataset { return m })
	return nil
}