mDoc/admin/db.go

package admin

import (
	"os"
	"strings"

	"github.com/qor5/x/login"
	"gorm.io/driver/postgres"
	"gorm.io/driver/sqlite"
	"gorm.io/gorm"
	"gorm.io/gorm/logger"
)

var IsSqlite bool // 是否为Sqlite数据库

func ConnectDB() (db *gorm.DB) {
	var err error
	// Ease 修改
	// 自动判断是否为sqlite3数据库： 以.db结尾，则为sqlite3数据库连接; 否则为postgresql数据库
	dbConnStr := os.Getenv("DB_PARAMS")
	if strings.HasSuffix(dbConnStr, ".db") {
		IsSqlite = true
		db, err = gorm.Open(sqlite.Open("test.db"), &gorm.Config{
			DisableForeignKeyConstraintWhenMigrating: true, // 禁止gorm自动建立外键约束
			SkipDefaultTransaction:                   true, // 跳过默认事务
		})
		db.Exec("PRAGMA journal_mode = OFF;")  // 设置数据库的日志模式: WAL/
		db.Exec("PRAGMA synchronous = OFF;")   // 关闭写同步
		db.Exec("PRAGMA cache_size = 8000;")   // 缓存大小(数据库文件页数): 1000000为页面缓存分配最大4GB的RAM
		db.Exec("PRAGMA temp_store = MEMORY;") // 临时存储放到内存中(临时表和索引会存储在内存中)
		// PRAGMA max_page_count 设置数据库文件最大的页数
		// PRAGMA page_size 数据库的页面大小
		sqlDB, err := db.DB()
		if err != nil {
			panic(err)
		}
		sqlDB.SetMaxIdleConns(10)  // 设置空闲连接池中连接的最大数量
		sqlDB.SetMaxOpenConns(100) // 设置打开数据库连接的最大数量

	} else {
		db, err = gorm.Open(postgres.Open(dbConnStr)) // 通过环境变量连接数据库
	}
	if err != nil {
		panic(err)
	}

	autoMigrate(db)

	db.Logger = db.Logger.LogMode(logger.Info)

	return
}

func autoMigrate(db *gorm.DB) {
	err := db.AutoMigrate(
		&User{}, &UserFavorite{}, &UserFollow{}, &UserCategorize{},
		&DocumentCategorize{}, &Document{},
	)
	if err != nil { // 生成数据表
		panic(err)
	}
}

func initWebsiteData(db *gorm.DB) {
	var cnt int64
	if err := db.Table("page_builder_pages").Count(&cnt).Error; err != nil {
		panic(err)
	}

	if cnt == 0 {
		if err := db.Exec(initWebsiteSQL).Error; err != nil {
			panic(err)
		}

	}

	createDefaultManager(db)
	return
}

// 创建默认管理员
func createDefaultManager(db *gorm.DB) {
	user := &User{ // 添加一个记录
		UserPass: login.UserPass{
			Account:  "ease@scwy.net",
			Password: "yjz129",
		},
	}
	user.EncryptPassword() // 对密码加密
	db.Create(user)
}

// 创建缺省数据
func initMediaLibraryData(db *gorm.DB) {
	var cnt int64
	if err := db.Table("media_libraries").Count(&cnt).Error; err != nil {
		panic(err)
	}

	if cnt == 0 {
		if err := db.Exec(initMediaLibrarySQL).Error; err != nil {
			panic(err)
		}
	}

	return
}