alt/services/worker/internal/storage/postgres/store.go

package postgres

import (
	"context"
	"errors"
	"fmt"
	"time"

	"github.com/jackc/pgx/v5"
	"github.com/jackc/pgx/v5/pgtype"
	"github.com/jackc/pgx/v5/pgxpool"

	"git.toki-labs.com/toki/alt/packages/domain/backtest"
	"git.toki-labs.com/toki/alt/packages/domain/market"
	"git.toki-labs.com/toki/alt/services/worker/internal/storage"
	"git.toki-labs.com/toki/alt/services/worker/internal/storage/postgres/sqlc"
)

var _ storage.InstrumentStore = (*Store)(nil)
var _ storage.BarStore = (*Store)(nil)
var _ storage.BacktestRunStore = (*Store)(nil)
var _ storage.BacktestResultStore = (*Store)(nil)
var _ storage.BacktestAnalysisStore = (*Store)(nil)

type Store struct {
	pool    *pgxpool.Pool
	queries *sqlc.Queries
}

func NewStore(pool *pgxpool.Pool) *Store {
	return &Store{
		pool:    pool,
		queries: sqlc.New(pool),
	}
}

func (s *Store) UpsertInstrument(ctx context.Context, inst market.Instrument) error {
	params, err := mapInstrumentToParams(inst)
	if err != nil {
		return fmt.Errorf("failed to map instrument: %w", err)
	}
	if err := s.queries.UpsertInstrument(ctx, params); err != nil {
		return fmt.Errorf("failed to upsert instrument: %w", err)
	}
	return nil
}

func (s *Store) GetInstrument(ctx context.Context, id market.InstrumentID) (market.Instrument, error) {
	row, err := s.queries.GetInstrument(ctx, string(id))
	if err != nil {
		if errors.Is(err, pgx.ErrNoRows) {
			return market.Instrument{}, storage.ErrInstrumentNotFound
		}
		return market.Instrument{}, fmt.Errorf("failed to get instrument: %w", err)
	}
	inst, err := mapRowToInstrument(row)
	if err != nil {
		return market.Instrument{}, fmt.Errorf("failed to map instrument row: %w", err)
	}
	return inst, nil
}

func (s *Store) ListInstruments(ctx context.Context) ([]market.Instrument, error) {
	rows, err := s.queries.ListInstruments(ctx)
	if err != nil {
		return nil, fmt.Errorf("failed to list instruments: %w", err)
	}
	insts := make([]market.Instrument, len(rows))
	for i, row := range rows {
		inst, err := mapRowToInstrument(row)
		if err != nil {
			return nil, fmt.Errorf("failed to map instrument row at index %d: %w", i, err)
		}
		insts[i] = inst
	}
	return insts, nil
}

func (s *Store) UpsertBar(ctx context.Context, bar market.Bar) error {
	params, err := mapBarToParams(bar)
	if err != nil {
		return fmt.Errorf("failed to map bar: %w", err)
	}
	if err := s.queries.UpsertBar(ctx, params); err != nil {
		return fmt.Errorf("failed to upsert bar: %w", err)
	}
	return nil
}

func (s *Store) GetBars(ctx context.Context, id market.InstrumentID, timeframe market.Timeframe, from, to time.Time) ([]market.Bar, error) {
	// First, fetch the instrument to know its currency
	instRow, err := s.queries.GetInstrument(ctx, string(id))
	if err != nil {
		if errors.Is(err, pgx.ErrNoRows) {
			return nil, storage.ErrInstrumentNotFound
		}
		return nil, fmt.Errorf("failed to get instrument for bar currency: %w", err)
	}

	params := sqlc.GetBarsParams{
		InstrumentID: string(id),
		Timeframe:    string(timeframe),
		Timestamp:    pgtype.Timestamptz{Time: from, Valid: !from.IsZero()},
		Timestamp_2:  pgtype.Timestamptz{Time: to, Valid: !to.IsZero()},
	}
	rows, err := s.queries.GetBars(ctx, params)
	if err != nil {
		return nil, fmt.Errorf("failed to get bars: %w", err)
	}

	bars := make([]market.Bar, len(rows))
	for i, row := range rows {
		bar, err := mapRowToBar(row, market.Currency(instRow.Currency))
		if err != nil {
			return nil, fmt.Errorf("failed to map bar row at index %d: %w", i, err)
		}
		bars[i] = bar
	}
	return bars, nil
}

func (s *Store) UpsertRun(ctx context.Context, run backtest.Run) error {
	params, err := mapRunToParams(run)
	if err != nil {
		return fmt.Errorf("failed to map run: %w", err)
	}
	if err := s.queries.UpsertRun(ctx, params); err != nil {
		return fmt.Errorf("failed to upsert run: %w", err)
	}
	return nil
}

func (s *Store) GetRun(ctx context.Context, id backtest.RunID) (backtest.Run, error) {
	row, err := s.queries.GetRun(ctx, string(id))
	if err != nil {
		if errors.Is(err, pgx.ErrNoRows) {
			return backtest.Run{}, storage.ErrRunNotFound
		}
		return backtest.Run{}, fmt.Errorf("failed to get run: %w", err)
	}
	run, err := mapRowToRun(row)
	if err != nil {
		return backtest.Run{}, fmt.Errorf("failed to map run row: %w", err)
	}
	return run, nil
}

func (s *Store) ListRuns(ctx context.Context, status backtest.RunStatus) ([]backtest.Run, error) {
	rows, err := s.queries.ListRuns(ctx, string(status))
	if err != nil {
		return nil, fmt.Errorf("failed to list runs: %w", err)
	}
	runs := make([]backtest.Run, len(rows))
	for i, row := range rows {
		run, err := mapRowToRun(row)
		if err != nil {
			return nil, fmt.Errorf("failed to map run row at index %d: %w", i, err)
		}
		runs[i] = run
	}
	return runs, nil
}

func (s *Store) GetRunDetail(ctx context.Context, id backtest.RunID) (storage.RunDetail, error) {
	run, err := s.GetRun(ctx, id)
	if err != nil {
		return storage.RunDetail{}, err
	}

	result, err := s.GetResult(ctx, id)
	if err != nil {
		if errors.Is(err, storage.ErrResultNotFound) {
			return storage.RunDetail{Run: run, HasResult: false}, nil
		}
		return storage.RunDetail{}, err
	}
	return storage.RunDetail{Run: run, Result: result, HasResult: true}, nil
}

func (s *Store) CompareResults(ctx context.Context, ids []backtest.RunID) ([]backtest.Result, error) {
	results := make([]backtest.Result, 0, len(ids))
	for _, id := range ids {
		result, err := s.GetResult(ctx, id)
		if err != nil {
			return nil, fmt.Errorf("failed to get result for run %s: %w", id, err)
		}
		results = append(results, result)
	}
	return results, nil
}

func (s *Store) UpsertResult(ctx context.Context, result backtest.Result) error {
	params, err := mapResultToParams(result)
	if err != nil {
		return fmt.Errorf("failed to map result: %w", err)
	}
	if err := s.queries.UpsertResult(ctx, params); err != nil {
		return fmt.Errorf("failed to upsert result: %w", err)
	}
	return nil
}

func (s *Store) GetResult(ctx context.Context, id backtest.RunID) (backtest.Result, error) {
	row, err := s.queries.GetResult(ctx, string(id))
	if err != nil {
		if errors.Is(err, pgx.ErrNoRows) {
			return backtest.Result{}, storage.ErrResultNotFound
		}
		return backtest.Result{}, fmt.Errorf("failed to get result: %w", err)
	}
	res, err := mapRowToResult(row)
	if err != nil {
		return backtest.Result{}, fmt.Errorf("failed to map result row: %w", err)
	}
	return res, nil
}