280 lines
8.0 KiB
Go
280 lines
8.0 KiB
Go
package storage
|
|
|
|
import (
|
|
"fmt"
|
|
"log"
|
|
"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"
|
|
)
|
|
|
|
// Storage 数据库存储器
|
|
type Storage struct {
|
|
db *gorm.DB
|
|
}
|
|
|
|
// NewStorage 创建新的数据库连接
|
|
func NewStorage(connStr string) (*Storage, error) {
|
|
// 配置GORM日志
|
|
newLogger := logger.New(
|
|
log.New(log.Writer(), "\r\n", log.LstdFlags),
|
|
logger.Config{
|
|
SlowThreshold: time.Second,
|
|
LogLevel: logger.Warn,
|
|
IgnoreRecordNotFoundError: true,
|
|
Colorful: true,
|
|
},
|
|
)
|
|
|
|
db, err := gorm.Open(postgres.Open(connStr), &gorm.Config{
|
|
Logger: newLogger,
|
|
})
|
|
if err != nil {
|
|
return nil, fmt.Errorf("打开数据库连接失败: %w", err)
|
|
}
|
|
|
|
// 获取底层的sql.DB以设置连接池
|
|
sqlDB, err := db.DB()
|
|
if err != nil {
|
|
return nil, fmt.Errorf("获取数据库实例失败: %w", err)
|
|
}
|
|
|
|
// 设置连接池参数
|
|
sqlDB.SetMaxOpenConns(25)
|
|
sqlDB.SetMaxIdleConns(5)
|
|
sqlDB.SetConnMaxLifetime(5 * time.Minute)
|
|
|
|
log.Println("数据库连接成功")
|
|
return &Storage{db: db}, nil
|
|
}
|
|
|
|
// Close 关闭数据库连接
|
|
func (s *Storage) Close() error {
|
|
if s.db != nil {
|
|
sqlDB, err := s.db.DB()
|
|
if err != nil {
|
|
return err
|
|
}
|
|
return sqlDB.Close()
|
|
}
|
|
return nil
|
|
}
|
|
|
|
// AutoMigrate 自动迁移数据库表结构
|
|
func (s *Storage) AutoMigrate() error {
|
|
log.Println("开始自动迁移数据库表结构...")
|
|
|
|
err := s.db.AutoMigrate(
|
|
&models.CollectorAssets{},
|
|
&models.CollectorOrder{},
|
|
&models.CollectorPosition{},
|
|
&models.CollectorTick{},
|
|
&models.CollectorLog{},
|
|
)
|
|
|
|
if err != nil {
|
|
return fmt.Errorf("自动迁移失败: %w", err)
|
|
}
|
|
|
|
log.Println("数据库表结构迁移完成")
|
|
return nil
|
|
}
|
|
|
|
// SaveStatus 保存完整状态数据(使用事务)
|
|
func (s *Storage) SaveStatus(status *types.Status, dataHash string) error {
|
|
// 验证必要字段 - AccountID是所有Upsert的共同条件
|
|
if status.Data.Assets.AccountID == "" {
|
|
return fmt.Errorf("账户ID不能为空")
|
|
}
|
|
|
|
return s.db.Transaction(func(tx *gorm.DB) error {
|
|
// 计算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: now,
|
|
}
|
|
|
|
// 使用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.CollectorOrder, 0, len(status.Data.Orders))
|
|
for _, order := range status.Data.Orders {
|
|
// 验证Upsert必要条件: OrderID和StockCode
|
|
if order.OrderID == 0 {
|
|
continue
|
|
}
|
|
if order.StockCode == "" {
|
|
continue
|
|
}
|
|
|
|
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,
|
|
TradedVolume: order.TradedVolume,
|
|
OrderStatus: order.OrderStatus,
|
|
OrderTime: order.OrderTime,
|
|
OrderRemark: order.OrderRemark,
|
|
DataHash: dataHash,
|
|
CollectedAt: now,
|
|
})
|
|
}
|
|
// 使用Upsert逻辑: 以 account_id, order_id, ymd 为条件
|
|
if len(orders) > 0 {
|
|
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.CollectorPosition, 0, len(status.Data.Positions))
|
|
for _, pos := range status.Data.Positions {
|
|
// 验证Upsert必要条件: Code
|
|
if pos.Code == "" {
|
|
continue
|
|
}
|
|
|
|
positions = append(positions, models.CollectorPosition{
|
|
AccountID: status.Data.Assets.AccountID,
|
|
Code: pos.Code,
|
|
Ymd: ymd,
|
|
Volume: pos.Volume,
|
|
CanUseVolume: pos.CanUseVolume,
|
|
FrozenVolume: pos.FrozenVolume,
|
|
AvgPrice: pos.AvgPrice,
|
|
OpenPrice: pos.OpenPrice,
|
|
CurrentPrice: pos.CurrentPrice,
|
|
MarketValue: pos.MarketValue,
|
|
Profit: pos.Profit,
|
|
ProfitRate: pos.ProfitRate,
|
|
MinProfitRate: pos.MinProfitRate,
|
|
DataHash: dataHash,
|
|
CollectedAt: now,
|
|
})
|
|
}
|
|
// 使用Upsert逻辑: 以 account_id, code, ymd 为条件
|
|
if len(positions) > 0 {
|
|
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.CollectorTick, 0, len(status.Data.TickData))
|
|
for code, tick := range status.Data.TickData {
|
|
// 验证Upsert必要条件: StockCode
|
|
if code == "" {
|
|
continue
|
|
}
|
|
|
|
ticks = append(ticks, models.CollectorTick{
|
|
StockCode: code,
|
|
Ymd: ymd,
|
|
LastPrice: tick.LastPrice,
|
|
Open: tick.Open,
|
|
High: tick.High,
|
|
Low: tick.Low,
|
|
LastClose: tick.LastClose,
|
|
Volume: tick.Volume,
|
|
Amount: tick.Amount,
|
|
PVolume: tick.PVolume,
|
|
BidPrices: tick.BidPrice,
|
|
BidVolumes: tick.BidVol,
|
|
AskPrices: tick.AskPrice,
|
|
AskVolumes: tick.AskVol,
|
|
Time: tick.Time,
|
|
TimeTag: tick.TimeTag,
|
|
StockStatus: tick.StockStatus,
|
|
DataHash: dataHash,
|
|
CollectedAt: now,
|
|
})
|
|
}
|
|
// 使用Upsert逻辑: 以 stock_code, ymd 为条件
|
|
if len(ticks) > 0 {
|
|
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)
|
|
}
|
|
}
|
|
}
|
|
|
|
return nil
|
|
})
|
|
}
|
|
|
|
// SaveCollectionLog 保存采集日志
|
|
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(),
|
|
}
|
|
|
|
if err := s.db.Create(&log).Error; err != nil {
|
|
return fmt.Errorf("保存采集日志失败: %w", err)
|
|
}
|
|
|
|
return nil
|
|
}
|
|
|
|
// GetDB 获取GORM DB实例(用于高级查询)
|
|
func (s *Storage) GetDB() *gorm.DB {
|
|
return s.db
|
|
}
|