92 lines
2.2 KiB
Go
92 lines
2.2 KiB
Go
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)
|
|
}
|