hpds_control_center/model/index.go

package model

import (
	"fmt"
	"go.uber.org/zap"
	"os"

	_ "github.com/go-sql-driver/mysql"
	"xorm.io/xorm"
	"xorm.io/xorm/dialects"
)

var (
	DB *xorm.Engine
)

func New(driveName, dsn string, showSql bool) {
	DB, _ = NewDbConnection(driveName, dsn)
	DB.ShowSQL(showSql)
	DB.Dialect().SetQuotePolicy(dialects.QuotePolicyReserved)
	err := DB.Sync2(
		&IssueModel{},
		&Model{},
		&Node{},
		&NodeState{},
		&Task{},
		&TaskLog{},
		&TaskResult{},
		&TrainingDataset{},
		&TrainingDatasetDetail{},
		&TrainTask{},
		&TrainTaskLog{},
		&TrainTaskResult{},
	)
	if err != nil {
		fmt.Println("同步数据库表结构", err)
		zap.L().Error("同步数据库表结构", zap.Error(err))
		os.Exit(1)
	}
}

func NewDbConnection(driveName, dsn string) (db *xorm.Engine, err error) {
	db, err = xorm.NewEngine(driveName, dsn)
	if err != nil {
		zap.L().Error("创建数据库连接", zap.Error(err))
		os.Exit(-1)
	}
	db.SetMaxOpenConns(300)
	return
}