Refactoring of check category connections (#905)

#### Description

This PR refactors the replication category. Now all checks share the
same PostgreSQL connection. The new `postgres.LazyConn` is not
thread-safe. It is used for executing checks sequentially.

Once the check finishes, the cleanup function must be called.

#### Type of Change

Please select the relevant option(s):

- [ ] 🐛 Bug fix (non-breaking change that fixes an issue)
- [ ] ✨ New feature (non-breaking change that adds functionality)
- [ ] 💥 Breaking change (fix or feature that would cause existing
functionality to not work as expected)
- [ ] 📚 Documentation update
- [x] 🔧 Refactoring (no functional changes)
- [ ] ⚡ Performance improvement
- [ ] 🧪 Test coverage improvement
- [ ] 🔨 Build/CI changes
- [ ] 🧹 Code cleanup

#### Testing

- [ ] Unit tests added/updated
- [ ] Integration tests added/updated
- [ ] Manual testing performed
- [x] All existing tests pass

#### Checklist

- [x] Code follows project style guidelines
- [x] Self-review completed
- [x] Code is well-commented
- [ ] Documentation updated where necessary

Author: kvch

cmd/check_cmd.go

@@ -36,13 +36,19 @@ var checkCmd = &cobra.Command{
 	RunE: func(cmd *cobra.Command, args []string) error {
 		sp, _ := pterm.DefaultSpinner.WithText("running pgstream checks...").Start()
 
-		err := func() error {
+		err := func() (retErr error) {
 			streamConfig, err := config.ParseStreamConfig()
 			if err != nil {
 				return fmt.Errorf("parsing stream config: %w", err)
 			}
 
-			checks := preflight.BuildChecks(streamConfig, selectedCategories(cmd))
+			checks, cleanup := preflight.BuildChecks(streamConfig, selectedCategories(cmd))
+			defer func() {
+				if cerr := cleanup(context.Background()); cerr != nil && retErr == nil {
+					retErr = fmt.Errorf("releasing check resources: %w", cerr)
+				}
+			}()
+
 			if len(checks) == 0 {
 				sp.Success("no checks to run")
 				return nil

internal/postgres/pg_lazy_conn.go

@@ -0,0 +1,47 @@
+// SPDX-License-Identifier: Apache-2.0
+
+package postgres
+
+import "context"
+
+// AcquireFunc lazily yields a Postgres connection. Useful when a set of
+// related callers want to share a single TCP connection without each one
+// having to think about lifecycle.
+type AcquireFunc func(ctx context.Context) (Querier, error)
+
+// LazyConn memoises a single *Conn (or its dial error) for a URL. Cheap to
+// construct: nothing is opened until Acquire is called for the first time.
+// Not safe for concurrent use — designed for the sequential case (e.g. a
+// preflight check engine that runs checks one after another).
+type LazyConn struct {
+	url  string
+	conn *Conn
+	err  error
+}
+
+// NewLazyConn returns a LazyConn that will open a connection to url on first
+// Acquire.
+func NewLazyConn(url string) *LazyConn {
+	return &LazyConn{url: url}
+}
+
+// Acquire returns the cached conn, opening it on the first call. A dial
+// failure is cached too — subsequent calls return the same error without
+// retrying.
+func (l *LazyConn) Acquire(ctx context.Context) (Querier, error) {
+	if l.conn != nil || l.err != nil {
+		return l.conn, l.err
+	}
+	l.conn, l.err = NewConn(ctx, l.url)
+	return l.conn, l.err
+}
+
+// Close releases the underlying connection if one was opened.
+func (l *LazyConn) Close(ctx context.Context) error {
+	if l.conn == nil {
+		return nil
+	}
+	c := l.conn
+	l.conn = nil
+	return c.Close(ctx)
+}

pkg/stream/preflight/CLAUDE.md

@@ -6,9 +6,11 @@ Guidance for Claude Code when working inside `pkg/stream/preflight`. The planned
 
 - `preflight.go` — `Check` interface (`Name()` + `Run(ctx) ([]Finding, error)`), `Finding`, `CheckResult`, `Report`, `Run(ctx, []Check, ...RunOption)` engine.
 - `printer.go` — `ReportPrinter{Report}` is the only thing that formats reports. The `Report` struct itself stays pure data.
-- `builder.go` — `Builder` struct, `Builders` registry slice, per-category builder functions (`BuildConnectivityChecks`, …), `BuildChecks(cfg, selected)`.
+- `builder.go` — `Builder` struct (returns `[]Check` + optional cleanup), `Builders` registry slice, per-category builder functions (`BuildConnectivityChecks`, …), `BuildChecks(cfg, selected)`.
 - One file per category of concrete checks (`connectivity.go`, `replication.go`, …).
 
+The shared-conn primitive lives one floor down at `internal/postgres.LazyConn` so other callers can reuse it.
+
 ## Adding a new check
 
 Adding a check is meant to be a small, mechanical edit. Keep it that way.
@@ -21,6 +23,7 @@ Adding a check is meant to be a small, mechanical edit. Keep it that way.
    - **Return `error` only when the check itself couldn't run** (timeout, internal bug, malformed input). A detected problem is a `Finding`, not an error.
    - **Put remediation in `Finding.Message`** — the user should be able to act on it without reading source.
 3. **Materialise instances in the category builder** (e.g. `BuildConnectivityChecks`). The builder is the applicability gate: it reads `*stream.Config` and decides which instances are relevant. Inapplicable checks are silently omitted today; an explicit "skipped: <reason>" mechanism is deferred (see `docs/migration_preflight_issue.md` "Architecture decisions" #6).
+   - **If checks in the category share a Postgres connection**, call `postgres.NewLazyConn(url)` in the builder, hand `src.Acquire` (a `postgres.AcquireFunc`) to every check, and return `src.Close` as the cleanup. See `BuildReplicationChecks` for the pattern. The engine runs sequentially, so the first check to call `Source(ctx)` opens the conn and the rest reuse it. A failed dial is memoised too — only one connection attempt happens, even if every check reports its own check error.
 4. **Tests.** Unit-test the check directly against mocked dependencies (`internal/postgres/mocks` has the postgres conn mock). For new categories, exercise the builder selection path through the cmd layer too.
 
 ## Do not

pkg/stream/preflight/builder.go

@@ -2,15 +2,25 @@
 
 package preflight
 
-import "github.com/xataio/pgstream/pkg/stream"
+import (
+	"context"
 
-// Builder turns a stream.Config into the concrete checks for a category. Each
-// new category adds an entry to Builders and a matching CLI flag in
-// cmd/root_cmd.go.
+	"github.com/xataio/pgstream/internal/postgres"
+	"github.com/xataio/pgstream/pkg/stream"
+)
+
+// CleanupFunc releases any resources a builder set up (e.g. a shared Postgres
+// connection). Builders return nil when there's nothing to clean up.
+type CleanupFunc func(context.Context) error
+
+// Builder turns a stream.Config into the concrete checks for a category, plus
+// an optional cleanup function that releases resources the checks share (e.g.
+// a Postgres connection). Each new category adds an entry to Builders and a
+// matching CLI flag in cmd/root_cmd.go.
 type Builder struct {
 	Category Category
 	Flag     string
-	Build    func(*stream.Config) []Check
+	Build    func(*stream.Config) ([]Check, CleanupFunc)
 }
 
 // Builders is the registry of category builders. Adding a new category = one
@@ -22,8 +32,9 @@ var Builders = []Builder{
 
 // BuildConnectivityChecks returns the connectivity checks applicable to cfg.
 // A source check is added when a source postgres URL is configured; a target
-// check is added when a postgres target is configured.
-func BuildConnectivityChecks(cfg *stream.Config) []Check {
+// check is added when a postgres target is configured. Each check opens its
+// own conn (to its own URL), so no shared cleanup is needed.
+func BuildConnectivityChecks(cfg *stream.Config) ([]Check, CleanupFunc) {
 	checks := []Check{}
 	if url := cfg.SourcePostgresURL(); url != "" {
 		checks = append(checks, &ConnectivityCheck{Label: "source", URL: url})
@@ -33,42 +44,58 @@ func BuildConnectivityChecks(cfg *stream.Config) []Check {
 			checks = append(checks, &ConnectivityCheck{Label: "target", URL: url})
 		}
 	}
-	return checks
+	return checks, nil
 }
 
 // BuildReplicationChecks returns the replication-preflight checks applicable
-// to cfg. Replication checks only apply when the source is configured with a
-// replication slot (i.e. the run is doing logical replication, not a one-shot
-// snapshot).
-func BuildReplicationChecks(cfg *stream.Config) []Check {
+// to cfg, plus a cleanup function that closes the shared source connection.
+// Replication checks only apply when the source is configured with a
+// replication slot.
+func BuildReplicationChecks(cfg *stream.Config) ([]Check, CleanupFunc) {
 	if cfg.PostgresReplicationSlot() == "" {
-		return nil
+		return nil, nil
 	}
 	url := cfg.SourcePostgresURL()
 	if url == "" {
-		return nil
+		return nil, nil
 	}
+	src := postgres.NewLazyConn(url)
 	return []Check{
-		&WALLevelCheck{URL: url},
-		&WAL2JSONCheck{URL: url},
-		&ReplicationSlotHeadroomCheck{URL: url},
-		&ReplicationRoleAttrCheck{URL: url},
-	}
+		&WALLevelCheck{Source: src.Acquire},
+		&WAL2JSONCheck{Source: src.Acquire},
+		&ReplicationSlotHeadroomCheck{Source: src.Acquire},
+		&ReplicationRoleAttrCheck{Source: src.Acquire},
+	}, src.Close
 }
 
 // BuildChecks returns the concrete checks for the selected categories,
-// preserving the registration order in Builders. An empty selection runs every
-// registered category.
-func BuildChecks(cfg *stream.Config, selected []Category) []Check {
+// preserving the registration order in Builders, plus a single cleanup
+// function that releases every category's resources. The returned cleanup is
+// always non-nil; callers can defer it unconditionally. An empty selection
+// runs every registered category.
+func BuildChecks(cfg *stream.Config, selected []Category) ([]Check, CleanupFunc) {
 	want := make(map[Category]bool, len(selected))
 	for _, c := range selected {
 		want[c] = true
 	}
 	checks := []Check{}
+	cleanups := []CleanupFunc{}
 	for _, b := range Builders {
 		if len(want) == 0 || want[b.Category] {
-			checks = append(checks, b.Build(cfg)...)
+			cs, cleanup := b.Build(cfg)
+			checks = append(checks, cs...)
+			if cleanup != nil {
+				cleanups = append(cleanups, cleanup)
+			}
+		}
+	}
+	return checks, func(ctx context.Context) error {
+		var firstErr error
+		for _, c := range cleanups {
+			if err := c(ctx); err != nil && firstErr == nil {
+				firstErr = err
+			}
 		}
+		return firstErr
 	}
-	return checks
 }

pkg/stream/preflight/replication.go

@@ -12,17 +12,16 @@ import (
 // WALLevelCheck verifies the source Postgres has `wal_level=logical`, which
 // pgstream's replication path requires.
 type WALLevelCheck struct {
-	URL string
+	Source postgres.AcquireFunc
 }
 
 func (c *WALLevelCheck) Name() string { return "wal_level" }
 
 func (c *WALLevelCheck) Run(ctx context.Context) ([]Finding, error) {
-	conn, err := postgres.NewConn(ctx, c.URL)
+	conn, err := c.Source(ctx)
 	if err != nil {
 		return nil, fmt.Errorf("connecting to source: %w", err)
 	}
-	defer conn.Close(ctx)
 
 	var level string
 	if err := conn.QueryRow(ctx, []any{&level}, "SHOW wal_level"); err != nil {
@@ -39,17 +38,16 @@ func (c *WALLevelCheck) Run(ctx context.Context) ([]Finding, error) {
 // WAL2JSONCheck verifies that the wal2json output plugin is installed and
 // loadable on the source. pgstream decodes WAL through wal2json.
 type WAL2JSONCheck struct {
-	URL string
+	Source postgres.AcquireFunc
 }
 
 func (c *WAL2JSONCheck) Name() string { return "wal2json" }
 
 func (c *WAL2JSONCheck) Run(ctx context.Context) ([]Finding, error) {
-	conn, err := postgres.NewConn(ctx, c.URL)
+	conn, err := c.Source(ctx)
 	if err != nil {
 		return nil, fmt.Errorf("connecting to source: %w", err)
 	}
-	defer conn.Close(ctx)
 
 	var present int
 	if err := conn.QueryRow(ctx, []any{&present}, "SELECT count(*)::int FROM pg_available_extensions WHERE name = 'wal2json'"); err != nil {
@@ -66,17 +64,16 @@ func (c *WAL2JSONCheck) Run(ctx context.Context) ([]Finding, error) {
 // ReplicationSlotHeadroomCheck reports whether the source has at least one
 // slot still available before max_replication_slots is reached.
 type ReplicationSlotHeadroomCheck struct {
-	URL string
+	Source postgres.AcquireFunc
 }
 
 func (c *ReplicationSlotHeadroomCheck) Name() string { return "replication_slot_headroom" }
 
 func (c *ReplicationSlotHeadroomCheck) Run(ctx context.Context) ([]Finding, error) {
-	conn, err := postgres.NewConn(ctx, c.URL)
+	conn, err := c.Source(ctx)
 	if err != nil {
 		return nil, fmt.Errorf("connecting to source: %w", err)
 	}
-	defer conn.Close(ctx)
 
 	var maxSlots, usedSlots int
 	err = conn.QueryRow(ctx, []any{&maxSlots, &usedSlots}, `
@@ -98,17 +95,16 @@ func (c *ReplicationSlotHeadroomCheck) Run(ctx context.Context) ([]Finding, erro
 // ReplicationRoleAttrCheck verifies the current source role has the
 // REPLICATION attribute, which is required to open a logical replication slot.
 type ReplicationRoleAttrCheck struct {
-	URL string
+	Source postgres.AcquireFunc
 }
 
 func (c *ReplicationRoleAttrCheck) Name() string { return "replication_role_attr" }
 
 func (c *ReplicationRoleAttrCheck) Run(ctx context.Context) ([]Finding, error) {
-	conn, err := postgres.NewConn(ctx, c.URL)
+	conn, err := c.Source(ctx)
 	if err != nil {
 		return nil, fmt.Errorf("connecting to source: %w", err)
 	}
-	defer conn.Close(ctx)
 
 	var roleName string
 	var hasReplication bool