feat(coderd/database/dbtestutil): add ability to dump database on failure (#9704)

johnstcn · web-flow · commit 1df75891055b · 2023-09-18T11:50:15.000+01:00
Adds dbtestutil.DumpOnFailure() to allow dumping the entire test database contents upon test failure.
This does nothing for dbfake currently.
diff --git a/.gitignore b/.gitignore
@@ -64,3 +64,6 @@ scaletest/terraform/secrets.tfvars
 
 # Nix
 result
+
+# Data dumps from unit tests
+**/*.test.sql
diff --git a/.prettierignore b/.prettierignore
@@ -67,6 +67,9 @@ scaletest/terraform/secrets.tfvars
 
 # Nix
 result
+
+# Data dumps from unit tests
+**/*.test.sql
 # .prettierignore.include:
 # Helm templates contain variables that are invalid YAML and can't be formatted
 # by Prettier.
diff --git a/coderd/database/dbtestutil/db.go b/coderd/database/dbtestutil/db.go
@@ -1,15 +1,21 @@
 package dbtestutil
 
 import (
+	"bytes"
 	"context"
 	"database/sql"
 	"fmt"
 	"net/url"
 	"os"
+	"os/exec"
+	"path/filepath"
+	"regexp"
 	"strings"
 	"testing"
+	"time"
 
 	"github.com/stretchr/testify/require"
+	"golang.org/x/xerrors"
 
 	"github.com/coder/coder/v2/coderd/database"
 	"github.com/coder/coder/v2/coderd/database/dbfake"
@@ -24,6 +30,7 @@ func WillUsePostgres() bool {
 
 type options struct {
 	fixedTimezone string
+	dumpOnFailure bool
 }
 
 type Option func(*options)
@@ -35,6 +42,13 @@ func WithTimezone(tz string) Option {
 	}
 }
 
+// WithDumpOnFailure will dump the entire database on test failure.
+func WithDumpOnFailure() Option {
+	return func(o *options) {
+		o.dumpOnFailure = true
+	}
+}
+
 func NewDB(t testing.TB, opts ...Option) (database.Store, pubsub.Pubsub) {
 	t.Helper()
 
@@ -74,6 +88,9 @@ func NewDB(t testing.TB, opts ...Option) (database.Store, pubsub.Pubsub) {
 		t.Cleanup(func() {
 			_ = sqlDB.Close()
 		})
+		if o.dumpOnFailure {
+			t.Cleanup(func() { DumpOnFailure(t, connectionURL) })
+		}
 		db = database.New(sqlDB)
 
 		ps, err = pubsub.New(context.Background(), sqlDB, connectionURL)
@@ -110,3 +127,87 @@ func dbNameFromConnectionURL(t testing.TB, connectionURL string) string {
 	require.NoError(t, err)
 	return strings.TrimPrefix(u.Path, "/")
 }
+
+// DumpOnFailure exports the database referenced by connectionURL to a file
+// corresponding to the current test, with a suffix indicating the time the
+// test was run.
+// To import this into a new database (assuming you have already run make test-postgres-docker):
+//   - Create a new test database:
+//     go run ./scripts/migrate-ci/main.go and note the database name it outputs
+//   - Import the file into the above database:
+//     psql 'postgres://postgres:postgres@127.0.0.1:5432/<dbname>?sslmode=disable' -f <path to file.test.sql>
+//   - Run a dev server against that database:
+//     ./scripts/coder-dev.sh server --postgres-url='postgres://postgres:postgres@127.0.0.1:5432/<dbname>?sslmode=disable'
+func DumpOnFailure(t testing.TB, connectionURL string) {
+	if !t.Failed() {
+		return
+	}
+	cwd, err := filepath.Abs(".")
+	if err != nil {
+		t.Errorf("dump on failure: cannot determine current working directory")
+		return
+	}
+	snakeCaseName := regexp.MustCompile("[^a-zA-Z0-9-_]+").ReplaceAllString(t.Name(), "_")
+	now := time.Now()
+	timeSuffix := fmt.Sprintf("%d%d%d%d%d%d", now.Year(), now.Month(), now.Day(), now.Hour(), now.Minute(), now.Second())
+	outPath := filepath.Join(cwd, snakeCaseName+"."+timeSuffix+".test.sql")
+	dump, err := pgDump(connectionURL)
+	if err != nil {
+		t.Errorf("dump on failure: failed to run pg_dump")
+		return
+	}
+	if err := os.WriteFile(outPath, filterDump(dump), 0o600); err != nil {
+		t.Errorf("dump on failure: failed to write: %s", err.Error())
+		return
+	}
+	t.Logf("Dumped database to %q due to failed test. I hope you find what you're looking for!", outPath)
+}
+
+// pgDump runs pg_dump against dbURL and returns the output.
+func pgDump(dbURL string) ([]byte, error) {
+	if _, err := exec.LookPath("pg_dump"); err != nil {
+		return nil, xerrors.Errorf("could not find pg_dump in path: %w", err)
+	}
+	cmdArgs := []string{
+		"pg_dump",
+		dbURL,
+		"--data-only",
+		"--column-inserts",
+		"--no-comments",
+		"--no-privileges",
+		"--no-publication",
+		"--no-security-labels",
+		"--no-subscriptions",
+		"--no-tablespaces",
+		// "--no-unlogged-table-data", // some tables are unlogged and may contain data of interest
+		"--no-owner",
+		"--exclude-table=schema_migrations",
+	}
+	cmd := exec.Command(cmdArgs[0], cmdArgs[1:]...) // nolint:gosec
+	cmd.Env = []string{
+		// "PGTZ=UTC", // This is probably not going to be useful if tz has been changed.
+		"PGCLIENTENCODINDG=UTF8",
+		"PGDATABASE=", // we should always specify the database name in the connection string
+	}
+	var stdout bytes.Buffer
+	cmd.Stdout = &stdout
+	if err := cmd.Run(); err != nil {
+		return nil, xerrors.Errorf("exec pg_dump: %w", err)
+	}
+	return stdout.Bytes(), nil
+}
+
+func filterDump(dump []byte) []byte {
+	lines := bytes.Split(dump, []byte{'\n'})
+	var buf bytes.Buffer
+	for _, line := range lines {
+		// We dump in column-insert format, so these are the only lines
+		// we care about
+		if !bytes.HasPrefix(line, []byte("INSERT")) {
+			continue
+		}
+		_, _ = buf.Write(line)
+		_, _ = buf.WriteRune('\n')
+	}
+	return buf.Bytes()
+}
diff --git a/enterprise/cli/server_dbcrypt_test.go b/enterprise/cli/server_dbcrypt_test.go
@@ -36,6 +36,7 @@ func TestServerDBCrypt(t *testing.T) {
 	connectionURL, closePg, err := postgres.Open()
 	require.NoError(t, err)
 	t.Cleanup(closePg)
+	t.Cleanup(func() { dbtestutil.DumpOnFailure(t, connectionURL) })
 
 	sqlDB, err := sql.Open("postgres", connectionURL)
 	require.NoError(t, err)
@@ -44,13 +45,6 @@ func TestServerDBCrypt(t *testing.T) {
 	})
 	db := database.New(sqlDB)
 
-	t.Cleanup(func() {
-		if t.Failed() {
-			t.Logf("Dumping data due to failed test. I hope you find what you're looking for!")
-			dumpUsers(t, sqlDB)
-		}
-	})
-
 	// Populate the database with some unencrypted data.
 	t.Logf("Generating unencrypted data")
 	users := genData(t, db)
@@ -250,50 +244,6 @@ func genData(t *testing.T, db database.Store) []database.User {
 	return users
 }
 
-func dumpUsers(t *testing.T, db *sql.DB) {
-	t.Helper()
-	rows, err := db.QueryContext(context.Background(), `SELECT
-	u.id,
-	u.login_type,
-	u.status,
-	u.deleted,
-	ul.oauth_access_token_key_id AS uloatkid,
-	ul.oauth_refresh_token_key_id AS ulortkid,
-	gal.oauth_access_token_key_id AS galoatkid,
-	gal.oauth_refresh_token_key_id AS galortkid
-FROM users u
-LEFT OUTER JOIN user_links ul ON u.id = ul.user_id
-LEFT OUTER JOIN git_auth_links gal ON u.id = gal.user_id
-ORDER BY u.created_at ASC;`)
-	require.NoError(t, err)
-	defer rows.Close()
-	for rows.Next() {
-		var (
-			id        string
-			loginType string
-			status    string
-			deleted   bool
-			UlOatKid  sql.NullString
-			UlOrtKid  sql.NullString
-			GalOatKid sql.NullString
-			GalOrtKid sql.NullString
-		)
-		require.NoError(t, rows.Scan(
-			&id,
-			&loginType,
-			&status,
-			&deleted,
-			&UlOatKid,
-			&UlOrtKid,
-			&GalOatKid,
-			&GalOrtKid,
-		))
-		t.Logf("user: id:%s login_type:%-8s status:%-9s deleted:%-5t ul_kids{at:%-7s rt:%-7s} gal_kids{at:%-7s rt:%-7s}",
-			id, loginType, status, deleted, UlOatKid.String, UlOrtKid.String, GalOatKid.String, GalOrtKid.String,
-		)
-	}
-}
-
 func mustString(t *testing.T, n int) string {
 	t.Helper()
 	s, err := cryptorand.String(n)
diff --git a/site/.eslintignore b/site/.eslintignore
@@ -67,6 +67,9 @@ stats/
 
 # Nix
 result
+
+# Data dumps from unit tests
+**/*.test.sql
 # .prettierignore.include:
 # Helm templates contain variables that are invalid YAML and can't be formatted
 # by Prettier.
diff --git a/site/.prettierignore b/site/.prettierignore
@@ -67,6 +67,9 @@ stats/
 
 # Nix
 result
+
+# Data dumps from unit tests
+**/*.test.sql
 # .prettierignore.include:
 # Helm templates contain variables that are invalid YAML and can't be formatted
 # by Prettier.
