package dbtestutil

import (

"context"

"crypto/sha256"

"database/sql"

"encoding/hex"

"errors"

"fmt"

"net"

"os"

"path/filepath"

"strconv"

"strings"

"sync"

"time"

"github.com/cenkalti/backoff/v4"

"github.com/gofrs/flock"

"github.com/ory/dockertest/v3"

"github.com/ory/dockertest/v3/docker"

"golang.org/x/xerrors"

"github.com/coder/coder/v2/coderd/database/migrations"

"github.com/coder/retry"

)

const postgresImage = "us-docker.pkg.dev/coder-v2-images-public/public/postgres"

type ConnectionParams struct {

Username string

Password string

Host string

Port string

DBName string

}

func (p ConnectionParams) DSN() string {

return fmt.Sprintf("postgres://%s:%s@%s:%s/%s?sslmode=disable", p.Username, p.Password, p.Host, p.Port, p.DBName)

}

// These variables are global because all tests share them.

var (

connectionParamsInitOnce sync.Once

defaultConnectionParams ConnectionParams

errDefaultConnectionParamsInit error

retryableErrSubstrings = []string{

"connection reset by peer",

}

noPostgresRunningErrSubstrings = []string{

"connection refused", // nothing is listening on the port

"No connection could be made", // Windows variant of the above

}

DefaultBroker = Broker{}

)

// initDefaultConnection initializes the default postgres connection parameters.

// It first checks if the database is running at localhost:5432. If it is, it will

// use that database. If it's not, it will start a new container and use that.

func initDefaultConnection(t TBSubset) error {

params := ConnectionParams{

Username: "postgres",

Password: "postgres",

Host: "127.0.0.1",

Port: "5432",

DBName: "postgres",

}

dsn := params.DSN()

// Helper closure to try opening and pinging the default Postgres instance.

// Used within a single retry loop that handles both retryable and permanent errors.

attemptConn := func() error {

db, err := sql.Open("postgres", dsn)

if err == nil {

err = db.Ping()

if closeErr := db.Close(); closeErr != nil {

return xerrors.Errorf("close db: %w", closeErr)

}

}

return err

}

var dbErr error

// Retry up to 10 seconds for temporary errors.

ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)

defer cancel()

for r := retry.New(10*time.Millisecond, 500*time.Millisecond); r.Wait(ctx); {

dbErr = attemptConn()

if dbErr == nil {

break

}

errString := dbErr.Error()

if !containsAnySubstring(errString, retryableErrSubstrings) {

break

}

t.Logf("%s failed to connect to postgres, retrying: %s", time.Now().Format(time.StampMilli), errString)

}

// After the loop dbErr is the last connection error (if any).

if dbErr != nil && containsAnySubstring(dbErr.Error(), noPostgresRunningErrSubstrings) {

// If there's no database running on the default port, we'll start a

// postgres container. We won't be cleaning it up so it can be reused

// by subsequent tests. It'll keep on running until the user terminates

// it manually.

container, _, err := openContainer(t, DBContainerOptions{

Name: "coder-test-postgres",

Port: 5432,

})

if err != nil {

return xerrors.Errorf("open container: %w", err)

}

params.Host = container.Host

params.Port = container.Port

dsn = params.DSN()

// Retry connecting for at most 10 seconds.

// The fact that openContainer succeeded does not

// mean that port forwarding is ready.

for r := retry.New(100*time.Millisecond, 10*time.Second); r.Wait(context.Background()); {

db, connErr := sql.Open("postgres", dsn)

if connErr == nil {

connErr = db.Ping()

if closeErr := db.Close(); closeErr != nil {

return xerrors.Errorf("close db, container: %w", closeErr)

}

}

if connErr == nil {

break

}

t.Logf("failed to connect to postgres after starting container, may retry: %s", connErr.Error())

}

} else if dbErr != nil {

return xerrors.Errorf("open postgres connection: %w", dbErr)

}

defaultConnectionParams = params

return nil

}

type OpenOptions struct {

DBFrom *string

LogDSN bool

}

type OpenOption func(*OpenOptions)

// WithDBFrom sets the template database to use when creating a new database.

// Overrides the DB_FROM environment variable.

func WithDBFrom(dbFrom string) OpenOption {

return func(o *OpenOptions) {

o.DBFrom = &dbFrom

}

}

// WithLogDSN sets whether the DSN should be logged during testing.

// This provides an ergonomic way to connect to test databases during debugging.

func WithLogDSN(logDSN bool) OpenOption {

return func(o *OpenOptions) {

o.LogDSN = logDSN

}

}

// TBSubset is a subset of the testing.TB interface.

// It allows to use dbtestutil.Open outside of tests.

type TBSubset interface {

Name() string

Cleanup(func())

Helper()

Logf(format string, args ...any)

TempDir() string

}

// Open creates a new PostgreSQL database instance.

// If there's a database running at localhost:5432, it will use that.

// Otherwise, it will start a new postgres container.

func Open(t TBSubset, opts ...OpenOption) (string, error) {

t.Helper()

params, err := DefaultBroker.Create(t, opts...)

if err != nil {

return "", err

}

return params.DSN(), nil

}

// createDatabaseFromTemplate creates a new database from a template database.

// If templateDBName is empty, it will create a new template database based on

// the current migrations, and name it "tpl_<migrations_hash>". Or if it's

// already been created, it will use that.

func createDatabaseFromTemplate(t TBSubset, connParams ConnectionParams, db *sql.DB, newDBName string, templateDBName string) error {

t.Helper()

emptyTemplateDBName := templateDBName == ""

if emptyTemplateDBName {

templateDBName = fmt.Sprintf("tpl_%s", migrations.GetMigrationsHash()[:32])

}

_, err := db.Exec("CREATE DATABASE " + newDBName + " WITH TEMPLATE " + templateDBName)

if err == nil {

// Template database already exists and we successfully created the new database.

return nil

}

tplDbDoesNotExistOccurred := strings.Contains(err.Error(), "template database") && strings.Contains(err.Error(), "does not exist")

if (tplDbDoesNotExistOccurred && !emptyTemplateDBName) || !tplDbDoesNotExistOccurred {

// First and case: user passed a templateDBName that doesn't exist.

// Second and case: some other error.

return xerrors.Errorf("create db with template: %w", err)

}

if !emptyTemplateDBName {

// sanity check

panic("templateDBName is not empty. there's a bug in the code above")

}

// The templateDBName is empty, so we need to create the template database.

err = createAndInitDatabase(t, connParams, db, templateDBName, func(tplDb *sql.DB) error {

if err := migrations.Up(tplDb); err != nil {

return xerrors.Errorf("migrate template db: %w", err)

}

return nil

})

if err != nil {

return xerrors.Errorf("create template database: %w", err)

}

// Try to create the database again now that a template exists.

if _, err = db.Exec("CREATE DATABASE " + newDBName + " WITH TEMPLATE " + templateDBName); err != nil {

return xerrors.Errorf("create db with template after migrations: %w", err)

}

return nil

}

func createAndInitDatabase(t TBSubset, connParams ConnectionParams, db *sql.DB, name string, initialize func(*sql.DB) error) error {

// We will use a tx to obtain a lock, so another test or process doesn't race with us.

tx, err := db.BeginTx(context.Background(), nil)

if err != nil {

return xerrors.Errorf("begin tx: %w", err)

}

// we only use the transaction for locking and querying, so it's fine to always roll it back.

defer func() {

err := tx.Rollback()

if err != nil && !errors.Is(err, sql.ErrTxDone) {

t.Logf("create database: failed to rollback tx: %s\n", err.Error())

}

}()

// 2137 is an arbitrary number. We just need a lock that is unique to creating

// the database.

_, err = tx.Exec("SELECT pg_advisory_xact_lock(2137)")

if err != nil {

return xerrors.Errorf("acquire lock: %w", err)

}

// Someone else might have created the db while we were waiting.

dbExistsRes, err := tx.Query("SELECT 1 FROM pg_database WHERE datname = $1", name)

if err != nil {

return xerrors.Errorf("check if db exists: %w", err)

}

dbAlreadyExists := dbExistsRes.Next()

if err := dbExistsRes.Close(); err != nil {

return xerrors.Errorf("close tpl db exists res: %w", err)

}

if dbAlreadyExists {

return nil

}

// We will use a temporary database to avoid race conditions. We will

// rename it to the real database name after we're sure it was fully

// initialized.

// It's dropped here to ensure that if a previous run of this function failed

// midway, we don't encounter issues with the temporary database still existing.

tmpDBName := "tmp_" + name

// We're using db instead of tx here because you can't run `DROP DATABASE` inside

// a transaction.

if _, err := db.Exec("DROP DATABASE IF EXISTS " + tmpDBName); err != nil {

return xerrors.Errorf("drop tmp db: %w", err)

}

if _, err := db.Exec("CREATE DATABASE " + tmpDBName); err != nil {

return xerrors.Errorf("create tmp db: %w", err)

}

tmpDbURL := ConnectionParams{

Username: connParams.Username,

Password: connParams.Password,

Host: connParams.Host,

Port: connParams.Port,

DBName: tmpDBName,

}.DSN()

tmpDb, err := sql.Open("postgres", tmpDbURL)

if err != nil {

return xerrors.Errorf("connect to template db: %w", err)

}

defer func() {

if err := tmpDb.Close(); err != nil {

t.Logf("failed to close temp db: %s\n", err.Error())

}

}()

if err := initialize(tmpDb); err != nil {

return xerrors.Errorf("initialize: %w", err)

}

if err := tmpDb.Close(); err != nil {

return xerrors.Errorf("close template db: %w", err)

}

if _, err := db.Exec("ALTER DATABASE " + tmpDBName + " RENAME TO " + name); err != nil {

return xerrors.Errorf("rename tmp db: %w", err)

}

return nil

}

type DBContainerOptions struct {

Port int

Name string

}

type container struct {

Resource *dockertest.Resource

Pool *dockertest.Pool

Host string

Port string

}

// OpenContainer creates a new PostgreSQL server using a Docker container. If port is nonzero, forward host traffic

// to that port to the database. If port is zero, allocate a free port from the OS.

// If name is set, we'll ensure that only one container is started with that name. If it's already running, we'll use that.

// Otherwise, we'll start a new container.

func openContainer(t TBSubset, opts DBContainerOptions) (container, func(), error) {

if opts.Name != "" {

// We only want to start the container once per unique name,

// so we take an inter-process lock to avoid concurrent test runs

// racing with us.

nameHash := sha256.Sum256([]byte(opts.Name))

nameHashStr := hex.EncodeToString(nameHash[:])

lock := flock.New(filepath.Join(os.TempDir(), "coder-postgres-container-"+nameHashStr[:8]))

if err := lock.Lock(); err != nil {

return container{}, nil, xerrors.Errorf("lock: %w", err)

}

defer func() {

err := lock.Unlock()

if err != nil {

t.Logf("create database from template: failed to unlock: %s\n", err.Error())

}

}()

}

pool, err := dockertest.NewPool("")

if err != nil {

return container{}, nil, xerrors.Errorf("create pool: %w", err)

}

var resource *dockertest.Resource

var tempDir string

if opts.Name != "" {

// If the container already exists, we'll use it.

resource, _ = pool.ContainerByName(opts.Name)

}

if resource == nil {

tempDir, err = os.MkdirTemp(os.TempDir(), "postgres")

if err != nil {

return container{}, nil, xerrors.Errorf("create tempdir: %w", err)

}

runOptions := dockertest.RunOptions{

Repository: postgresImage,

Tag: strconv.Itoa(minimumPostgreSQLVersion),

Env: []string{

"POSTGRES_PASSWORD=postgres",

"POSTGRES_USER=postgres",

"POSTGRES_DB=postgres",

// The location for temporary database files!

"PGDATA=/tmp",

"listen_addresses = '*'",

},

PortBindings: map[docker.Port][]docker.PortBinding{

"5432/tcp": {{

// Manually specifying a host IP tells Docker just to use an IPV4 address.

// If we don't do this, we hit a fun bug:

// https://github.com/moby/moby/issues/42442

// where the ipv4 and ipv6 ports might be _different_ and collide with other running docker containers.

HostIP: "0.0.0.0",

HostPort: strconv.FormatInt(int64(opts.Port), 10),

}},

},

Mounts: []string{

// The postgres image has a VOLUME parameter in it's image.

// If we don't mount at this point, Docker will allocate a

// volume for this directory.

//

// This isn't used anyways, since we override PGDATA.

fmt.Sprintf("%s:/var/lib/postgresql/data", tempDir),

},

Cmd: []string{"-c", "max_connections=1000"},

}

if opts.Name != "" {

runOptions.Name = opts.Name

}

resource, err = pool.RunWithOptions(&runOptions, func(config *docker.HostConfig) {

// set AutoRemove to true so that stopped container goes away by itself

config.AutoRemove = true

config.RestartPolicy = docker.RestartPolicy{Name: "no"}

config.Tmpfs = map[string]string{

"/tmp": "rw",

}

})

if err != nil {

return container{}, nil, xerrors.Errorf("could not start resource: %w", err)

}

}

hostAndPort := resource.GetHostPort("5432/tcp")

host, port, err := net.SplitHostPort(hostAndPort)

if err != nil {

return container{}, nil, xerrors.Errorf("split host and port: %w", err)

}

for r := retry.New(50*time.Millisecond, 15*time.Second); r.Wait(context.Background()); {

stdout := &strings.Builder{}

stderr := &strings.Builder{}

_, err = resource.Exec([]string{"pg_isready", "-h", "127.0.0.1"}, dockertest.ExecOptions{

StdOut: stdout,

StdErr: stderr,

})

if err == nil {

break

}

}

if err != nil {

return container{}, nil, xerrors.Errorf("pg_isready: %w", err)

}

return container{

Host: host,

Port: port,

Resource: resource,

Pool: pool,

}, func() {

_ = pool.Purge(resource)

if tempDir != "" {

_ = os.RemoveAll(tempDir)

}

}, nil

}

// OpenContainerized creates a new PostgreSQL server using a Docker container. If port is nonzero, forward host traffic

// to that port to the database. If port is zero, allocate a free port from the OS.

// The user is responsible for calling the returned cleanup function.

func OpenContainerized(t TBSubset, opts DBContainerOptions) (string, func(), error) {

container, containerCleanup, err := openContainer(t, opts)

if err != nil {

return "", nil, xerrors.Errorf("open container: %w", err)

}

defer func() {

if err != nil {

containerCleanup()

}

}()

dbURL := ConnectionParams{

Username: "postgres",

Password: "postgres",

Host: container.Host,

Port: container.Port,

DBName: "postgres",

}.DSN()

// Docker should hard-kill the container after 120 seconds.

err = container.Resource.Expire(120)

if err != nil {

return "", nil, xerrors.Errorf("expire resource: %w", err)

}

container.Pool.MaxWait = 120 * time.Second

// Record the error that occurs during the retry.

// The 'pool' pkg hardcodes a deadline error devoid

// of any useful context.

var retryErr error

err = container.Pool.Retry(func() error {

db, err := sql.Open("postgres", dbURL)

if err != nil {

retryErr = xerrors.Errorf("open postgres: %w", err)

return retryErr

}

defer db.Close()

err = db.Ping()

if err != nil {

retryErr = xerrors.Errorf("ping postgres: %w", err)

return retryErr

}

err = migrations.Up(db)

if err != nil {

retryErr = xerrors.Errorf("migrate db: %w", err)

// Only try to migrate once.

return backoff.Permanent(retryErr)

}

return nil

})

if err != nil {

return "", nil, retryErr

}

return dbURL, containerCleanup, nil

}

func containsAnySubstring(s string, substrings []string) bool {

for _, substr := range substrings {

if strings.Contains(s, substr) {

return true

}

}

return false

}