pages/internal/store/db.go

// Package store owns the SQLite database: connection management, migrations and
// every query the server runs. It is server-side only — the CLI must never reach
// it, and cmd/pages/deps_test.go enforces that.
package store

import (
	"context"
	"database/sql"
	"errors"
	"fmt"
	"log/slog"
	"math/rand/v2"
	"net/url"
	"runtime"
	"time"

	"modernc.org/sqlite"
	sqlite3 "modernc.org/sqlite/lib"
)

const driverName = "sqlite"

// maxTxAttempts bounds the BEGIN-level retry loop. busy_timeout already handles
// contention inside a transaction; these retries exist for the cases it cannot
// cover (see Tx).
const maxTxAttempts = 5

// DB holds the two connection pools SQLite wants under concurrent load.
//
// SQLite permits exactly one writer at a time. Handing database/sql a single
// pool means readers and the writer compete for the same connections and every
// static request can end up queued behind a deployment commit. Two pools make
// the rule explicit instead: W is capped at one connection so writes serialise
// in Go (where waiting is cheap and fair) rather than in SQLite (where it
// surfaces as SQLITE_BUSY), and R holds the concurrent readers that WAL mode
// lets run undisturbed alongside the writer.
type DB struct {
	w    *sql.DB
	r    *sql.DB
	path string
	log  *slog.Logger
}

// Open connects to the database at path, applying the pragmas both pools need,
// and runs any outstanding migrations.
func Open(ctx context.Context, path string, log *slog.Logger) (*DB, error) {
	w, err := sql.Open(driverName, dsn(path, false))
	if err != nil {
		return nil, fmt.Errorf("open %s: %w", path, err)
	}
	// One writer, by construction. Idle == open and no lifetime cap: rebuilding a
	// SQLite connection means replaying every pragma, and a server-side
	// connection has nothing to go stale against.
	w.SetMaxOpenConns(1)
	w.SetMaxIdleConns(1)
	w.SetConnMaxLifetime(0)
	w.SetConnMaxIdleTime(0)

	// The writer opens first so the database file, and its WAL, exist before any
	// read-only connection tries to attach to them.
	if err := w.PingContext(ctx); err != nil {
		w.Close()
		return nil, fmt.Errorf("open %s: %w", path, err)
	}

	readers := max(4, runtime.NumCPU())
	r, err := sql.Open(driverName, dsn(path, true))
	if err != nil {
		w.Close()
		return nil, fmt.Errorf("open %s (read pool): %w", path, err)
	}
	r.SetMaxOpenConns(readers)
	r.SetMaxIdleConns(readers)
	r.SetConnMaxLifetime(0)
	r.SetConnMaxIdleTime(0)
	if err := r.PingContext(ctx); err != nil {
		w.Close()
		r.Close()
		return nil, fmt.Errorf("open %s (read pool): %w", path, err)
	}

	d := &DB{w: w, r: r, path: path, log: log}
	if err := d.migrate(ctx); err != nil {
		d.Close()
		return nil, err
	}
	return d, nil
}

// dsn builds the connection string. The pragmas are per-connection state, so
// every connection in both pools must carry them.
func dsn(path string, readOnly bool) string {
	q := url.Values{}
	// Ordering inside the driver is fixed (busy_timeout first), so the list here
	// is grouped by intent rather than by application order.
	q.Add("_pragma", "busy_timeout(10000)") // wait for a lock instead of failing
	q.Add("_pragma", "journal_mode(WAL)")   // readers do not block the writer
	q.Add("_pragma", "synchronous(NORMAL)") // the correct pairing with WAL
	q.Add("_pragma", "foreign_keys(ON)")    // off by default in SQLite
	q.Add("_pragma", "recursive_triggers(ON)")
	q.Add("_pragma", "temp_store(MEMORY)")
	q.Add("_pragma", "wal_autocheckpoint(1000)")

	if readOnly {
		// A guard rail, not a security boundary: it turns "this query was meant to
		// be a read" from a silent lock contention bug into an immediate error.
		q.Set("_query_only", "true")
	} else {
		// Every transaction takes the write lock at BEGIN. Without this a
		// transaction that starts with a SELECT and later writes must upgrade its
		// lock, and an upgrade that finds another writer fails with SQLITE_BUSY
		// immediately — busy_timeout does not apply to it, because waiting could
		// only ever deadlock.
		q.Set("_txlock", "immediate")
	}
	return "file:" + path + "?" + q.Encode()
}

// Reader returns the read-only pool. Queries that must observe uncommitted
// changes belong inside the writing transaction instead.
func (d *DB) Reader() *sql.DB { return d.r }

// Path is the database file's location.
func (d *DB) Path() string { return d.path }

// Tx runs fn inside a single write transaction and commits it, retrying from
// the top when SQLite reports contention.
//
// fn may be called more than once, so it must not have side effects outside the
// transaction — no file writes, no atomic pointer stores, no channel sends.
// Publishing a change to the rest of the process is the caller's job, after Tx
// returns nil.
func (d *DB) Tx(ctx context.Context, fn func(*sql.Tx) error) error {
	var err error
	for attempt := range maxTxAttempts {
		if attempt > 0 {
			// Exponential backoff with jitter, so two contending writers do not
			// retry in lockstep.
			backoff := time.Duration(1<<attempt) * time.Millisecond
			select {
			case <-time.After(backoff/2 + time.Duration(rand.Int64N(int64(backoff)))):
			case <-ctx.Done():
				return ctx.Err()
			}
		}
		err = d.tx(ctx, fn)
		if err == nil {
			return nil
		}
		if !IsBusy(err) || ctx.Err() != nil {
			return err
		}
		if d.log != nil {
			d.log.Debug("write transaction contended, retrying",
				"attempt", attempt+1, "err", err)
		}
	}
	return fmt.Errorf("write transaction still contended after %d attempts: %w", maxTxAttempts, err)
}

func (d *DB) tx(ctx context.Context, fn func(*sql.Tx) error) (err error) {
	tx, err := d.w.BeginTx(ctx, nil)
	if err != nil {
		return err
	}
	defer func() {
		if p := recover(); p != nil {
			tx.Rollback()
			panic(p)
		}
		if err != nil {
			tx.Rollback()
		}
	}()
	if err = fn(tx); err != nil {
		return err
	}
	return tx.Commit()
}

// Close checkpoints the WAL and shuts both pools down.
func (d *DB) Close() error {
	var first error
	if d.r != nil {
		// Readers go first: a TRUNCATE checkpoint needs the WAL to itself.
		if err := d.r.Close(); err != nil {
			first = err
		}
	}
	if d.w != nil {
		// Fold the WAL back into the database file so a restart — or an operator
		// copying the directory — sees one self-contained file.
		if _, err := d.w.ExecContext(context.Background(), "PRAGMA wal_checkpoint(TRUNCATE)"); err != nil && d.log != nil {
			d.log.Warn("wal checkpoint on close failed", "err", err)
		}
		if err := d.w.Close(); err != nil && first == nil {
			first = err
		}
	}
	return first
}

// IsBusy reports whether err is SQLite telling us the database was locked.
func IsBusy(err error) bool {
	var serr *sqlite.Error
	if !errors.As(err, &serr) {
		return false
	}
	switch serr.Code() {
	case sqlite3.SQLITE_BUSY, sqlite3.SQLITE_BUSY_SNAPSHOT,
		sqlite3.SQLITE_LOCKED, sqlite3.SQLITE_LOCKED_SHAREDCACHE:
		return true
	}
	return false
}

// IsConstraint reports whether err is a constraint violation, optionally of a
// specific extended kind (sqlite3.SQLITE_CONSTRAINT_UNIQUE, ...).
func IsConstraint(err error, extended ...int) bool {
	var serr *sqlite.Error
	if !errors.As(err, &serr) {
		return false
	}
	code := serr.Code()
	if code&0xff != sqlite3.SQLITE_CONSTRAINT {
		return false
	}
	if len(extended) == 0 {
		return true
	}
	for _, e := range extended {
		if code == e {
			return true
		}
	}
	return false
}

// nowFunc is the clock the store stamps rows with. Tests replace it.
var nowFunc = time.Now

// unixNow is the timestamp format every *_at column uses: whole seconds, which
// is what the SQL triggers' unixepoch() produces.
func unixNow() int64 { return nowFunc().Unix() }