mDoc/admin/db.go

package admin

import (
	"os"
	"strings"

	"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)
	}

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

	return
}

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)
		}
	}

	return
}

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
}