mesh-controller/internal/store/fortest.go

package store

import (
	"context"
	"crypto/sha256"
	"encoding/binary"
	"encoding/hex"
	"errors"
	"fmt"
	"strings"
	"sync"
	"testing"
	"time"

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

// CreateMigratedForTest creates the database name on the test PostgreSQL that admin reaches, holding
// every one of migrations already applied and recorded, as a fresh database migrated by Migrate would.
//
// **A copy of a template, migrated once** (novox/hq issue 591): a store test used to create an empty
// database and run every migration into it, some 90 for the inventory, and with hundreds of store tests
// that alone was most of the controller's check. The template is migrated by Migrate itself, once for
// each set of migrations, and every test's database is copied from it by PostgreSQL. The caller still
// runs Migrate on the copy, which finds nothing pending only if the copy's record matches the migrations
// exactly, so a template that disagreed with them would fail the test rather than pass it unseen.
//
// The template is named after a digest of the migrations, so a changed or added migration is a new
// template, never an old one reused. Packages tested at once in separate processes share one: it is
// built under an advisory lock, and a template a crashed build left half done is not marked a template
// and is built again. It takes a *testing.T, so nothing that is not a test can call it.
func CreateMigratedForTest(t *testing.T, admin, name string, migrations []Migration) {
	t.Helper()
	template, err := templateFor(admin, migrations)
	if err != nil {
		t.Fatalf("cannot raise the migrated template: %v", err)
	}
	conn, err := pgx.Connect(t.Context(), admin)
	if err != nil {
		t.Fatalf("cannot reach the test PostgreSQL: %v", err)
	}
	defer conn.Close(context.Background())
	// A copy is refused while anything is connected to its template; the template allows no connection
	// once built, so only the builder's last connection closing can still be in the way.
	for deadline := time.Now().Add(10 * time.Second); ; {
		_, err = conn.Exec(t.Context(), "create database "+name+" template "+template)
		var pgErr *pgconn.PgError
		if err == nil || !errors.As(err, &pgErr) || pgErr.Code != "55006" || time.Now().After(deadline) {
			break
		}
		time.Sleep(20 * time.Millisecond)
	}
	if err != nil {
		t.Fatalf("cannot create %s from %s: %v", name, template, err)
	}
}

// templates are the templates this process has raised or found, by name.
var templates sync.Map // name → *templateOnce

type templateOnce struct {
	once sync.Once
	err  error
}

// templateFor is the name of the template holding migrations, raised if no process has yet.
func templateFor(admin string, migrations []Migration) (string, error) {
	digest := sha256.New()
	for _, m := range migrations {
		fmt.Fprintf(digest, "%d %s %s\n", m.Number, m.Name, m.Checksum)
	}
	sum := digest.Sum(nil)
	name := "tpl_" + hex.EncodeToString(sum[:12])
	held, _ := templates.LoadOrStore(name, &templateOnce{})
	once := held.(*templateOnce)
	once.once.Do(func() {
		// Not the asking test's context: the template outlives that test, and is raised once for all.
		ctx, cancel := context.WithTimeout(context.Background(), 2*time.Minute)
		defer cancel()
		once.err = raiseTemplate(ctx, admin, name, int64(binary.BigEndian.Uint64(sum[12:20])), migrations)
	})
	return name, once.err
}

// raiseTemplate builds the template name under the advisory lock key, unless it stands built.
func raiseTemplate(ctx context.Context, admin, name string, key int64, migrations []Migration) error {
	conn, err := pgx.Connect(ctx, admin)
	if err != nil {
		return err
	}
	defer conn.Close(context.Background())
	if _, err := conn.Exec(ctx, `select pg_advisory_lock($1)`, key); err != nil {
		return err
	}
	defer func() { _, _ = conn.Exec(context.Background(), `select pg_advisory_unlock($1)`, key) }()

	var built bool
	if err := conn.QueryRow(ctx, `select exists (select 1 from pg_database where datname = $1 and datistemplate)`,
		name).Scan(&built); err != nil {
		return err
	}
	if built {
		return nil
	}
	if _, err := conn.Exec(ctx, "drop database if exists "+name+" with (force)"); err != nil {
		return err
	}
	if _, err := conn.Exec(ctx, "create database "+name); err != nil {
		return err
	}
	cut := strings.LastIndex(admin, "/")
	rest := ""
	if q := strings.Index(admin[cut:], "?"); q >= 0 {
		rest = admin[cut+q:]
	}
	pool, err := pgxpool.New(ctx, admin[:cut]+"/"+name+rest)
	if err != nil {
		return err
	}
	_, err = (&Store{context: "template", pool: pool}).Migrate(ctx, migrations)
	pool.Close()
	if err != nil {
		return err
	}
	// Marked a template only once whole, and closed to connections, so a copy is never refused for one.
	for _, sql := range []string{
		"alter database " + name + " with allow_connections false",
		"select pg_terminate_backend(pid) from pg_stat_activity where datname = '" + name + "'",
		"alter database " + name + " with is_template true",
	} {
		if _, err := conn.Exec(ctx, sql); err != nil {
			return err
		}
	}
	return nil
}