waitfordb/main.go

// waitfordb blocks until a target database answers SELECT 1 under the given
// credentials, then exits 0. It is designed to run as a Kubernetes
// initContainer, configured entirely by environment variables.
//
// Exit codes: 0 ready, 1 timeout, 2 configuration error.
package main

import (
	"context"
	"fmt"
	"os"
	"os/signal"
	"strings"
	"syscall"
	"time"

	"git.unkin.net/unkin/waitfordb/internal/config"
	"git.unkin.net/unkin/waitfordb/internal/driver"
	"git.unkin.net/unkin/waitfordb/internal/wait"
)

// version is overridden at build time via -ldflags "-X main.version=...".
var version = "dev"

const (
	exitReady   = 0
	exitTimeout = 1
	exitConfig  = 2
)

func main() {
	os.Exit(run(os.Args[1:]))
}

func run(args []string) int {
	for _, a := range args {
		switch a {
		case "version", "--version", "-v":
			fmt.Println(version)
			return exitReady
		case "help", "--help", "-h":
			usage()
			return exitReady
		default:
			fmt.Fprintf(os.Stderr, "waitfordb: unknown argument %q\n\n", a)
			usage()
			return exitConfig
		}
	}

	logger := newLogger()

	cfg, err := config.Load(os.Getenv)
	if err != nil {
		logger.printf("config error: %v", err)
		return exitConfig
	}

	drv, err := driver.Get(cfg.Driver)
	if err != nil {
		logger.printf("config error: %v", err)
		return exitConfig
	}

	logger.printf("waitfordb %s: waiting for database (%s)", version, cfg.Redacted())

	pinger, err := drv.Open(cfg)
	if err != nil {
		logger.printf("config error: opening %s driver: %v", cfg.Driver, err)
		return exitConfig
	}
	defer pinger.Close()

	ctx, stop := signal.NotifyContext(context.Background(), syscall.SIGINT, syscall.SIGTERM)
	defer stop()

	target := targetName(cfg)

	res := wait.Run(ctx,
		wait.Params{Timeout: cfg.Timeout, Interval: cfg.Interval, ConnectTimeout: cfg.ConnectTimeout},
		pinger.Ping,
		func(attempt int, err error, elapsed, timeout, retryIn time.Duration) {
			logger.printf("attempt %d: %s; retrying in %s (elapsed %s/%s)",
				attempt, shortReason(err), roundDur(retryIn), roundDur(elapsed), timeoutStr(timeout))
		},
		wait.RealClock{},
	)

	switch {
	case res.OK:
		logger.printf("database %s ready after %s (%d attempt%s)",
			target, roundDur(res.Elapsed), res.Attempts, plural(res.Attempts))
		return exitReady
	case res.Cancelled:
		logger.printf("interrupted after %s waiting for %s (%d attempt%s): %s",
			roundDur(res.Elapsed), target, res.Attempts, plural(res.Attempts), shortReason(res.LastErr))
		return exitTimeout
	default: // timed out
		logger.printf("timed out after %s waiting for %s: %s",
			roundDur(res.Elapsed), target, shortReason(res.LastErr))
		return exitTimeout
	}
}

func targetName(cfg config.Config) string {
	if cfg.Database != "" {
		return cfg.Database
	}
	if cfg.DSN != "" {
		return "database"
	}
	return cfg.Host
}

// shortReason collapses a (possibly multi-line) driver error into one
// informative line. Driver connection errors include host/user/database but
// never the password, so this is safe to log.
func shortReason(err error) string {
	if err == nil {
		return "unknown error"
	}
	fields := strings.Fields(err.Error())
	return strings.TrimRight(strings.Join(fields, " "), ":")
}

func roundDur(d time.Duration) time.Duration {
	if d >= time.Second {
		return d.Round(100 * time.Millisecond)
	}
	return d.Round(time.Millisecond)
}

func timeoutStr(d time.Duration) string {
	if d == 0 {
		return "forever"
	}
	return d.String()
}

func plural(n int) string {
	if n == 1 {
		return ""
	}
	return "s"
}

func usage() {
	fmt.Fprint(os.Stderr, `waitfordb — block until a database is ready (SELECT 1 succeeds).

Usage:
  waitfordb            wait using the WAITFORDB_*/PG* environment variables
  waitfordb version    print the version
  waitfordb help       print this help

Environment:
  WAITFORDB_DRIVER           postgres (default) or mysql
  WAITFORDB_HOST             database host           (PGHOST fallback)
  WAITFORDB_PORT             database port           (PGPORT fallback)
  WAITFORDB_USER             username                (PGUSER fallback)
  WAITFORDB_PASSWORD         password                (PGPASSWORD fallback)
  WAITFORDB_DATABASE         database name           (PGDATABASE fallback)
  WAITFORDB_SSLMODE          postgres sslmode        (PGSSLMODE fallback)
  WAITFORDB_DSN              full connection string (overrides the fields above)
  WAITFORDB_TIMEOUT          total wait budget, Go duration; 0 = forever (default 0)
  WAITFORDB_INTERVAL         gap between retries    (default 2s)
  WAITFORDB_CONNECT_TIMEOUT  per-attempt connect timeout (default 5s)

Precedence for connection parameters: WAITFORDB_DSN > WAITFORDB_* > PG*.
Exit codes: 0 ready, 1 timeout/interrupted, 2 configuration error.
`)
}

type logger struct{ out *os.File }

func newLogger() logger { return logger{out: os.Stderr} }

func (l logger) printf(format string, a ...any) {
	ts := time.Now().Format("15:04:05")
	fmt.Fprintf(l.out, "%s %s\n", ts, fmt.Sprintf(format, a...))
}