quant/collector
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

optz

Browse Source
This commit is contained in:
yanweidong
2026-04-07 14:46:49 +08:00
parent aaa8308a90
commit 35e52bc5ea
14 changed files with 234 additions and 620 deletions
Download Patch File Download Diff File Expand all files Collapse all files

97
storage/storage.go
View File

@@ -6,8 +6,10 @@ import (
"time"
"git.apinb.com/quant/collector/models"
"git.apinb.com/quant/collector/types"
"gorm.io/driver/postgres"
"gorm.io/gorm"
"gorm.io/gorm/clause"
"gorm.io/gorm/logger"
)
@@ -68,11 +70,11 @@ func (s *Storage) AutoMigrate() error {
log.Println("开始自动迁移数据库表结构...")
err := s.db.AutoMigrate(
&models.AssetSnapshot{},
&models.OrderRecord{},
&models.PositionRecord{},
&models.TickRecord{},
&models.CollectionLog{},
&models.CollectorAssets{},
&models.CollectorOrder{},
&models.CollectorPosition{},
&models.CollectorTick{},
&models.CollectorLog{},
)
if err != nil {
@@ -84,31 +86,42 @@ func (s *Storage) AutoMigrate() error {
}
// SaveStatus 保存完整状态数据(使用事务)
func (s *Storage) SaveStatus(status *models.Status, dataHash string) error {
func (s *Storage) SaveStatus(status *types.Status, dataHash string) error {
return s.db.Transaction(func(tx *gorm.DB) error {
// 保存资产快照
asset := models.AssetSnapshot{
// 计算Ymd (年月日数字格式,如20260407)
now := time.Now()
ymd := now.Year()*10000 + int(now.Month())*100 + now.Day()
// 保存资产快照 (Upsert: 存在则更新,不存在则插入)
asset := models.CollectorAssets{
AccountID: status.Data.Assets.AccountID,
Ymd: ymd,
Cash: status.Data.Assets.Cash,
FrozenCash: status.Data.Assets.FrozenCash,
MarketValue: status.Data.Assets.MarketValue,
Profit: status.Data.Assets.Profit,
TotalAsset: status.Data.Assets.TotalAsset,
DataHash: dataHash,
CollectedAt: time.Now(),
CollectedAt: now,
}
if err := tx.Create(&asset).Error; err != nil {
// 使用GORM的Clauses实现Upsert
if err := tx.Clauses(clause.OnConflict{
Columns: []clause.Column{{Name: "account_id"}, {Name: "ymd"}},
DoUpdates: clause.AssignmentColumns([]string{"cash", "frozen_cash", "market_value", "profit", "total_asset", "data_hash", "collected_at"}),
}).Create(&asset).Error; err != nil {
return fmt.Errorf("保存资产快照失败: %w", err)
}
// 批量保存订单
// 批量保存订单 (Upsert: 存在则更新,不存在则插入)
if len(status.Data.Orders) > 0 {
orders := make([]models.OrderRecord, 0, len(status.Data.Orders))
orders := make([]models.CollectorOrder, 0, len(status.Data.Orders))
for _, order := range status.Data.Orders {
orders = append(orders, models.OrderRecord{
orders = append(orders, models.CollectorOrder{
OrderID: order.OrderID,
AccountID: status.Data.Assets.AccountID,
StockCode: order.StockCode,
Ymd: ymd,
Price: order.Price,
Volume: order.Volume,
TradedPrice: order.TradedPrice,
@@ -117,21 +130,30 @@ func (s *Storage) SaveStatus(status *models.Status, dataHash string) error {
OrderTime: order.OrderTime,
OrderRemark: order.OrderRemark,
DataHash: dataHash,
CollectedAt: time.Now(),
CollectedAt: now,
})
}
if err := tx.CreateInBatches(orders, 100).Error; err != nil {
// 使用Upsert逻辑: 以 account_id, order_id, ymd 为条件
if err := tx.Clauses(clause.OnConflict{
Columns: []clause.Column{{Name: "account_id"}, {Name: "order_id"}, {Name: "ymd"}},
DoUpdates: clause.AssignmentColumns([]string{
"stock_code", "price", "volume",
"traded_price", "traded_volume", "order_status",
"order_time", "order_remark", "data_hash", "collected_at",
}),
}).Create(&orders).Error; err != nil {
return fmt.Errorf("保存订单失败: %w", err)
}
}
// 批量保存持仓
// 批量保存持仓 (Upsert: 存在则更新,不存在则插入)
if len(status.Data.Positions) > 0 {
positions := make([]models.PositionRecord, 0, len(status.Data.Positions))
positions := make([]models.CollectorPosition, 0, len(status.Data.Positions))
for _, pos := range status.Data.Positions {
positions = append(positions, models.PositionRecord{
positions = append(positions, models.CollectorPosition{
AccountID: status.Data.Assets.AccountID,
Code: pos.Code,
Ymd: ymd,
Volume: pos.Volume,
CanUseVolume: pos.CanUseVolume,
FrozenVolume: pos.FrozenVolume,
@@ -143,20 +165,30 @@ func (s *Storage) SaveStatus(status *models.Status, dataHash string) error {
ProfitRate: pos.ProfitRate,
MinProfitRate: pos.MinProfitRate,
DataHash: dataHash,
CollectedAt: time.Now(),
CollectedAt: now,
})
}
if err := tx.CreateInBatches(positions, 100).Error; err != nil {
// 使用Upsert逻辑: 以 account_id, code, ymd 为条件
if err := tx.Clauses(clause.OnConflict{
Columns: []clause.Column{{Name: "account_id"}, {Name: "code"}, {Name: "ymd"}},
DoUpdates: clause.AssignmentColumns([]string{
"volume", "can_use_volume", "frozen_volume",
"avg_price", "open_price", "current_price",
"market_value", "profit", "profit_rate",
"min_profit_rate", "data_hash", "collected_at",
}),
}).Create(&positions).Error; err != nil {
return fmt.Errorf("保存持仓失败: %w", err)
}
}
// 批量保存行情数据
// 批量保存行情数据 (Upsert: 存在则更新,不存在则插入)
if len(status.Data.TickData) > 0 {
ticks := make([]models.TickRecord, 0, len(status.Data.TickData))
ticks := make([]models.CollectorTick, 0, len(status.Data.TickData))
for code, tick := range status.Data.TickData {
ticks = append(ticks, models.TickRecord{
ticks = append(ticks, models.CollectorTick{
StockCode: code,
Ymd: ymd,
LastPrice: tick.LastPrice,
Open: tick.Open,
High: tick.High,
@@ -173,10 +205,20 @@ func (s *Storage) SaveStatus(status *models.Status, dataHash string) error {
TimeTag: tick.TimeTag,
StockStatus: tick.StockStatus,
DataHash: dataHash,
CollectedAt: time.Now(),
CollectedAt: now,
})
}
if err := tx.CreateInBatches(ticks, 100).Error; err != nil {
// 使用Upsert逻辑: 以 stock_code, ymd 为条件
if err := tx.Clauses(clause.OnConflict{
Columns: []clause.Column{{Name: "stock_code"}, {Name: "ymd"}},
DoUpdates: clause.AssignmentColumns([]string{
"last_price", "open", "high", "low", "last_close",
"volume", "amount", "pvolume",
"bid_prices", "bid_volumes", "ask_prices", "ask_volumes",
"time", "timetag", "stock_status",
"data_hash", "collected_at",
}),
}).Create(&ticks).Error; err != nil {
return fmt.Errorf("保存行情数据失败: %w", err)
}
}
@@ -186,9 +228,10 @@ func (s *Storage) SaveStatus(status *models.Status, dataHash string) error {
}
// SaveCollectionLog 保存采集日志
func (s *Storage) SaveCollectionLog(dataHash string, hasChanged bool, statusMessage string) error {
log := models.CollectionLog{
func (s *Storage) SaveCollectionLog(dataHash string, ymd int, hasChanged bool, statusMessage string) error {
log := models.CollectorLog{
DataHash: dataHash,
Ymd: ymd,
HasChanged: hasChanged,
StatusMessage: statusMessage,
CollectedAt: time.Now(),