194 lines
5.7 KiB
Go
194 lines
5.7 KiB
Go
package initcmd
|
||
|
||
import (
|
||
"bufio"
|
||
"fmt"
|
||
"os"
|
||
"path/filepath"
|
||
"strings"
|
||
|
||
"sunhpc/internal/auth"
|
||
"sunhpc/internal/db"
|
||
"sunhpc/internal/log"
|
||
|
||
"github.com/spf13/cobra"
|
||
"github.com/spf13/viper"
|
||
)
|
||
|
||
const (
|
||
appName string = "sunhpc"
|
||
defaultDBPath string = "/var/lib/sunhpc"
|
||
defaultDBName string = "sunhpc.db"
|
||
)
|
||
|
||
var (
|
||
forceDB bool
|
||
dbPath string
|
||
dbName string
|
||
dbFullPath string
|
||
)
|
||
|
||
func initDBPathWithViper() error {
|
||
/*
|
||
从 Viper 配置文件获取数据库路径
|
||
配置文件里的键要和 Viper.GetXXX 的键对应
|
||
配置文件格式：
|
||
db:
|
||
path: "/tmp/sunhpc" # 自定义数据库路径
|
||
name: "my_sunhpc.db" # 自定义数据库名
|
||
*/
|
||
|
||
log.Infof("从 Viper 配置文件获取数据库路径...")
|
||
|
||
// ========== 第一步：设置 Viper 配置文件规则（核心） ==========
|
||
// 1. 设置Viper基础规则
|
||
viper.SetConfigType("yaml") // 配置文件类型
|
||
viper.SetConfigName("config") // 配置文件名（不带后缀）
|
||
viper.SetEnvPrefix(appName) // 环境变量前缀：SUNHPC_
|
||
viper.AutomaticEnv() // 自动读取环境变量（可选，增强兼容性）
|
||
|
||
// 2. 添加配置文件搜索目录（Viper 会按顺序查找，找到第一个就停止）
|
||
// 优先级：当前目录 → 用户级目录 → 系统级目录
|
||
|
||
// ① 当前目录（开发/测试常用）
|
||
viper.AddConfigPath(".")
|
||
|
||
// ② Linux/macOS 用户级目录（~/.config/sunhpc/）
|
||
if homeDir, err := os.UserHomeDir(); err == nil {
|
||
viper.AddConfigPath(filepath.Join(homeDir, ".config", appName))
|
||
}
|
||
// ③ Linux/macOS 系统级目录（/etc/sunhpc/）
|
||
viper.AddConfigPath(filepath.Join("/etc", appName))
|
||
|
||
// ========== 第二步：设置默认值（最低优先级） ==========
|
||
viper.SetDefault("db.path", defaultDBPath)
|
||
viper.SetDefault("db.name", defaultDBName)
|
||
|
||
// ========== 第三步：绑定环境变量（优先级高于默认值，低于配置文件） ==========
|
||
viper.BindEnv("db.path", "DB_PATH") // 绑定 SUNHPC_DB_PATH → db.path
|
||
viper.BindEnv("db.name", "DB_NAME") // 绑定 SUNHPC_DB_NAME → db.name
|
||
|
||
// ========== 第四步：读取配置文件（优先级高于环境变量，低于默认值） ==========
|
||
if err := viper.ReadInConfig(); err != nil {
|
||
if _, ok := err.(viper.ConfigFileNotFoundError); ok {
|
||
log.Info("未找到配置文件,将使用环境变量/默认值")
|
||
return nil // 配置文件存在但格式错误，返回错误
|
||
}
|
||
log.Warnf("读取配置文件失败: %v", err)
|
||
return fmt.Errorf("读取配置文件失败: %w", err)
|
||
}
|
||
|
||
log.Infof("成功加载配置文件: %s", viper.ConfigFileUsed())
|
||
return nil
|
||
}
|
||
|
||
func initDBPath() error {
|
||
|
||
// 1. 从 Viper 配置文件获取数据库路径(加载配置文件->环境变量->默认值)
|
||
if err := initDBPathWithViper(); err != nil {
|
||
return fmt.Errorf("Viper初始化数据库失败: %v", err)
|
||
}
|
||
|
||
// 2. 从Viper获取数据库路径
|
||
dbPath = viper.GetString("db.path")
|
||
dbName = viper.GetString("db.name")
|
||
|
||
// 3. 拼接数据库路径
|
||
dbFullPath = filepath.Join(dbPath, dbName)
|
||
log.Infof("数据库完整路径: %s", dbFullPath)
|
||
|
||
// 3. 检查数据库文件是否存在
|
||
dir := filepath.Dir(dbFullPath)
|
||
// 4. 检查数据库目录是否存在
|
||
if _, err := os.Stat(dir); os.IsNotExist(err) {
|
||
log.Infof("数据库目录不存在,创建目录: %s", dir)
|
||
|
||
if err := os.MkdirAll(dir, 0755); err != nil {
|
||
return fmt.Errorf("创建数据库目录失败: %v", err)
|
||
}
|
||
}
|
||
return nil
|
||
}
|
||
|
||
var databaseCmd = &cobra.Command{
|
||
Use: "database",
|
||
Short: "初始化数据库",
|
||
Long: `初始化SQLite数据库,创建所有表结构和默认数据。
|
||
|
||
示例:
|
||
sunhpc init database # 初始化数据库
|
||
sunhpc init database --force # 强制重新初始化`,
|
||
|
||
RunE: func(cmd *cobra.Command, args []string) error {
|
||
if err := auth.RequireRoot(); err != nil {
|
||
return err
|
||
}
|
||
|
||
log.Debug("初始化数据库...")
|
||
|
||
if err := initDBPath(); err != nil {
|
||
return fmt.Errorf("初始化数据库路径失败: %w", err)
|
||
}
|
||
|
||
log.Debugf("数据库目录存在: %s", dbPath)
|
||
|
||
// 强制模式：用户确认
|
||
if forceDB {
|
||
log.Warn("⚠️ 警告：强制重新初始化将清空数据库中的所有数据！")
|
||
fmt.Printf("数据库路径: %s\n", dbFullPath)
|
||
fmt.Print("确认要重新初始化数据库吗？这将删除所有现有数据。(Y/yes): ")
|
||
|
||
reader := bufio.NewReader(os.Stdin)
|
||
input, err := reader.ReadString('\n')
|
||
if err != nil {
|
||
return fmt.Errorf("读取用户输入失败: %v", err)
|
||
}
|
||
|
||
input = strings.TrimSpace(strings.ToLower(input))
|
||
if input != "y" && input != "yes" {
|
||
log.Info("操作已取消")
|
||
return nil
|
||
}
|
||
|
||
log.Info("用户确认重新初始化数据库")
|
||
}
|
||
|
||
// 数据库存在且不是强制模式则跳过初始化
|
||
if _, err := os.Stat(dbFullPath); err == nil && !forceDB {
|
||
log.Infof("数据库文件已存在: %s", dbFullPath)
|
||
return nil
|
||
}
|
||
|
||
// 初始化数据库（使用配置的路径）
|
||
database, err := db.GetInstanceWithConfig(dbPath, dbName)
|
||
if err != nil {
|
||
return fmt.Errorf("初始化数据库失败: %v", err)
|
||
}
|
||
defer database.Close()
|
||
|
||
// 如果是强制模式，设置强制重新初始化标志
|
||
if forceDB {
|
||
database.SetForceInit(true)
|
||
log.Info("强制重新初始化数据库表...")
|
||
|
||
// 关闭现有连接以触发重新连接
|
||
if err := database.CloseConnection(); err != nil {
|
||
return fmt.Errorf("关闭现有数据库连接失败: %v", err)
|
||
}
|
||
|
||
// 重新连接并初始化
|
||
if err := database.Connect(); err != nil {
|
||
return fmt.Errorf("强制重新初始化数据库失败: %v", err)
|
||
}
|
||
}
|
||
|
||
log.Infof("数据库初始化成功: %s", dbFullPath)
|
||
return nil
|
||
},
|
||
}
|
||
|
||
func init() {
|
||
databaseCmd.Flags().BoolVarP(&forceDB, "force", "f", false, "强制重新初始化，删除现有数据库")
|
||
Cmd.AddCommand(databaseCmd)
|
||
}
|