diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..0e63a18 --- /dev/null +++ b/.gitignore @@ -0,0 +1,4 @@ +# built binary (repo root) +/waitfordb +# build output +dist/ diff --git a/.woodpecker/build.yaml b/.woodpecker/build.yaml new file mode 100644 index 0000000..4126eb7 --- /dev/null +++ b/.woodpecker/build.yaml @@ -0,0 +1,43 @@ +when: + - event: [pull_request, push] + +steps: + - name: test + image: golang:1.25 + commands: + - go test -race ./... + backend_options: + kubernetes: + serviceAccountName: default + resources: + requests: + memory: 512Mi + cpu: 1 + limits: + memory: 2Gi + cpu: 2 + + # Validate the image builds without pushing. The CA-baked buildx plugin trusts + # the internal registry's CA; the plain woodpeckerci plugin fails x509 here. + - name: build-check + image: artifactapi.k8s.syd1.au.unkin.net/docker-internal/plugin-docker-buildx:latest + settings: + repo: artifactapi.k8s.syd1.au.unkin.net/docker-internal/waitfordb + dockerfile: Dockerfile + dry_run: true + platforms: linux/amd64 + build_args: + VERSION: ${CI_COMMIT_SHA} + buildkit_config: | + [registry."artifactapi.k8s.syd1.au.unkin.net"] + ca = ["/etc/docker/certs.d/artifactapi.k8s.syd1.au.unkin.net/ca.crt"] + backend_options: + kubernetes: + serviceAccountName: default + resources: + requests: + memory: 1Gi + cpu: 1 + limits: + memory: 4Gi + cpu: 2 diff --git a/.woodpecker/docker.yaml b/.woodpecker/docker.yaml new file mode 100644 index 0000000..155baaf --- /dev/null +++ b/.woodpecker/docker.yaml @@ -0,0 +1,30 @@ +when: + - event: tag + ref: refs/tags/v* + +steps: + - name: docker + image: artifactapi.k8s.syd1.au.unkin.net/docker-internal/plugin-docker-buildx:latest + settings: + registry: artifactapi.k8s.syd1.au.unkin.net + repo: artifactapi.k8s.syd1.au.unkin.net/docker-internal/waitfordb + dockerfile: Dockerfile + platforms: linux/amd64 + build_args: + VERSION: ${CI_COMMIT_TAG} + buildkit_config: | + [registry."artifactapi.k8s.syd1.au.unkin.net"] + ca = ["/etc/docker/certs.d/artifactapi.k8s.syd1.au.unkin.net/ca.crt"] + tags: + - ${CI_COMMIT_TAG} + - latest + backend_options: + kubernetes: + serviceAccountName: default + resources: + requests: + memory: 1Gi + cpu: 1 + limits: + memory: 4Gi + cpu: 2 diff --git a/Dockerfile b/Dockerfile new file mode 100644 index 0000000..17c8b2c --- /dev/null +++ b/Dockerfile @@ -0,0 +1,18 @@ +FROM golang:1.25-alpine AS builder + +WORKDIR /build + +COPY go.mod go.sum ./ +RUN go mod download + +COPY . . + +ARG VERSION=dev +RUN CGO_ENABLED=0 go build -ldflags="-s -w -X main.version=${VERSION}" -o waitfordb . + +# distroless static ships ca-certificates and runs as an unprivileged user. +FROM gcr.io/distroless/static-debian12:nonroot + +COPY --from=builder /build/waitfordb /usr/local/bin/waitfordb + +ENTRYPOINT ["waitfordb"] diff --git a/Makefile b/Makefile new file mode 100644 index 0000000..fe8acbf --- /dev/null +++ b/Makefile @@ -0,0 +1,58 @@ +BINARY := waitfordb +DIST := dist +IMAGE := artifactapi.k8s.syd1.au.unkin.net/docker-internal/waitfordb +VERSION := $(shell git describe --tags --always --dirty 2>/dev/null || echo dev) +GOFLAGS := -ldflags="-s -w -X main.version=$(VERSION)" +OS ?= $(shell go env GOOS) +ARCH ?= $(shell go env GOARCH) + +.PHONY: all build test test-integration lint fmt clean docker patch minor major _tag + +all: build + +build: + CGO_ENABLED=0 GOOS=$(OS) GOARCH=$(ARCH) go build $(GOFLAGS) -o $(DIST)/$(BINARY) . + +test: + go test -race ./... + +# Integration test spins up a real postgres container (podman); skipped when no +# container runtime is present. +test-integration: + go test -tags integration -race -run TestIntegration -v . + +lint: + golangci-lint run ./... + +fmt: + gofmt -w . + +clean: + rm -rf $(DIST) + +# Build the container image locally for a smoke test. +docker: + docker build --build-arg VERSION=$(VERSION) -t $(IMAGE):$(VERSION) . + +# Bump helpers — read the latest semver tag and create the next one, then push +# it so the Woodpecker tag pipeline builds and publishes the image. +_LATEST := $(shell git tag --sort=-v:refname | grep -E '^v[0-9]+\.[0-9]+\.[0-9]+$$' | head -1) +_BASE := $(if $(_LATEST),$(_LATEST),v0.0.0) +_MAJ := $(shell echo $(_BASE) | sed 's/^v//' | cut -d. -f1) +_MIN := $(shell echo $(_BASE) | sed 's/^v//' | cut -d. -f2) +_PAT := $(shell echo $(_BASE) | sed 's/^v//' | cut -d. -f3) + +patch: + @NEW=v$(_MAJ).$(_MIN).$(shell expr $(_PAT) + 1); \ + git tag $$NEW && echo "Tagged $$NEW" && $(MAKE) _tag TAG=$$NEW + +minor: + @NEW=v$(_MAJ).$(shell expr $(_MIN) + 1).0; \ + git tag $$NEW && echo "Tagged $$NEW" && $(MAKE) _tag TAG=$$NEW + +major: + @NEW=v$(shell expr $(_MAJ) + 1).0.0; \ + git tag $$NEW && echo "Tagged $$NEW" && $(MAKE) _tag TAG=$$NEW + +_tag: + git push origin $(TAG) diff --git a/README.md b/README.md index ab52993..a2f4827 100644 --- a/README.md +++ b/README.md @@ -1,3 +1,121 @@ # waitfordb -A small Go tool, shipped as a container image, used as a Kubernetes initContainer to block an app from starting until its database is ready. \ No newline at end of file +A tiny, stateless Go tool that blocks until a database is ready, then exits `0`. +It is designed to run as a Kubernetes **initContainer** so an app never starts +before its database can serve queries. It replaces the hand-written +`psql`-in-a-shell initContainers we used on the arrstack (sonarr/radarr/prowlarr). + +Readiness means a trivial liveness query (`SELECT 1`) succeeds under the given +credentials and database — proving the **server is up, auth works, and the +target database/role exist**. On timeout it exits non-zero so the pod fails +fast and Kubernetes restarts it. + +`postgres` is the default driver; `mysql` is also supported. Configuration is +entirely via environment variables. + +## Environment variables + +| Variable | Default | Description | +| --------------------------- | ------------ | ------------------------------------------------------------------ | +| `WAITFORDB_DRIVER` | `postgres` | Database driver: `postgres` or `mysql`. | +| `WAITFORDB_HOST` | `localhost` | Database host. Postgres falls back to `PGHOST`. | +| `WAITFORDB_PORT` | `5432`/`3306`| Database port (driver default). Postgres falls back to `PGPORT`. | +| `WAITFORDB_USER` | — | Username. Postgres falls back to `PGUSER`. Required (unless DSN). | +| `WAITFORDB_PASSWORD` | — | Password. Postgres falls back to `PGPASSWORD`. | +| `WAITFORDB_DATABASE` | — | Database name. Postgres falls back to `PGDATABASE`. Required (unless DSN). | +| `WAITFORDB_SSLMODE` | driver default | Postgres sslmode (e.g. `disable`). Falls back to `PGSSLMODE`. | +| `WAITFORDB_DSN` | — | Full driver-native connection string. Overrides all fields above. | +| `WAITFORDB_TIMEOUT` | `0` (forever)| Total wait budget as a Go duration (e.g. `5m`). `0` waits forever. | +| `WAITFORDB_INTERVAL` | `2s` | Gap between retries. | +| `WAITFORDB_CONNECT_TIMEOUT` | `5s` | Per-attempt connect timeout, so a black-holed host cannot hang. | + +**Precedence for connection parameters:** `WAITFORDB_DSN` > `WAITFORDB_*` > `PG*`. +The `PG*` (libpq) fallback applies to the **postgres** driver only, so the tool +is a drop-in replacement anywhere those variables are already set. + +### Exit codes + +| Code | Meaning | +| ---- | ------------------------------------------- | +| `0` | Database ready. | +| `1` | Timed out (or interrupted by SIGTERM/SIGINT). | +| `2` | Configuration error. | + +`waitfordb version` prints the version; `waitfordb help` prints usage. + +### Example logs + +``` +12:49:50 waitfordb v0.1.0: waiting for database (driver=postgres addr=db:5432 database=sonarr-main user=sonarr password=*** timeout=5m interval=2s connect_timeout=5s) +12:49:52 attempt 1: failed to connect to `user=sonarr database=sonarr-main`: dial tcp ... connection refused; retrying in 2s (elapsed 2s/5m) +12:49:56 database sonarr-main ready after 4.2s (3 attempts) +``` + +The password is never printed — not in the startup line, not in any error line. + +## Kubernetes initContainer example + +Drop this initContainer into a Deployment/StatefulSet. It reuses the same +`Secret` the app consumes, so there is no second place to keep credentials. + +```yaml +initContainers: + - name: wait-for-db + image: artifactapi.k8s.syd1.au.unkin.net/docker-internal/waitfordb:latest + env: + - name: WAITFORDB_HOST + value: arrstack-postgres-rw.arrstack.svc.cluster.local + - name: WAITFORDB_PORT + value: "5432" + - name: WAITFORDB_DATABASE + value: sonarr-main + - name: WAITFORDB_SSLMODE + value: disable + - name: WAITFORDB_TIMEOUT + value: 5m + - name: WAITFORDB_USER + valueFrom: + secretKeyRef: + name: sonarr-db + key: username + - name: WAITFORDB_PASSWORD + valueFrom: + secretKeyRef: + name: sonarr-db + key: password + resources: + requests: + cpu: 10m + memory: 16Mi + limits: + cpu: 100m + memory: 64Mi + securityContext: + allowPrivilegeEscalation: false + readOnlyRootFilesystem: true + runAsNonRoot: true + capabilities: + drop: ["ALL"] +``` + +The libpq fallback means you can instead pass a single set of `PG*` variables +(e.g. via `envFrom` a shared ConfigMap/Secret) with no `WAITFORDB_*` at all. + +## Development + +```sh +make build # build ./dist/waitfordb +make test # unit tests (fake driver + fake clock, deterministic) +make test-integration # spins up postgres:16 via podman, asserts wait-then-succeed +make docker # build the container image locally +``` + +Releases are cut by tagging: `make patch|minor|major` bumps the semver tag and +pushes it, which triggers the Woodpecker pipeline to build and push +`artifactapi.k8s.syd1.au.unkin.net/docker-internal/waitfordb:` (and `latest`). + +## Adding a driver + +Implement `driver.Driver` (`Name()` + `Open(config.Config) (Pinger, error)`) and +register it from an `init()`. Engines that plug into `database/sql` can reuse the +shared `sqlPinger` (see `internal/driver/postgres.go` and `mysql.go`). diff --git a/go.mod b/go.mod new file mode 100644 index 0000000..de27ca7 --- /dev/null +++ b/go.mod @@ -0,0 +1,18 @@ +module git.unkin.net/unkin/waitfordb + +go 1.25 + +require ( + github.com/go-sql-driver/mysql v1.9.3 + github.com/jackc/pgx/v5 v5.7.6 +) + +require ( + filippo.io/edwards25519 v1.1.0 // indirect + github.com/jackc/pgpassfile v1.0.0 // indirect + github.com/jackc/pgservicefile v0.0.0-20240606120523-5a60cdf6a761 // indirect + github.com/jackc/puddle/v2 v2.2.2 // indirect + golang.org/x/crypto v0.37.0 // indirect + golang.org/x/sync v0.13.0 // indirect + golang.org/x/text v0.24.0 // indirect +) diff --git a/go.sum b/go.sum new file mode 100644 index 0000000..5e1fe02 --- /dev/null +++ b/go.sum @@ -0,0 +1,32 @@ +filippo.io/edwards25519 v1.1.0 h1:FNf4tywRC1HmFuKW5xopWpigGjJKiJSV0Cqo0cJWDaA= +filippo.io/edwards25519 v1.1.0/go.mod h1:BxyFTGdWcka3PhytdK4V28tE5sGfRvvvRV7EaN4VDT4= +github.com/davecgh/go-spew v1.1.0/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38= +github.com/davecgh/go-spew v1.1.1 h1:vj9j/u1bqnvCEfJOwUhtlOARqs3+rkHYY13jYWTU97c= +github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38= +github.com/go-sql-driver/mysql v1.9.3 h1:U/N249h2WzJ3Ukj8SowVFjdtZKfu9vlLZxjPXV1aweo= +github.com/go-sql-driver/mysql v1.9.3/go.mod h1:qn46aNg1333BRMNU69Lq93t8du/dwxI64Gl8i5p1WMU= +github.com/jackc/pgpassfile v1.0.0 h1:/6Hmqy13Ss2zCq62VdNG8tM1wchn8zjSGOBJ6icpsIM= +github.com/jackc/pgpassfile v1.0.0/go.mod h1:CEx0iS5ambNFdcRtxPj5JhEz+xB6uRky5eyVu/W2HEg= +github.com/jackc/pgservicefile v0.0.0-20240606120523-5a60cdf6a761 h1:iCEnooe7UlwOQYpKFhBabPMi4aNAfoODPEFNiAnClxo= +github.com/jackc/pgservicefile v0.0.0-20240606120523-5a60cdf6a761/go.mod h1:5TJZWKEWniPve33vlWYSoGYefn3gLQRzjfDlhSJ9ZKM= +github.com/jackc/pgx/v5 v5.7.6 h1:rWQc5FwZSPX58r1OQmkuaNicxdmExaEz5A2DO2hUuTk= +github.com/jackc/pgx/v5 v5.7.6/go.mod h1:aruU7o91Tc2q2cFp5h4uP3f6ztExVpyVv88Xl/8Vl8M= +github.com/jackc/puddle/v2 v2.2.2 h1:PR8nw+E/1w0GLuRFSmiioY6UooMp6KJv0/61nB7icHo= +github.com/jackc/puddle/v2 v2.2.2/go.mod h1:vriiEXHvEE654aYKXXjOvZM39qJ0q+azkZFrfEOc3H4= +github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM= +github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4= +github.com/stretchr/objx v0.1.0/go.mod h1:HFkY916IF+rwdDfMAkV7OtwuqBVzrE8GR6GFx+wExME= +github.com/stretchr/testify v1.3.0/go.mod h1:M5WIy9Dh21IEIfnGCwXGc5bZfKNJtfHm1UVUgZn+9EI= +github.com/stretchr/testify v1.7.0/go.mod h1:6Fq8oRcR53rry900zMqJjRRixrwX3KX962/h/Wwjteg= +github.com/stretchr/testify v1.8.1 h1:w7B6lhMri9wdJUVmEZPGGhZzrYTPvgJArz7wNPgYKsk= +github.com/stretchr/testify v1.8.1/go.mod h1:w2LPCIKwWwSfY2zedu0+kehJoqGctiVI29o6fzry7u4= +golang.org/x/crypto v0.37.0 h1:kJNSjF/Xp7kU0iB2Z+9viTPMW4EqqsrywMXLJOOsXSE= +golang.org/x/crypto v0.37.0/go.mod h1:vg+k43peMZ0pUMhYmVAWysMK35e6ioLh3wB8ZCAfbVc= +golang.org/x/sync v0.13.0 h1:AauUjRAJ9OSnvULf/ARrrVywoJDy0YS2AwQ98I37610= +golang.org/x/sync v0.13.0/go.mod h1:1dzgHSNfp02xaA81J2MS99Qcpr2w7fw1gpm99rleRqA= +golang.org/x/text v0.24.0 h1:dd5Bzh4yt5KYA8f9CJHCP4FB4D51c2c6JvN37xJJkJ0= +golang.org/x/text v0.24.0/go.mod h1:L8rBsPeo2pSS+xqN0d5u2ikmjtmoJbDBT1b7nHvFCdU= +gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0= +gopkg.in/yaml.v3 v3.0.0-20200313102051-9f266ea9e77c/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM= +gopkg.in/yaml.v3 v3.0.1 h1:fxVm/GzAzEWqLHuvctI91KS9hhNmmWOoWu0XTYJS7CA= +gopkg.in/yaml.v3 v3.0.1/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM= diff --git a/integration_test.go b/integration_test.go new file mode 100644 index 0000000..6943c0a --- /dev/null +++ b/integration_test.go @@ -0,0 +1,94 @@ +//go:build integration + +// Integration test: runs a real postgres container that accepts connections a +// few seconds after start and asserts waitfordb retries then succeeds. Run with: +// +// go test -tags integration -run TestIntegration -v . +// +// Requires podman (or docker via WAITFORDB_TEST_RUNTIME). Skipped otherwise. +package main + +import ( + "context" + "os" + "os/exec" + "strings" + "testing" + "time" + + "git.unkin.net/unkin/waitfordb/internal/config" + "git.unkin.net/unkin/waitfordb/internal/driver" + "git.unkin.net/unkin/waitfordb/internal/wait" +) + +func TestIntegrationPostgresComesUpLate(t *testing.T) { + runtime := envOr("WAITFORDB_TEST_RUNTIME", "podman") + if _, err := exec.LookPath(runtime); err != nil { + t.Skipf("%s not available: %v", runtime, err) + } + image := envOr("WAITFORDB_TEST_PG_IMAGE", "docker.io/library/postgres:16") + + name := "waitfordb-it-" + time.Now().Format("150405") + // Publish on a fixed host port; map to a random-ish high port. + hostPort := "55432" + args := []string{ + "run", "-d", "--rm", "--name", name, + "-e", "POSTGRES_PASSWORD=secret", + "-e", "POSTGRES_USER=app", + "-e", "POSTGRES_DB=appdb", + "-p", hostPort + ":5432", + image, + } + if out, err := exec.Command(runtime, args...).CombinedOutput(); err != nil { + t.Skipf("cannot start container (%s): %v\n%s", runtime, err, out) + } + t.Cleanup(func() { + _ = exec.Command(runtime, "rm", "-f", name).Run() + }) + + t.Setenv("WAITFORDB_HOST", "127.0.0.1") + t.Setenv("WAITFORDB_PORT", hostPort) + t.Setenv("WAITFORDB_USER", "app") + t.Setenv("WAITFORDB_PASSWORD", "secret") + t.Setenv("WAITFORDB_DATABASE", "appdb") + t.Setenv("WAITFORDB_SSLMODE", "disable") + t.Setenv("WAITFORDB_TIMEOUT", "60s") + t.Setenv("WAITFORDB_INTERVAL", "1s") + + cfg, err := config.Load(os.Getenv) + if err != nil { + t.Fatalf("config: %v", err) + } + drv, err := driver.Get(cfg.Driver) + if err != nil { + t.Fatalf("driver: %v", err) + } + pinger, err := drv.Open(cfg) + if err != nil { + t.Fatalf("open: %v", err) + } + defer pinger.Close() + + res := wait.Run(context.Background(), + wait.Params{Timeout: cfg.Timeout, Interval: cfg.Interval, ConnectTimeout: cfg.ConnectTimeout}, + pinger.Ping, + func(n int, err error, _, _, _ time.Duration) { + t.Logf("attempt %d failed as expected: %v", n, err) + }, + wait.RealClock{}) + + if !res.OK { + t.Fatalf("expected readiness, got %+v", res) + } + t.Logf("ready after %s (%d attempts)", res.Elapsed, res.Attempts) + if res.Attempts < 1 { + t.Errorf("expected at least one attempt") + } +} + +func envOr(k, def string) string { + if v := strings.TrimSpace(os.Getenv(k)); v != "" { + return v + } + return def +} diff --git a/internal/config/config.go b/internal/config/config.go new file mode 100644 index 0000000..f4dfd22 --- /dev/null +++ b/internal/config/config.go @@ -0,0 +1,236 @@ +// Package config resolves the waitfordb runtime configuration from environment +// variables and formats a secret-free summary for logging. +package config + +import ( + "fmt" + "strings" + "time" +) + +// Defaults applied when the corresponding env var is unset. +const ( + DefaultDriver = "postgres" + DefaultHost = "localhost" + DefaultInterval = 2 * time.Second + DefaultConnectTimeout = 5 * time.Second +) + +// defaultPort maps a driver to the port used when none is configured. +var defaultPort = map[string]string{ + "postgres": "5432", + "mysql": "3306", +} + +// Config is the resolved, validated configuration for a single run. +type Config struct { + Driver string + + // DSN, when set, is a full driver-native connection string that overrides + // the discrete Host/Port/User/Password/Database fields. + DSN string + + Host string + Port string + User string + Password string + Database string + SSLMode string + + // Timeout is the total budget to wait for readiness. Zero means wait + // forever. + Timeout time.Duration + // Interval is the gap between retries. + Interval time.Duration + // ConnectTimeout bounds a single connect+ping attempt. + ConnectTimeout time.Duration +} + +// Error is a configuration error; main maps it to exit code 2. +type Error struct{ msg string } + +func (e *Error) Error() string { return e.msg } + +func errf(format string, a ...any) *Error { return &Error{msg: fmt.Sprintf(format, a...)} } + +// Getenv matches os.Getenv; injected in tests. +type Getenv func(string) string + +// Load resolves configuration from env. Connection-parameter precedence is +// DSN > WAITFORDB_* > PG* (the libpq fallback applies to the postgres driver +// only). +func Load(get Getenv) (Config, error) { + c := Config{ + Driver: firstNonEmpty(get("WAITFORDB_DRIVER"), DefaultDriver), + DSN: get("WAITFORDB_DSN"), + ConnectTimeout: DefaultConnectTimeout, + Interval: DefaultInterval, + } + c.Driver = strings.ToLower(strings.TrimSpace(c.Driver)) + + pg := c.Driver == "postgres" + + // WAITFORDB_* first, then the libpq PG* fallback for postgres. + c.Host = pick(get, pg, "WAITFORDB_HOST", "PGHOST") + c.Port = pick(get, pg, "WAITFORDB_PORT", "PGPORT") + c.User = pick(get, pg, "WAITFORDB_USER", "PGUSER") + c.Password = pick(get, pg, "WAITFORDB_PASSWORD", "PGPASSWORD") + c.Database = pick(get, pg, "WAITFORDB_DATABASE", "PGDATABASE") + c.SSLMode = pick(get, pg, "WAITFORDB_SSLMODE", "PGSSLMODE") + + if c.Host == "" { + c.Host = DefaultHost + } + if c.Port == "" { + c.Port = defaultPort[c.Driver] + } + + var err error + if c.Timeout, err = parseDuration(get("WAITFORDB_TIMEOUT"), 0); err != nil { + return Config{}, errf("WAITFORDB_TIMEOUT: %v", err) + } + if c.Interval, err = parseDuration(get("WAITFORDB_INTERVAL"), DefaultInterval); err != nil { + return Config{}, errf("WAITFORDB_INTERVAL: %v", err) + } + if c.ConnectTimeout, err = parseDuration(get("WAITFORDB_CONNECT_TIMEOUT"), DefaultConnectTimeout); err != nil { + return Config{}, errf("WAITFORDB_CONNECT_TIMEOUT: %v", err) + } + + if err := c.validate(); err != nil { + return Config{}, err + } + return c, nil +} + +func (c Config) validate() error { + if _, ok := defaultPort[c.Driver]; !ok { + return errf("unsupported WAITFORDB_DRIVER %q (supported: postgres, mysql)", c.Driver) + } + if c.Interval <= 0 { + return errf("WAITFORDB_INTERVAL must be > 0") + } + if c.ConnectTimeout <= 0 { + return errf("WAITFORDB_CONNECT_TIMEOUT must be > 0") + } + if c.Timeout < 0 { + return errf("WAITFORDB_TIMEOUT must be >= 0") + } + // With a DSN the discrete fields are optional (the DSN carries them). + if c.DSN == "" { + if c.Database == "" { + return errf("no database configured: set WAITFORDB_DATABASE (or PGDATABASE for postgres) or WAITFORDB_DSN") + } + if c.User == "" { + return errf("no user configured: set WAITFORDB_USER (or PGUSER for postgres) or WAITFORDB_DSN") + } + } + return nil +} + +// Redacted returns a single-line, password-free summary of the resolved +// configuration suitable for the startup log line. +func (c Config) Redacted() string { + var b strings.Builder + fmt.Fprintf(&b, "driver=%s", c.Driver) + if c.DSN != "" { + fmt.Fprintf(&b, " dsn=%s", redactDSN(c.DSN)) + } else { + fmt.Fprintf(&b, " addr=%s:%s database=%s user=%s password=%s", + c.Host, c.Port, c.Database, c.User, redactSecret(c.Password)) + if c.SSLMode != "" { + fmt.Fprintf(&b, " sslmode=%s", c.SSLMode) + } + } + fmt.Fprintf(&b, " timeout=%s interval=%s connect_timeout=%s", + timeoutStr(c.Timeout), c.Interval, c.ConnectTimeout) + return b.String() +} + +func redactSecret(s string) string { + if s == "" { + return "(unset)" + } + return "***" +} + +func timeoutStr(d time.Duration) string { + if d == 0 { + return "forever" + } + return d.String() +} + +// redactDSN masks the password in either a URL-style or keyword-style DSN so it +// never reaches a log line. +func redactDSN(dsn string) string { + // URL form: scheme://user:password@host/... + if i := strings.Index(dsn, "://"); i >= 0 { + rest := dsn[i+3:] + if at := strings.Index(rest, "@"); at >= 0 { + creds := rest[:at] + if colon := strings.Index(creds, ":"); colon >= 0 { + return dsn[:i+3] + creds[:colon] + ":***@" + rest[at+1:] + } + } + return dsn + } + // Keyword form: key=value pairs and mysql user:pass@tcp(...) form. + out := dsn + for _, key := range []string{"password", "passwd"} { + out = redactKeyword(out, key) + } + // mysql DSN: user:pass@tcp(host)/db + if at := strings.Index(out, "@tcp("); at >= 0 { + if colon := strings.LastIndex(out[:at], ":"); colon >= 0 { + out = out[:colon+1] + "***" + out[at:] + } + } + return out +} + +func redactKeyword(dsn, key string) string { + lower := strings.ToLower(dsn) + idx := strings.Index(lower, key+"=") + if idx < 0 { + return dsn + } + valStart := idx + len(key) + 1 + valEnd := valStart + for valEnd < len(dsn) && dsn[valEnd] != ' ' { + valEnd++ + } + return dsn[:valStart] + "***" + dsn[valEnd:] +} + +// pick returns the WAITFORDB_* value, falling back to the PG* value only when +// fallback is true (postgres). +func pick(get Getenv, fallback bool, primary, secondary string) string { + if v := get(primary); v != "" { + return v + } + if fallback { + return get(secondary) + } + return "" +} + +func firstNonEmpty(vals ...string) string { + for _, v := range vals { + if v != "" { + return v + } + } + return "" +} + +func parseDuration(s string, def time.Duration) (time.Duration, error) { + s = strings.TrimSpace(s) + if s == "" { + return def, nil + } + d, err := time.ParseDuration(s) + if err != nil { + return 0, fmt.Errorf("invalid duration %q", s) + } + return d, nil +} diff --git a/internal/config/config_test.go b/internal/config/config_test.go new file mode 100644 index 0000000..f74a409 --- /dev/null +++ b/internal/config/config_test.go @@ -0,0 +1,173 @@ +package config + +import ( + "strings" + "testing" + "time" +) + +// envFrom returns a Getenv backed by a map. +func envFrom(m map[string]string) Getenv { + return func(k string) string { return m[k] } +} + +func TestLoadDefaults(t *testing.T) { + c, err := Load(envFrom(map[string]string{ + "WAITFORDB_USER": "sonarr", + "WAITFORDB_DATABASE": "sonarr-main", + "WAITFORDB_HOST": "db", + })) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + if c.Driver != "postgres" { + t.Errorf("driver = %q, want postgres", c.Driver) + } + if c.Port != "5432" { + t.Errorf("port = %q, want default 5432", c.Port) + } + if c.Interval != 2*time.Second { + t.Errorf("interval = %v, want 2s", c.Interval) + } + if c.ConnectTimeout != 5*time.Second { + t.Errorf("connect_timeout = %v, want 5s", c.ConnectTimeout) + } + if c.Timeout != 0 { + t.Errorf("timeout = %v, want 0 (forever)", c.Timeout) + } +} + +func TestPrecedenceDSNWins(t *testing.T) { + c, err := Load(envFrom(map[string]string{ + "WAITFORDB_DSN": "postgres://u:p@h:5432/d", + "WAITFORDB_HOST": "ignored", + "WAITFORDB_USER": "ignored", + "WAITFORDB_DATABASE": "ignored", + })) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + if c.DSN == "" { + t.Fatal("DSN should be set") + } +} + +func TestPrecedenceWaitfordbOverPG(t *testing.T) { + c, err := Load(envFrom(map[string]string{ + "WAITFORDB_HOST": "native-host", + "PGHOST": "pg-host", + "WAITFORDB_USER": "native-user", + "PGUSER": "pg-user", + "PGDATABASE": "pg-db", // only PG* set for database -> fallback used + })) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + if c.Host != "native-host" { + t.Errorf("host = %q, want native-host (WAITFORDB_* wins)", c.Host) + } + if c.User != "native-user" { + t.Errorf("user = %q, want native-user", c.User) + } + if c.Database != "pg-db" { + t.Errorf("database = %q, want pg-db (PG* fallback)", c.Database) + } +} + +func TestPGFallbackPostgresOnly(t *testing.T) { + // For a non-postgres driver the PG* fallback must not apply. + _, err := Load(envFrom(map[string]string{ + "WAITFORDB_DRIVER": "mysql", + "PGUSER": "pg-user", + "PGDATABASE": "pg-db", + // No WAITFORDB_USER/DATABASE -> must be a config error, PG* ignored. + })) + if err == nil { + t.Fatal("expected config error: PG* must not satisfy mysql config") + } +} + +func TestMySQLDefaultPort(t *testing.T) { + c, err := Load(envFrom(map[string]string{ + "WAITFORDB_DRIVER": "mysql", + "WAITFORDB_USER": "u", + "WAITFORDB_DATABASE": "d", + })) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + if c.Port != "3306" { + t.Errorf("port = %q, want 3306", c.Port) + } +} + +func TestValidationErrors(t *testing.T) { + cases := map[string]map[string]string{ + "missing database": {"WAITFORDB_USER": "u"}, + "missing user": {"WAITFORDB_DATABASE": "d"}, + "bad driver": {"WAITFORDB_DRIVER": "oracle", "WAITFORDB_USER": "u", "WAITFORDB_DATABASE": "d"}, + "bad timeout": {"WAITFORDB_USER": "u", "WAITFORDB_DATABASE": "d", "WAITFORDB_TIMEOUT": "nope"}, + "zero interval": {"WAITFORDB_USER": "u", "WAITFORDB_DATABASE": "d", "WAITFORDB_INTERVAL": "0"}, + } + for name, env := range cases { + t.Run(name, func(t *testing.T) { + if _, err := Load(envFrom(env)); err == nil { + t.Fatalf("expected error for %s", name) + } else if _, ok := err.(*Error); !ok { + t.Fatalf("expected *config.Error, got %T", err) + } + }) + } +} + +func TestRedactedHidesPassword(t *testing.T) { + c, err := Load(envFrom(map[string]string{ + "WAITFORDB_HOST": "h", + "WAITFORDB_USER": "u", + "WAITFORDB_PASSWORD": "sup3r-s3cret", + "WAITFORDB_DATABASE": "d", + "WAITFORDB_TIMEOUT": "5m", + })) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + got := c.Redacted() + if strings.Contains(got, "sup3r-s3cret") { + t.Fatalf("redacted output leaked password: %q", got) + } + for _, want := range []string{"driver=postgres", "database=d", "user=u", "password=***", "timeout=5m", "interval=2s"} { + if !strings.Contains(got, want) { + t.Errorf("redacted output missing %q: %q", want, got) + } + } +} + +func TestRedactedForeverTimeout(t *testing.T) { + c, _ := Load(envFrom(map[string]string{"WAITFORDB_USER": "u", "WAITFORDB_DATABASE": "d"})) + if !strings.Contains(c.Redacted(), "timeout=forever") { + t.Errorf("want timeout=forever, got %q", c.Redacted()) + } +} + +func TestRedactedDSNMasksPassword(t *testing.T) { + cases := []struct { + name string + dsn string + }{ + {"url", "postgres://user:topsecret@host:5432/db?sslmode=disable"}, + {"keyword", "host=h user=u password=topsecret dbname=d"}, + {"mysql", "user:topsecret@tcp(h:3306)/d?timeout=5s"}, + } + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + c := Config{DSN: tc.dsn, Timeout: 0, Interval: time.Second, ConnectTimeout: time.Second, Driver: "postgres"} + got := c.Redacted() + if strings.Contains(got, "topsecret") { + t.Fatalf("leaked password: %q", got) + } + if !strings.Contains(got, "***") { + t.Errorf("expected redaction marker in %q", got) + } + }) + } +} diff --git a/internal/driver/driver.go b/internal/driver/driver.go new file mode 100644 index 0000000..5f619c0 --- /dev/null +++ b/internal/driver/driver.go @@ -0,0 +1,49 @@ +// Package driver abstracts the per-database connection details behind a small +// interface so new engines can be added without touching the wait loop. +package driver + +import ( + "context" + "fmt" + "sort" + "strings" + + "git.unkin.net/unkin/waitfordb/internal/config" +) + +// Pinger is a live handle to a database that can answer a trivial liveness +// query. A successful Ping proves the server is up, auth succeeded, and the +// target database/role exist. +type Pinger interface { + Ping(ctx context.Context) error + Close() error +} + +// Driver knows how to open a Pinger for one database engine. +type Driver interface { + Name() string + Open(cfg config.Config) (Pinger, error) +} + +var registry = map[string]Driver{} + +func register(d Driver) { registry[d.Name()] = d } + +// Get returns the registered driver for name. +func Get(name string) (Driver, error) { + d, ok := registry[strings.ToLower(name)] + if !ok { + return nil, fmt.Errorf("unsupported driver %q (supported: %s)", name, strings.Join(Names(), ", ")) + } + return d, nil +} + +// Names lists the registered driver names, sorted. +func Names() []string { + out := make([]string, 0, len(registry)) + for n := range registry { + out = append(out, n) + } + sort.Strings(out) + return out +} diff --git a/internal/driver/driver_test.go b/internal/driver/driver_test.go new file mode 100644 index 0000000..1844335 --- /dev/null +++ b/internal/driver/driver_test.go @@ -0,0 +1,61 @@ +package driver + +import ( + "strings" + "testing" + + "git.unkin.net/unkin/waitfordb/internal/config" +) + +func TestGetRegisteredDrivers(t *testing.T) { + for _, name := range []string{"postgres", "mysql"} { + if _, err := Get(name); err != nil { + t.Errorf("Get(%q) failed: %v", name, err) + } + } + if _, err := Get("oracle"); err == nil { + t.Error("Get(oracle) should fail") + } +} + +func TestPostgresDSN(t *testing.T) { + cfg := config.Config{ + Host: "db.example", Port: "5432", User: "sonarr", + Password: "p@ss/w:rd", Database: "sonarr-main", + SSLMode: "disable", ConnectTimeout: 5e9, + } + dsn := postgresDSN(cfg) + if !strings.HasPrefix(dsn, "postgres://sonarr:") { + t.Errorf("unexpected prefix: %s", dsn) + } + // The special-character password must be percent-escaped, not raw. + if strings.Contains(dsn, "p@ss/w:rd") { + t.Errorf("password not escaped in DSN: %s", dsn) + } + if !strings.Contains(dsn, "db.example:5432") { + t.Errorf("missing host:port: %s", dsn) + } + if !strings.Contains(dsn, "sslmode=disable") { + t.Errorf("missing sslmode: %s", dsn) + } + if !strings.Contains(dsn, "connect_timeout=5") { + t.Errorf("missing connect_timeout: %s", dsn) + } + if !strings.Contains(dsn, "/sonarr-main") { + t.Errorf("missing dbname: %s", dsn) + } +} + +func TestMySQLDSN(t *testing.T) { + cfg := config.Config{ + Host: "db", Port: "3306", User: "u", Password: "pw", + Database: "app", ConnectTimeout: 5e9, + } + dsn := mysqlDSN(cfg) + if !strings.Contains(dsn, "@tcp(db:3306)/app") { + t.Errorf("unexpected mysql dsn: %s", dsn) + } + if !strings.Contains(dsn, "timeout=5s") { + t.Errorf("missing timeout: %s", dsn) + } +} diff --git a/internal/driver/mysql.go b/internal/driver/mysql.go new file mode 100644 index 0000000..76b7308 --- /dev/null +++ b/internal/driver/mysql.go @@ -0,0 +1,35 @@ +package driver + +import ( + "fmt" + + "git.unkin.net/unkin/waitfordb/internal/config" + "github.com/go-sql-driver/mysql" +) + +func init() { register(mysqlDriver{}) } + +type mysqlDriver struct{} + +func (mysqlDriver) Name() string { return "mysql" } + +func (mysqlDriver) Open(cfg config.Config) (Pinger, error) { + dsn := cfg.DSN + if dsn == "" { + dsn = mysqlDSN(cfg) + } + return openSQL("mysql", dsn) +} + +// mysqlDSN builds a driver-native DSN via mysql.Config so credentials and the +// address are escaped correctly. +func mysqlDSN(cfg config.Config) string { + c := mysql.NewConfig() + c.User = cfg.User + c.Passwd = cfg.Password + c.Net = "tcp" + c.Addr = fmt.Sprintf("%s:%s", cfg.Host, cfg.Port) + c.DBName = cfg.Database + c.Timeout = cfg.ConnectTimeout + return c.FormatDSN() +} diff --git a/internal/driver/postgres.go b/internal/driver/postgres.go new file mode 100644 index 0000000..21bb958 --- /dev/null +++ b/internal/driver/postgres.go @@ -0,0 +1,48 @@ +package driver + +import ( + "fmt" + "net/url" + "strconv" + + "git.unkin.net/unkin/waitfordb/internal/config" + _ "github.com/jackc/pgx/v5/stdlib" // registers the "pgx" database/sql driver +) + +func init() { register(postgres{}) } + +type postgres struct{} + +func (postgres) Name() string { return "postgres" } + +func (postgres) Open(cfg config.Config) (Pinger, error) { + dsn := cfg.DSN + if dsn == "" { + dsn = postgresDSN(cfg) + } + return openSQL("pgx", dsn) +} + +// postgresDSN builds a URL-style DSN. net/url escapes the userinfo and query so +// passwords with special characters are handled safely. +func postgresDSN(cfg config.Config) string { + u := url.URL{ + Scheme: "postgres", + Host: fmt.Sprintf("%s:%s", cfg.Host, cfg.Port), + Path: "/" + cfg.Database, + } + if cfg.User != "" { + u.User = url.UserPassword(cfg.User, cfg.Password) + } + q := url.Values{} + if cfg.SSLMode != "" { + q.Set("sslmode", cfg.SSLMode) + } + // connect_timeout is a per-attempt safety net in addition to the context + // deadline the wait loop applies; it is in whole seconds. + if secs := int(cfg.ConnectTimeout.Seconds()); secs > 0 { + q.Set("connect_timeout", strconv.Itoa(secs)) + } + u.RawQuery = q.Encode() + return u.String() +} diff --git a/internal/driver/sql.go b/internal/driver/sql.go new file mode 100644 index 0000000..dc1315a --- /dev/null +++ b/internal/driver/sql.go @@ -0,0 +1,31 @@ +package driver + +import ( + "context" + "database/sql" +) + +// sqlPinger runs the liveness query against a database/sql handle. It is shared +// by every engine whose Go driver plugs into database/sql. +type sqlPinger struct { + db *sql.DB +} + +func openSQL(driverName, dsn string) (*sqlPinger, error) { + db, err := sql.Open(driverName, dsn) + if err != nil { + return nil, err + } + // A readiness check only ever needs one connection; keep the pool tiny so a + // failed attempt does not leave idle half-open connections behind. + db.SetMaxOpenConns(1) + db.SetMaxIdleConns(0) + return &sqlPinger{db: db}, nil +} + +func (p *sqlPinger) Ping(ctx context.Context) error { + var one int + return p.db.QueryRowContext(ctx, "SELECT 1").Scan(&one) +} + +func (p *sqlPinger) Close() error { return p.db.Close() } diff --git a/internal/wait/wait.go b/internal/wait/wait.go new file mode 100644 index 0000000..f2c2196 --- /dev/null +++ b/internal/wait/wait.go @@ -0,0 +1,113 @@ +// Package wait implements the retry/timeout loop that polls a database until a +// liveness attempt succeeds. The clock and the attempt are injected so the loop +// is fully testable without a real database or real time. +package wait + +import ( + "context" + "time" +) + +// Clock abstracts time so tests can advance it instantly. +type Clock interface { + Now() time.Time + // Sleep blocks for d or until ctx is done, returning ctx.Err() if it was + // cancelled first. + Sleep(ctx context.Context, d time.Duration) error +} + +// AttemptFunc performs one connect+ping. The ctx carries the per-attempt +// connect timeout. +type AttemptFunc func(ctx context.Context) error + +// OnFailure is invoked after each failed attempt that will be retried. +type OnFailure func(attempt int, err error, elapsed, timeout, retryIn time.Duration) + +// Params configures the loop. +type Params struct { + Timeout time.Duration // 0 = wait forever + Interval time.Duration + ConnectTimeout time.Duration +} + +// Result reports how a run ended. +type Result struct { + OK bool + TimedOut bool + Cancelled bool + Attempts int + Elapsed time.Duration + LastErr error +} + +// Run polls attempt until it succeeds, the timeout is exhausted, or ctx is +// cancelled. It always makes at least one attempt. +func Run(ctx context.Context, p Params, attempt AttemptFunc, onFail OnFailure, clk Clock) Result { + start := clk.Now() + var deadline time.Time + if p.Timeout > 0 { + deadline = start.Add(p.Timeout) + } + + res := Result{} + for { + res.Attempts++ + + actx, cancel := context.WithTimeout(ctx, p.ConnectTimeout) + err := attempt(actx) + cancel() + + now := clk.Now() + res.Elapsed = now.Sub(start) + + if err == nil { + res.OK = true + return res + } + res.LastErr = err + + // A cancelled parent context (SIGTERM/SIGINT) wins over a retry. + if ctx.Err() != nil { + res.Cancelled = true + return res + } + + // No time budget left for another attempt. + if p.Timeout > 0 && !now.Before(deadline) { + res.TimedOut = true + return res + } + + sleep := p.Interval + if p.Timeout > 0 { + if remaining := deadline.Sub(now); remaining < sleep { + sleep = remaining + } + } + onFail(res.Attempts, err, res.Elapsed, p.Timeout, sleep) + + if serr := clk.Sleep(ctx, sleep); serr != nil { + res.Cancelled = true + return res + } + } +} + +// RealClock is the production Clock backed by the wall clock. +type RealClock struct{} + +func (RealClock) Now() time.Time { return time.Now() } + +func (RealClock) Sleep(ctx context.Context, d time.Duration) error { + if d <= 0 { + return ctx.Err() + } + t := time.NewTimer(d) + defer t.Stop() + select { + case <-ctx.Done(): + return ctx.Err() + case <-t.C: + return nil + } +} diff --git a/internal/wait/wait_test.go b/internal/wait/wait_test.go new file mode 100644 index 0000000..d6cb7ec --- /dev/null +++ b/internal/wait/wait_test.go @@ -0,0 +1,125 @@ +package wait + +import ( + "context" + "errors" + "testing" + "time" +) + +// fakeClock advances instantly on Sleep so the loop runs with no real delay. +type fakeClock struct { + t time.Time + cancelAt time.Duration // if >0, cancel the run once elapsed reaches this + cancel context.CancelFunc + start time.Time + sleepCall int +} + +func newFakeClock() *fakeClock { + start := time.Unix(0, 0) + return &fakeClock{t: start, start: start} +} + +func (c *fakeClock) Now() time.Time { return c.t } + +func (c *fakeClock) Sleep(ctx context.Context, d time.Duration) error { + c.sleepCall++ + c.t = c.t.Add(d) + if c.cancelAt > 0 && c.t.Sub(c.start) >= c.cancelAt && c.cancel != nil { + c.cancel() + } + return ctx.Err() +} + +var errDown = errors.New("connection refused") + +// failNThenOK returns an AttemptFunc that fails the first n calls then succeeds. +func failNThenOK(n int, calls *int) AttemptFunc { + return func(ctx context.Context) error { + *calls++ + if *calls <= n { + return errDown + } + return nil + } +} + +func noFail(int, error, time.Duration, time.Duration, time.Duration) {} + +func TestSucceedsFirstAttempt(t *testing.T) { + calls := 0 + res := Run(context.Background(), + Params{Timeout: time.Minute, Interval: 2 * time.Second, ConnectTimeout: time.Second}, + failNThenOK(0, &calls), noFail, newFakeClock()) + if !res.OK || res.Attempts != 1 { + t.Fatalf("want OK after 1 attempt, got %+v", res) + } +} + +func TestWaitsThenSucceeds(t *testing.T) { + calls := 0 + clk := newFakeClock() + failures := 0 + res := Run(context.Background(), + Params{Timeout: time.Minute, Interval: 2 * time.Second, ConnectTimeout: time.Second}, + failNThenOK(3, &calls), + func(int, error, time.Duration, time.Duration, time.Duration) { failures++ }, + clk) + if !res.OK { + t.Fatalf("want OK, got %+v", res) + } + if res.Attempts != 4 { + t.Errorf("attempts = %d, want 4", res.Attempts) + } + if failures != 3 { + t.Errorf("onFail called %d times, want 3", failures) + } + // 3 sleeps of 2s each. + if got := res.Elapsed; got != 6*time.Second { + t.Errorf("elapsed = %v, want 6s", got) + } +} + +func TestTimesOut(t *testing.T) { + calls := 0 + alwaysFail := func(ctx context.Context) error { calls++; return errDown } + res := Run(context.Background(), + Params{Timeout: 10 * time.Second, Interval: 3 * time.Second, ConnectTimeout: time.Second}, + alwaysFail, noFail, newFakeClock()) + if res.OK || !res.TimedOut { + t.Fatalf("want timeout, got %+v", res) + } + if !errors.Is(res.LastErr, errDown) { + t.Errorf("LastErr = %v, want errDown", res.LastErr) + } + // Deadline 10s, interval 3s: attempts at 0,3,6,9, then next check at ~12s > deadline. + if res.Attempts < 3 { + t.Errorf("attempts = %d, want several before timeout", res.Attempts) + } +} + +func TestWaitForeverEventuallySucceeds(t *testing.T) { + calls := 0 + res := Run(context.Background(), + Params{Timeout: 0, Interval: time.Second, ConnectTimeout: time.Second}, + failNThenOK(100, &calls), noFail, newFakeClock()) + if !res.OK || res.Attempts != 101 { + t.Fatalf("want OK after 101 attempts with no timeout, got %+v", res) + } +} + +func TestCancelledDuringSleep(t *testing.T) { + ctx, cancel := context.WithCancel(context.Background()) + clk := newFakeClock() + clk.cancelAt = 4 * time.Second + clk.cancel = cancel + calls := 0 + alwaysFail := func(ctx context.Context) error { calls++; return errDown } + res := Run(ctx, + Params{Timeout: time.Hour, Interval: 2 * time.Second, ConnectTimeout: time.Second}, + alwaysFail, noFail, clk) + if !res.Cancelled { + t.Fatalf("want Cancelled, got %+v", res) + } +} diff --git a/main.go b/main.go new file mode 100644 index 0000000..9118722 --- /dev/null +++ b/main.go @@ -0,0 +1,180 @@ +// 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...)) +}