coder/coderd/database/gen/dump/main.go

package main

import (
	"database/sql"
	"fmt"
	"net"
	"os"
	"os/exec"
	"os/signal"
	"path/filepath"
	"runtime"
	"strconv"
	"strings"
	"sync"
	"syscall"

	embeddedpostgres "github.com/fergusstrange/embedded-postgres"
	"golang.org/x/xerrors"

	"github.com/coder/coder/v2/coderd/database/dbtestutil"
	"github.com/coder/coder/v2/coderd/database/migrations"
	"github.com/coder/coder/v2/scripts/atomicwrite"
)

var preamble = []byte("-- Code generated by 'make coderd/database/generate'. DO NOT EDIT.")

type mockTB struct {
	cleanup []func()
}

func (*mockTB) Name() string {
	return "mockTB"
}

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 (*mockTB) TempDir() string {
	panic("not implemented")
}

func main() {
	t := &mockTB{}

	// Ensure cleanups run on both normal exit and SIGINT/SIGTERM.
	// Go's default signal handlers call os.Exit, which skips deferred
	// funcs and would leave an embedded-postgres daemon orphaned.
	var cleanupOnce sync.Once
	runCleanup := func() {
		cleanupOnce.Do(func() {
			for _, f := range t.cleanup {
				f()
			}
		})
	}
	defer runCleanup()

	sigCh := make(chan os.Signal, 1)
	signal.Notify(sigCh, os.Interrupt, syscall.SIGTERM)
	go func() {
		<-sigCh
		runCleanup()
		os.Exit(130)
	}()

	connection := os.Getenv("DB_DUMP_CONNECTION_URL")
	if connection == "" {
		var cleanup func()
		var err error
		connection, cleanup, err = dbtestutil.OpenContainerized(t, dbtestutil.DBContainerOptions{})
		if err != nil {
			_, _ = fmt.Fprintf(os.Stderr, "containerized postgres unavailable (%s); falling back to embedded postgres\n", err)
			connection, cleanup, err = openEmbeddedPostgres()
			if err != nil {
				panic(err)
			}
		}
		t.Cleanup(cleanup)
	}

	db, err := sql.Open("postgres", connection)
	if err != nil {
		err = xerrors.Errorf("open database failed: %w", err)
		panic(err)
	}
	defer db.Close()

	err = migrations.Up(db)
	if err != nil {
		err = xerrors.Errorf("run migrations failed: %w", err)
		panic(err)
	}

	dumpBytes, err := dbtestutil.PGDumpSchemaOnly(connection)
	if err != nil {
		if !pgDumpUsable() {
			_, _ = fmt.Fprintf(os.Stderr,
				"\nThis step needs pg_dump (PostgreSQL v13 or later) on PATH OR a Docker-compatible daemon.\n"+
					"Install pg_dump locally to avoid Docker:\n"+
					"  mise:  mise use -g postgres@13\n"+
					"  brew:  brew install libpq && brew link --force libpq\n"+
					"  apt:   sudo apt-get install -y postgresql-client\n\n")
		}
		err = xerrors.Errorf("dump schema failed: %w", err)
		panic(err)
	}

	_, mainPath, _, ok := runtime.Caller(0)
	if !ok {
		panic("couldn't get caller path")
	}
	err = atomicwrite.File(filepath.Join(mainPath, "..", "..", "..", "dump.sql"), append(preamble, dumpBytes...))
	if err != nil {
		err = xerrors.Errorf("write dump failed: %w", err)
		panic(err)
	}
}

// pgDumpUsable mirrors PGDumpSchemaOnly's requirement (pg_dump on PATH at
// v13 or later). PGDumpSchemaOnly silently falls back to `docker run` when
// either condition fails, so we only show the install hint here when the
// local pg_dump is genuinely unusable. Otherwise an old pg_dump would
// produce a misleading Docker-not-found message.
func pgDumpUsable() bool {
	path, err := exec.LookPath("pg_dump")
	if err != nil {
		return false
	}
	out, err := exec.Command(path, "--version").Output()
	if err != nil {
		return false
	}
	// Output format: "pg_dump (PostgreSQL) 14.5 ..."
	parts := strings.Fields(string(out))
	if len(parts) < 3 {
		return false
	}
	major, err := strconv.Atoi(strings.SplitN(parts[2], ".", 2)[0])
	if err != nil {
		return false
	}
	return major >= 13
}

func openEmbeddedPostgres() (string, func(), error) {
	listener, err := net.Listen("tcp4", "127.0.0.1:0")
	if err != nil {
		return "", nil, xerrors.Errorf("find ephemeral port: %w", err)
	}
	tcpAddr, ok := listener.Addr().(*net.TCPAddr)
	if !ok {
		_ = listener.Close()
		return "", nil, xerrors.New("listener returned non-TCP addr")
	}
	port := tcpAddr.Port
	_ = listener.Close()

	cacheRoot, err := os.UserCacheDir()
	if err != nil {
		cacheRoot = os.TempDir()
	}
	cacheDir := filepath.Join(cacheRoot, "coder", "dbdump-postgres")

	runtimeDir, err := os.MkdirTemp("", "coder-dbdump-postgres-")
	if err != nil {
		return "", nil, xerrors.Errorf("create runtime dir: %w", err)
	}

	const password = "postgres"
	ep := embeddedpostgres.NewDatabase(
		embeddedpostgres.DefaultConfig().
			Version(embeddedpostgres.V13).
			// repo1.maven.org is flaky; matches cli/server.go and scripts/embedded-pg/main.go.
			BinaryRepositoryURL("https://repo.maven.apache.org/maven2").
			BinariesPath(filepath.Join(cacheDir, "bin")).
			CachePath(filepath.Join(cacheDir, "cache")).
			DataPath(filepath.Join(runtimeDir, "data")).
			RuntimePath(filepath.Join(runtimeDir, "runtime")).
			Port(uint32(port)). //nolint:gosec // port from listener, fits uint32.
			Username("postgres").
			Password(password).
			Database("postgres").
			// Postgres canonicalizes timestamptz DEFAULT expressions at
			// parse time using the server timezone GUC, then stores the
			// canonical form in pg_attrdef. Without UTC, the host's TZ
			// leaks into dump.sql as values like '0001-12-31 23:06:32+00 BC'.
			StartParameters(map[string]string{"timezone": "UTC"}).
			Logger(nil),
	)

	_, _ = fmt.Fprintln(os.Stderr, "starting embedded postgres (first run may download binaries)...")
	if err := ep.Start(); err != nil {
		_ = os.RemoveAll(runtimeDir)
		return "", nil, xerrors.Errorf("start embedded postgres: %w", err)
	}

	dsn := dbtestutil.ConnectionParams{
		Username: "postgres",
		Password: password,
		Host:     "127.0.0.1",
		Port:     strconv.Itoa(port),
		DBName:   "postgres",
	}.DSN()

	cleanup := func() {
		if stopErr := ep.Stop(); stopErr != nil {
			_, _ = fmt.Fprintf(os.Stderr, "failed to stop embedded postgres: %s\n", stopErr)
		}
		_ = os.RemoveAll(runtimeDir)
	}
	return dsn, cleanup, nil
}