package cmd import ( "fmt" "log" "os" "time" "github.com/spf13/cobra" "gorm.io/driver/mysql" "gorm.io/gorm" "gorm.io/gorm/logger" "my-spiders/internal/config" ) var ( DB *gorm.DB cfgFile string debug bool // 命令行 debug 标志 rootCmd = &cobra.Command{ Use: "spider-cli", Short: "自动化多平台爬虫集成系统", } ) func Execute() { if err := rootCmd.Execute(); err != nil { fmt.Println(err) os.Exit(1) } } func init() { // 支持通过 --config 参数指定其他的配置文件路径 rootCmd.PersistentFlags().StringVarP(&cfgFile, "config", "c", "", "配置文件路径 (默认是 ./config.yaml)") // 新增全局 -d / --debug 参数 rootCmd.PersistentFlags().BoolVarP(&debug, "debug", "d", false, "调试模式 (仅输出不入库)") cobra.OnInitialize(initConfigAndDB) } func initConfigAndDB() { // 1. 初始化配置文件 if err := config.InitConfig(cfgFile); err != nil { log.Fatalf("[致命错误] %v", err) } // 命令行 -d 参数优先于 YAML 配置 if debug { config.GlobalConfig.App.Debug = true } isDebug := config.GlobalConfig.App.Debug // 2. 初始化数据库连接 dbCfg := config.GlobalConfig.Database var err error DB, err = gorm.Open(mysql.Open(dbCfg.DSN), &gorm.Config{ Logger: logger.Default.LogMode(logger.Warn), }) if err != nil { if isDebug { log.Println("[Warning] 数据库连接失败,当前处于 Debug 模式,将跳过数据库操作...") DB = nil return } log.Fatalf("[致命错误] 数据库连接失败: %v", err) } // 3. 设置连接池参数 sqlDB, err := DB.DB() if err != nil || sqlDB.Ping() != nil { if isDebug { log.Println("[Warning] 数据库 Ping 失败,当前处于 Debug 模式,将跳过数据库操作...") DB = nil return } log.Fatalf("[致命错误] 数据库不可用: %v", err) } sqlDB.SetMaxOpenConns(dbCfg.MaxOpenConns) sqlDB.SetMaxIdleConns(dbCfg.MaxIdleConns) sqlDB.SetConnMaxLifetime(time.Duration(dbCfg.ConnMaxLifetime) * time.Minute) sqlDB.SetConnMaxIdleTime(time.Duration(dbCfg.ConnMaxIdleTime) * time.Minute) log.Printf("[Info] MySQL 连接池初始化完成 | MaxOpen: %d | MaxIdle: %d", dbCfg.MaxOpenConns, dbCfg.MaxIdleConns) }