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

johnstcn · johnstcn · commit d06e9c1d869f · 2023-09-15T17:49:32.000+01:00
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/coderd/database/dbtestutil/db.go b/coderd/database/dbtestutil/db.go
@@ -1,15 +1,20 @@
 package dbtestutil
 
 import (
+	"bytes"
 	"context"
 	"database/sql"
 	"fmt"
 	"net/url"
 	"os"
+	"os/exec"
+	"path/filepath"
+	"regexp"
 	"strings"
 	"testing"
 
 	"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 +29,7 @@ func WillUsePostgres() bool {
 
 type options struct {
 	fixedTimezone string
+	dumpOnFailure bool
 }
 
 type Option func(*options)
@@ -35,6 +41,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 +87,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 +126,85 @@ func dbNameFromConnectionURL(t testing.TB, connectionURL string) string {
 	require.NoError(t, err)
 	return strings.TrimPrefix(u.Path, "/")
 }
+
+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(), "_")
+	outPath := filepath.Join(cwd, snakeCaseName+".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: %w", err)
+		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",
+		"--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 {
+		// Skip blank lines
+		if len(line) == 0 {
+			continue
+		}
+		// Skip comments
+		if bytes.HasPrefix(line, []byte("--")) {
+			continue
+		}
+		// Skip SELECT or SET statements
+		if bytes.HasPrefix(line, []byte("SELECT")) {
+			continue
+		}
+		if bytes.HasPrefix(line, []byte("SET")) {
+			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)
