Veyra/internal/repository/backtest.go

package repository

import (
	"context"
	"crypto/sha256"
	"encoding/hex"
	"encoding/json"
	"git.sundynix.cn/Blizzard/Veyra/internal/apperror"
	"git.sundynix.cn/Blizzard/Veyra/internal/model"
	"github.com/google/uuid"
	"github.com/jackc/pgx/v5"
	"github.com/jackc/pgx/v5/pgxpool"
)

type BacktestRepository struct{ pool *pgxpool.Pool }

func NewBacktestRepository(pool *pgxpool.Pool) *BacktestRepository { return &BacktestRepository{pool} }

const backtestColumns = `id,strategy_version_id,status,input,input_hash,result,failure,lease_id,created_at,finished_at`

func scanBacktest(row pgx.Row) (model.BacktestRun, error) {
	var v model.BacktestRun
	err := row.Scan(&v.ID, &v.StrategyVersionID, &v.Status, &v.Input, &v.InputHash, &v.Result, &v.Failure, &v.LeaseID, &v.CreatedAt, &v.FinishedAt)
	return v, mapError(err)
}
func (r *BacktestRepository) Detail(ctx context.Context, id string) (model.BacktestRun, error) {
	return scanBacktest(r.pool.QueryRow(ctx, `SELECT `+backtestColumns+` FROM veyra_backtest_run WHERE id=$1`, id))
}
func (r *BacktestRepository) List(ctx context.Context) ([]model.BacktestRun, error) {
	rows, e := r.pool.Query(ctx, `SELECT `+backtestColumns+` FROM veyra_backtest_run ORDER BY created_at DESC,id DESC LIMIT 100`)
	if e != nil {
		return nil, mapError(e)
	}
	defer rows.Close()
	list := make([]model.BacktestRun, 0)
	for rows.Next() {
		v, e := scanBacktest(rows)
		if e != nil {
			return nil, e
		}
		v.Input = nil
		v.Result = nil
		list = append(list, v)
	}
	return list, mapError(rows.Err())
}
func (r *BacktestRepository) Create(ctx context.Context, input model.BacktestInput, actor, requestID string) (model.BacktestRun, error) {
	tx, e := r.pool.Begin(ctx)
	if e != nil {
		return model.BacktestRun{}, mapError(e)
	}
	defer tx.Rollback(context.Background())
	var p model.StrategyParameters
	var raw []byte
	if e := tx.QueryRow(ctx, `SELECT parameters FROM veyra_strategy_version WHERE id=$1`, input.StrategyVersionID).Scan(&raw); e != nil {
		return model.BacktestRun{}, mapError(e)
	}
	if json.Unmarshal(raw, &p) != nil {
		return model.BacktestRun{}, apperror.ErrUnavailable
	}
	if len(input.Bars) < p.LongWindow+2 {
		return model.BacktestRun{}, apperror.ErrInvalidInput
	}
	data, e := json.Marshal(model.ResearchInput{BacktestInput: input, Parameters: p})
	if e != nil {
		return model.BacktestRun{}, apperror.ErrInvalidInput
	}
	hash := sha256.Sum256(data)
	id := uuid.NewString()
	_, e = tx.Exec(ctx, `INSERT INTO veyra_backtest_run(id,strategy_version_id,input,input_hash,author_id) VALUES($1,$2,$3,$4,$5)`, id, input.StrategyVersionID, data, hex.EncodeToString(hash[:]), actor)
	if e != nil {
		return model.BacktestRun{}, mapError(e)
	}
	if e := appendAudit(ctx, tx, actor, "BACKTEST_QUEUED", "backtest_run", id, requestID, map[string]any{"strategy_version_id": input.StrategyVersionID, "input_hash": hex.EncodeToString(hash[:]), "bars": len(input.Bars)}); e != nil {
		return model.BacktestRun{}, e
	}
	v, e := scanBacktest(tx.QueryRow(ctx, `SELECT `+backtestColumns+` FROM veyra_backtest_run WHERE id=$1`, id))
	if e != nil {
		return v, e
	}
	return v, mapError(tx.Commit(ctx))
}
func (r *BacktestRepository) Claim(ctx context.Context) (model.BacktestRun, error) {
	return scanBacktest(r.pool.QueryRow(ctx, `WITH next AS (SELECT id FROM veyra_backtest_run WHERE status='PENDING' OR (status='RUNNING' AND lease_expires_at<now()) ORDER BY created_at,id LIMIT 1 FOR UPDATE SKIP LOCKED) UPDATE veyra_backtest_run r SET status='RUNNING',lease_id=$1,lease_expires_at=now()+interval '60 seconds' FROM next WHERE r.id=next.id RETURNING r.id,r.strategy_version_id,r.status,r.input,r.input_hash,r.result,r.failure,r.lease_id,r.created_at,r.finished_at`, uuid.NewString()))
}
func (r *BacktestRepository) Finish(ctx context.Context, job model.BacktestRun, result json.RawMessage, failure string) error {
	status := "COMPLETED"
	if failure != "" {
		status = "FAILED"
		result = nil
	}
	tag, e := r.pool.Exec(ctx, `UPDATE veyra_backtest_run SET status=$3,result=$4,failure=$5,lease_id=NULL,lease_expires_at=NULL,finished_at=now() WHERE id=$1 AND lease_id=$2 AND status='RUNNING' AND lease_expires_at>now()`, job.ID, job.LeaseID, status, result, failure)
	if e != nil {
		return mapError(e)
	}
	if tag.RowsAffected() != 1 {
		return apperror.ErrConflict
	}
	return nil
}