From da4e2d547b212f6fc48516bb9863cc321c41a509 Mon Sep 17 00:00:00 2001
From: Hugo Dutka <hugo@coder.com>
Date: Thu, 31 Oct 2024 17:03:33 +0000
Subject: [PATCH 01/10] utilities for creating postgres databases from
 templates

---
 coderd/database/dbtestutil/postgres.go      | 481 ++++++++++++++++----
 coderd/database/migrations/migrate.go       |  55 +++
 coderd/database/pubsub/pubsub_linux_test.go |   4 +-
 3 files changed, 456 insertions(+), 84 deletions(-)

diff --git a/coderd/database/dbtestutil/postgres.go b/coderd/database/dbtestutil/postgres.go
index 3a559778b6968..8ed947c2bc7b7 100644
--- a/coderd/database/dbtestutil/postgres.go
+++ b/coderd/database/dbtestutil/postgres.go
@@ -1,13 +1,22 @@
 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"
@@ -16,119 +25,430 @@ import (
 	"github.com/coder/coder/v2/cryptorand"
 )
 
-// Open creates a new PostgreSQL database instance.  With DB_FROM environment variable set, it clones a database
-// from the provided template.  With the environment variable unset, it creates a new Docker container running 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)
+}
+
+var (
+	connectionParamsInitOnce       sync.Once
+	defaultConnectionParams        ConnectionParams
+	errDefaultConnectionParamsInit error
+)
+
+// initDefaultConnectionParams initializes the default connection parameters
+// by checking if the database is running locally. If it is, it will use the
+// local database. If it's not, it will start a new container and use that.
+func initDefaultConnectionParams() error {
+	params := ConnectionParams{
+		Username: "postgres",
+		Password: "postgres",
+		Host:     "127.0.0.1",
+		Port:     "5432",
+		DBName:   "postgres",
+	}
+	dsn := params.DSN()
+	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)
+		}
+	}
+	shouldOpenContainer := false
+	if err != nil {
+		errSubstrings := []string{
+			"connection refused",          // this happens on Linux when there's nothing listening on the port
+			"No connection could be made", // like above but Windows
+		}
+		errString := err.Error()
+		for _, errSubstring := range errSubstrings {
+			if strings.Contains(errString, errSubstring) {
+				shouldOpenContainer = true
+				break
+			}
+		}
+	}
+	if err != nil && shouldOpenContainer {
+		// 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 themselves.
+		container, _, err := openContainer(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 a cumulative 5 seconds.
+		// The fact that openContainer succeeded does not
+		// mean that port forwarding is ready.
+		for i := 0; i < 20; i++ {
+			db, err = sql.Open("postgres", dsn)
+			if err == nil {
+				err = db.Ping()
+				if closeErr := db.Close(); closeErr != nil {
+					return xerrors.Errorf("close db, container: %w", closeErr)
+				}
+			}
+			if err == nil {
+				break
+			}
+			time.Sleep(250 * time.Millisecond)
+		}
+	} else if err != nil {
+		return xerrors.Errorf("open postgres connection: %w", err)
+	}
+	defaultConnectionParams = params
+	return nil
+}
+
+// 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() (string, func(), error) {
-	if os.Getenv("DB_FROM") != "" {
-		// In CI, creating a Docker container for each test is slow.
-		// This expects a PostgreSQL instance with the hardcoded credentials
-		// available.
-		dbURL := "postgres://postgres:postgres@127.0.0.1:5432/postgres?sslmode=disable"
-		db, err := sql.Open("postgres", dbURL)
+	connectionParamsInitOnce.Do(func() {
+		errDefaultConnectionParamsInit = initDefaultConnectionParams()
+	})
+	if errDefaultConnectionParamsInit != nil {
+		return "", func() {}, xerrors.Errorf("init default connection params: %w", errDefaultConnectionParamsInit)
+	}
+
+	var (
+		username = defaultConnectionParams.Username
+		password = defaultConnectionParams.Password
+		host     = defaultConnectionParams.Host
+		port     = defaultConnectionParams.Port
+	)
+
+	// Use a time-based prefix to make it easier to find the database
+	// when debugging.
+	now := time.Now().Format("test_2006_01_02_15_04_05")
+	dbSuffix, err := cryptorand.StringCharset(cryptorand.Lower, 10)
+	if err != nil {
+		return "", func() {}, xerrors.Errorf("generate db suffix: %w", err)
+	}
+	dbName := now + "_" + dbSuffix
+
+	// if empty createDatabaseFromTemplate will create a new template db
+	templateDBName := os.Getenv("DB_FROM")
+	if err = createDatabaseFromTemplate(ConnectionParams{
+		Username: username,
+		Password: password,
+		Host:     host,
+		Port:     port,
+		DBName:   dbName,
+	}, dbName, templateDBName); err != nil {
+		return "", func() {}, xerrors.Errorf("create database: %w", err)
+	}
+
+	cleanup := func() {
+		cleanupDbURL := defaultConnectionParams.DSN()
+		cleanupConn, err := sql.Open("postgres", cleanupDbURL)
 		if err != nil {
-			return "", nil, xerrors.Errorf("connect to ci postgres: %w", err)
+			_, _ = fmt.Fprintf(os.Stderr, "cleanup database %q: failed to connect to postgres: %s\n", dbName, err.Error())
 		}
+		defer cleanupConn.Close()
+		_, err = cleanupConn.Exec("DROP DATABASE " + dbName + ";")
+		if err != nil {
+			_, _ = fmt.Fprintf(os.Stderr, "failed to clean up database %q: %s\n", dbName, err.Error())
+		}
+	}
 
-		defer db.Close()
+	dsn := ConnectionParams{
+		Username: username,
+		Password: password,
+		Host:     host,
+		Port:     port,
+		DBName:   dbName,
+	}.DSN()
+	return dsn, cleanup, nil
+}
 
-		dbName, err := cryptorand.StringCharset(cryptorand.Lower, 10)
+// 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(connParams ConnectionParams, newDBName string, templateDBName string) error {
+	dbURL := connParams.DSN()
+	db, err := sql.Open("postgres", dbURL)
+	if err != nil {
+		return xerrors.Errorf("connect to postgres: %w", err)
+	}
+	defer db.Close()
+
+	if templateDBName == "" {
+		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
+	}
+	tplDbDoesNotExist := strings.Contains(err.Error(), "template database") && strings.Contains(err.Error(), "does not exist")
+	if (tplDbDoesNotExist && templateDBName != "") || (!tplDbDoesNotExist) {
+		// 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 templateDBName != "" {
+		// 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.
+	// 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)
+	}
+	defer func() {
+		err := tx.Rollback()
+		if err != nil && !errors.Is(err, sql.ErrTxDone) {
+			panic(err)
+		}
+	}()
+	_, 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 template db while we were waiting.
+	tplDbExistsRes, err := tx.Query("SELECT 1 FROM pg_database WHERE datname = $1", templateDBName)
+	if err != nil {
+		return xerrors.Errorf("check if db exists: %w", err)
+	}
+	tplDbAlreadyExists := tplDbExistsRes.Next()
+	if !tplDbAlreadyExists {
+		// We will use a temporary template database to avoid race conditions. We will
+		// rename it to the real template 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.
+		tmpTemplateDBName := "tmp_" + templateDBName
+		_, err = db.Exec("DROP DATABASE IF EXISTS " + tmpTemplateDBName)
 		if err != nil {
-			return "", nil, xerrors.Errorf("generate db name: %w", err)
+			return xerrors.Errorf("drop tmp template db: %w", err)
 		}
 
-		dbName = "ci" + dbName
-		_, err = db.Exec("CREATE DATABASE " + dbName + " WITH TEMPLATE " + os.Getenv("DB_FROM"))
+		_, err = db.Exec("CREATE DATABASE " + tmpTemplateDBName)
+		if err != nil {
+			return xerrors.Errorf("create tmp template db: %w", err)
+		}
+		tplDbURL := ConnectionParams{
+			Username: connParams.Username,
+			Password: connParams.Password,
+			Host:     connParams.Host,
+			Port:     connParams.Port,
+			DBName:   tmpTemplateDBName,
+		}.DSN()
+		tplDb, err := sql.Open("postgres", tplDbURL)
+		if err != nil {
+			return xerrors.Errorf("connect to template db: %w", err)
+		}
+		defer tplDb.Close()
+		if err := migrations.Up(tplDb); err != nil {
+			return xerrors.Errorf("migrate template db: %w", err)
+		}
+		if err := tplDb.Close(); err != nil {
+			return xerrors.Errorf("close template db: %w", err)
+		}
+		_, err = db.Exec("ALTER DATABASE " + tmpTemplateDBName + " RENAME TO " + templateDBName)
 		if err != nil {
-			return "", nil, xerrors.Errorf("create db with template: %w", err)
+			return xerrors.Errorf("rename tmp template db: %w", err)
 		}
+	}
 
-		dsn := "postgres://postgres:postgres@127.0.0.1:5432/" + dbName + "?sslmode=disable"
-		// Normally this would get cleaned up by removing the container but if we
-		// reuse the same container for multiple tests we run the risk of filling
-		// up our disk. Avoid this!
-		cleanup := func() {
-			cleanupConn, err := sql.Open("postgres", dbURL)
-			if err != nil {
-				_, _ = fmt.Fprintf(os.Stderr, "cleanup database %q: failed to connect to postgres: %s\n", dbName, err.Error())
-			}
-			defer cleanupConn.Close()
-			_, err = cleanupConn.Exec("DROP DATABASE " + dbName + ";")
+	// Try to create the database again now that a template exists.
+	_, err = db.Exec("CREATE DATABASE " + newDBName + " WITH TEMPLATE " + templateDBName)
+	if err != nil {
+		return xerrors.Errorf("create db with template after migrations: %w", err)
+	}
+
+	err = tx.Commit()
+	if err != nil {
+		return xerrors.Errorf("commit tx: %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(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 {
-				_, _ = fmt.Fprintf(os.Stderr, "failed to clean up database %q: %s\n", dbName, err.Error())
+				panic(err)
 			}
-		}
-		return dsn, cleanup, nil
+		}()
 	}
-	return OpenContainerized(0)
-}
 
-// 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.
-func OpenContainerized(port int) (string, func(), error) {
 	pool, err := dockertest.NewPool("")
 	if err != nil {
-		return "", nil, xerrors.Errorf("create pool: %w", err)
+		return container{}, nil, xerrors.Errorf("create pool: %w", err)
 	}
 
 	tempDir, err := os.MkdirTemp(os.TempDir(), "postgres")
 	if err != nil {
-		return "", nil, xerrors.Errorf("create tempdir: %w", err)
-	}
-
-	resource, err := pool.RunWithOptions(&dockertest.RunOptions{
-		Repository: "gcr.io/coder-dev-1/postgres",
-		Tag:        "13",
-		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(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),
-		},
-	}, 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"}
-	})
-	if err != nil {
-		return "", nil, xerrors.Errorf("could not start resource: %w", err)
+		return container{}, nil, xerrors.Errorf("create tempdir: %w", err)
+	}
+
+	var resource *dockertest.Resource
+	if opts.Name != "" {
+		// If the container already exists, we'll use it.
+		resource, _ = pool.ContainerByName(opts.Name)
+	}
+	if resource == nil {
+		runOptions := dockertest.RunOptions{
+			Repository: "gcr.io/coder-dev-1/postgres",
+			Tag:        "13",
+			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")
-	dbURL := fmt.Sprintf("postgres://postgres:postgres@%s/postgres?sslmode=disable", hostAndPort)
+	host, port, err := net.SplitHostPort(hostAndPort)
+	if err != nil {
+		return container{}, nil, xerrors.Errorf("split host and port: %w", err)
+	}
+
+	// wait for a cumulative 60 * 250ms = 15 seconds for the database to start
+	for i := 0; i < 60; i++ {
+		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
+		}
+		time.Sleep(250 * time.Millisecond)
+	}
+	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)
+			_ = 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.
+func OpenContainerized(opts DBContainerOptions) (string, func(), error) {
+	container, containerCleanup, err := openContainer(opts)
+	defer func() {
+		if err != nil {
+			containerCleanup()
+		}
+	}()
+	if err != nil {
+		return "", nil, xerrors.Errorf("open container: %w", err)
+	}
+	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 = resource.Expire(120)
+	err = container.Resource.Expire(120)
 	if err != nil {
 		return "", nil, xerrors.Errorf("expire resource: %w", err)
 	}
 
-	pool.MaxWait = 120 * time.Second
+	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 = pool.Retry(func() error {
+	err = container.Pool.Retry(func() error {
 		db, err := sql.Open("postgres", dbURL)
 		if err != nil {
 			retryErr = xerrors.Errorf("open postgres: %w", err)
@@ -155,8 +475,5 @@ func OpenContainerized(port int) (string, func(), error) {
 		return "", nil, retryErr
 	}
 
-	return dbURL, func() {
-		_ = pool.Purge(resource)
-		_ = os.RemoveAll(tempDir)
-	}, nil
+	return dbURL, containerCleanup, nil
 }
diff --git a/coderd/database/migrations/migrate.go b/coderd/database/migrations/migrate.go
index 213408bbadd8c..c6c1b5740f873 100644
--- a/coderd/database/migrations/migrate.go
+++ b/coderd/database/migrations/migrate.go
@@ -2,11 +2,16 @@ package migrations
 
 import (
 	"context"
+	"crypto/sha256"
 	"database/sql"
 	"embed"
 	"errors"
+	"fmt"
 	"io/fs"
 	"os"
+	"sort"
+	"strings"
+	"sync"
 
 	"github.com/golang-migrate/migrate/v4"
 	"github.com/golang-migrate/migrate/v4/source"
@@ -17,6 +22,56 @@ import (
 //go:embed *.sql
 var migrations embed.FS
 
+var (
+	migrationsHash     string
+	migrationsHashOnce sync.Once
+)
+
+// A migrations hash is a sha256 hash of the contents and names
+// of the migrations sorted by filename.
+func calculateMigrationsHash(migrationsFs embed.FS) (string, error) {
+	files, err := migrationsFs.ReadDir(".")
+	if err != nil {
+		return "", xerrors.Errorf("read migrations directory: %w", err)
+	}
+	sortedFiles := make([]fs.DirEntry, len(files))
+	copy(sortedFiles, files)
+	sort.Slice(sortedFiles, func(i, j int) bool {
+		return sortedFiles[i].Name() < sortedFiles[j].Name()
+	})
+
+	var builder strings.Builder
+	for _, file := range sortedFiles {
+		if _, err := builder.WriteString(file.Name()); err != nil {
+			return "", xerrors.Errorf("write migration file name %q: %w", file.Name(), err)
+		}
+		content, err := migrationsFs.ReadFile(file.Name())
+		if err != nil {
+			return "", xerrors.Errorf("read migration file %q: %w", file.Name(), err)
+		}
+		if _, err := builder.Write(content); err != nil {
+			return "", xerrors.Errorf("write migration file content %q: %w", file.Name(), err)
+		}
+	}
+
+	hash := sha256.New()
+	if _, err := hash.Write([]byte(builder.String())); err != nil {
+		return "", xerrors.Errorf("write to hash: %w", err)
+	}
+	return fmt.Sprintf("%x", hash.Sum(nil)), nil
+}
+
+func GetMigrationsHash() string {
+	migrationsHashOnce.Do(func() {
+		hash, err := calculateMigrationsHash(migrations)
+		if err != nil {
+			panic(err)
+		}
+		migrationsHash = hash
+	})
+	return migrationsHash
+}
+
 func setup(db *sql.DB, migs fs.FS) (source.Driver, *migrate.Migrate, error) {
 	if migs == nil {
 		migs = migrations
diff --git a/coderd/database/pubsub/pubsub_linux_test.go b/coderd/database/pubsub/pubsub_linux_test.go
index f208af921b441..a4be5927a4725 100644
--- a/coderd/database/pubsub/pubsub_linux_test.go
+++ b/coderd/database/pubsub/pubsub_linux_test.go
@@ -167,7 +167,7 @@ const disconnectTestPort = 26892
 func TestPubsub_Disconnect(t *testing.T) {
 	// we always use a Docker container for this test, even in CI, since we need to be able to kill
 	// postgres and bring it back on the same port.
-	connectionURL, closePg, err := dbtestutil.OpenContainerized(disconnectTestPort)
+	connectionURL, closePg, err := dbtestutil.OpenContainerized(dbtestutil.DBContainerOptions{Port: disconnectTestPort})
 	require.NoError(t, err)
 	defer closePg()
 	db, err := sql.Open("postgres", connectionURL)
@@ -238,7 +238,7 @@ func TestPubsub_Disconnect(t *testing.T) {
 
 	// restart postgres on the same port --- since we only use LISTEN/NOTIFY it doesn't
 	// matter that the new postgres doesn't have any persisted state from before.
-	_, closeNewPg, err := dbtestutil.OpenContainerized(disconnectTestPort)
+	_, closeNewPg, err := dbtestutil.OpenContainerized(dbtestutil.DBContainerOptions{Port: disconnectTestPort})
 	require.NoError(t, err)
 	defer closeNewPg()
 

From 9635fc78cb0c455a5140aa0f4752c9baa6c4a340 Mon Sep 17 00:00:00 2001
From: Hugo Dutka <hugo@coder.com>
Date: Thu, 31 Oct 2024 18:59:30 +0000
Subject: [PATCH 02/10] tests

---
 coderd/database/dbtestutil/postgres.go      | 74 +++++++++++-------
 coderd/database/dbtestutil/postgres_test.go | 84 ++++++++++++++++++---
 2 files changed, 122 insertions(+), 36 deletions(-)

diff --git a/coderd/database/dbtestutil/postgres.go b/coderd/database/dbtestutil/postgres.go
index 8ed947c2bc7b7..25ccfd0fd11ae 100644
--- a/coderd/database/dbtestutil/postgres.go
+++ b/coderd/database/dbtestutil/postgres.go
@@ -115,10 +115,24 @@ func initDefaultConnectionParams() error {
 	return nil
 }
 
+type OpenOptions struct {
+	DBFrom *string
+}
+
+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
+	}
+}
+
 // 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() (string, func(), error) {
+func Open(opts ...OpenOption) (string, func(), error) {
 	connectionParamsInitOnce.Do(func() {
 		errDefaultConnectionParamsInit = initDefaultConnectionParams()
 	})
@@ -126,6 +140,11 @@ func Open() (string, func(), error) {
 		return "", func() {}, xerrors.Errorf("init default connection params: %w", errDefaultConnectionParamsInit)
 	}
 
+	openOptions := OpenOptions{}
+	for _, opt := range opts {
+		opt(&openOptions)
+	}
+
 	var (
 		username = defaultConnectionParams.Username
 		password = defaultConnectionParams.Password
@@ -144,13 +163,10 @@ func Open() (string, func(), error) {
 
 	// if empty createDatabaseFromTemplate will create a new template db
 	templateDBName := os.Getenv("DB_FROM")
-	if err = createDatabaseFromTemplate(ConnectionParams{
-		Username: username,
-		Password: password,
-		Host:     host,
-		Port:     port,
-		DBName:   dbName,
-	}, dbName, templateDBName); err != nil {
+	if openOptions.DBFrom != nil {
+		templateDBName = *openOptions.DBFrom
+	}
+	if err = createDatabaseFromTemplate(defaultConnectionParams, dbName, templateDBName); err != nil {
 		return "", func() {}, xerrors.Errorf("create database: %w", err)
 	}
 
@@ -187,9 +203,14 @@ func createDatabaseFromTemplate(connParams ConnectionParams, newDBName string, t
 	if err != nil {
 		return xerrors.Errorf("connect to postgres: %w", err)
 	}
-	defer db.Close()
+	defer func() {
+		if err := db.Close(); err != nil {
+			panic(err)
+		}
+	}()
 
-	if templateDBName == "" {
+	emptyTemplateDBName := templateDBName == ""
+	if emptyTemplateDBName {
 		templateDBName = fmt.Sprintf("tpl_%s", migrations.GetMigrationsHash()[:32])
 	}
 	_, err = db.Exec("CREATE DATABASE " + newDBName + " WITH TEMPLATE " + templateDBName)
@@ -197,13 +218,13 @@ func createDatabaseFromTemplate(connParams ConnectionParams, newDBName string, t
 		// Template database already exists and we successfully created the new database.
 		return nil
 	}
-	tplDbDoesNotExist := strings.Contains(err.Error(), "template database") && strings.Contains(err.Error(), "does not exist")
-	if (tplDbDoesNotExist && templateDBName != "") || (!tplDbDoesNotExist) {
+	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 templateDBName != "" {
+	if !emptyTemplateDBName {
 		// sanity check
 		panic("templateDBName is not empty. there's a bug in the code above")
 	}
@@ -230,6 +251,9 @@ func createDatabaseFromTemplate(connParams ConnectionParams, newDBName string, t
 		return xerrors.Errorf("check if db exists: %w", err)
 	}
 	tplDbAlreadyExists := tplDbExistsRes.Next()
+	if err := tplDbExistsRes.Close(); err != nil {
+		return xerrors.Errorf("close tpl db exists res: %w", err)
+	}
 	if !tplDbAlreadyExists {
 		// We will use a temporary template database to avoid race conditions. We will
 		// rename it to the real template database name after we're sure it was fully
@@ -237,13 +261,10 @@ func createDatabaseFromTemplate(connParams ConnectionParams, newDBName string, t
 		// 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.
 		tmpTemplateDBName := "tmp_" + templateDBName
-		_, err = db.Exec("DROP DATABASE IF EXISTS " + tmpTemplateDBName)
-		if err != nil {
+		if _, err := db.Exec("DROP DATABASE IF EXISTS " + tmpTemplateDBName); err != nil {
 			return xerrors.Errorf("drop tmp template db: %w", err)
 		}
-
-		_, err = db.Exec("CREATE DATABASE " + tmpTemplateDBName)
-		if err != nil {
+		if _, err := db.Exec("CREATE DATABASE " + tmpTemplateDBName); err != nil {
 			return xerrors.Errorf("create tmp template db: %w", err)
 		}
 		tplDbURL := ConnectionParams{
@@ -257,30 +278,29 @@ func createDatabaseFromTemplate(connParams ConnectionParams, newDBName string, t
 		if err != nil {
 			return xerrors.Errorf("connect to template db: %w", err)
 		}
-		defer tplDb.Close()
+		defer func() {
+			if err := tplDb.Close(); err != nil {
+				panic(err)
+			}
+		}()
 		if err := migrations.Up(tplDb); err != nil {
 			return xerrors.Errorf("migrate template db: %w", err)
 		}
 		if err := tplDb.Close(); err != nil {
 			return xerrors.Errorf("close template db: %w", err)
 		}
-		_, err = db.Exec("ALTER DATABASE " + tmpTemplateDBName + " RENAME TO " + templateDBName)
-		if err != nil {
+		if _, err := db.Exec("ALTER DATABASE " + tmpTemplateDBName + " RENAME TO " + templateDBName); err != nil {
 			return xerrors.Errorf("rename tmp template db: %w", err)
 		}
 	}
 
 	// Try to create the database again now that a template exists.
-	_, err = db.Exec("CREATE DATABASE " + newDBName + " WITH TEMPLATE " + templateDBName)
-	if err != nil {
+	if _, err = db.Exec("CREATE DATABASE " + newDBName + " WITH TEMPLATE " + templateDBName); err != nil {
 		return xerrors.Errorf("create db with template after migrations: %w", err)
 	}
-
-	err = tx.Commit()
-	if err != nil {
+	if err = tx.Commit(); err != nil {
 		return xerrors.Errorf("commit tx: %w", err)
 	}
-
 	return nil
 }
 
diff --git a/coderd/database/dbtestutil/postgres_test.go b/coderd/database/dbtestutil/postgres_test.go
index ec500d824a9ba..a09c87843e9be 100644
--- a/coderd/database/dbtestutil/postgres_test.go
+++ b/coderd/database/dbtestutil/postgres_test.go
@@ -11,25 +11,20 @@ import (
 	"go.uber.org/goleak"
 
 	"github.com/coder/coder/v2/coderd/database/dbtestutil"
+	"github.com/coder/coder/v2/coderd/database/migrations"
 )
 
 func TestMain(m *testing.M) {
 	goleak.VerifyTestMain(m)
 }
 
-// nolint:paralleltest
-func TestPostgres(t *testing.T) {
-	// postgres.Open() seems to be creating race conditions when run in parallel.
-	// t.Parallel()
-
-	if testing.Short() {
-		t.SkipNow()
-		return
-	}
+func TestOpen(t *testing.T) {
+	t.Parallel()
 
 	connect, closePg, err := dbtestutil.Open()
 	require.NoError(t, err)
 	defer closePg()
+
 	db, err := sql.Open("postgres", connect)
 	require.NoError(t, err)
 	err = db.Ping()
@@ -37,3 +32,74 @@ func TestPostgres(t *testing.T) {
 	err = db.Close()
 	require.NoError(t, err)
 }
+
+func TestOpen_InvalidDBFrom(t *testing.T) {
+	t.Parallel()
+
+	_, _, err := dbtestutil.Open(dbtestutil.WithDBFrom("__invalid__"))
+	require.Error(t, err)
+	require.ErrorContains(t, err, "template database")
+	require.ErrorContains(t, err, "does not exist")
+}
+
+func TestOpen_ValidDBFrom(t *testing.T) {
+	t.Parallel()
+
+	// first check if we can create a new template db
+	// and a db from that template db
+	dsn, cleanup, err := dbtestutil.Open(dbtestutil.WithDBFrom(""))
+	require.NoError(t, err)
+	t.Cleanup(cleanup)
+
+	db, err := sql.Open("postgres", dsn)
+	require.NoError(t, err)
+	t.Cleanup(func() {
+		require.NoError(t, db.Close())
+	})
+
+	err = db.Ping()
+	require.NoError(t, err)
+
+	templateDBName := "tpl_" + migrations.GetMigrationsHash()[:32]
+	tplDbExistsRes, err := db.Query("SELECT 1 FROM pg_database WHERE datname = $1", templateDBName)
+	if err != nil {
+		require.NoError(t, err)
+	}
+	require.True(t, tplDbExistsRes.Next())
+	require.NoError(t, tplDbExistsRes.Close())
+
+	// now populate the db with some data and use it as a new template db
+	// to verify that we can create a new db from it
+	_, err = db.Exec("CREATE TABLE my_wonderful_table (id serial PRIMARY KEY, name text)")
+	require.NoError(t, err)
+	_, err = db.Exec("INSERT INTO my_wonderful_table (name) VALUES ('test')")
+	require.NoError(t, err)
+
+	rows, err := db.Query("SELECT current_database()")
+	require.NoError(t, err)
+	require.True(t, rows.Next())
+	var freshTemplateDBName string
+	require.NoError(t, rows.Scan(&freshTemplateDBName))
+	require.NoError(t, rows.Close())
+	require.NoError(t, db.Close())
+
+	for i := 0; i < 10; i++ {
+		db, err := sql.Open("postgres", dsn)
+		require.NoError(t, err)
+		require.NoError(t, db.Ping())
+		require.NoError(t, db.Close())
+	}
+
+	// now create a new db from the template db
+	newDsn, newCleanup, err := dbtestutil.Open(dbtestutil.WithDBFrom(freshTemplateDBName))
+	require.NoError(t, err)
+	defer newCleanup()
+
+	newDb, err := sql.Open("postgres", newDsn)
+	require.NoError(t, err)
+	defer newDb.Close()
+
+	rows, err = newDb.Query("SELECT 1 FROM my_wonderful_table WHERE name = 'test'")
+	require.NoError(t, err)
+	require.True(t, rows.Next())
+}

From 72c2ce1fe18f778e8f028696969a83d5f93fcd2b Mon Sep 17 00:00:00 2001
From: Hugo Dutka <hugo@coder.com>
Date: Fri, 1 Nov 2024 12:08:12 +0000
Subject: [PATCH 03/10] update comments

---
 coderd/database/dbtestutil/postgres.go | 10 ++++++----
 1 file changed, 6 insertions(+), 4 deletions(-)

diff --git a/coderd/database/dbtestutil/postgres.go b/coderd/database/dbtestutil/postgres.go
index 25ccfd0fd11ae..f9ee7b593da9a 100644
--- a/coderd/database/dbtestutil/postgres.go
+++ b/coderd/database/dbtestutil/postgres.go
@@ -43,9 +43,9 @@ var (
 	errDefaultConnectionParamsInit error
 )
 
-// initDefaultConnectionParams initializes the default connection parameters
-// by checking if the database is running locally. If it is, it will use the
-// local database. If it's not, it will start a new container and use that.
+// initDefaultConnectionParams 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 initDefaultConnectionParams() error {
 	params := ConnectionParams{
 		Username: "postgres",
@@ -80,7 +80,7 @@ func initDefaultConnectionParams() error {
 		// 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 themselves.
+		// it manually.
 		container, _, err := openContainer(DBContainerOptions{
 			Name: "coder-test-postgres",
 			Port: 5432,
@@ -261,6 +261,8 @@ func createDatabaseFromTemplate(connParams ConnectionParams, newDBName string, t
 		// 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.
 		tmpTemplateDBName := "tmp_" + templateDBName
+		// 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 " + tmpTemplateDBName); err != nil {
 			return xerrors.Errorf("drop tmp template db: %w", err)
 		}

From ce9bc2736cac6529332b27c58e1feb33bc708d87 Mon Sep 17 00:00:00 2001
From: Hugo Dutka <hugo@coder.com>
Date: Fri, 1 Nov 2024 12:17:51 +0000
Subject: [PATCH 04/10] update comments

---
 coderd/database/dbtestutil/postgres_test.go | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/coderd/database/dbtestutil/postgres_test.go b/coderd/database/dbtestutil/postgres_test.go
index a09c87843e9be..3768a94dc035a 100644
--- a/coderd/database/dbtestutil/postgres_test.go
+++ b/coderd/database/dbtestutil/postgres_test.go
@@ -46,7 +46,6 @@ func TestOpen_ValidDBFrom(t *testing.T) {
 	t.Parallel()
 
 	// first check if we can create a new template db
-	// and a db from that template db
 	dsn, cleanup, err := dbtestutil.Open(dbtestutil.WithDBFrom(""))
 	require.NoError(t, err)
 	t.Cleanup(cleanup)
@@ -69,7 +68,7 @@ func TestOpen_ValidDBFrom(t *testing.T) {
 	require.NoError(t, tplDbExistsRes.Close())
 
 	// now populate the db with some data and use it as a new template db
-	// to verify that we can create a new db from it
+	// to verify that dbtestutil.Open respects WithDBFrom
 	_, err = db.Exec("CREATE TABLE my_wonderful_table (id serial PRIMARY KEY, name text)")
 	require.NoError(t, err)
 	_, err = db.Exec("INSERT INTO my_wonderful_table (name) VALUES ('test')")
@@ -102,4 +101,5 @@ func TestOpen_ValidDBFrom(t *testing.T) {
 	rows, err = newDb.Query("SELECT 1 FROM my_wonderful_table WHERE name = 'test'")
 	require.NoError(t, err)
 	require.True(t, rows.Next())
+	require.NoError(t, rows.Close())
 }

From 99f13c2a8a6c90382d8cd311bf6b9b53f4046b4e Mon Sep 17 00:00:00 2001
From: Hugo Dutka <hugo@coder.com>
Date: Fri, 1 Nov 2024 12:26:30 +0000
Subject: [PATCH 05/10] remove DB_FROM from test-postgres

---
 Makefile | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/Makefile b/Makefile
index 084e8bb77e5f0..fe6e527fe712e 100644
--- a/Makefile
+++ b/Makefile
@@ -765,7 +765,7 @@ sqlc-vet: test-postgres-docker
 test-postgres: test-postgres-docker
 	# The postgres test is prone to failure, so we limit parallelism for
 	# more consistent execution.
-	$(GIT_FLAGS)  DB=ci DB_FROM=$(shell go run scripts/migrate-ci/main.go) gotestsum \
+	$(GIT_FLAGS)  DB=ci gotestsum \
 		--junitfile="gotests.xml" \
 		--jsonfile="gotests.json" \
 		--packages="./..." -- \

From 1e484f7f99b69cfeac3ca58ac9cbd4d3b9265c66 Mon Sep 17 00:00:00 2001
From: Hugo Dutka <hugo@coder.com>
Date: Fri, 1 Nov 2024 12:26:45 +0000
Subject: [PATCH 06/10] only initialize dbmem if needed

---
 coderd/database/dbtestutil/db.go | 7 +++++--
 1 file changed, 5 insertions(+), 2 deletions(-)

diff --git a/coderd/database/dbtestutil/db.go b/coderd/database/dbtestutil/db.go
index 327d880f69648..5a15329da1bbd 100644
--- a/coderd/database/dbtestutil/db.go
+++ b/coderd/database/dbtestutil/db.go
@@ -95,8 +95,8 @@ func NewDB(t testing.TB, opts ...Option) (database.Store, pubsub.Pubsub) {
 		opt(&o)
 	}
 
-	db := dbmem.New()
-	ps := pubsub.NewInMemory()
+	var db database.Store
+	var ps pubsub.Pubsub
 	if WillUsePostgres() {
 		connectionURL := os.Getenv("CODER_PG_CONNECTION_URL")
 		if connectionURL == "" && o.url != "" {
@@ -142,6 +142,9 @@ func NewDB(t testing.TB, opts ...Option) (database.Store, pubsub.Pubsub) {
 		t.Cleanup(func() {
 			_ = ps.Close()
 		})
+	} else {
+		db = dbmem.New()
+		ps = pubsub.NewInMemory()
 	}
 
 	return db, ps

From ef47a3339933980f396fb44d3086421c28bb59b6 Mon Sep 17 00:00:00 2001
From: Hugo Dutka <hugo@coder.com>
Date: Fri, 1 Nov 2024 14:57:08 +0000
Subject: [PATCH 07/10] address review comments

---
 coderd/database/dbtestutil/postgres.go | 55 +++++++++++++++-----------
 1 file changed, 31 insertions(+), 24 deletions(-)

diff --git a/coderd/database/dbtestutil/postgres.go b/coderd/database/dbtestutil/postgres.go
index f9ee7b593da9a..c26cdad4a0497 100644
--- a/coderd/database/dbtestutil/postgres.go
+++ b/coderd/database/dbtestutil/postgres.go
@@ -23,6 +23,7 @@ import (
 
 	"github.com/coder/coder/v2/coderd/database/migrations"
 	"github.com/coder/coder/v2/cryptorand"
+	"github.com/coder/retry"
 )
 
 type ConnectionParams struct {
@@ -37,16 +38,17 @@ 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
 )
 
-// initDefaultConnectionParams initializes the default postgres connection parameters.
+// 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 initDefaultConnectionParams() error {
+func initDefaultConnection() error {
 	params := ConnectionParams{
 		Username: "postgres",
 		Password: "postgres",
@@ -55,20 +57,20 @@ func initDefaultConnectionParams() error {
 		DBName:   "postgres",
 	}
 	dsn := params.DSN()
-	db, err := sql.Open("postgres", dsn)
-	if err == nil {
-		err = db.Ping()
+	db, dbErr := sql.Open("postgres", dsn)
+	if dbErr == nil {
+		dbErr = db.Ping()
 		if closeErr := db.Close(); closeErr != nil {
 			return xerrors.Errorf("close db: %w", closeErr)
 		}
 	}
 	shouldOpenContainer := false
-	if err != nil {
+	if dbErr != nil {
 		errSubstrings := []string{
 			"connection refused",          // this happens on Linux when there's nothing listening on the port
 			"No connection could be made", // like above but Windows
 		}
-		errString := err.Error()
+		errString := dbErr.Error()
 		for _, errSubstring := range errSubstrings {
 			if strings.Contains(errString, errSubstring) {
 				shouldOpenContainer = true
@@ -76,7 +78,7 @@ func initDefaultConnectionParams() error {
 			}
 		}
 	}
-	if err != nil && shouldOpenContainer {
+	if dbErr != nil && shouldOpenContainer {
 		// 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
@@ -92,24 +94,23 @@ func initDefaultConnectionParams() error {
 		params.Port = container.Port
 		dsn = params.DSN()
 
-		// Retry connecting for a cumulative 5 seconds.
+		// Retry connecting for at most 10 seconds.
 		// The fact that openContainer succeeded does not
 		// mean that port forwarding is ready.
-		for i := 0; i < 20; i++ {
-			db, err = sql.Open("postgres", dsn)
-			if err == nil {
-				err = db.Ping()
+		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 err == nil {
+			if connErr == nil {
 				break
 			}
-			time.Sleep(250 * time.Millisecond)
 		}
-	} else if err != nil {
-		return xerrors.Errorf("open postgres connection: %w", err)
+	} else if dbErr != nil {
+		return xerrors.Errorf("open postgres connection: %w", dbErr)
 	}
 	defaultConnectionParams = params
 	return nil
@@ -134,7 +135,7 @@ func WithDBFrom(dbFrom string) OpenOption {
 // Otherwise, it will start a new postgres container.
 func Open(opts ...OpenOption) (string, func(), error) {
 	connectionParamsInitOnce.Do(func() {
-		errDefaultConnectionParamsInit = initDefaultConnectionParams()
+		errDefaultConnectionParamsInit = initDefaultConnection()
 	})
 	if errDefaultConnectionParamsInit != nil {
 		return "", func() {}, xerrors.Errorf("init default connection params: %w", errDefaultConnectionParamsInit)
@@ -175,11 +176,13 @@ func Open(opts ...OpenOption) (string, func(), error) {
 		cleanupConn, err := sql.Open("postgres", cleanupDbURL)
 		if err != nil {
 			_, _ = fmt.Fprintf(os.Stderr, "cleanup database %q: failed to connect to postgres: %s\n", dbName, err.Error())
+			return
 		}
 		defer cleanupConn.Close()
 		_, err = cleanupConn.Exec("DROP DATABASE " + dbName + ";")
 		if err != nil {
 			_, _ = fmt.Fprintf(os.Stderr, "failed to clean up database %q: %s\n", dbName, err.Error())
+			return
 		}
 	}
 
@@ -240,6 +243,8 @@ func createDatabaseFromTemplate(connParams ConnectionParams, newDBName string, t
 			panic(err)
 		}
 	}()
+	// 2137 is an arbitrary number. We just need a lock that is unique to creating
+	// the template database.
 	_, err = tx.Exec("SELECT pg_advisory_xact_lock(2137)")
 	if err != nil {
 		return xerrors.Errorf("acquire lock: %w", err)
@@ -346,17 +351,17 @@ func openContainer(opts DBContainerOptions) (container, func(), error) {
 		return container{}, nil, xerrors.Errorf("create pool: %w", err)
 	}
 
-	tempDir, err := os.MkdirTemp(os.TempDir(), "postgres")
-	if err != nil {
-		return container{}, nil, xerrors.Errorf("create tempdir: %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: "gcr.io/coder-dev-1/postgres",
 			Tag:        "13",
@@ -434,7 +439,9 @@ func openContainer(opts DBContainerOptions) (container, func(), error) {
 			Pool:     pool,
 		}, func() {
 			_ = pool.Purge(resource)
-			_ = os.RemoveAll(tempDir)
+			if tempDir != "" {
+				_ = os.RemoveAll(tempDir)
+			}
 		}, nil
 }
 

From c11253db08499f82c7a3d1240a4f2a38bf557830 Mon Sep 17 00:00:00 2001
From: Hugo Dutka <hugo@coder.com>
Date: Fri, 1 Nov 2024 15:40:46 +0000
Subject: [PATCH 08/10] pass testing.T to dbtestutil.Open

---
 cli/resetpassword_test.go                   |  3 +-
 cli/server_createadminuser_test.go          | 12 ++--
 cli/server_test.go                          |  6 +-
 coderd/database/db_test.go                  |  3 +-
 coderd/database/dbtestutil/db.go            |  6 +-
 coderd/database/dbtestutil/postgres.go      | 61 +++++++++++++--------
 coderd/database/dbtestutil/postgres_test.go | 15 +++--
 coderd/database/gen/dump/main.go            | 27 ++++++++-
 coderd/database/migrations/migrate_test.go  |  3 +-
 coderd/database/pubsub/pubsub_linux_test.go | 19 +++----
 coderd/database/pubsub/pubsub_test.go       |  6 +-
 enterprise/cli/server_dbcrypt_test.go       |  3 +-
 enterprise/tailnet/pgcoord_test.go          |  3 +-
 13 files changed, 93 insertions(+), 74 deletions(-)

diff --git a/cli/resetpassword_test.go b/cli/resetpassword_test.go
index 0cd90f5b4cd00..de712874f3f07 100644
--- a/cli/resetpassword_test.go
+++ b/cli/resetpassword_test.go
@@ -32,9 +32,8 @@ func TestResetPassword(t *testing.T) {
 	const newPassword = "MyNewPassword!"
 
 	// start postgres and coder server processes
-	connectionURL, closeFunc, err := dbtestutil.Open()
+	connectionURL, err := dbtestutil.Open(t)
 	require.NoError(t, err)
-	defer closeFunc()
 	ctx, cancelFunc := context.WithCancel(context.Background())
 	serverDone := make(chan struct{})
 	serverinv, cfg := clitest.New(t,
diff --git a/cli/server_createadminuser_test.go b/cli/server_createadminuser_test.go
index 17c02b6548c09..7660d71e89d99 100644
--- a/cli/server_createadminuser_test.go
+++ b/cli/server_createadminuser_test.go
@@ -85,9 +85,8 @@ func TestServerCreateAdminUser(t *testing.T) {
 			// Skip on non-Linux because it spawns a PostgreSQL instance.
 			t.SkipNow()
 		}
-		connectionURL, closeFunc, err := dbtestutil.Open()
+		connectionURL, err := dbtestutil.Open(t)
 		require.NoError(t, err)
-		defer closeFunc()
 
 		sqlDB, err := sql.Open("postgres", connectionURL)
 		require.NoError(t, err)
@@ -151,9 +150,8 @@ func TestServerCreateAdminUser(t *testing.T) {
 			// Skip on non-Linux because it spawns a PostgreSQL instance.
 			t.SkipNow()
 		}
-		connectionURL, closeFunc, err := dbtestutil.Open()
+		connectionURL, err := dbtestutil.Open(t)
 		require.NoError(t, err)
-		defer closeFunc()
 
 		ctx, cancel := context.WithTimeout(context.Background(), testutil.WaitMedium)
 		defer cancel()
@@ -185,9 +183,8 @@ func TestServerCreateAdminUser(t *testing.T) {
 			// Skip on non-Linux because it spawns a PostgreSQL instance.
 			t.SkipNow()
 		}
-		connectionURL, closeFunc, err := dbtestutil.Open()
+		connectionURL, err := dbtestutil.Open(t)
 		require.NoError(t, err)
-		defer closeFunc()
 
 		ctx, cancel := context.WithTimeout(context.Background(), testutil.WaitMedium)
 		defer cancel()
@@ -225,9 +222,8 @@ func TestServerCreateAdminUser(t *testing.T) {
 			// Skip on non-Linux because it spawns a PostgreSQL instance.
 			t.SkipNow()
 		}
-		connectionURL, closeFunc, err := dbtestutil.Open()
+		connectionURL, err := dbtestutil.Open(t)
 		require.NoError(t, err)
-		defer closeFunc()
 		ctx, cancelFunc := context.WithCancel(context.Background())
 		defer cancelFunc()
 
diff --git a/cli/server_test.go b/cli/server_test.go
index ad6a98038c7bb..83a7f7171c6f5 100644
--- a/cli/server_test.go
+++ b/cli/server_test.go
@@ -1598,9 +1598,8 @@ func TestServer_Production(t *testing.T) {
 		// Skip on non-Linux because it spawns a PostgreSQL instance.
 		t.SkipNow()
 	}
-	connectionURL, closeFunc, err := dbtestutil.Open()
+	connectionURL, err := dbtestutil.Open(t)
 	require.NoError(t, err)
-	defer closeFunc()
 
 	// Postgres + race detector + CI = slow.
 	ctx, cancelFunc := context.WithTimeout(context.Background(), testutil.WaitSuperLong*3)
@@ -1803,9 +1802,8 @@ func TestConnectToPostgres(t *testing.T) {
 
 	log := slogtest.Make(t, nil)
 
-	dbURL, closeFunc, err := dbtestutil.Open()
+	dbURL, err := dbtestutil.Open(t)
 	require.NoError(t, err)
-	t.Cleanup(closeFunc)
 
 	sqlDB, err := cli.ConnectToPostgres(ctx, log, "postgres", dbURL)
 	require.NoError(t, err)
diff --git a/coderd/database/db_test.go b/coderd/database/db_test.go
index a6df18fcbb8c8..b4580527c843a 100644
--- a/coderd/database/db_test.go
+++ b/coderd/database/db_test.go
@@ -87,9 +87,8 @@ func TestNestedInTx(t *testing.T) {
 func testSQLDB(t testing.TB) *sql.DB {
 	t.Helper()
 
-	connection, closeFn, err := dbtestutil.Open()
+	connection, err := dbtestutil.Open(t)
 	require.NoError(t, err)
-	t.Cleanup(closeFn)
 
 	db, err := sql.Open("postgres", connection)
 	require.NoError(t, err)
diff --git a/coderd/database/dbtestutil/db.go b/coderd/database/dbtestutil/db.go
index 5a15329da1bbd..398e9ae9f0393 100644
--- a/coderd/database/dbtestutil/db.go
+++ b/coderd/database/dbtestutil/db.go
@@ -104,12 +104,10 @@ func NewDB(t testing.TB, opts ...Option) (database.Store, pubsub.Pubsub) {
 		}
 		if connectionURL == "" {
 			var (
-				err     error
-				closePg func()
+				err error
 			)
-			connectionURL, closePg, err = Open()
+			connectionURL, err = Open(t)
 			require.NoError(t, err)
-			t.Cleanup(closePg)
 		}
 
 		if o.fixedTimezone == "" {
diff --git a/coderd/database/dbtestutil/postgres.go b/coderd/database/dbtestutil/postgres.go
index c26cdad4a0497..3169ec0bb68df 100644
--- a/coderd/database/dbtestutil/postgres.go
+++ b/coderd/database/dbtestutil/postgres.go
@@ -48,7 +48,7 @@ var (
 // 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() error {
+func initDefaultConnection(t TBSubset) error {
 	params := ConnectionParams{
 		Username: "postgres",
 		Password: "postgres",
@@ -83,7 +83,7 @@ func initDefaultConnection() error {
 		// 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(DBContainerOptions{
+		container, _, err := openContainer(t, DBContainerOptions{
 			Name: "coder-test-postgres",
 			Port: 5432,
 		})
@@ -130,15 +130,25 @@ func WithDBFrom(dbFrom string) OpenOption {
 	}
 }
 
+// TBSubset is a subset of the testing.TB interface.
+// It allows to use dbtestutil.Open outside of tests.
+type TBSubset interface {
+	Cleanup(func())
+	Helper()
+	Logf(format string, args ...any)
+}
+
 // 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(opts ...OpenOption) (string, func(), error) {
+func Open(t TBSubset, opts ...OpenOption) (string, error) {
+	t.Helper()
+
 	connectionParamsInitOnce.Do(func() {
-		errDefaultConnectionParamsInit = initDefaultConnection()
+		errDefaultConnectionParamsInit = initDefaultConnection(t)
 	})
 	if errDefaultConnectionParamsInit != nil {
-		return "", func() {}, xerrors.Errorf("init default connection params: %w", errDefaultConnectionParamsInit)
+		return "", xerrors.Errorf("init default connection params: %w", errDefaultConnectionParamsInit)
 	}
 
 	openOptions := OpenOptions{}
@@ -158,7 +168,7 @@ func Open(opts ...OpenOption) (string, func(), error) {
 	now := time.Now().Format("test_2006_01_02_15_04_05")
 	dbSuffix, err := cryptorand.StringCharset(cryptorand.Lower, 10)
 	if err != nil {
-		return "", func() {}, xerrors.Errorf("generate db suffix: %w", err)
+		return "", xerrors.Errorf("generate db suffix: %w", err)
 	}
 	dbName := now + "_" + dbSuffix
 
@@ -167,24 +177,28 @@ func Open(opts ...OpenOption) (string, func(), error) {
 	if openOptions.DBFrom != nil {
 		templateDBName = *openOptions.DBFrom
 	}
-	if err = createDatabaseFromTemplate(defaultConnectionParams, dbName, templateDBName); err != nil {
-		return "", func() {}, xerrors.Errorf("create database: %w", err)
+	if err = createDatabaseFromTemplate(t, defaultConnectionParams, dbName, templateDBName); err != nil {
+		return "", xerrors.Errorf("create database: %w", err)
 	}
 
-	cleanup := func() {
+	t.Cleanup(func() {
 		cleanupDbURL := defaultConnectionParams.DSN()
 		cleanupConn, err := sql.Open("postgres", cleanupDbURL)
 		if err != nil {
-			_, _ = fmt.Fprintf(os.Stderr, "cleanup database %q: failed to connect to postgres: %s\n", dbName, err.Error())
+			t.Logf("cleanup database %q: failed to connect to postgres: %s\n", dbName, err.Error())
 			return
 		}
-		defer cleanupConn.Close()
+		defer func() {
+			if err := cleanupConn.Close(); err != nil {
+				t.Logf("cleanup database %q: failed to close connection: %s\n", dbName, err.Error())
+			}
+		}()
 		_, err = cleanupConn.Exec("DROP DATABASE " + dbName + ";")
 		if err != nil {
-			_, _ = fmt.Fprintf(os.Stderr, "failed to clean up database %q: %s\n", dbName, err.Error())
+			t.Logf("failed to clean up database %q: %s\n", dbName, err.Error())
 			return
 		}
-	}
+	})
 
 	dsn := ConnectionParams{
 		Username: username,
@@ -193,14 +207,16 @@ func Open(opts ...OpenOption) (string, func(), error) {
 		Port:     port,
 		DBName:   dbName,
 	}.DSN()
-	return dsn, cleanup, nil
+	return 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(connParams ConnectionParams, newDBName string, templateDBName string) error {
+func createDatabaseFromTemplate(t TBSubset, connParams ConnectionParams, newDBName string, templateDBName string) error {
+	t.Helper()
+
 	dbURL := connParams.DSN()
 	db, err := sql.Open("postgres", dbURL)
 	if err != nil {
@@ -208,7 +224,7 @@ func createDatabaseFromTemplate(connParams ConnectionParams, newDBName string, t
 	}
 	defer func() {
 		if err := db.Close(); err != nil {
-			panic(err)
+			t.Logf("create database from template: failed to close connection: %s\n", err.Error())
 		}
 	}()
 
@@ -240,7 +256,7 @@ func createDatabaseFromTemplate(connParams ConnectionParams, newDBName string, t
 	defer func() {
 		err := tx.Rollback()
 		if err != nil && !errors.Is(err, sql.ErrTxDone) {
-			panic(err)
+			t.Logf("create database from template: failed to rollback tx: %s\n", err.Error())
 		}
 	}()
 	// 2137 is an arbitrary number. We just need a lock that is unique to creating
@@ -287,7 +303,7 @@ func createDatabaseFromTemplate(connParams ConnectionParams, newDBName string, t
 		}
 		defer func() {
 			if err := tplDb.Close(); err != nil {
-				panic(err)
+				t.Logf("create database from template: failed to close template db: %s\n", err.Error())
 			}
 		}()
 		if err := migrations.Up(tplDb); err != nil {
@@ -327,7 +343,7 @@ type container struct {
 // 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(opts DBContainerOptions) (container, func(), error) {
+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
@@ -341,7 +357,7 @@ func openContainer(opts DBContainerOptions) (container, func(), error) {
 		defer func() {
 			err := lock.Unlock()
 			if err != nil {
-				panic(err)
+				t.Logf("create database from template: failed to unlock: %s\n", err.Error())
 			}
 		}()
 	}
@@ -447,8 +463,9 @@ func openContainer(opts DBContainerOptions) (container, func(), error) {
 
 // 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.
-func OpenContainerized(opts DBContainerOptions) (string, func(), error) {
-	container, containerCleanup, err := openContainer(opts)
+// 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)
 	defer func() {
 		if err != nil {
 			containerCleanup()
diff --git a/coderd/database/dbtestutil/postgres_test.go b/coderd/database/dbtestutil/postgres_test.go
index 3768a94dc035a..9cae9411289ad 100644
--- a/coderd/database/dbtestutil/postgres_test.go
+++ b/coderd/database/dbtestutil/postgres_test.go
@@ -21,9 +21,8 @@ func TestMain(m *testing.M) {
 func TestOpen(t *testing.T) {
 	t.Parallel()
 
-	connect, closePg, err := dbtestutil.Open()
+	connect, err := dbtestutil.Open(t)
 	require.NoError(t, err)
-	defer closePg()
 
 	db, err := sql.Open("postgres", connect)
 	require.NoError(t, err)
@@ -36,7 +35,7 @@ func TestOpen(t *testing.T) {
 func TestOpen_InvalidDBFrom(t *testing.T) {
 	t.Parallel()
 
-	_, _, err := dbtestutil.Open(dbtestutil.WithDBFrom("__invalid__"))
+	_, err := dbtestutil.Open(t, dbtestutil.WithDBFrom("__invalid__"))
 	require.Error(t, err)
 	require.ErrorContains(t, err, "template database")
 	require.ErrorContains(t, err, "does not exist")
@@ -46,9 +45,8 @@ func TestOpen_ValidDBFrom(t *testing.T) {
 	t.Parallel()
 
 	// first check if we can create a new template db
-	dsn, cleanup, err := dbtestutil.Open(dbtestutil.WithDBFrom(""))
+	dsn, err := dbtestutil.Open(t, dbtestutil.WithDBFrom(""))
 	require.NoError(t, err)
-	t.Cleanup(cleanup)
 
 	db, err := sql.Open("postgres", dsn)
 	require.NoError(t, err)
@@ -90,13 +88,14 @@ func TestOpen_ValidDBFrom(t *testing.T) {
 	}
 
 	// now create a new db from the template db
-	newDsn, newCleanup, err := dbtestutil.Open(dbtestutil.WithDBFrom(freshTemplateDBName))
+	newDsn, err := dbtestutil.Open(t, dbtestutil.WithDBFrom(freshTemplateDBName))
 	require.NoError(t, err)
-	defer newCleanup()
 
 	newDb, err := sql.Open("postgres", newDsn)
 	require.NoError(t, err)
-	defer newDb.Close()
+	t.Cleanup(func() {
+		require.NoError(t, newDb.Close())
+	})
 
 	rows, err = newDb.Query("SELECT 1 FROM my_wonderful_table WHERE name = 'test'")
 	require.NoError(t, err)
diff --git a/coderd/database/gen/dump/main.go b/coderd/database/gen/dump/main.go
index f563e1142619e..e3e80c528144e 100644
--- a/coderd/database/gen/dump/main.go
+++ b/coderd/database/gen/dump/main.go
@@ -2,6 +2,7 @@ package main
 
 import (
 	"database/sql"
+	"fmt"
 	"os"
 	"path/filepath"
 	"runtime"
@@ -12,12 +13,34 @@ import (
 
 var preamble = []byte("-- Code generated by 'make coderd/database/generate'. DO NOT EDIT.")
 
+type mockTB struct {
+	cleanup []func()
+}
+
+func (t *mockTB) Cleanup(f func()) {
+	t.cleanup = append(t.cleanup, f)
+}
+
+func (*mockTB) Helper() {
+	// noop
+}
+
+func (*mockTB) Logf(format string, args ...any) {
+	_, _ = fmt.Printf(format, args...)
+}
+
 func main() {
-	connection, closeFn, err := dbtestutil.Open()
+	t := &mockTB{}
+	defer func() {
+		for _, f := range t.cleanup {
+			f()
+		}
+	}()
+
+	connection, err := dbtestutil.Open(t)
 	if err != nil {
 		panic(err)
 	}
-	defer closeFn()
 
 	db, err := sql.Open("postgres", connection)
 	if err != nil {
diff --git a/coderd/database/migrations/migrate_test.go b/coderd/database/migrations/migrate_test.go
index 51e7fcc86cb03..c64c2436da18d 100644
--- a/coderd/database/migrations/migrate_test.go
+++ b/coderd/database/migrations/migrate_test.go
@@ -95,9 +95,8 @@ func TestMigrate(t *testing.T) {
 func testSQLDB(t testing.TB) *sql.DB {
 	t.Helper()
 
-	connection, closeFn, err := dbtestutil.Open()
+	connection, err := dbtestutil.Open(t)
 	require.NoError(t, err)
-	t.Cleanup(closeFn)
 
 	db, err := sql.Open("postgres", connection)
 	require.NoError(t, err)
diff --git a/coderd/database/pubsub/pubsub_linux_test.go b/coderd/database/pubsub/pubsub_linux_test.go
index a4be5927a4725..819de0a71ba52 100644
--- a/coderd/database/pubsub/pubsub_linux_test.go
+++ b/coderd/database/pubsub/pubsub_linux_test.go
@@ -40,9 +40,8 @@ func TestPubsub(t *testing.T) {
 		defer cancelFunc()
 		logger := slogtest.Make(t, nil).Leveled(slog.LevelDebug)
 
-		connectionURL, closePg, err := dbtestutil.Open()
+		connectionURL, err := dbtestutil.Open(t)
 		require.NoError(t, err)
-		defer closePg()
 		db, err := sql.Open("postgres", connectionURL)
 		require.NoError(t, err)
 		defer db.Close()
@@ -69,9 +68,8 @@ func TestPubsub(t *testing.T) {
 		ctx, cancelFunc := context.WithCancel(context.Background())
 		defer cancelFunc()
 		logger := slogtest.Make(t, nil).Leveled(slog.LevelDebug)
-		connectionURL, closePg, err := dbtestutil.Open()
+		connectionURL, err := dbtestutil.Open(t)
 		require.NoError(t, err)
-		defer closePg()
 		db, err := sql.Open("postgres", connectionURL)
 		require.NoError(t, err)
 		defer db.Close()
@@ -85,9 +83,8 @@ func TestPubsub(t *testing.T) {
 		ctx, cancel := context.WithCancel(context.Background())
 		defer cancel()
 		logger := slogtest.Make(t, nil).Leveled(slog.LevelDebug)
-		connectionURL, closePg, err := dbtestutil.Open()
+		connectionURL, err := dbtestutil.Open(t)
 		require.NoError(t, err)
-		defer closePg()
 		db, err := sql.Open("postgres", connectionURL)
 		require.NoError(t, err)
 		defer db.Close()
@@ -122,9 +119,8 @@ func TestPubsub_ordering(t *testing.T) {
 	defer cancelFunc()
 	logger := slogtest.Make(t, nil).Leveled(slog.LevelDebug)
 
-	connectionURL, closePg, err := dbtestutil.Open()
+	connectionURL, err := dbtestutil.Open(t)
 	require.NoError(t, err)
-	defer closePg()
 	db, err := sql.Open("postgres", connectionURL)
 	require.NoError(t, err)
 	defer db.Close()
@@ -167,7 +163,7 @@ const disconnectTestPort = 26892
 func TestPubsub_Disconnect(t *testing.T) {
 	// we always use a Docker container for this test, even in CI, since we need to be able to kill
 	// postgres and bring it back on the same port.
-	connectionURL, closePg, err := dbtestutil.OpenContainerized(dbtestutil.DBContainerOptions{Port: disconnectTestPort})
+	connectionURL, closePg, err := dbtestutil.OpenContainerized(t, dbtestutil.DBContainerOptions{Port: disconnectTestPort})
 	require.NoError(t, err)
 	defer closePg()
 	db, err := sql.Open("postgres", connectionURL)
@@ -238,7 +234,7 @@ func TestPubsub_Disconnect(t *testing.T) {
 
 	// restart postgres on the same port --- since we only use LISTEN/NOTIFY it doesn't
 	// matter that the new postgres doesn't have any persisted state from before.
-	_, closeNewPg, err := dbtestutil.OpenContainerized(dbtestutil.DBContainerOptions{Port: disconnectTestPort})
+	_, closeNewPg, err := dbtestutil.OpenContainerized(t, dbtestutil.DBContainerOptions{Port: disconnectTestPort})
 	require.NoError(t, err)
 	defer closeNewPg()
 
@@ -305,7 +301,7 @@ func TestMeasureLatency(t *testing.T) {
 	newPubsub := func() (pubsub.Pubsub, func()) {
 		ctx, cancel := context.WithCancel(context.Background())
 		logger := slogtest.Make(t, nil).Leveled(slog.LevelDebug)
-		connectionURL, closePg, err := dbtestutil.Open()
+		connectionURL, err := dbtestutil.Open(t)
 		require.NoError(t, err)
 		db, err := sql.Open("postgres", connectionURL)
 		require.NoError(t, err)
@@ -315,7 +311,6 @@ func TestMeasureLatency(t *testing.T) {
 		return ps, func() {
 			_ = ps.Close()
 			_ = db.Close()
-			closePg()
 			cancel()
 		}
 	}
diff --git a/coderd/database/pubsub/pubsub_test.go b/coderd/database/pubsub/pubsub_test.go
index 21b4b1d54c171..6b8181ea7d834 100644
--- a/coderd/database/pubsub/pubsub_test.go
+++ b/coderd/database/pubsub/pubsub_test.go
@@ -24,9 +24,8 @@ func TestPGPubsub_Metrics(t *testing.T) {
 	}
 
 	logger := slogtest.Make(t, nil).Leveled(slog.LevelDebug)
-	connectionURL, closePg, err := dbtestutil.Open()
+	connectionURL, err := dbtestutil.Open(t)
 	require.NoError(t, err)
-	defer closePg()
 	db, err := sql.Open("postgres", connectionURL)
 	require.NoError(t, err)
 	defer db.Close()
@@ -132,9 +131,8 @@ func TestPGPubsubDriver(t *testing.T) {
 		IgnoreErrors: true,
 	}).Leveled(slog.LevelDebug)
 
-	connectionURL, closePg, err := dbtestutil.Open()
+	connectionURL, err := dbtestutil.Open(t)
 	require.NoError(t, err)
-	defer closePg()
 
 	// use a separate subber and pubber so we can keep track of listener connections
 	db, err := sql.Open("postgres", connectionURL)
diff --git a/enterprise/cli/server_dbcrypt_test.go b/enterprise/cli/server_dbcrypt_test.go
index b1767889d9c33..070f172bcbe7b 100644
--- a/enterprise/cli/server_dbcrypt_test.go
+++ b/enterprise/cli/server_dbcrypt_test.go
@@ -32,9 +32,8 @@ func TestServerDBCrypt(t *testing.T) {
 	t.Cleanup(cancel)
 
 	// Setup a postgres database.
-	connectionURL, closePg, err := dbtestutil.Open()
+	connectionURL, err := dbtestutil.Open(t)
 	require.NoError(t, err)
-	t.Cleanup(closePg)
 	t.Cleanup(func() { dbtestutil.DumpOnFailure(t, connectionURL) })
 
 	sqlDB, err := sql.Open("postgres", connectionURL)
diff --git a/enterprise/tailnet/pgcoord_test.go b/enterprise/tailnet/pgcoord_test.go
index c0d122aa74992..49248e636f04b 100644
--- a/enterprise/tailnet/pgcoord_test.go
+++ b/enterprise/tailnet/pgcoord_test.go
@@ -798,9 +798,8 @@ func TestPGCoordinatorDual_FailedHeartbeat(t *testing.T) {
 		t.Skip("test only with postgres")
 	}
 
-	dburl, closeFn, err := dbtestutil.Open()
+	dburl, err := dbtestutil.Open(t)
 	require.NoError(t, err)
-	t.Cleanup(closeFn)
 
 	store1, ps1, sdb1 := dbtestutil.NewDBWithSQLDB(t, dbtestutil.WithURL(dburl))
 	defer sdb1.Close()

From 215cd81e2c59a9d2777390bcedbb9e8e10d28f47 Mon Sep 17 00:00:00 2001
From: Hugo Dutka <hugo@coder.com>
Date: Fri, 1 Nov 2024 15:45:13 +0000
Subject: [PATCH 09/10] format

---
 coderd/database/dbtestutil/db.go | 4 +---
 1 file changed, 1 insertion(+), 3 deletions(-)

diff --git a/coderd/database/dbtestutil/db.go b/coderd/database/dbtestutil/db.go
index 398e9ae9f0393..966efd8386643 100644
--- a/coderd/database/dbtestutil/db.go
+++ b/coderd/database/dbtestutil/db.go
@@ -103,9 +103,7 @@ func NewDB(t testing.TB, opts ...Option) (database.Store, pubsub.Pubsub) {
 			connectionURL = o.url
 		}
 		if connectionURL == "" {
-			var (
-				err error
-			)
+			var err error
 			connectionURL, err = Open(t)
 			require.NoError(t, err)
 		}

From 0aed3bdfc20b7c3a33a71a9c266b9091cb0ef449 Mon Sep 17 00:00:00 2001
From: Hugo Dutka <hugo@coder.com>
Date: Sat, 2 Nov 2024 00:08:45 +0000
Subject: [PATCH 10/10] convert one more loop to use retry

---
 coderd/database/dbtestutil/postgres.go | 4 +---
 1 file changed, 1 insertion(+), 3 deletions(-)

diff --git a/coderd/database/dbtestutil/postgres.go b/coderd/database/dbtestutil/postgres.go
index 3169ec0bb68df..a58ffb570763f 100644
--- a/coderd/database/dbtestutil/postgres.go
+++ b/coderd/database/dbtestutil/postgres.go
@@ -431,8 +431,7 @@ func openContainer(t TBSubset, opts DBContainerOptions) (container, func(), erro
 		return container{}, nil, xerrors.Errorf("split host and port: %w", err)
 	}
 
-	// wait for a cumulative 60 * 250ms = 15 seconds for the database to start
-	for i := 0; i < 60; i++ {
+	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{
@@ -442,7 +441,6 @@ func openContainer(t TBSubset, opts DBContainerOptions) (container, func(), erro
 		if err == nil {
 			break
 		}
-		time.Sleep(250 * time.Millisecond)
 	}
 	if err != nil {
 		return container{}, nil, xerrors.Errorf("pg_isready: %w", err)