package main import ( "fmt" "log" "os" "os/signal" "strconv" "syscall" "time" "git.apinb.com/quant/collector/collector" "git.apinb.com/quant/collector/storage" "github.com/robfig/cron/v3" ) func main() { log.Println("=== QMT数据采集器启动 ===") // 从环境变量获取配置 collectorURL := getEnv("COLLECTOR_URL", "http://localhost:5000/status") dbConnStr := getEnv("DATABASE_URL", "host=139.224.247.176 user=postgres password=Stock0310~! dbname=stock_prod port=19432 sslmode=disable TimeZone=Asia/Shanghai") interval := getEnvAsInt("COLLECTION_INTERVAL", 60) // 默认60秒 log.Printf("采集地址: %s", collectorURL) log.Printf("采集间隔: %d秒", interval) // 创建采集器 coll := collector.NewCollector(collectorURL) // 创建数据库存储 store, err := storage.NewStorage(dbConnStr) if err != nil { log.Fatalf("数据库连接失败: %v", err) } defer store.Close() // 自动迁移数据库表结构 if err := store.AutoMigrate(); err != nil { log.Fatalf("数据库迁移失败: %v", err) } // 创建cron调度器 c := cron.New(cron.WithSeconds()) // 构建cron表达式 (每N秒执行一次) cronSpec := fmt.Sprintf("@every %ds", interval) log.Printf("定时任务表达式: %s", cronSpec) // 添加定时任务 _, err = c.AddFunc(cronSpec, func() { runCollection(coll, store) }) if err != nil { log.Fatalf("添加定时任务失败: %v", err) } // 启动调度器 c.Start() log.Println("定时任务已启动") // 等待退出信号 quit := make(chan os.Signal, 1) signal.Notify(quit, syscall.SIGINT, syscall.SIGTERM) <-quit log.Println("收到退出信号,正在关闭...") c.Stop() log.Println("采集器已停止") } // runCollection 执行一次数据采集和存储 func runCollection(coll *collector.Collector, store *storage.Storage) { log.Println("开始采集...") // 计算Ymd now := time.Now() ymd := now.Year()*10000 + int(now.Month())*100 + now.Day() // 采集数据并检查变化 status, dataHash, changed, err := coll.CollectAndCheck() if err != nil { log.Printf("采集失败: %v", err) // 记录失败的日志 if err := store.SaveCollectionLog("", ymd, false, err.Error()); err != nil { log.Printf("保存采集日志失败: %v", err) } return } log.Printf("数据哈希: %s", dataHash) log.Printf("数据是否变化: %v", changed) // 如果数据没有变化,只记录日志 if !changed { log.Println("数据未变化,跳过存储") if err := store.SaveCollectionLog(dataHash, ymd, false, "数据未变化"); err != nil { log.Printf("保存采集日志失败: %v", err) } return } // 数据有变化,保存到数据库 log.Println("数据已变化,开始存储到数据库...") if err := store.SaveStatus(status, dataHash); err != nil { log.Printf("保存数据失败: %v", err) // 记录失败的日志 if err := store.SaveCollectionLog(dataHash, ymd, true, err.Error()); err != nil { log.Printf("保存采集日志失败: %v", err) } return } // 记录成功的日志 if err := store.SaveCollectionLog(dataHash, ymd, true, "数据保存成功"); err != nil { log.Printf("保存采集日志失败: %v", err) } log.Printf("数据存储成功 - 资产账户: %s, 订单数: %d, 持仓数: %d, 行情数: %d", status.Data.Assets.AccountID, len(status.Data.Orders), len(status.Data.Positions), len(status.Data.TickData)) } // getEnv 获取环境变量,如果不存在则返回默认值 func getEnv(key, defaultValue string) string { value := os.Getenv(key) if value == "" { return defaultValue } return value } // getEnvAsInt 获取环境变量并转换为整数 func getEnvAsInt(key string, defaultValue int) int { value := os.Getenv(key) if value == "" { return defaultValue } result, err := strconv.Atoi(value) if err != nil { log.Printf("环境变量 %s 转换失败: %v,使用默认值 %d", key, err, defaultValue) return defaultValue } return result }