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

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

package dbtest

import (
	"context"
	"database/sql"
	"fmt"
	"maps"
	"slices"
	"strings"
	"sync"
	"testing"

	"github.com/jackc/pgx/v5"
	"github.com/jackc/pgx/v5/stdlib"

	"git.coopcloud.tech/wiki-cafe/member-console/internal/testkit"
)

// Counted opens TEST_DATABASE_URL as Open does, on a connection whose every
// statement is recorded, so a test can count what one request runs. Hand the
// handle to the code under test only, and build fixtures on a handle from
// Open, so the record holds the code's statements and nothing else. It
// starts empty.
func Counted(tb testing.TB) (*sql.DB, *Statements) {
	tb.Helper()
	return CountedAt(tb, DSN(tb))
}

// CountedAt is Counted over the database at dsn, a scratch database for
// instance. Its core ledger must be at the highest version, as Open requires
// of the test database.
func CountedAt(tb testing.TB, dsn string) (*sql.DB, *Statements) {
	tb.Helper()
	cfg, err := pgx.ParseConfig(dsn)
	if err != nil {
		tb.Fatalf("parse the database address: %v", err)
	}
	s := &Statements{}
	cfg.Tracer = s
	d := stdlib.OpenDB(*cfg)
	tb.Cleanup(func() { testkit.Close(tb, d) })
	if err := ledgerBehind(d); err != nil {
		tb.Fatalf("%v", err)
	}
	s.Reset()
	return d, s
}

// Statements records the statements run on a Counted handle, each by its
// name: the sqlc query name for a generated query, whose text starts
// "-- name: <Name> :<kind>", and otherwise the first line of the text with
// its spaces collapsed (a raw query, BEGIN, COMMIT). Pings are not
// statements and are not recorded.
type Statements struct {
	mu    sync.Mutex
	names []string
	notes []note
}

// note is a line a test added to the record, placed after the first at
// statements.
type note struct {
	at   int
	text string
}

// TraceQueryStart records one statement. It implements pgx.QueryTracer.
func (s *Statements) TraceQueryStart(ctx context.Context, _ *pgx.Conn, data pgx.TraceQueryStartData) context.Context {
	s.mu.Lock()
	defer s.mu.Unlock()
	s.names = append(s.names, statementName(data.SQL))
	return ctx
}

// TraceQueryEnd implements pgx.QueryTracer; the end of a statement adds
// nothing.
func (s *Statements) TraceQueryEnd(context.Context, *pgx.Conn, pgx.TraceQueryEndData) {}

// Reset forgets every statement recorded so far.
func (s *Statements) Reset() {
	s.mu.Lock()
	defer s.mu.Unlock()
	s.names = nil
	s.notes = nil
}

// Note adds text to the record after the statements recorded so far, so
// Names places an event the database does not see, a call to Stripe for
// instance, among the statements around it. A note is not a statement:
// Total and ByName leave it out.
func (s *Statements) Note(text string) {
	s.mu.Lock()
	defer s.mu.Unlock()
	s.notes = append(s.notes, note{at: len(s.names), text: text})
}

// Total is the number of statements recorded since the last Reset.
func (s *Statements) Total() int {
	s.mu.Lock()
	defer s.mu.Unlock()
	return len(s.names)
}

// Names returns the names of the statements recorded since the last Reset,
// in the order they ran, with each note at its place among them. A
// transaction's begin, commit and rollback are among them.
func (s *Statements) Names() []string {
	s.mu.Lock()
	defer s.mu.Unlock()
	out := make([]string, 0, len(s.names)+len(s.notes))
	next := 0
	for i, n := range s.names {
		for next < len(s.notes) && s.notes[next].at == i {
			out = append(out, s.notes[next].text)
			next++
		}
		out = append(out, n)
	}
	for _, n := range s.notes[next:] {
		out = append(out, n.text)
	}
	return out
}

// ByName counts the statements recorded since the last Reset by name.
func (s *Statements) ByName() map[string]int {
	s.mu.Lock()
	defer s.mu.Unlock()
	out := map[string]int{}
	for _, n := range s.names {
		out[n]++
	}
	return out
}

// Diff lists, one "name: before -> after" line per name in sorted order, the
// names whose counts differ between two ByName results.
func Diff(before, after map[string]int) string {
	names := maps.Clone(before)
	maps.Copy(names, after)
	var b strings.Builder
	for _, n := range slices.Sorted(maps.Keys(names)) {
		if before[n] != after[n] {
			fmt.Fprintf(&b, "%s: %d -> %d\n", n, before[n], after[n])
		}
	}
	return b.String()
}

func statementName(text string) string {
	for line := range strings.Lines(text) {
		line = strings.Join(strings.Fields(line), " ")
		if line == "" {
			continue
		}
		if rest, ok := strings.CutPrefix(line, "-- name: "); ok {
			name, _, _ := strings.Cut(rest, " ")
			return name
		}
		return line
	}
	return ""
}