数据库模块完成-1

2026-02-18 17:09:52 +08:00
parent 8a7bf8a39c
commit bf1df59bc8
7 changed files with 609 additions and 127 deletions
--- a/internal/db/db.go
+++ b/internal/db/db.go
@@ -1,10 +1,12 @@
 package db

 import (
+	"bufio"
 	"database/sql"
 	"fmt"
 	"os"
 	"path/filepath"
+	"strings"
 	"sync"

 	_ "github.com/go-sql-driver/mysql"
@@ -17,31 +19,127 @@ import (
 // DB wraps the sql.DB connection pool.
 type DB struct {
 	engine *sql.DB
+	config *config.DBConfig // 保存配置
 }

+/*
 // Engine returns the underlying *sql.DB.
 func (d *DB) Engine() *sql.DB {
 	return d.engine
 }
+*/
+
+func ConfirmWithRetry(prompt string, maxAttempts int) bool {
+	reader := bufio.NewReader(os.Stdin)
+
+	for attempt := 1; attempt <= maxAttempts; attempt++ {
+		log.Infof("%s [y/n]", prompt)
+
+		response, err := reader.ReadString('\n')
+		if err != nil {
+			continue
+		}
+
+		response = strings.ToLower(strings.TrimSpace(response))
+
+		switch response {
+		case "y", "yes":
+			return true
+		case "n", "no", "":
+			return false
+		default:
+			if attempt < maxAttempts {
+				log.Warnf(
+					"⚠️ 无效输入、请输入 'y' 或 'n'(剩余尝试次数: %d)",
+					maxAttempts-attempt)
+			}
+		}
+	}
+
+	log.Warn("⚠️  警告：尝试次数过多、操作已取消")
+	return false
+}

 // InitSchema initializes the database schema.
 // If force is true, drops existing tables before recreating them.
 func (d *DB) InitSchema(force bool) error {
-	db := d.engine
+	fullPath := filepath.Join(d.config.Path, d.config.Name)

-	if force {
-		if err := dropTables(db); err != nil {
-			return fmt.Errorf("failed to drop tables: %w", err)
+	// 检查文件是否存在
+	_, err := os.Stat(fullPath)
+	fileExists := err == nil
+
+	// 处理不同的场景
+	switch {
+	case !fileExists:
+		// 场景1：文件不存在，连接并创建(allowCreate = true).
+		log.Infof("数据库文件不存在，将创建: %s", fullPath)
+		if err := d.Connect(true); err != nil {
+			return err
 		}
+		return createTables(d.engine)
+
+	case fileExists && !force:
+		// 场景2：文件存在、无 force 参数、提示友好退出.
+		log.Warnf("数据库文件已存在: %s", fullPath)
+		log.Warn("如果需要强制重新初始化，请添加 --force 参数")
+		log.Warn("数据库已存在、退出初始化操作.")
+		os.Exit(1)
+
+	case fileExists && force:
+		// 场景3：文件存在、force 参数 -> 需要用户确认并重建.
+		log.Warn("警告：强制重新初始化将清空数据库中的所有数据！")
+		if !ConfirmWithRetry("是否继续？", 3) {
+			return fmt.Errorf("用户取消操作")
+		}
+
+		// 连接现有数据库(allowCreate = true, 因为文件已经存在)
+		if err := d.Connect(true); err != nil {
+			return err
+		}
+
+		// 清空现有数据.
+		if err := dropTables(d.engine); err != nil {
+			return err
+		}
+
+		// 清空表
+		if err := dropTriggers(d.engine); err != nil {
+			return err
+		}
+
+		log.Info("已清空现有数据库触发器")
+		return createTables(d.engine)
 	}

+	log.Info("数据库创建成功")
+	return nil
+}
+
+// 辅助函数：检查文件是否存在
+func fileExists(path string) bool {
+	_, err := os.Stat(path)
+	if err == nil {
+		return true
+	}
+	if os.IsNotExist(err) {
+		return false
+	}
+	// 其他问题（如权限问题）也视为文件不存在，但应该记录日志
+	log.Debugf("检查数据库文件状态失败: %v", err)
+	return false
+}
+
+// 辅助函数: 创建数据库表
+func createTables(db *sql.DB) error {
 	// ✅ 调用 schema.go 中的函数
 	for _, ddl := range CreateTableStatements() {
+		log.Debugf("执行: %s", ddl)
 		if _, err := db.Exec(ddl); err != nil {
-			return fmt.Errorf("failed to create table: %w", err)
+			return fmt.Errorf("数据表创建失败: %w", err)
 		}
 	}
-
+	log.Info("数据库表创建成功")
 	return nil
 }

@@ -55,6 +153,16 @@ func dropTables(db *sql.DB) error {
 	return nil
 }

+func dropTriggers(db *sql.DB) error {
+	// ✅ 调用 schema.go 中的函数
+	for _, trigger := range DropTriggerStatements() {
+		if _, err := db.Exec(fmt.Sprintf("DROP TRIGGER IF EXISTS `%s`", trigger)); err != nil {
+			return err
+		}
+	}
+	return nil
+}
+
 // --- Singleton DB Instance ---
 var (
 	globalDB *DB
@@ -70,52 +178,77 @@ func GetDB() (*DB, error) {
 			return
 		}

-		if _, err := os.Stat(cfg.DB.Path); err != nil {
-			// 创建数据库目录
-			if err := os.MkdirAll(cfg.DB.Path, 0755); err != nil {
-				log.Fatalf("创建数据库目录失败: %v", err)
-			}
-			log.Infof("数据库目录创建成功: %s", cfg.DB.Path)
+		globalDB = &DB{
+			config: &cfg.DB,
+		}
+	})
+	return globalDB, initErr
+}
+
+func (d *DB) Connect(allowCreate bool) error {
+	// 如果已经连接，直接返回
+	if d.engine != nil {
+		return nil
+	}
+
+	switch d.config.Type {
+	case "sqlite":
+		fullPath := filepath.Join(d.config.Path, d.config.Name)
+
+		// 检查文件是否存在
+		_, err := os.Stat(fullPath)
+		fileExists := err == nil
+
+		// 如果文件不存在且不允许创建，返回错误
+		if !fileExists && !allowCreate {
+			return fmt.Errorf("数据库文件不存在: %s, 请先初始化.", fullPath)
 		}

-		var dsn string
-		var driver string
-
-		switch cfg.DB.Type {
-		case "sqlite":
-			driver = "sqlite3"
-			fullPath := filepath.Join(cfg.DB.Path, cfg.DB.Name)
-			dsn = fmt.Sprintf("%s?_foreign_keys=on&_journal_mode=WAL&_timeout=5000", fullPath)
-		case "mysql":
-			driver = "mysql"
-			if cfg.DB.Socket != "" {
-				dsn = fmt.Sprintf("%s:%s@unix(%s)/%s?parseTime=true&loc=Local",
-					cfg.DB.User, cfg.DB.Password, cfg.DB.Socket, cfg.DB.Name)
-			} else {
-				dsn = fmt.Sprintf("%s:%s@tcp(%s)/%s?parseTime=true&loc=Local",
-					cfg.DB.User, cfg.DB.Password, cfg.DB.Host, cfg.DB.Name)
-			}
-		default:
-			initErr = fmt.Errorf("unsupported database type: %s", cfg.DB.Type)
-			return
+		// 确保目录存在
+		if err := os.MkdirAll(d.config.Path, 0755); err != nil {
+			return fmt.Errorf("创建数据库目录失败: %w", err)
 		}

-		engine, err := sql.Open(driver, dsn)
+		// 连接参数
+		dsn := fmt.Sprintf("%s?_foreign_keys=on&_journal_mode=WAL&_timeout=5000",
+			fullPath)
+
+		engine, err := sql.Open("sqlite3", dsn)
 		if err != nil {
-			initErr = fmt.Errorf("failed to open database: %w", err)
-			return
+			return fmt.Errorf("数据库打开失败: %w", err)
 		}

 		if err := engine.Ping(); err != nil {
 			engine.Close()
-			initErr = fmt.Errorf("failed to ping database: %w", err)
-			return
+			return fmt.Errorf("数据库连接失败: %w", err)
 		}

-		globalDB = &DB{engine: engine}
-	})
+		d.engine = engine
+	case "mysql":
+		// TODO: 实现 MySQL 连接逻辑
+		return fmt.Errorf("mysql 数据库连接未实现")
+	}

-	return globalDB, initErr
+	return nil
+}
+
+// Close 关闭数据库连接
+func (d *DB) Close() error {
+	if d.engine != nil {
+		return d.engine.Close()
+	}
+	return nil
+}
+
+// GetEngine 获取数据库引擎(自动连接)
+func (d *DB) GetEngine() (*sql.DB, error) {
+	// 如果还没有连接，自动连接（但不创建新文件)
+	if d.engine == nil {
+		if err := d.Connect(false); err != nil {
+			return nil, err
+		}
+	}
+	return d.engine, nil
 }

 // MustGetDB is a helper that panics on error (use in main/init only).
@@ -126,3 +259,48 @@ func MustGetDB() *DB {
 	}
 	return db
 }
+
+func GetDBConfig() (*config.DBConfig, error) {
+	cfg, err := config.LoadConfig()
+	if err != nil {
+		return nil, fmt.Errorf("数据库配置文件加载失败: %w", err)
+	}
+	return &cfg.DB, nil
+}
+
+func CheckDB() (*config.Config, error) {
+	cfg, err := config.LoadConfig()
+	if err != nil {
+		log.Warnf("加载配置失败: %v", err)
+	}
+
+	// 统一转为小写，避免用户输入错误
+	dbType := strings.ToLower(cfg.DB.Type)
+
+	// 打印配置（调试用）
+	log.Debugf("数据库类型: %s", dbType)
+	log.Debugf("数据库名称: %s", cfg.DB.Name)
+	log.Debugf("数据库路径: %s", cfg.DB.Path)
+	log.Debugf("数据库用户: %s", cfg.DB.User)
+	log.Debugf("数据库主机: %s", cfg.DB.Host)
+	log.Debugf("数据库套接字: %s", cfg.DB.Socket)
+	log.Debugf("数据库详细日志: %v", cfg.DB.Verbose)
+
+	// 支持 sqlite，mysql的常见别名
+	isSQLite := dbType == "sqlite" || dbType == "sqlite3"
+	isMySQL := dbType == "mysql"
+
+	// 检查数据库类型,只允许 sqlite 和 mysql
+	if !isSQLite && !isMySQL {
+		log.Fatalf("不支持的数据库类型: %s(仅支持 sqlite、sqlite3、mysql)", dbType)
+	}
+
+	// 检查数据库路径是否存在
+	if isSQLite {
+		if _, err := os.Stat(cfg.DB.Path); os.IsNotExist(err) {
+			log.Warnf("SQLite 数据库路径 %s 不存在", cfg.DB.Path)
+			log.Warn("必须先执行 'sunhpc init database' 初始化数据库")
+		}
+	}
+	return cfg, nil
+}