LocalAI/core/services/testutil/testdb.go

package testutil

import (
	"context"
	"fmt"
	"net/url"
	"runtime"
	"sync"
	"sync/atomic"
	"time"

	"github.com/testcontainers/testcontainers-go"
	tcpostgres "github.com/testcontainers/testcontainers-go/modules/postgres"
	"github.com/testcontainers/testcontainers-go/wait"
	"gorm.io/driver/postgres"
	"gorm.io/gorm"
	"gorm.io/gorm/logger"

	. "github.com/onsi/ginkgo/v2"
	. "github.com/onsi/gomega"
)

// One PostgreSQL container per test PROCESS, not per spec, with a database per
// SetupTestDB call.
//
// Starting a container per spec was both slow and flaky. Slow because a
// postgres:16 start is seconds and the packages behind this helper hold several
// hundred specs; flaky because every start was a fresh chance to miss the
// readiness deadline, and a miss lands in the caller's BeforeEach as a failure
// of whichever spec happened to be running. That is the exact shape of the
// intermittent single-spec failure seen twice in this package and never
// reproduced: one spec of many, no pattern, never twice in the same place.
// Starting the container once per process leaves one chance to miss it instead
// of one per spec, and moves that chance onto a deadline that only has to be met
// while nothing else is competing for the machine.
//
// Isolation is unchanged and is what callers actually depend on: each call still
// hands back an empty database that no other spec can see. The database is
// dropped when the spec that asked for it ends. Advisory locks, sequences and
// extensions are all per-database in PostgreSQL, so nothing the packages behind
// this helper rely on leaks between specs.
//
// This mirrors the pattern already proven in tests/e2e/distributed
// (testhelpers_test.go), which is where the argument and the measurements come
// from.
//
// One container per process rather than one shared across `ginkgo -p` workers is
// deliberate: parallel Ginkgo processes are separate OS processes, each gets its
// own container, and nothing has to coordinate database names across them.
var (
	sharedOnce sync.Once
	sharedPG   *tcpostgres.PostgresContainer
	sharedDSN  string
	sharedErr  error

	// dbCounter makes each database name unique within this process. The
	// container is not shared across processes, so a process-local counter is
	// enough.
	dbCounter atomic.Int64
)

// The container outlives every spec, so its teardown belongs to the suite. This
// registers one AfterSuite in every suite that imports this package, which is
// every suite that could have started a container; it is a no-op in the ones
// that never call SetupTestDB.
//
// Package-level rather than something callers have to remember: a helper whose
// cleanup depends on 56 test files each declaring a hook is a helper that leaks
// containers the first time someone forgets. Registration happens during package
// initialisation, which is before RunSpecs, so Ginkgo is still building its tree.
var _ = AfterSuite(func() {
	if sharedPG == nil {
		return
	}
	// Best-effort: a failed terminate must not fail a suite whose specs all
	// passed. Testcontainers' reaper removes it in that case.
	_ = sharedPG.Terminate(context.Background())
})

// sharedPostgres returns the DSN of this process's PostgreSQL container,
// starting it on first use.
//
// The error is remembered rather than only asserted inside the sync.Once: an
// assertion there fails the one spec that happened to be first, and every later
// spec would then find a nil container and fail for some unrelated-looking
// reason. Re-asserting the stored error makes every affected spec say the same
// true thing.
func sharedPostgres() string {
	GinkgoHelper()

	sharedOnce.Do(func() {
		ctx := context.Background()
		sharedPG, sharedErr = tcpostgres.Run(ctx, "postgres:16",
			tcpostgres.WithDatabase("testdb"),
			tcpostgres.WithUsername("test"),
			tcpostgres.WithPassword("test"),
			// The deadline is per process now, not per spec, so it is generous
			// on purpose: it is paid once, and the cost of missing it is a
			// whole suite rather than one spec.
			testcontainers.WithWaitStrategyAndDeadline(120*time.Second,
				wait.ForLog("database system is ready to accept connections").WithOccurrence(2)),
		)
		if sharedErr != nil {
			return
		}
		sharedDSN, sharedErr = sharedPG.ConnectionString(ctx, "sslmode=disable")
	})

	Expect(sharedErr).ToNot(HaveOccurred(), "the suite's PostgreSQL container could not be started")
	return sharedDSN
}

// SetupTestDB returns a gorm.DB on a PostgreSQL database created for the calling
// spec. The database is dropped, and its connection pool closed, when the spec
// ends.
//
// It is a wrapper rather than a second implementation, so a caller that also
// needs the DSN cannot end up on a database created differently from the one
// every other suite gets.
func SetupTestDB() *gorm.DB {
	GinkgoHelper()
	db, _ := SetupTestDBWithDSN()
	return db
}

// SetupTestDBWithDSN is SetupTestDB, and additionally returns the DSN of the
// database it created.
//
// The DSN exists for callers that have to open a connection of their own rather
// than borrow one from the pool. A PostgreSQL LISTEN connection is the case
// that motivated it: LISTEN registrations belong to one backend session, so a
// pooled handle would register them on whichever connection it happened to hand
// out and lose them on the next one.
func SetupTestDBWithDSN() (*gorm.DB, string) {
	GinkgoHelper()
	if runtime.GOOS == "darwin" {
		Skip("testcontainers requires Docker, not available on macOS CI")
	}

	dsn := sharedPostgres()
	name := fmt.Sprintf("testdb_%d", dbCounter.Add(1))

	// Scoped so a failed CREATE cannot leak the pool: the assertion panics out
	// of this function, and a leaked pool per failing spec exhausts the
	// server's connection limit for every spec after it.
	//
	// CREATE and DROP DATABASE cannot run against the target database itself,
	// so both go through a short-lived connection to the container's own
	// maintenance database.
	func() {
		admin := openPool(dsn)
		defer closePool(admin)
		Expect(admin.Exec(fmt.Sprintf("CREATE DATABASE %q", name)).Error).To(Succeed())
	}()

	specDSN := replaceDBName(dsn, name)
	db, err := gorm.Open(postgres.Open(specDSN), &gorm.Config{
		Logger: logger.Default.LogMode(logger.Silent),
	})
	Expect(err).ToNot(HaveOccurred())

	DeferCleanup(func() {
		// The caller's own DeferCleanups were registered later and so run
		// first, which is what lets a spec keep using this handle in its
		// teardown.
		closePool(db)

		drop, err := openTolerantPool(dsn)
		if err != nil {
			// Reported, never asserted. A cleanup that fails the spec turns one
			// database hiccup into a failure that buries whatever the spec was
			// actually about.
			AddReportEntry("drop test database skipped", fmt.Sprintf("%s: %v", name, err))
			return
		}
		defer closePool(drop)
		// FORCE terminates whatever connections the spec left open, including
		// any a background goroutine is still holding (PostgreSQL 13+).
		if err := drop.Exec(fmt.Sprintf("DROP DATABASE IF EXISTS %q WITH (FORCE)", name)).Error; err != nil {
			AddReportEntry("drop test database failed", fmt.Sprintf("%s: %v", name, err))
		}
	})

	return db, specDSN
}

// maintenanceDSN is dsn with every server-side timeout disabled as a CONNECTION
// STARTUP OPTION rather than as a statement.
//
// The timeouts have to go because CREATE DATABASE and DROP DATABASE must not be
// bounded by anything a spec configured. A spec that sets a short
// statement_timeout on ITS own database cannot reach this connection, but a spec
// that names the maintenance database by mistake can, and that is not
// hypothetical: two advisory-lock specs did exactly that.
//
// Clearing it with `SET statement_timeout = 0` on an already-open connection is
// circular and was a real defect here: that connection has already inherited the
// database's bound, so the statement that clears the bound runs under it and can
// be aborted by it with SQLSTATE 57014. It failed roughly once in fifty at
// 8-way concurrency, which is the same invisible load-dependent single-spec
// flake this helper exists to remove. A startup option removes the circularity
// instead of buying headroom against it: the value is delivered in the startup
// packet, so the connection is already unbounded before it can run anything.
//
// The route is verified in the driver rather than assumed. pgx puts every URL
// query parameter into settings (pgconn/config.go:614), `options` is absent from
// notRuntimeParams (pgconn/config.go:340-362) so it becomes a runtime parameter
// (pgconn/config.go:374-378), and runtime parameters are copied into the startup
// message (pgconn/pgconn.go:382-388). PostgreSQL treats `options` as backend
// command-line switches, so `-c statement_timeout=0` is applied before the
// session accepts a query.
func maintenanceDSN(dsn string) (string, error) {
	u, err := url.Parse(dsn)
	if err != nil {
		return "", err
	}
	q := u.Query()
	// Percent-encoded by Encode, and pgx decodes query values before they reach
	// settings, so the server receives the switches with their spaces intact.
	q.Set("options", "-c statement_timeout=0 -c lock_timeout=0")
	u.RawQuery = q.Encode()
	return u.String(), nil
}

// openPool connects to the maintenance database with logging off and no
// server-side timeouts. Used for the short-lived maintenance connections only;
// the database a spec is handed keeps gorm's silent logger and the server's
// defaults, because setting timeouts on it is a thing specs do on purpose.
func openPool(dsn string) *gorm.DB {
	GinkgoHelper()
	db, err := openTolerantPool(dsn)
	Expect(err).ToNot(HaveOccurred())
	return db
}

// openTolerantPool is openPool for the cleanup path, which must report a
// failure rather than assert one: an assertion here would fail a spec that had
// already passed, and bury whatever the next real failure was.
//
// It carries the same startup options, and the DROP is the statement that most
// needs them: FORCE waits on terminating other sessions, measured at up to 169ms
// against the 300ms bound that used to leak here, and a DROP aborted mid-way is
// swallowed and leaks a database.
func openTolerantPool(dsn string) (*gorm.DB, error) {
	maintenance, err := maintenanceDSN(dsn)
	if err != nil {
		return nil, err
	}
	db, err := gorm.Open(postgres.Open(maintenance), &gorm.Config{Logger: logger.Discard})
	if err != nil {
		return nil, err
	}
	return db, nil
}

func closePool(db *gorm.DB) {
	if db == nil {
		return
	}
	if sqlDB, err := db.DB(); err == nil {
		_ = sqlDB.Close()
	}
}

// replaceDBName swaps the database component of a DSN, preserving credentials,
// host, port and query parameters.
func replaceDBName(dsn, name string) string {
	GinkgoHelper()
	u, err := url.Parse(dsn)
	Expect(err).ToNot(HaveOccurred())
	u.Path = "/" + name
	return u.String()
}