feat: migrate store layer to GORM with PostgreSQL support

- Migrate all store packages from raw database/sql to GORM ORM - Add PostgreSQL support alongside SQLite - Move EncryptedString type to crypto package for cleaner architecture - Add automatic encryption/decryption for sensitive fields (API keys, secrets) - Fix PostgreSQL AutoMigrate conflicts by skipping existing tables - Fix duplicate /klines route registration - Update tests to use GORM database connections - Add database configuration support in config package
2026-07-16 09:11:03 +08:00 · 2026-01-01 19:32:49 +08:00
parent d547863ebb
commit 2d272bb7b8
32 changed files with 2573 additions and 3771 deletions
--- a/store/equity.go
+++ b/store/equity.go
@@ -1,55 +1,48 @@
 package store

 import (
-	"database/sql"
 	"fmt"
 	"time"
+
+	"gorm.io/gorm"
 )

 // EquityStore account equity storage (for plotting return curves)
 type EquityStore struct {
-	db *sql.DB
+	db *gorm.DB
 }

 // EquitySnapshot equity snapshot
 type EquitySnapshot struct {
-	ID            int64     `json:"id"`
-	TraderID      string    `json:"trader_id"`
-	Timestamp     time.Time `json:"timestamp"`
-	TotalEquity   float64   `json:"total_equity"`    // Account equity (balance + unrealized PnL)
-	Balance       float64   `json:"balance"`         // Account balance
-	UnrealizedPnL float64   `json:"unrealized_pnl"`  // Unrealized profit and loss
-	PositionCount int       `json:"position_count"`  // Position count
-	MarginUsedPct float64   `json:"margin_used_pct"` // Margin usage percentage
+	ID            int64     `gorm:"primaryKey;autoIncrement" json:"id"`
+	TraderID      string    `gorm:"column:trader_id;not null;index:idx_equity_trader_time" json:"trader_id"`
+	Timestamp     time.Time `gorm:"not null;index:idx_equity_trader_time,sort:desc;index:idx_equity_timestamp,sort:desc" json:"timestamp"`
+	TotalEquity   float64   `gorm:"column:total_equity;not null;default:0" json:"total_equity"`
+	Balance       float64   `gorm:"not null;default:0" json:"balance"`
+	UnrealizedPnL float64   `gorm:"column:unrealized_pnl;not null;default:0" json:"unrealized_pnl"`
+	PositionCount int       `gorm:"column:position_count;default:0" json:"position_count"`
+	MarginUsedPct float64   `gorm:"column:margin_used_pct;default:0" json:"margin_used_pct"`
+	CreatedAt     time.Time `json:"created_at"`
+}
+
+func (EquitySnapshot) TableName() string { return "trader_equity_snapshots" }
+
+// NewEquityStore creates a new EquityStore
+func NewEquityStore(db *gorm.DB) *EquityStore {
+	return &EquityStore{db: db}
 }

 // initTables initializes equity tables
 func (s *EquityStore) initTables() error {
-	queries := []string{
-		// Equity snapshot table - specifically for return curves
-		`CREATE TABLE IF NOT EXISTS trader_equity_snapshots (
-			id INTEGER PRIMARY KEY AUTOINCREMENT,
-			trader_id TEXT NOT NULL,
-			timestamp DATETIME NOT NULL,
-			total_equity REAL NOT NULL DEFAULT 0,
-			balance REAL NOT NULL DEFAULT 0,
-			unrealized_pnl REAL NOT NULL DEFAULT 0,
-			position_count INTEGER DEFAULT 0,
-			margin_used_pct REAL DEFAULT 0,
-			created_at DATETIME DEFAULT CURRENT_TIMESTAMP
-		)`,
-		// Indexes
-		`CREATE INDEX IF NOT EXISTS idx_equity_trader_time ON trader_equity_snapshots(trader_id, timestamp DESC)`,
-		`CREATE INDEX IF NOT EXISTS idx_equity_timestamp ON trader_equity_snapshots(timestamp DESC)`,
-	}
-
-	for _, query := range queries {
-		if _, err := s.db.Exec(query); err != nil {
-			return fmt.Errorf("failed to execute SQL: %w", err)
+	// For PostgreSQL with existing table, skip AutoMigrate
+	if s.db.Dialector.Name() == "postgres" {
+		var tableExists int64
+		s.db.Raw(`SELECT COUNT(*) FROM information_schema.tables WHERE table_name = 'trader_equity_snapshots'`).Scan(&tableExists)
+		if tableExists > 0 {
+			return nil
 		}
 	}
-
-	return nil
+	return s.db.AutoMigrate(&EquitySnapshot{})
 }

 // Save saves equity snapshot
@@ -60,58 +53,22 @@ func (s *EquityStore) Save(snapshot *EquitySnapshot) error {
 		snapshot.Timestamp = snapshot.Timestamp.UTC()
 	}

-	result, err := s.db.Exec(`
-		INSERT INTO trader_equity_snapshots (
-			trader_id, timestamp, total_equity, balance,
-			unrealized_pnl, position_count, margin_used_pct
-		) VALUES (?, ?, ?, ?, ?, ?, ?)
-	`,
-		snapshot.TraderID,
-		snapshot.Timestamp.Format(time.RFC3339),
-		snapshot.TotalEquity,
-		snapshot.Balance,
-		snapshot.UnrealizedPnL,
-		snapshot.PositionCount,
-		snapshot.MarginUsedPct,
-	)
-	if err != nil {
+	if err := s.db.Create(snapshot).Error; err != nil {
 		return fmt.Errorf("failed to save equity snapshot: %w", err)
 	}
-
-	id, _ := result.LastInsertId()
-	snapshot.ID = id
 	return nil
 }

 // GetLatest gets the latest N equity records for specified trader (sorted in ascending chronological order: old to new)
 func (s *EquityStore) GetLatest(traderID string, limit int) ([]*EquitySnapshot, error) {
-	rows, err := s.db.Query(`
-		SELECT id, trader_id, timestamp, total_equity, balance,
-		       unrealized_pnl, position_count, margin_used_pct
-		FROM trader_equity_snapshots
-		WHERE trader_id = ?
-		ORDER BY timestamp DESC
-		LIMIT ?
-	`, traderID, limit)
+	var snapshots []*EquitySnapshot
+	err := s.db.Where("trader_id = ?", traderID).
+		Order("timestamp DESC").
+		Limit(limit).
+		Find(&snapshots).Error
 	if err != nil {
 		return nil, fmt.Errorf("failed to query equity records: %w", err)
 	}
-	defer rows.Close()
-
-	var snapshots []*EquitySnapshot
-	for rows.Next() {
-		snap := &EquitySnapshot{}
-		var timestampStr string
-		err := rows.Scan(
-			&snap.ID, &snap.TraderID, &timestampStr, &snap.TotalEquity,
-			&snap.Balance, &snap.UnrealizedPnL, &snap.PositionCount, &snap.MarginUsedPct,
-		)
-		if err != nil {
-			continue
-		}
-		snap.Timestamp, _ = time.Parse(time.RFC3339, timestampStr)
-		snapshots = append(snapshots, snap)
-	}

 	// Reverse the array to sort time from old to new (suitable for plotting curves)
 	for i, j := 0, len(snapshots)-1; i < j; i, j = i+1, j-1 {
@@ -123,116 +80,81 @@ func (s *EquityStore) GetLatest(traderID string, limit int) ([]*EquitySnapshot,

 // GetByTimeRange gets equity records within specified time range
 func (s *EquityStore) GetByTimeRange(traderID string, start, end time.Time) ([]*EquitySnapshot, error) {
-	rows, err := s.db.Query(`
-		SELECT id, trader_id, timestamp, total_equity, balance,
-		       unrealized_pnl, position_count, margin_used_pct
-		FROM trader_equity_snapshots
-		WHERE trader_id = ? AND timestamp >= ? AND timestamp <= ?
-		ORDER BY timestamp ASC
-	`, traderID, start.Format(time.RFC3339), end.Format(time.RFC3339))
+	var snapshots []*EquitySnapshot
+	err := s.db.Where("trader_id = ? AND timestamp >= ? AND timestamp <= ?", traderID, start, end).
+		Order("timestamp ASC").
+		Find(&snapshots).Error
 	if err != nil {
 		return nil, fmt.Errorf("failed to query equity records: %w", err)
 	}
-	defer rows.Close()
-
-	var snapshots []*EquitySnapshot
-	for rows.Next() {
-		snap := &EquitySnapshot{}
-		var timestampStr string
-		err := rows.Scan(
-			&snap.ID, &snap.TraderID, &timestampStr, &snap.TotalEquity,
-			&snap.Balance, &snap.UnrealizedPnL, &snap.PositionCount, &snap.MarginUsedPct,
-		)
-		if err != nil {
-			continue
-		}
-		snap.Timestamp, _ = time.Parse(time.RFC3339, timestampStr)
-		snapshots = append(snapshots, snap)
-	}
-
 	return snapshots, nil
 }

 // GetAllTradersLatest gets latest equity for all traders (for leaderboards)
 func (s *EquityStore) GetAllTradersLatest() (map[string]*EquitySnapshot, error) {
-	rows, err := s.db.Query(`
+	// Use raw SQL for this complex query with subquery
+	var snapshots []*EquitySnapshot
+	err := s.db.Raw(`
 		SELECT e.id, e.trader_id, e.timestamp, e.total_equity, e.balance,
-		       e.unrealized_pnl, e.position_count, e.margin_used_pct
+		       e.unrealized_pnl, e.position_count, e.margin_used_pct, e.created_at
 		FROM trader_equity_snapshots e
 		INNER JOIN (
 			SELECT trader_id, MAX(timestamp) as max_ts
 			FROM trader_equity_snapshots
 			GROUP BY trader_id
 		) latest ON e.trader_id = latest.trader_id AND e.timestamp = latest.max_ts
-	`)
+	`).Scan(&snapshots).Error
 	if err != nil {
 		return nil, fmt.Errorf("failed to query latest equity: %w", err)
 	}
-	defer rows.Close()

 	result := make(map[string]*EquitySnapshot)
-	for rows.Next() {
-		snap := &EquitySnapshot{}
-		var timestampStr string
-		err := rows.Scan(
-			&snap.ID, &snap.TraderID, &timestampStr, &snap.TotalEquity,
-			&snap.Balance, &snap.UnrealizedPnL, &snap.PositionCount, &snap.MarginUsedPct,
-		)
-		if err != nil {
-			continue
-		}
-		snap.Timestamp, _ = time.Parse(time.RFC3339, timestampStr)
+	for _, snap := range snapshots {
 		result[snap.TraderID] = snap
 	}
-
 	return result, nil
 }

 // CleanOldRecords cleans old records from N days ago
 func (s *EquityStore) CleanOldRecords(traderID string, days int) (int64, error) {
-	cutoffTime := time.Now().AddDate(0, 0, -days).Format(time.RFC3339)
+	cutoffTime := time.Now().AddDate(0, 0, -days)

-	result, err := s.db.Exec(`
-		DELETE FROM trader_equity_snapshots
-		WHERE trader_id = ? AND timestamp < ?
-	`, traderID, cutoffTime)
-	if err != nil {
-		return 0, fmt.Errorf("failed to clean old records: %w", err)
+	result := s.db.Where("trader_id = ? AND timestamp < ?", traderID, cutoffTime).
+		Delete(&EquitySnapshot{})
+	if result.Error != nil {
+		return 0, fmt.Errorf("failed to clean old records: %w", result.Error)
 	}
-
-	return result.RowsAffected()
+	return result.RowsAffected, nil
 }

 // GetCount gets record count for specified trader
 func (s *EquityStore) GetCount(traderID string) (int, error) {
-	var count int
-	err := s.db.QueryRow(`
-		SELECT COUNT(*) FROM trader_equity_snapshots WHERE trader_id = ?
-	`, traderID).Scan(&count)
-	return count, err
+	var count int64
+	err := s.db.Model(&EquitySnapshot{}).Where("trader_id = ?", traderID).Count(&count).Error
+	return int(count), err
 }

 // MigrateFromDecision migrates data from old decision_account_snapshots table
 func (s *EquityStore) MigrateFromDecision() (int64, error) {
 	// Check if migration is needed (whether new table is empty)
-	var count int
-	s.db.QueryRow(`SELECT COUNT(*) FROM trader_equity_snapshots`).Scan(&count)
+	var count int64
+	s.db.Model(&EquitySnapshot{}).Count(&count)
 	if count > 0 {
 		return 0, nil // Already has data, skip migration
 	}

-	// Check if old table exists
+	// Check if old table exists (SQLite specific check, but works for migration)
 	var tableName string
-	err := s.db.QueryRow(`
+	err := s.db.Raw(`
 		SELECT name FROM sqlite_master
 		WHERE type='table' AND name='decision_account_snapshots'
-	`).Scan(&tableName)
-	if err != nil {
+	`).Scan(&tableName).Error
+	if err != nil || tableName == "" {
 		return 0, nil // Old table doesn't exist, skip
 	}

 	// Migrate data: join query from decision_records + decision_account_snapshots
-	result, err := s.db.Exec(`
+	result := s.db.Exec(`
 		INSERT INTO trader_equity_snapshots (
 			trader_id, timestamp, total_equity, balance,
 			unrealized_pnl, position_count, margin_used_pct
@@ -249,9 +171,9 @@ func (s *EquityStore) MigrateFromDecision() (int64, error) {
 		JOIN decision_account_snapshots das ON dr.id = das.decision_id
 		ORDER BY dr.timestamp ASC
 	`)
-	if err != nil {
-		return 0, fmt.Errorf("failed to migrate data: %w", err)
+	if result.Error != nil {
+		return 0, fmt.Errorf("failed to migrate data: %w", result.Error)
 	}

-	return result.RowsAffected()
+	return result.RowsAffected, nil
 }