market-data-service/internal/repository/stock.go

package repository

import (
	"context"
	"database/sql"
	"fmt"
	"strings"
	"time"

	"market-data-service/api"
)

// DB PostgreSQL连接
type DB struct {
	*sql.DB
}

// NewDB 创建数据库连接
func NewDB(connStr string) (*DB, error) {
	db, err := sql.Open("postgres", connStr)
	if err != nil {
		return nil, err
	}

	if err := db.Ping(); err != nil {
		return nil, err
	}

	return &DB{db}, nil
}

// ============================================
// 股票Repository
// ============================================

// StockRepository 股票数据仓库
type StockRepository struct {
	db *DB
}

// NewStockRepository 创建股票Repository
func NewStockRepository(db *DB) *StockRepository {
	return &StockRepository{db: db}
}

// GetKLines 获取K线数据
func (r *StockRepository) GetKLines(ctx context.Context, symbol string, freq api.Frequency, start, end time.Time, adjust api.AdjustType) ([]api.KLineItem, error) {
	tableName := fmt.Sprintf("stock.klines_%s", freq)
	
	query := fmt.Sprintf(`
		SELECT ts, open, high, low, close, volume, amount
		FROM %s
		WHERE symbol_id = $1 AND ts >= $2 AND ts <= $3
		ORDER BY ts ASC
	`, tableName)

	rows, err := r.db.QueryContext(ctx, query, symbol, start, end)
	if err != nil {
		return nil, err
	}
	defer rows.Close()

	var items []api.KLineItem
	for rows.Next() {
		var item api.KLineItem
		if err := rows.Scan(&item.Time, &item.Open, &item.High, &item.Low, &item.Close, &item.Volume, &item.Amount); err != nil {
			return nil, err
		}
		items = append(items, item)
	}

	return items, rows.Err()
}

// SaveKLines 保存K线数据
func (r *StockRepository) SaveKLines(ctx context.Context, freq api.Frequency, items []api.KLineItem) error {
	if len(items) == 0 {
		return nil
	}

	tableName := fmt.Sprintf("stock.klines_%s", freq)
	
	// 使用批量插入
	valueStrs := make([]string, 0, len(items))
	args := make([]interface{}, 0, len(items)*7)
	argIdx := 1

	for _, item := range items {
		valueStrs = append(valueStrs, fmt.Sprintf("($%d, $%d, $%d, $%d, $%d, $%d, $%d)", 
			argIdx, argIdx+1, argIdx+2, argIdx+3, argIdx+4, argIdx+5, argIdx+6))
		args = append(args, item.Symbol, item.Time, item.Open, item.High, item.Low, item.Close, item.Volume, item.Amount)
		argIdx += 7
	}

	query := fmt.Sprintf(`
		INSERT INTO %s (symbol_id, ts, open, high, low, close, volume, amount)
		VALUES %s
		ON CONFLICT (symbol_id, ts) DO UPDATE SET
			open = EXCLUDED.open,
			high = EXCLUDED.high,
			low = EXCLUDED.low,
			close = EXCLUDED.close,
			volume = EXCLUDED.volume,
			amount = EXCLUDED.amount
	`, tableName, strings.Join(valueStrs, ","))

	_, err := r.db.ExecContext(ctx, query, args...)
	return err
}

// ListSymbols 查询标的列表
func (r *StockRepository) ListSymbols(ctx context.Context, req *api.SymbolListRequest) ([]api.Symbol, int, error) {
	whereClause := "WHERE 1=1"
	args := []interface{}{}
	argIdx := 1

	if req.Exchange != "" {
		whereClause += fmt.Sprintf(" AND exchange = $%d", argIdx)
		args = append(args, req.Exchange)
		argIdx++
	}

	if req.Keyword != "" {
		whereClause += fmt.Sprintf(" AND (symbol_id ILIKE $%d OR name ILIKE $%d)", argIdx, argIdx)
		args = append(args, "%"+req.Keyword+"%")
		argIdx++
	}

	// 查询总数
	countQuery := "SELECT COUNT(*) FROM stock.symbols " + whereClause
	var total int
	if err := r.db.QueryRowContext(ctx, countQuery, args...).Scan(&total); err != nil {
		return nil, 0, err
	}

	// 查询数据
	query := fmt.Sprintf(`
		SELECT symbol_id, symbol_type, exchange, name, name_en, list_date, delist_date, industry, status
		FROM stock.symbols
		%s
		ORDER BY symbol_id
		LIMIT $%d OFFSET $%d
	`, whereClause, argIdx, argIdx+1)
	args = append(args, req.Size, (req.Page-1)*req.Size)

	rows, err := r.db.QueryContext(ctx, query, args...)
	if err != nil {
		return nil, 0, err
	}
	defer rows.Close()

	var symbols []api.Symbol
	for rows.Next() {
		var s api.Symbol
		var listDate, delistDate sql.NullTime
		if err := rows.Scan(&s.SymbolID, &s.SymbolType, &s.Exchange, &s.Name, &s.NameEN,
			&listDate, &delistDate, &s.Industry, &s.Status); err != nil {
			return nil, 0, err
		}
		if listDate.Valid {
			s.ListDate = &listDate.Time
		}
		if delistDate.Valid {
			s.DelistDate = &delistDate.Time
		}
		symbols = append(symbols, s)
	}

	return symbols, total, rows.Err()
}

// GetTradingDates 获取交易日历
func (r *StockRepository) GetTradingDates(ctx context.Context, start, end string) (*api.TradingDatesData, error) {
	query := `
		SELECT trade_date
		FROM stock.trading_calendar
		WHERE trade_date >= $1 AND trade_date <= $2 AND is_trading_day = true
		ORDER BY trade_date ASC
	`

	rows, err := r.db.QueryContext(ctx, query, start, end)
	if err != nil {
		return nil, err
	}
	defer rows.Close()

	var dates []string
	for rows.Next() {
		var date string
		if err := rows.Scan(&date); err != nil {
			return nil, err
		}
		dates = append(dates, date)
	}

	// 计算总天数
	startDate, _ := time.Parse("20060102", start)
	endDate, _ := time.Parse("20060102", end)
	totalDays := int(endDate.Sub(startDate).Hours()/24) + 1

	return &api.TradingDatesData{
		Start:        start,
		End:          end,
		TotalDays:    totalDays,
		TradingDays:  len(dates),
		TradingDates: dates,
	}, rows.Err()
}

// SaveSymbols 保存标的列表
func (r *StockRepository) SaveSymbols(ctx context.Context, symbols []api.Symbol) error {
	if len(symbols) == 0 {
		return nil
	}

	tx, err := r.db.BeginTx(ctx, nil)
	if err != nil {
		return err
	}
	defer tx.Rollback()

	stmt, err := tx.PrepareContext(ctx, `
		INSERT INTO stock.symbols (symbol_id, symbol_type, exchange, name, name_en, list_date, delist_date, industry, status)
		VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9)
		ON CONFLICT (symbol_id) DO UPDATE SET
			name = EXCLUDED.name,
			name_en = EXCLUDED.name_en,
			list_date = EXCLUDED.list_date,
			delist_date = EXCLUDED.delist_date,
			industry = EXCLUDED.industry,
			status = EXCLUDED.status,
			updated_at = NOW()
	`)
	if err != nil {
		return err
	}
	defer stmt.Close()

	for _, s := range symbols {
		var listDate, delistDate interface{}
		if s.ListDate != nil {
			listDate = *s.ListDate
		}
		if s.DelistDate != nil {
			delistDate = *s.DelistDate
		}

		_, err := stmt.ExecContext(ctx, s.SymbolID, s.SymbolType, s.Exchange, s.Name, s.NameEN,
			listDate, delistDate, s.Industry, s.Status)
		if err != nil {
			return err
		}
	}

	return tx.Commit()
}

// SaveTradingCalendar 保存交易日历
func (r *StockRepository) SaveTradingCalendar(ctx context.Context, dates []api.TradeCalData) error {
	if len(dates) == 0 {
		return nil
	}

	tx, err := r.db.BeginTx(ctx, nil)
	if err != nil {
		return err
	}
	defer tx.Rollback()

	stmt, err := tx.PrepareContext(ctx, `
		INSERT INTO stock.trading_calendar (trade_date, is_trading_day, week_day)
		VALUES ($1, $2, $3)
		ON CONFLICT (trade_date) DO UPDATE SET
			is_trading_day = EXCLUDED.is_trading_day,
			week_day = EXCLUDED.week_day,
			updated_at = NOW()
	`)
	if err != nil {
		return err
	}
	defer stmt.Close()

	for _, d := range dates {
		_, err := stmt.ExecContext(ctx, d.Date.Format("2006-01-02"), d.IsTradingDay, int(d.Date.Weekday())+1)
		if err != nil {
			return err
		}
	}

	return tx.Commit()
}