llmgw/internal/store/postgres.go

package store

import (
	"context"
	"fmt"
	"net"
	"time"

	// The pure-Go PostgreSQL driver. Pure-Go is not a preference here: the
	// whole binary is built with CGO_ENABLED=0 so that it runs on a node whose
	// libc is older than the builder's, and a cgo driver would take that away
	// from every deployment, including the ones still on SQLite.
	"github.com/jackc/pgx/v5"
	"github.com/jackc/pgx/v5/stdlib"
)

// How soon a store machine that is gone is given up on. A replicated store's
// address names every member, and a new connection tries them in order: with
// no timeout, one to a machine that had lost its power waited for the
// kernel's own — half a minute and more — before the next member was asked,
// and every request whose key was not cached waited with it.
const (
	// storeConnectTimeout is per member: the next is tried after it.
	storeConnectTimeout = 3 * time.Second
	// storeDeadAfter is how long a connection's sent data may go unanswered,
	// or an idle one's probes, before it is closed: a connection to the
	// primary that was lost fails, and the next opens to the new one.
	storeDeadAfter = 10 * time.Second
)

// failFast gives a connection configuration those limits, where its address
// gave none of its own.
func failFast(cfg *pgx.ConnConfig) {
	if cfg.ConnectTimeout == 0 {
		cfg.ConnectTimeout = storeConnectTimeout
	}
	d := &net.Dialer{
		KeepAliveConfig: net.KeepAliveConfig{Enable: true, Idle: 5 * time.Second, Interval: time.Second,
			Count: int((storeDeadAfter - 5*time.Second) / time.Second)},
		Control: userTimeout(storeDeadAfter),
	}
	cfg.DialFunc = d.DialContext
}

// OpenPostgres opens the store against PostgreSQL.
//
// The same queries run here as on SQLite; only the placeholders differ, and
// the dialect handles that. What differs operationally is concurrency: SQLite
// is held to one connection because more would only produce lock contention
// that reads as a deadlock, while PostgreSQL is the reason someone moves to
// it and wants a pool.
//
// The DSN carries a password, so it is taken as a string the caller read from
// somewhere protected rather than as a path this package knows how to find —
// and it is never logged, which is why errors below name the database and not
// the DSN.
func OpenPostgres(ctx context.Context, dsn string) (*DB, error) {
	if dsn == "" {
		return nil, fmt.Errorf("store: no PostgreSQL address configured")
	}
	cfg, err := pgx.ParseConfig(dsn)
	if err != nil {
		// Deliberately not wrapped with the DSN: it holds a password, and an
		// error string ends up in logs, tickets and screenshots.
		return nil, fmt.Errorf("opening the PostgreSQL store: %w", err)
	}
	failFast(cfg)
	db := stdlib.OpenDB(*cfg)

	// Sized for a control plane rather than for a fleet of application
	// servers. The request path barely touches the store — usage is batched,
	// budgets are counted in memory, authentication is cached — so the pool
	// exists for conversation turns and reporting queries, not for inference.
	db.SetMaxOpenConns(16)
	db.SetMaxIdleConns(4)
	// Long enough that a quiet cluster keeps its connections, short enough
	// that a restarted database is not answered with dead ones for an hour.
	db.SetConnMaxLifetime(30 * time.Minute)
	db.SetConnMaxIdleTime(5 * time.Minute)

	if err := db.PingContext(ctx); err != nil {
		_ = db.Close()
		return nil, fmt.Errorf("reaching the PostgreSQL store: %w", err)
	}
	if err := migrate(ctx, db, string(PostgresDialect)); err != nil {
		_ = db.Close()
		return nil, err
	}
	s := &DB{
		raw:     db,
		dialect: PostgresDialect,
		db:      bound{c: PostgresDialect, to: db},
	}
	// Not fatal: a store without pgvector serves everything but knowledge
	// bases, and says why when somebody asks for one.
	s.kbErr = s.ensureKnowledge(ctx)
	return s, nil
}