llmgw/internal/cli/serverinit_postgres.go

package cli

import (
	"context"
	"crypto/rand"
	"encoding/hex"
	"fmt"
	"log/slog"
	"os"
	"os/exec"
	"path/filepath"
	"strconv"
	"strings"
	"time"

	"gitea.mynodes.xyz/ilker/llmgw/internal/agent/discovery"
	"gitea.mynodes.xyz/ilker/llmgw/internal/agent/provision"
	"gitea.mynodes.xyz/ilker/llmgw/internal/config"
	"gitea.mynodes.xyz/ilker/llmgw/internal/store"
)

// pgRole and pgDatabase are what init creates in a PostgreSQL it manages.
const (
	pgRole     = "llmgw"
	pgDatabase = "llmgw"
)

// dsnPath is where the PostgreSQL address is kept, out of server.yaml, so the
// password in it is not in a file that gets copied, diffed and pasted.
func (in *serverInit) dsnPath() string { return filepath.Join(in.etcDir, "secrets", "store-dsn") }

// storeAddress is what auth.store says for the store init manages.
func (in *serverInit) storeAddress() string {
	if in.postgres {
		return "@" + in.dsnPath()
	}
	return filepath.Join(in.stateDir, "llmgw.db")
}

// ensureDatabase puts the store in PostgreSQL when asked to, moving whatever
// the SQLite store already holds.
//
// The move is the point. Keys are stored hashed, so starting PostgreSQL empty
// means reissuing every key and being unable to say which was which; and the
// administrator init created a minute ago would be gone too.
func (in *serverInit) ensureDatabase(ctx context.Context) (err error) {
	if len(in.patroni) > 0 || len(in.patroniNodes) > 0 {
		if in.postgres {
			return fmt.Errorf("--postgres keeps the store on this machine and --patroni on others: give one")
		}
		return in.ensurePatroni(ctx)
	}
	if !in.postgres {
		return nil
	}
	setup := in.pgSetup
	if setup == nil {
		setup = in.setupPostgres
	}
	target, err := setup(ctx)
	if err != nil {
		return err
	}

	cfg, err := config.Load(in.configPath())
	if err != nil {
		return err
	}
	current := cfg.Auth.Store
	if current == target {
		in.kept("store in PostgreSQL (%s)", in.dsnPath())
		return nil
	}
	if isPostgresAddress(current) {
		// Two PostgreSQL stores, and a guess about which is meant would lose
		// one of them.
		return fmt.Errorf("server.yaml already names a PostgreSQL store (%s); "+
			"init moves SQLite into PostgreSQL, not one PostgreSQL into another", current)
	}

	// Nothing may write the store while it is copied: a row landing after its
	// table has been passed would be absent, and nothing would say so.
	stopped, err := in.stopService(ctx)
	if err != nil {
		return err
	}
	defer func() {
		// A move that failed leaves the configuration naming SQLite, which
		// still holds everything: start the service on it again rather than
		// leave the gateway down over a database that did not work out.
		if err != nil && stopped {
			_ = exec.Command("systemctl", "start", "llmgw-server").Run()
			fmt.Fprintf(in.out, "  ..    llmgw-server started again on its SQLite store\n")
		}
	}()

	if current != "" {
		if _, err := os.Stat(current); err == nil {
			if err := in.copyStore(ctx, current, target); err != nil {
				return err
			}
		}
	}

	if err := setAuthStore(in.configPath(), target); err != nil {
		return err
	}
	in.done("server.yaml now keeps the store in PostgreSQL")
	return nil
}

func (in *serverInit) copyStore(ctx context.Context, from, to string) error {
	src, err := store.Open(ctx, from)
	if err != nil {
		return fmt.Errorf("opening %s: %w", from, err)
	}
	defer src.Close()
	dst, err := store.Open(ctx, to)
	if err != nil {
		return fmt.Errorf("opening PostgreSQL: %w", err)
	}
	defer dst.Close()

	empty, err := dst.Empty(ctx)
	if err != nil {
		return err
	}
	if !empty {
		// A previous run that copied and then stopped before changing the
		// configuration. The service was stopped for that copy too, so what
		// is there is what SQLite held.
		in.kept("PostgreSQL already holds a store; %s is left as it is", from)
		return nil
	}
	res, err := src.CopyTo(ctx, dst)
	if err != nil {
		return fmt.Errorf("copying %s into PostgreSQL: %w", from, err)
	}
	var rows int64
	for _, r := range res {
		rows += r.Rows
	}
	in.done("copied %d rows from %s into PostgreSQL (the SQLite file is kept)", rows, from)
	return nil
}

func isPostgresAddress(a string) bool {
	return strings.HasPrefix(a, "@") || strings.HasPrefix(a, "postgres://") ||
		strings.HasPrefix(a, "postgresql://")
}

// setAuthStore changes auth.store in server.yaml and nothing else.
//
// Edited as text rather than unmarshalled and written back, because writing
// it back would drop every comment in a file people annotate. The old file is
// kept beside it.
func setAuthStore(path, address string) error {
	b, err := os.ReadFile(path)
	if err != nil {
		return err
	}
	lines := strings.Split(string(b), "\n")
	value := fmt.Sprintf("%q", address)

	authAt, storeAt, indent := -1, -1, "  "
	for i, l := range lines {
		trimmed := strings.TrimSpace(l)
		top := l != "" && l[0] != ' ' && l[0] != '\t' && !strings.HasPrefix(trimmed, "#")
		if top {
			if authAt >= 0 && i > authAt {
				break // past the auth block
			}
			if strings.HasPrefix(trimmed, "auth:") {
				authAt = i
			}
			continue
		}
		if authAt >= 0 && strings.HasPrefix(trimmed, "store:") {
			storeAt = i
			indent = l[:len(l)-len(strings.TrimLeft(l, " \t"))]
			break
		}
	}
	switch {
	case storeAt >= 0:
		lines[storeAt] = indent + "store: " + value
	case authAt >= 0:
		lines = append(lines[:authAt+1], append([]string{indent + "store: " + value}, lines[authAt+1:]...)...)
	default:
		lines = append(lines, "auth:", "  enabled: true", "  store: "+value)
	}

	backup := fmt.Sprintf("%s.bak-%s", path, time.Now().Format("20060102-150405"))
	if err := os.WriteFile(backup, b, 0o640); err != nil {
		return err
	}
	info, err := os.Stat(path)
	if err != nil {
		return err
	}
	return os.WriteFile(path, []byte(strings.Join(lines, "\n")), info.Mode().Perm())
}

// pgServer is the PostgreSQL installation on this machine, as systemd and the
// file system describe it.
//
// Found rather than assumed. A distribution's own package calls its unit
// postgresql.service and puts initdb and psql on PATH; the PostgreSQL
// project's packages (PGDG) call it postgresql-16.service and keep them in
// /usr/pgsql-16/bin — and they provide "postgresql-server" too, so asking dnf
// for that name can install or find one of them. Assuming the first found a
// Rocky machine with the second unable to initialise the database it had.
type pgServer struct {
	unit   string // postgresql.service, postgresql-16.service
	pgdata string // "" when the unit does not say, which is Debian's case
	bin    string // directory holding initdb and psql; "" means PATH
}

func (p pgServer) tool(name string) string {
	if p.bin != "" {
		return filepath.Join(p.bin, name)
	}
	return name
}

// pickPostgresUnit chooses the server unit from systemctl's list of
// postgresql* unit files: the plain one when there is one, else the highest
// version. Templates (postgresql@.service) are not a server by themselves.
func pickPostgresUnit(list string) string {
	best, bestVersion := "", -1
	for _, line := range strings.Split(list, "\n") {
		f := strings.Fields(line)
		if len(f) == 0 {
			continue
		}
		unit := f[0]
		if !strings.HasSuffix(unit, ".service") || strings.Contains(unit, "@") {
			continue
		}
		name := strings.TrimSuffix(unit, ".service")
		if name == "postgresql" {
			return unit
		}
		v := strings.TrimPrefix(name, "postgresql-")
		if n, err := strconv.Atoi(strings.SplitN(v, ".", 2)[0]); err == nil && n > bestVersion {
			best, bestVersion = unit, n
		}
	}
	return best
}

// pgdataFrom reads PGDATA out of `systemctl show -p Environment --value`.
func pgdataFrom(env string) string {
	for _, kv := range strings.Fields(env) {
		if v, ok := strings.CutPrefix(kv, "PGDATA="); ok {
			return v
		}
	}
	return ""
}

func findPostgres(ctx context.Context) (pgServer, bool) {
	out, _ := exec.CommandContext(ctx, "systemctl", "list-unit-files", "--type=service",
		"--no-legend", "--plain", "postgresql*").Output()
	unit := pickPostgresUnit(string(out))
	if unit == "" {
		return pgServer{}, false
	}
	p := pgServer{unit: unit}
	env, _ := exec.CommandContext(ctx, "systemctl", "show", "-p", "Environment", "--value", unit).Output()
	p.pgdata = pgdataFrom(string(env))

	if v, ok := strings.CutPrefix(strings.TrimSuffix(unit, ".service"), "postgresql-"); ok {
		if dir := "/usr/pgsql-" + v + "/bin"; isDir(dir) {
			p.bin = dir
		}
		if p.pgdata == "" && isDir("/var/lib/pgsql/"+v) {
			p.pgdata = "/var/lib/pgsql/" + v + "/data"
		}
	}
	if p.bin == "" {
		if _, err := exec.LookPath("psql"); err != nil {
			// Debian keeps the server's tools per version and puts only
			// wrappers on PATH; without even those, look where it keeps them.
			if dirs, _ := filepath.Glob("/usr/lib/postgresql/*/bin"); len(dirs) > 0 {
				p.bin = dirs[len(dirs)-1]
			}
		}
	}
	return p, true
}

func isDir(p string) bool {
	st, err := os.Stat(p)
	return err == nil && st.IsDir()
}

// setupPostgres makes sure a PostgreSQL on this machine holds a database for
// llmgw, installing the server if there is none, and returns the store
// address.
func (in *serverInit) setupPostgres(ctx context.Context) (string, error) {
	installed, err := in.installPostgres(ctx)
	if err != nil {
		return "", err
	}
	if err := in.waitForPostgres(ctx); err != nil {
		return "", err
	}

	exists, err := in.psqlBool(ctx, "SELECT 1 FROM pg_roles WHERE rolname = '"+pgRole+"'")
	if err != nil {
		return "", err
	}
	_, statErr := os.Stat(in.dsnPath())
	if exists && statErr == nil {
		in.kept("PostgreSQL role %s", pgRole)
	} else {
		// A role with no address written down is one whose password nobody
		// knows; giving it a new one is the only way to use it.
		password := make([]byte, 24)
		_, _ = rand.Read(password)
		pass := hex.EncodeToString(password)
		verb := "CREATE ROLE"
		if exists {
			verb = "ALTER ROLE"
		}
		if _, err := in.psql(ctx, verb+" "+pgRole+" WITH LOGIN PASSWORD '"+pass+"'"); err != nil {
			return "", err
		}
		dsn := fmt.Sprintf("postgres://%s:%s@127.0.0.1:5432/%s?sslmode=disable", pgRole, pass, pgDatabase)
		if err := in.writeDSN(dsn); err != nil {
			return "", err
		}
		in.done("PostgreSQL role %s, its address in %s", pgRole, in.dsnPath())
	}

	exists, err = in.psqlBool(ctx, "SELECT 1 FROM pg_database WHERE datname = '"+pgDatabase+"'")
	if err != nil {
		return "", err
	}
	if exists {
		in.kept("PostgreSQL database %s", pgDatabase)
	} else {
		if _, err := in.psql(ctx, "CREATE DATABASE "+pgDatabase+" OWNER "+pgRole); err != nil {
			return "", err
		}
		in.done("PostgreSQL database %s", pgDatabase)
	}

	if installed {
		// Its defaults assume it owns the machine. A gateway swapped out to
		// give the database cache is worse than a smaller cache, and a store
		// whose working set is megabytes needs little. Only a PostgreSQL init
		// installed itself: one that was already here may serve something
		// else that needs its connections.
		for _, q := range []string{
			"ALTER SYSTEM SET shared_buffers = '96MB'",
			"ALTER SYSTEM SET max_connections = 40",
		} {
			if _, err := in.psql(ctx, q); err != nil {
				return "", err
			}
		}
		if out, err := exec.CommandContext(ctx, "systemctl", "restart", in.pg.unit).CombinedOutput(); err != nil {
			return "", fmt.Errorf("restarting PostgreSQL: %v: %s", err, strings.TrimSpace(string(out)))
		}
		if err := in.waitForPostgres(ctx); err != nil {
			return "", err
		}
		in.done("PostgreSQL sized for a machine it shares (shared_buffers 96MB, 40 connections)")
	}

	in.setupPgvector(ctx, discovery.DiscoverHost())

	// Proof the address works, before anything is moved into it.
	address := "@" + in.dsnPath()
	s, err := store.Open(ctx, address)
	if err != nil {
		return "", fmt.Errorf("connecting with the address in %s: %w (PostgreSQL's pg_hba.conf "+
			"must allow password logins from 127.0.0.1: scram-sha-256 or md5, not ident)", in.dsnPath(), err)
	}
	_ = s.Close()
	return address, nil
}

func (in *serverInit) writeDSN(dsn string) error {
	dir := filepath.Dir(in.dsnPath())
	if err := os.MkdirAll(dir, 0o750); err != nil {
		return err
	}
	if err := in.chown(dir, false); err != nil {
		return err
	}
	// Root's, and the service's group may enter it. Its mode set too, not
	// only on creation: a directory made 0700 for the service user (as
	// `server add` made them) and handed to root here was one the service
	// could no longer open, and the gateway did not start.
	if in.user != "" {
		if err := os.Chmod(dir, 0o750); err != nil {
			return err
		}
	}
	if err := os.WriteFile(in.dsnPath(), []byte(dsn+"\n"), 0o600); err != nil {
		return err
	}
	// Read by the service, which runs as llmgw; root reads it anyway.
	return in.chown(in.dsnPath(), true)
}

// installPostgres installs the server when this machine has none, makes sure
// its database is initialised and running, and says whether it installed it.
func (in *serverInit) installPostgres(ctx context.Context) (bool, error) {
	host := discovery.DiscoverHost()
	pg, found := findPostgres(ctx)
	installed := false
	if found {
		in.kept("PostgreSQL server (%s)", pg.unit)
	} else {
		if host.Family() == "" {
			return false, fmt.Errorf("no PostgreSQL here and no package manager init knows (%s); "+
				"install PostgreSQL and run init again", host.OSID)
		}
		fmt.Fprintf(in.out, "  ..    installing PostgreSQL (a minute or two)\n")
		pi := &provision.Installer{Host: host, Log: slog.Default()}
		if err := pi.InstallPackages(ctx, map[string][]string{
			"deb": {"postgresql"},
			"rpm": {"postgresql-server", "postgresql"},
		}); err != nil {
			return false, err
		}
		if pg, found = findPostgres(ctx); !found {
			q := "rpm -qa 'postgresql*'"
			if host.Family() == "deb" {
				q = "dpkg -l 'postgresql*'"
			}
			have, _ := exec.CommandContext(ctx, "sh", "-c", q).CombinedOutput()
			return false, fmt.Errorf("PostgreSQL was installed but no postgresql service unit "+
				"appeared; what is installed:\n%s", strings.TrimSpace(string(have)))
		}
		installed = true
		in.done("PostgreSQL installed (%s)", pg.unit)
	}
	in.pg = pg

	// Debian's package creates and starts a cluster itself. The RHEL family's,
	// and PGDG's, leave the data directory empty until someone runs initdb.
	if pg.pgdata != "" {
		if _, err := os.Stat(filepath.Join(pg.pgdata, "PG_VERSION")); err != nil {
			// Directly rather than through postgresql-setup, which is named
			// differently by every packager. Password logins over TCP are
			// asked for: RHEL's default for them is ident, which no service
			// account passes.
			if err := os.MkdirAll(pg.pgdata, 0o700); err != nil {
				return false, err
			}
			if out, err := exec.CommandContext(ctx, "chown", "postgres:postgres", pg.pgdata).CombinedOutput(); err != nil {
				return false, fmt.Errorf("chown %s: %v: %s", pg.pgdata, err, strings.TrimSpace(string(out)))
			}
			cmd := exec.CommandContext(ctx, "runuser", "-u", "postgres", "--", pg.tool("initdb"),
				"-D", pg.pgdata, "--auth-host=scram-sha-256", "--auth-local=peer")
			cmd.Dir = "/"
			if out, err := cmd.CombinedOutput(); err != nil {
				return false, fmt.Errorf("initdb %s: %v: %s", pg.pgdata, err, lastLines(out, 5))
			}
			in.done("PostgreSQL database cluster initialised in %s", pg.pgdata)
		}
	}
	if exec.CommandContext(ctx, "systemctl", "is-active", "--quiet", pg.unit).Run() != nil {
		if out, err := exec.CommandContext(ctx, "systemctl", "enable", "--now", pg.unit).CombinedOutput(); err != nil {
			return false, fmt.Errorf("starting %s: %v: %s", pg.unit, err, strings.TrimSpace(string(out)))
		}
		in.done("%s started", pg.unit)
	}
	return installed, nil
}

func lastLines(b []byte, n int) string {
	lines := strings.Split(strings.TrimSpace(string(b)), "\n")
	if len(lines) > n {
		lines = lines[len(lines)-n:]
	}
	return strings.Join(lines, "\n")
}

// psql runs one statement as the postgres superuser, over the local socket.
func (in *serverInit) psql(ctx context.Context, query string) (string, error) {
	return in.psqlIn(ctx, "", query)
}

// psqlIn runs a statement as postgres in one database, or in postgres's
// default one when db is "".
func (in *serverInit) psqlIn(ctx context.Context, db, query string) (string, error) {
	args := []string{"-v", "ON_ERROR_STOP=1", "-qAtX", "-c", query}
	if db != "" {
		args = append(args, "-d", db)
	}
	cmd := exec.CommandContext(ctx, "runuser", append([]string{"-u", "postgres", "--", in.pg.tool("psql")}, args...)...)
	// postgres cannot enter root's home, and psql says so on every call.
	cmd.Dir = "/"
	out, err := cmd.CombinedOutput()
	if err != nil {
		// The statement can hold a password; the error names only its verb.
		verb := strings.Join(strings.Fields(query)[:2], " ")
		return "", fmt.Errorf("psql (%s): %v: %s", verb, err, strings.TrimSpace(string(out)))
	}
	return strings.TrimSpace(string(out)), nil
}

func (in *serverInit) psqlBool(ctx context.Context, query string) (bool, error) {
	out, err := in.psql(ctx, query)
	return out == "1", err
}

// waitForPostgres waits for a server that was just started or restarted to
// take connections.
func (in *serverInit) waitForPostgres(ctx context.Context) error {
	deadline := time.Now().Add(60 * time.Second)
	for {
		_, err := in.psql(ctx, "SELECT 1")
		if err == nil {
			return nil
		}
		if time.Now().After(deadline) {
			return fmt.Errorf("PostgreSQL is not taking connections after 60s: %w", err)
		}
		select {
		case <-ctx.Done():
			return ctx.Err()
		case <-time.After(2 * time.Second):
		}
	}
}

// stopService stops the control plane if it runs, so the store is still.
// It says whether it stopped it, so a failed move can start it again.
func (in *serverInit) stopService(ctx context.Context) (bool, error) {
	if in.unitPath == "" {
		return false, nil
	}
	if _, err := exec.LookPath("systemctl"); err != nil {
		return false, nil
	}
	if exec.CommandContext(ctx, "systemctl", "is-active", "--quiet", "llmgw-server").Run() != nil {
		return false, nil
	}
	if out, err := exec.CommandContext(ctx, "systemctl", "stop", "llmgw-server").CombinedOutput(); err != nil {
		return false, fmt.Errorf("stopping llmgw-server to move its store: %v: %s", err, strings.TrimSpace(string(out)))
	}
	in.done("llmgw-server stopped while its store moves")
	return true, nil
}

// pgvectorPackages are the package names of pgvector for a PostgreSQL major
// version: Debian's and Ubuntu's (and PGDG's apt repository) name it after the
// server, PGDG's rpm repository after the version alone.
func pgvectorPackages(major int) map[string][]string {
	return map[string][]string{
		"deb": {fmt.Sprintf("postgresql-%d-pgvector", major)},
		"rpm": {fmt.Sprintf("pgvector_%d", major)},
	}
}

// setupPgvector installs pgvector and enables it in llmgw's database, for
// knowledge bases. Not a reason to stop: without it everything else works,
// and the store says why knowledge bases are off when somebody asks for one.
func (in *serverInit) setupPgvector(ctx context.Context, h discovery.Host) {
	warn := func(format string, a ...any) {
		fmt.Fprintf(in.out, "  !!    pgvector: "+format+"; knowledge bases stay off\n", a...)
	}
	if on, err := in.psqlIn(ctx, pgDatabase, "SELECT 1 FROM pg_extension WHERE extname = 'vector'"); err == nil && on == "1" {
		in.kept("pgvector in database %s", pgDatabase)
		return
	}
	have, err := in.psqlBool(ctx, "SELECT 1 FROM pg_available_extensions WHERE name = 'vector'")
	if err != nil {
		warn("%v", err)
		return
	}
	if !have {
		num, err := in.psql(ctx, "SHOW server_version_num")
		n, convErr := strconv.Atoi(num)
		if err != nil || convErr != nil || h.Family() == "" {
			warn("cannot tell which package to install (version %q, %s)", num, h.OSID)
			return
		}
		pkgs := pgvectorPackages(n / 10000)
		pi := &provision.Installer{Host: h, Log: slog.Default()}
		if err := pi.InstallPackages(ctx, pkgs); err != nil {
			warn("installing %s: %v", strings.Join(pkgs[h.Family()], " "), err)
			return
		}
	}
	if _, err := in.psqlIn(ctx, pgDatabase, "CREATE EXTENSION IF NOT EXISTS vector"); err != nil {
		warn("%v", err)
		return
	}
	in.done("pgvector enabled in database %s", pgDatabase)
}