member-console/internal/testkit/dbtest/query.go

// SPDX-License-Identifier: AGPL-3.0-only OR LicenseRef-Commercial
// SPDX-FileCopyrightText: 2025-2026 Christian Galo

package dbtest

import (
	"context"
	"database/sql"
	"errors"
	"fmt"
	"strings"
	"sync/atomic"
	"testing"
	"time"

	"github.com/jackc/pgx/v5/pgconn"
)

// SQLSTATE codes the tests check for.
const (
	SQLStateCheck            = "23514"
	SQLStateForeignKey       = "23503"
	SQLStateUnique           = "23505"
	SQLStatePrivilege        = "42501"
	SQLStateOutOfRange       = "22003"
	SQLStateInvalidParameter = "22023"
	SQLStateNoDataFound      = "P0002"
	SQLStateNullValue        = "22004"
)

// Execer is what a Q runs its statements on: a *sql.DB, a *sql.Tx or a
// *sql.Conn.
type Execer interface {
	ExecContext(ctx context.Context, query string, args ...any) (sql.Result, error)
	QueryContext(ctx context.Context, query string, args ...any) (*sql.Rows, error)
	QueryRowContext(ctx context.Context, query string, args ...any) *sql.Row
}

// Q runs statements for one test and fails it on any error. It holds the
// testing.TB it was built with, so it belongs to the test or subtest that
// built it; a subtest that uses its parent's transaction rebinds with
// On(t, q.Tx()), so a failing statement fails the subtest and never calls the
// parent's FailNow from the subtest's goroutine.
type Q struct {
	tb  testing.TB
	ctx context.Context
	x   Execer
}

// On binds x to tb.
func On(tb testing.TB, x Execer) *Q {
	return &Q{tb: tb, ctx: context.Background(), x: x}
}

// Tx begins a transaction on d that the test rolls back when it ends.
func Tx(tb testing.TB, d *sql.DB) *Q {
	tb.Helper()
	tx, err := d.BeginTx(context.Background(), nil)
	if err != nil {
		tb.Fatalf("begin: %v", err)
	}
	tb.Cleanup(func() {
		if err := tx.Rollback(); err != nil && !errors.Is(err, sql.ErrTxDone) {
			tb.Errorf("rollback: %v", err)
		}
	})
	return On(tb, tx)
}

// Tx returns the transaction q runs on and fails the test when q runs on
// anything else.
func (q *Q) Tx() *sql.Tx {
	q.tb.Helper()
	tx, ok := q.x.(*sql.Tx)
	if !ok {
		q.tb.Fatalf("dbtest: Tx needs a Q over a *sql.Tx, this one runs on %T", q.x)
	}
	return tx
}

func fold(query string) string { return strings.Join(strings.Fields(query), " ") }

// Exec runs a statement.
func (q *Q) Exec(query string, args ...any) {
	q.tb.Helper()
	if _, err := q.x.ExecContext(q.ctx, query, args...); err != nil {
		q.tb.Fatalf("%s: %v", fold(query), err)
	}
}

// Scan runs a query that returns one row and scans its columns into dest.
func (q *Q) Scan(dest []any, query string, args ...any) {
	q.tb.Helper()
	if err := q.x.QueryRowContext(q.ctx, query, args...).Scan(dest...); err != nil {
		q.tb.Fatalf("%s: %v", fold(query), err)
	}
}

// ID runs a query that returns one text column in one row, a statement's
// RETURNING id for instance.
func (q *Q) ID(query string, args ...any) string {
	q.tb.Helper()
	return q.String(query, args...)
}

// Int runs a query that returns one integer.
func (q *Q) Int(query string, args ...any) int64 {
	q.tb.Helper()
	var n int64
	q.Scan([]any{&n}, query, args...)
	return n
}

// String runs a query that returns one text column in one row.
func (q *Q) String(query string, args ...any) string {
	q.tb.Helper()
	var s string
	q.Scan([]any{&s}, query, args...)
	return s
}

// Strings runs a query that selects one text column and returns its values in
// row order. The rows are closed on return, so the next statement can run on
// the same transaction.
func (q *Q) Strings(query string, args ...any) []string {
	q.tb.Helper()
	rows, err := q.x.QueryContext(q.ctx, query, args...)
	if err != nil {
		q.tb.Fatalf("%s: %v", fold(query), err)
	}
	defer func() { _ = rows.Close() }()
	var out []string
	for rows.Next() {
		var s string
		if err := rows.Scan(&s); err != nil {
			q.tb.Fatalf("%s: %v", fold(query), err)
		}
		out = append(out, s)
	}
	if err := rows.Err(); err != nil {
		q.tb.Fatalf("%s: %v", fold(query), err)
	}
	return out
}

// Rows runs query and returns its rows as text for a golden file: a line of
// the column names, then one line per row in the order the query returns
// them, the values separated by " | ". NULL reads "NULL", a time reads in
// UTC as RFC 3339 with its fraction, bytes read as text, and anything else as
// fmt prints it. The query names its own order; Rows adds none.
func (q *Q) Rows(query string, args ...any) string {
	q.tb.Helper()
	rows, err := q.x.QueryContext(q.ctx, query, args...)
	if err != nil {
		q.tb.Fatalf("%s: %v", fold(query), err)
	}
	defer func() { _ = rows.Close() }()
	cols, err := rows.Columns()
	if err != nil {
		q.tb.Fatalf("%s: columns: %v", fold(query), err)
	}
	var b strings.Builder
	b.WriteString(strings.Join(cols, " | "))
	b.WriteByte('\n')
	vals := make([]any, len(cols))
	ptrs := make([]any, len(cols))
	for i := range vals {
		ptrs[i] = &vals[i]
	}
	for rows.Next() {
		if err := rows.Scan(ptrs...); err != nil {
			q.tb.Fatalf("%s: scan: %v", fold(query), err)
		}
		cells := make([]string, len(vals))
		for i, v := range vals {
			cells[i] = cell(v)
		}
		b.WriteString(strings.Join(cells, " | "))
		b.WriteByte('\n')
	}
	if err := rows.Err(); err != nil {
		q.tb.Fatalf("%s: rows: %v", fold(query), err)
	}
	return b.String()
}

// cell is one value as Rows prints it.
func cell(v any) string {
	switch v := v.(type) {
	case nil:
		return "NULL"
	case []byte:
		return string(v)
	case time.Time:
		return v.UTC().Format(time.RFC3339Nano)
	default:
		return fmt.Sprint(v)
	}
}

// Attempt runs one statement inside a savepoint and returns its error,
// rolling back to the savepoint when it fails so the transaction goes on.
// Postgres refuses SAVEPOINT outside a transaction block (25P01), so q must
// run on a *sql.Tx.
func (q *Q) Attempt(query string, args ...any) error {
	q.tb.Helper()
	tx, ok := q.x.(*sql.Tx)
	if !ok {
		q.tb.Fatalf("dbtest: Attempt and Refuse need a Q over a *sql.Tx (a savepoint outside a transaction block is SQLSTATE 25P01), this one runs on %T", q.x)
	}
	q.Exec("SAVEPOINT dbtest_attempt")
	if _, err := tx.ExecContext(q.ctx, query, args...); err != nil {
		q.Exec("ROLLBACK TO SAVEPOINT dbtest_attempt")
		return err
	}
	q.Exec("RELEASE SAVEPOINT dbtest_attempt")
	return nil
}

// Refuse requires a statement to fail with the SQLSTATE and, when given, the
// constraint name.
func (q *Q) Refuse(sqlstate, constraint, query string, args ...any) {
	q.tb.Helper()
	WantPgError(q.tb, q.Attempt(query, args...), sqlstate, constraint)
}

// WantPgError requires err to be a *pgconn.PgError with the given SQLSTATE
// and, when constraint is non-empty, the given constraint name.
func WantPgError(tb testing.TB, err error, sqlstate, constraint string) {
	tb.Helper()
	if err == nil {
		tb.Fatalf("expected error (SQLSTATE %s, constraint %q), got nil", sqlstate, constraint)
	}
	var pgErr *pgconn.PgError
	if !errors.As(err, &pgErr) {
		tb.Fatalf("expected a *pgconn.PgError, got %T: %v", err, err)
	}
	if pgErr.Code != sqlstate {
		tb.Errorf("SQLSTATE = %s, want %s (err: %v)", pgErr.Code, sqlstate, err)
	}
	if constraint != "" && pgErr.ConstraintName != constraint {
		tb.Errorf("constraint = %q, want %q (err: %v)", pgErr.ConstraintName, constraint, err)
	}
}

var nameSeq atomic.Int64

// Name returns a lowercase alphanumeric name unique across runs, which a
// provider key's hygiene check accepts.
func Name(prefix string) string {
	return fmt.Sprintf("%s%d%d", prefix, time.Now().UnixNano(), nameSeq.Add(1))
}