llmgw/internal/cli/serverinit_postgres.go

package cli

import (
	"context"
	"crypto/rand"
	"encoding/hex"
	"fmt"
	"log/slog"
	"os"
	"os/exec"
	"path/filepath"
	"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 !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())
}

// 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 := waitForPostgres(ctx); err != nil {
		return "", err
	}

	exists, err := 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 := 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 = 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 := 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 := psql(ctx, q); err != nil {
				return "", err
			}
		}
		if out, err := exec.CommandContext(ctx, "systemctl", "restart", "postgresql").CombinedOutput(); err != nil {
			return "", fmt.Errorf("restarting PostgreSQL: %v: %s", err, strings.TrimSpace(string(out)))
		}
		if err := waitForPostgres(ctx); err != nil {
			return "", err
		}
		in.done("PostgreSQL sized for a machine it shares (shared_buffers 96MB, 40 connections)")
	}

	// 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 %s: %w", 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
	}
	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, and says
// whether it did.
func (in *serverInit) installPostgres(ctx context.Context) (bool, error) {
	if exec.CommandContext(ctx, "systemctl", "cat", "postgresql.service").Run() == nil {
		in.kept("PostgreSQL server")
		return false, nil
	}
	host := discovery.DiscoverHost()
	family := host.Family()
	if 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 family == "rpm" {
		// Debian's package creates and starts a cluster; RHEL's leaves an
		// empty data directory. Password authentication over TCP is asked
		// for, because RHEL's default for it is ident, which no service
		// account passes.
		if _, err := os.Stat("/var/lib/pgsql/data/PG_VERSION"); err != nil {
			cmd := exec.CommandContext(ctx, "postgresql-setup", "--initdb")
			cmd.Env = append(os.Environ(), "PGSETUP_INITDB_OPTIONS=--auth-host=scram-sha-256 --auth-local=peer")
			if out, err := cmd.CombinedOutput(); err != nil {
				return false, fmt.Errorf("postgresql-setup --initdb: %v: %s", err, strings.TrimSpace(string(out)))
			}
		}
	}
	if out, err := exec.CommandContext(ctx, "systemctl", "enable", "--now", "postgresql").CombinedOutput(); err != nil {
		return false, fmt.Errorf("starting PostgreSQL: %v: %s", err, strings.TrimSpace(string(out)))
	}
	in.done("PostgreSQL installed and started")
	return true, nil
}

// psql runs one statement as the postgres superuser, over the local socket.
func psql(ctx context.Context, query string) (string, error) {
	cmd := exec.CommandContext(ctx, "runuser", "-u", "postgres", "--",
		"psql", "-v", "ON_ERROR_STOP=1", "-qAtX", "-c", query)
	// 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 psqlBool(ctx context.Context, query string) (bool, error) {
	out, err := psql(ctx, query)
	return out == "1", err
}

// waitForPostgres waits for a server that was just started or restarted to
// take connections.
func waitForPostgres(ctx context.Context) error {
	deadline := time.Now().Add(60 * time.Second)
	for {
		_, err := 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
}