test: refactor dbtestutil to record database creation (#19843)

relates to https://github.com/coder/internal/issues/927

Refactors dbtestutil to use a `Broker` struct to create test databases. Additionally uses a `coder_testing` database to record test databases that are created and when they are dropped.

This is in preparation for the PR above this in the stack which adds a "cleaner" subprocess that cleans out any databases that were left when the test process ends.

coderd/database/dbtestutil/broker.go

@@ -0,0 +1,176 @@
+package dbtestutil
+
+import (
+	"database/sql"
+	_ "embed"
+	"fmt"
+	"os"
+	"sync"
+	"time"
+
+	"github.com/google/uuid"
+	"github.com/lib/pq"
+	"golang.org/x/xerrors"
+
+	"github.com/coder/coder/v2/cryptorand"
+)
+
+const CoderTestingDBName = "coder_testing"
+
+//go:embed coder_testing.sql
+var coderTestingSQLInit string
+
+type Broker struct {
+	sync.Mutex
+	uuid           uuid.UUID
+	coderTestingDB *sql.DB
+	refCount       int
+}
+
+func (b *Broker) Create(t TBSubset, opts ...OpenOption) (ConnectionParams, error) {
+	if err := b.init(t); err != nil {
+		return ConnectionParams{}, err
+	}
+	openOptions := OpenOptions{}
+	for _, opt := range opts {
+		opt(&openOptions)
+	}
+
+	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 ConnectionParams{}, xerrors.Errorf("generate db suffix: %w", err)
+	}
+	dbName := now + "_" + dbSuffix
+
+	// TODO: add package and test name
+	_, err = b.coderTestingDB.Exec(
+		"INSERT INTO test_databases (name, process_uuid) VALUES ($1, $2)", dbName, b.uuid)
+	if err != nil {
+		return ConnectionParams{}, xerrors.Errorf("insert test_database row: %w", err)
+	}
+
+	// if empty createDatabaseFromTemplate will create a new template db
+	templateDBName := os.Getenv("DB_FROM")
+	if openOptions.DBFrom != nil {
+		templateDBName = *openOptions.DBFrom
+	}
+	if err = createDatabaseFromTemplate(t, defaultConnectionParams, b.coderTestingDB, dbName, templateDBName); err != nil {
+		return ConnectionParams{}, xerrors.Errorf("create database: %w", err)
+	}
+
+	testDBParams := ConnectionParams{
+		Username: username,
+		Password: password,
+		Host:     host,
+		Port:     port,
+		DBName:   dbName,
+	}
+
+	// Optionally log the DSN to help connect to the test database.
+	if openOptions.LogDSN {
+		_, _ = fmt.Fprintf(os.Stderr, "Connect to the database for %s using: psql '%s'\n", t.Name(), testDBParams.DSN())
+	}
+	t.Cleanup(b.clean(t, dbName))
+	return testDBParams, nil
+}
+
+func (b *Broker) clean(t TBSubset, dbName string) func() {
+	return func() {
+		_, err := b.coderTestingDB.Exec("DROP DATABASE " + dbName + ";")
+		if err != nil {
+			t.Logf("failed to clean up database %q: %s\n", dbName, err.Error())
+			return
+		}
+		_, err = b.coderTestingDB.Exec("UPDATE test_databases SET dropped_at = CURRENT_TIMESTAMP WHERE name = $1", dbName)
+		if err != nil {
+			t.Logf("failed to mark test database '%s' dropped: %s\n", dbName, err.Error())
+		}
+	}
+}
+
+func (b *Broker) init(t TBSubset) error {
+	b.Lock()
+	defer b.Unlock()
+	b.refCount++
+	t.Cleanup(b.decRef)
+	if b.coderTestingDB != nil {
+		// already initialized
+		return nil
+	}
+
+	connectionParamsInitOnce.Do(func() {
+		errDefaultConnectionParamsInit = initDefaultConnection(t)
+	})
+	if errDefaultConnectionParamsInit != nil {
+		return xerrors.Errorf("init default connection params: %w", errDefaultConnectionParamsInit)
+	}
+	coderTestingParams := defaultConnectionParams
+	coderTestingParams.DBName = CoderTestingDBName
+	coderTestingDB, err := sql.Open("postgres", coderTestingParams.DSN())
+	if err != nil {
+		return xerrors.Errorf("open postgres connection: %w", err)
+	}
+
+	// creating the db can succeed even if the database doesn't exist. Ping it to find out.
+	err = coderTestingDB.Ping()
+	var pqErr *pq.Error
+	if xerrors.As(err, &pqErr) && pqErr.Code == "3D000" {
+		// database does not exist.
+		if closeErr := coderTestingDB.Close(); closeErr != nil {
+			return xerrors.Errorf("close postgres connection: %w", closeErr)
+		}
+		err = createCoderTestingDB(t)
+		if err != nil {
+			return xerrors.Errorf("create coder testing db: %w", err)
+		}
+		coderTestingDB, err = sql.Open("postgres", coderTestingParams.DSN())
+		if err != nil {
+			return xerrors.Errorf("open postgres connection: %w", err)
+		}
+	} else if err != nil {
+		_ = coderTestingDB.Close()
+		return xerrors.Errorf("ping '%s' database: %w", CoderTestingDBName, err)
+	}
+	b.coderTestingDB = coderTestingDB
+	b.uuid = uuid.New()
+	return nil
+}
+
+func createCoderTestingDB(t TBSubset) error {
+	db, err := sql.Open("postgres", defaultConnectionParams.DSN())
+	if err != nil {
+		return xerrors.Errorf("open postgres connection: %w", err)
+	}
+	defer func() {
+		_ = db.Close()
+	}()
+	err = createAndInitDatabase(t, defaultConnectionParams, db, CoderTestingDBName, func(testDB *sql.DB) error {
+		_, err := testDB.Exec(coderTestingSQLInit)
+		return err
+	})
+	if err != nil {
+		return xerrors.Errorf("create coder testing db: %w", err)
+	}
+	return nil
+}
+
+func (b *Broker) decRef() {
+	b.Lock()
+	defer b.Unlock()
+	b.refCount--
+	if b.refCount == 0 {
+		// ensures we don't leave go routines around for GoLeak to find.
+		_ = b.coderTestingDB.Close()
+		b.coderTestingDB = nil
+	}
+}