waitfordb: initial tool — env-configured wait-for-DB init container
Small Go tool + distroless container used as a K8s initContainer to block an app until its database (Postgres/MySQL) is reachable. Env-var configured (WAITFORDB_* + libpq PG* fallback), configurable timeout/interval, redacted logs, exit codes. Woodpecker CI publishes docker-internal/waitfordb on tag.
This commit is contained in:
@@ -0,0 +1,4 @@
|
||||
# built binary (repo root)
|
||||
/waitfordb
|
||||
# build output
|
||||
dist/
|
||||
@@ -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
|
||||
@@ -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
|
||||
+18
@@ -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"]
|
||||
@@ -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)
|
||||
@@ -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.
|
||||
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:<tag>` (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`).
|
||||
|
||||
@@ -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
|
||||
)
|
||||
@@ -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=
|
||||
@@ -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
|
||||
}
|
||||
@@ -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
|
||||
}
|
||||
@@ -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)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
@@ -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
|
||||
}
|
||||
@@ -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)
|
||||
}
|
||||
}
|
||||
@@ -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()
|
||||
}
|
||||
@@ -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()
|
||||
}
|
||||
@@ -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() }
|
||||
@@ -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
|
||||
}
|
||||
}
|
||||
@@ -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)
|
||||
}
|
||||
}
|
||||
@@ -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...))
|
||||
}
|
||||
Reference in New Issue
Block a user