Halcyon/store/store.go

// Package store persists players, ships and world snapshots in SQLite or
// PostgreSQL through database/sql.
package store

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

	_ "github.com/jackc/pgx/v5/stdlib"
	_ "modernc.org/sqlite"
)

type Dialect int

const (
	SQLite Dialect = iota
	Postgres
)

// Store is a database handle. Methods on the embedded Q run directly on the
// pool; use WithTx for atomic multi-step work.
type Store struct {
	*Q
	db *sql.DB
}

type execer interface {
	ExecContext(ctx context.Context, q string, args ...any) (sql.Result, error)
	QueryContext(ctx context.Context, q string, args ...any) (*sql.Rows, error)
	QueryRowContext(ctx context.Context, q string, args ...any) *sql.Row
}

// Q holds the query methods; it wraps either a pool or a transaction.
type Q struct {
	x  execer
	d  Dialect
	tx bool
}

// Open connects using dsn. A "postgres://" or "postgresql://" URL selects
// PostgreSQL; anything else is a SQLite file path (optionally prefixed with
// "sqlite:").
func Open(dsn string) (*Store, error) {
	var (
		driver string
		d      Dialect
	)
	switch {
	case strings.HasPrefix(dsn, "postgres://"), strings.HasPrefix(dsn, "postgresql://"):
		driver, d = "pgx", Postgres
	default:
		path := strings.TrimPrefix(dsn, "sqlite:")
		sep := "?"
		if strings.Contains(path, "?") {
			sep = "&"
		}
		dsn = "file:" + path + sep + "_txlock=immediate&_pragma=busy_timeout(10000)&_pragma=journal_mode(WAL)&_pragma=foreign_keys(1)"
		driver, d = "sqlite", SQLite
	}
	db, err := sql.Open(driver, dsn)
	if err != nil {
		return nil, err
	}
	if d == SQLite {
		db.SetMaxOpenConns(1) // SQLite allows a single writer; keep it simple
	}
	return &Store{Q: &Q{x: db, d: d}, db: db}, nil
}

func (s *Store) Close() error { return s.db.Close() }

// WithTx runs fn in a transaction, committing if it returns nil.
func (s *Store) WithTx(ctx context.Context, fn func(q *Q) error) error {
	tx, err := s.db.BeginTx(ctx, nil)
	if err != nil {
		return err
	}
	if err := fn(&Q{x: tx, d: s.d, tx: true}); err != nil {
		tx.Rollback()
		return err
	}
	return tx.Commit()
}

// rebind converts ? placeholders to $n for PostgreSQL.
func (q *Q) rebind(query string) string {
	if q.d != Postgres {
		return query
	}
	var b strings.Builder
	n := 0
	for _, r := range query {
		if r == '?' {
			n++
			fmt.Fprintf(&b, "$%d", n)
		} else {
			b.WriteRune(r)
		}
	}
	return b.String()
}

func (q *Q) exec(ctx context.Context, query string, args ...any) (sql.Result, error) {
	return q.x.ExecContext(ctx, q.rebind(query), args...)
}

func (q *Q) query(ctx context.Context, query string, args ...any) (*sql.Rows, error) {
	return q.x.QueryContext(ctx, q.rebind(query), args...)
}

func (q *Q) row(ctx context.Context, query string, args ...any) *sql.Row {
	return q.x.QueryRowContext(ctx, q.rebind(query), args...)
}

var schema = []string{
	`CREATE TABLE IF NOT EXISTS meta (key TEXT PRIMARY KEY, value TEXT NOT NULL)`,
	`CREATE TABLE IF NOT EXISTS players (
		id {ID}, name TEXT NOT NULL UNIQUE, key_hash TEXT NOT NULL UNIQUE,
		credits BIGINT NOT NULL DEFAULT 0)`,
	`CREATE TABLE IF NOT EXISTS ships (
		id {ID}, owner_id BIGINT NOT NULL REFERENCES players(id),
		status TEXT NOT NULL, program {BLOB}, uplink {BLOB}, downlink {BLOB},
		downlink_day BIGINT NOT NULL DEFAULT -1)`,
	`CREATE INDEX IF NOT EXISTS ships_owner ON ships(owner_id)`,
	`CREATE INDEX IF NOT EXISTS ships_status ON ships(status)`,
	`CREATE TABLE IF NOT EXISTS snapshots (
		day BIGINT PRIMARY KEY, hash TEXT NOT NULL, data {BLOB} NOT NULL)`,
	`CREATE TABLE IF NOT EXISTS events (
		id {ID}, day BIGINT NOT NULL, tick INTEGER NOT NULL,
		ship_id BIGINT NOT NULL, kind TEXT NOT NULL, detail TEXT NOT NULL)`,
	`CREATE INDEX IF NOT EXISTS events_day ON events(day)`,
	`CREATE TABLE IF NOT EXISTS day_runs (
		day BIGINT PRIMARY KEY, hash TEXT NOT NULL, ran_at TEXT NOT NULL)`,
	`INSERT INTO meta (key, value) VALUES ('lock', '') ON CONFLICT (key) DO NOTHING`,
}

// Migrate creates the schema if it does not exist.
func (s *Store) Migrate(ctx context.Context) error {
	id, blob := "INTEGER PRIMARY KEY AUTOINCREMENT", "BLOB"
	if s.d == Postgres {
		id, blob = "BIGSERIAL PRIMARY KEY", "BYTEA"
	}
	r := strings.NewReplacer("{ID}", id, "{BLOB}", blob)
	for _, stmt := range schema {
		if _, err := s.exec(ctx, r.Replace(stmt)); err != nil {
			return fmt.Errorf("migrate: %w", err)
		}
	}
	return nil
}

// LockWorld serialises daily runs. SQLite transactions are already exclusive
// (immediate mode); PostgreSQL takes a row lock held until commit.
func (q *Q) LockWorld(ctx context.Context) error {
	if q.d != Postgres {
		return nil
	}
	var v string
	return q.row(ctx, `SELECT value FROM meta WHERE key = 'lock' FOR UPDATE`).Scan(&v)
}

func (q *Q) GetMeta(ctx context.Context, key string) (string, bool, error) {
	var v string
	err := q.row(ctx, `SELECT value FROM meta WHERE key = ?`, key).Scan(&v)
	if err == sql.ErrNoRows {
		return "", false, nil
	}
	return v, err == nil, err
}

func (q *Q) SetMeta(ctx context.Context, key, value string) error {
	_, err := q.exec(ctx, `INSERT INTO meta (key, value) VALUES (?, ?)
		ON CONFLICT (key) DO UPDATE SET value = excluded.value`, key, value)
	return err
}