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 }