sunhpc-go/cmd/root.go

package cmd

import (
	"os"
	"strings"
	initcmd "sunhpc/cmd/init"
	"sunhpc/cmd/soft"
	"sunhpc/cmd/tmpl"
	"sunhpc/internal/auth"
	"sunhpc/internal/config"
	"sunhpc/internal/log"

	"github.com/spf13/cobra"
)

var (
	cfgFile string
	verbose bool
	noColor bool
)

func checkDB() error {
	cfg, err := config.LoadConfig()
	if err != nil {
		log.Fatal("加载配置失败: ", 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.Fatalf("必须先执行 'sunhpc init database' 初始化数据库")
		}
	}
	return nil
}

var rootCmd = &cobra.Command{
	Use:   "sunhpc",
	Short: "SunHPC - HPC集群一体化运维工具",
	PersistentPreRun: func(cmd *cobra.Command, args []string) {
		// 初始化日志（verbose=false 不显示调试信息）
		log.Init(verbose)

		// 是否禁用颜色
		log.EnableColor(!noColor)

		log.Debugf("当前命令 Annotations: %+v", cmd.Annotations)

		// 检查当前命令是否标记为跳过 DB 检查
		if cmd.Annotations["skip-db-check"] == "true" {
			log.Debugf("当前命令 %s 标记为跳过 DB 检查", cmd.Name())
			return
		} else {
			// 检查数据库
			if err := checkDB(); err != nil {
				log.Fatalf("数据库检查失败: %v", err)
			}
		}

		// 需要 root 权限
		if cmd.Annotations["require-root"] == "true" {
			if err := auth.RequireRoot(); err != nil {
				log.Fatalf("需要 root 权限: %v", err)
			}
		}

		log.Debugf("当前命令: %s", cmd.Name())
		log.Debugf("详细模式: %v", verbose)
		log.Debugf("禁用颜色: %v", noColor)
	},
	PersistentPostRun: func(cmd *cobra.Command, args []string) {
		// 同步日志
		log.Sync()
		log.Close()
	},
	Run: func(cmd *cobra.Command, args []string) {
		cmd.Help()
	},
}

func Execute() error {
	return rootCmd.Execute()
}

func init() {
	rootCmd.PersistentFlags().StringVar(&cfgFile, "config", "", "配置文件路径 (默认为 /etc/sunhpc/sunhpc.yaml)")
	rootCmd.PersistentFlags().BoolVarP(&verbose, "verbose", "v", false, "启用详细日志输出")
	rootCmd.PersistentFlags().BoolVar(&noColor, "no-color", false, "禁用彩色输出")

	// 注册一级子命令下的子命令树
	rootCmd.AddCommand(initcmd.Cmd)
	rootCmd.AddCommand(soft.Cmd)
	rootCmd.AddCommand(tmpl.Cmd)
}