scrabble-game/loadtest/internal/seed/seed.go

package seed

import (
	"context"
	"fmt"

	"github.com/google/uuid"
	"github.com/jackc/pgx/v5"
	"github.com/jackc/pgx/v5/pgxpool"
)

// Marker prefixes every display_name the harness writes. Cleanup matches on it, so
// the harness only ever deletes its own rows and never touches real accounts.
const Marker = "lt:"

// Schema-qualified targets so the seeder does not depend on the connection's
// search_path (the backend pins search_path=backend; we qualify explicitly).
var (
	accountsTbl   = pgx.Identifier{"backend", "accounts"}
	identitiesTbl = pgx.Identifier{"backend", "identities"}
	sessionsTbl   = pgx.Identifier{"backend", "sessions"}
)

// Account is one seeded player: its account id, marker display name and the
// plaintext bearer token the driver presents in the Authorization header. Guest
// marks a guest (no identity, accrues no statistics). Name is retained so a
// profile.update can resend the marker display name and keep the row findable by
// Cleanup.
type Account struct {
	ID    uuid.UUID
	Name  string
	Token string
	Guest bool
}

// Pool is the seeded population, split by durability.
type Pool struct {
	Guests   []Account
	Durables []Account
}

// All returns every seeded account, durables first.
func (p *Pool) All() []Account {
	out := make([]Account, 0, len(p.Durables)+len(p.Guests))
	out = append(out, p.Durables...)
	out = append(out, p.Guests...)
	return out
}

// Seeder writes and removes the harness population over a pgx pool against the
// backend Postgres schema.
type Seeder struct{ pool *pgxpool.Pool }

// New connects to dsn (the backend Postgres) and verifies the connection.
func New(ctx context.Context, dsn string) (*Seeder, error) {
	pool, err := pgxpool.New(ctx, dsn)
	if err != nil {
		return nil, fmt.Errorf("seed: connect: %w", err)
	}
	if err := pool.Ping(ctx); err != nil {
		pool.Close()
		return nil, fmt.Errorf("seed: ping: %w", err)
	}
	return &Seeder{pool: pool}, nil
}

// Close releases the pool.
func (s *Seeder) Close() { s.pool.Close() }

// Seed inserts nDurable durable accounts (each with a confirmed email identity) and
// nGuest guest accounts, an active session per account, and returns the population
// with the plaintext tokens. Rows go in over COPY in foreign-key order (accounts,
// then identities and sessions). Every row carries Marker in its display name /
// external id so Cleanup can find them.
func (s *Seeder) Seed(ctx context.Context, nDurable, nGuest int) (*Pool, error) {
	pool := &Pool{
		Durables: make([]Account, 0, nDurable),
		Guests:   make([]Account, 0, nGuest),
	}
	var acctRows, identRows, sessRows [][]any

	add := func(guest bool, i int) error {
		aid, err := uuid.NewV7()
		if err != nil {
			return err
		}
		sid, err := uuid.NewV7()
		if err != nil {
			return err
		}
		token, hash, err := GenerateToken()
		if err != nil {
			return err
		}
		lang := "en"
		if i%2 == 1 {
			lang = "ru"
		}
		kind := "d"
		if guest {
			kind = "g"
		}
		name := fmt.Sprintf("%s%s-%06d", Marker, kind, i)
		acctRows = append(acctRows, []any{aid, name, guest, lang})
		sessRows = append(sessRows, []any{sid, aid, hash, "active"})
		if !guest {
			iid, err := uuid.NewV7()
			if err != nil {
				return err
			}
			ext := fmt.Sprintf("%s%s@loadtest.invalid", Marker, aid)
			identRows = append(identRows, []any{iid, aid, "email", ext, true})
		}
		acc := Account{ID: aid, Name: name, Token: token, Guest: guest}
		if guest {
			pool.Guests = append(pool.Guests, acc)
		} else {
			pool.Durables = append(pool.Durables, acc)
		}
		return nil
	}

	for i := 0; i < nDurable; i++ {
		if err := add(false, i); err != nil {
			return nil, err
		}
	}
	for i := 0; i < nGuest; i++ {
		if err := add(true, i); err != nil {
			return nil, err
		}
	}

	if _, err := s.pool.CopyFrom(ctx, accountsTbl,
		[]string{"account_id", "display_name", "is_guest", "preferred_language"},
		pgx.CopyFromRows(acctRows)); err != nil {
		return nil, fmt.Errorf("seed: copy accounts: %w", err)
	}
	if len(identRows) > 0 {
		if _, err := s.pool.CopyFrom(ctx, identitiesTbl,
			[]string{"identity_id", "account_id", "kind", "external_id", "confirmed"},
			pgx.CopyFromRows(identRows)); err != nil {
			return nil, fmt.Errorf("seed: copy identities: %w", err)
		}
	}
	if _, err := s.pool.CopyFrom(ctx, sessionsTbl,
		[]string{"session_id", "account_id", "token_hash", "status"},
		pgx.CopyFromRows(sessRows)); err != nil {
		return nil, fmt.Errorf("seed: copy sessions: %w", err)
	}
	return pool, nil
}

// Cleanup removes everything the harness created: first the games any harness
// account is seated in (cascading game_players / game_moves / complaints / chat),
// then the harness accounts (cascading identities, sessions, stats, invitations,
// drafts and the rest). It is scoped by Marker, so it is safe to run against a
// contour that also holds real data. The authoritative hard reset remains the
// contour DB wipe (DROP SCHEMA backend CASCADE + backend restart). It returns the
// number of accounts removed.
func (s *Seeder) Cleanup(ctx context.Context) (int, error) {
	if _, err := s.pool.Exec(ctx, `
		DELETE FROM backend.games
		WHERE game_id IN (
			SELECT p.game_id FROM backend.game_players p
			JOIN backend.accounts a ON a.account_id = p.account_id
			WHERE a.display_name LIKE $1
		)`, Marker+"%"); err != nil {
		return 0, fmt.Errorf("seed: cleanup games: %w", err)
	}
	tag, err := s.pool.Exec(ctx,
		`DELETE FROM backend.accounts WHERE display_name LIKE $1`, Marker+"%")
	if err != nil {
		return 0, fmt.Errorf("seed: cleanup accounts: %w", err)
	}
	return int(tag.RowsAffected()), nil
}