gohub/pkg/database/database.go

// Package database 数据库连接
package database

import (
	"database/sql"
	"fmt"

	"gorm.io/gorm"
	gormlgger "gorm.io/gorm/logger"

	"github.com/runningwater/gohub/pkg/config"
)

// DB 数据库连接实例
var DB *gorm.DB
var SQLDB *sql.DB

func Connect(dbConfig gorm.Dialector, _logger gormlgger.Interface) {
	// 使用 gorm.Open 连接数据库
	var err error
	DB, err = gorm.Open(dbConfig, &gorm.Config{
		Logger: _logger,
	})
	// 处理错误
	if err != nil {
		fmt.Println("数据库连接失败", err.Error())
	}
	// 获取底层的 sql.DB 实例
	SQLDB, err = DB.DB()

	if err != nil {
		fmt.Println("数据库连接失败", err.Error())
	}
}

func CurrentDatabase() string {
	return DB.Migrator().CurrentDatabase()
}
func DeleteAllTables() error {
	var err error

	switch config.Get("database.connection") {
	case "mysql":
		err = deleteMySQLTables()
	case "sqlite":
		err = deleteAllSQLiteTables()
	default:
		err = fmt.Errorf("不支持的数据库类型: %s", config.Get("database.connection"))
	}

	return err
}

func deleteMySQLTables() error {
	dbname := CurrentDatabase()
	var tables []string

	err := DB.Table("information_schema.tables").
		Where("table_schema = ?", dbname).
		Pluck("table_name", &tables).Error
	if err != nil {
		return err
	}

	// 暂时关闭外键检测
	DB.Exec("SET FOREIGN_KEY_CHECKS = 0;")

	// 删除所有表
	for _, table := range tables {
		if err := DB.Migrator().DropTable(table); err != nil {
			return err
		}
	}
	// 重新启用外键检测
	DB.Exec("SET FOREIGN_KEY_CHECKS = 1;")

	return nil
}

func deleteAllSQLiteTables() error {
	var tables []string

	err := DB.Raw("SELECT name FROM sqlite_master WHERE type='table'").
		Pluck("name", &tables).Error
	if err != nil {
		return err
	}

	// 删除所有表
	for _, table := range tables {
		if err := DB.Migrator().DropTable(table); err != nil {
			return err
		}
	}

	return nil
}

// TableName 获取表名
func TableName(obj any) string {
	stmt := &gorm.Statement{DB: DB}
	err := stmt.Parse(obj)
	if err != nil {
		return ""
	}
	return stmt.Schema.Table
}