Agent-System/internal/localpg/localpg.go

package localpg

import (
	"context"
	"fmt"
	"os"
	"os/exec"
	"path/filepath"
	"strconv"
	"strings"
	"time"

	"chat-agent-system/internal/config"
)

// Mode selects how PostgreSQL is managed by the server process.
type Mode string

const (
	ModeOff     Mode = "off"     // external DSN only (compose / remote)
	ModeManaged Mode = "managed" // initdb + pg_ctl child process
	ModeAttach  Mode = "attach"  // connect only; do not start/stop
)

// Handle owns a managed postmaster child process.
type Handle struct {
	BinDir  string
	PGData  string
	Port    int
	Socket  string
	started bool
}

// ModeFromEnv reads LOCALPG_MODE (default off).
func ModeFromEnv() Mode {
	switch strings.ToLower(strings.TrimSpace(os.Getenv("LOCALPG_MODE"))) {
	case "managed":
		return ModeManaged
	case "attach":
		return ModeAttach
	default:
		return ModeOff
	}
}

// ResolveBinDir finds postgresql binaries (LOCALPG_BIN overrides).
func ResolveBinDir() (string, error) {
	if v := strings.TrimSpace(os.Getenv("LOCALPG_BIN")); v != "" {
		if _, err := os.Stat(filepath.Join(v, "pg_ctl")); err != nil {
			return "", fmt.Errorf("localpg: LOCALPG_BIN=%s missing pg_ctl: %w", v, err)
		}
		return v, nil
	}
	candidates := []string{
		"/home/linuxbrew/.linuxbrew/opt/postgresql@18/bin",
		"/opt/homebrew/opt/postgresql@18/bin",
		"/usr/local/opt/postgresql@18/bin",
		"/home/linuxbrew/.linuxbrew/opt/postgresql/bin",
	}
	for _, c := range candidates {
		if _, err := os.Stat(filepath.Join(c, "pg_ctl")); err == nil {
			return c, nil
		}
	}
	if p, err := exec.LookPath("pg_ctl"); err == nil {
		return filepath.Dir(p), nil
	}
	return "", fmt.Errorf("localpg: pg_ctl not found; install postgresql@18 or set LOCALPG_BIN")
}

// shortSocketDir returns a unix socket directory whose path stays under the
// Linux ~107-byte sun_path limit (including /.s.PGSQL.<port>).
func shortSocketDir(port int) string {
	if v := strings.TrimSpace(os.Getenv("LOCALPG_SOCKET_DIR")); v != "" {
		return v
	}
	base := os.Getenv("XDG_RUNTIME_DIR")
	if base == "" || len(base) > 40 {
		base = os.TempDir()
	}
	return filepath.Join(base, fmt.Sprintf("as-pg-%d", port))
}

// Ensure bootstraps a managed local PostgreSQL when LOCALPG_MODE=managed.
// Returns nil handle for off/attach (no process ownership).
func Ensure(ctx context.Context, cfg *config.Config) (*Handle, error) {
	mode := ModeFromEnv()
	switch mode {
	case ModeOff, ModeAttach:
		return nil, nil
	case ModeManaged:
		// continue
	default:
		return nil, fmt.Errorf("localpg: unknown LOCALPG_MODE %q", mode)
	}
	if cfg == nil {
		return nil, fmt.Errorf("localpg: config required")
	}

	binDir, err := ResolveBinDir()
	if err != nil {
		return nil, err
	}
	dataDir := strings.TrimSpace(cfg.DataDir)
	if dataDir == "" {
		dataDir = config.DefaultDataDir()
	}
	pgData := filepath.Join(dataDir, "pg")
	port := cfg.ContextDbPort
	if port <= 0 {
		port = 5433
	}
	// Socket dir must be short: Linux sun_path max is ~107 bytes; worktree
	// paths under .worktrees/... easily exceed that. Prefer XDG_RUNTIME_DIR
	// or /tmp with a stable short name keyed by port.
	socketDir := shortSocketDir(port)
	h := &Handle{BinDir: binDir, PGData: pgData, Port: port, Socket: socketDir}

	if err := os.MkdirAll(socketDir, 0o700); err != nil {
		return nil, fmt.Errorf("localpg: mkdir socket: %w", err)
	}

	if _, err := os.Stat(filepath.Join(pgData, "PG_VERSION")); err != nil {
		if err := os.MkdirAll(pgData, 0o700); err != nil {
			return nil, fmt.Errorf("localpg: mkdir PGDATA: %w", err)
		}
		if err := h.initdb(); err != nil {
			return nil, err
		}
	}

	if err := h.start(ctx); err != nil {
		return nil, err
	}
	if err := h.waitReady(ctx, 30*time.Second); err != nil {
		_ = h.Stop()
		return nil, err
	}
	if err := h.bootstrapRoleAndDB(cfg); err != nil {
		_ = h.Stop()
		return nil, err
	}
	return h, nil
}

func (h *Handle) bin(name string) string {
	return filepath.Join(h.BinDir, name)
}

func (h *Handle) initdb() error {
	cmd := exec.Command(h.bin("initdb"),
		"-D", h.PGData,
		"--locale=C",
		"--encoding=UTF8",
		"--auth-local=trust",
		"--auth-host=scram-sha-256",
		"-U", "postgres",
	)
	cmd.Env = os.Environ()
	out, err := cmd.CombinedOutput()
	if err != nil {
		return fmt.Errorf("localpg: initdb: %w\n%s", err, out)
	}
	// Allow TCP from localhost with password + unix trust for bootstrap.
	conf := filepath.Join(h.PGData, "postgresql.conf")
	appendConf := fmt.Sprintf("\n# agent-system localpg\nlisten_addresses = '127.0.0.1'\nport = %d\nunix_socket_directories = '%s'\n", h.Port, h.Socket)
	f, err := os.OpenFile(conf, os.O_APPEND|os.O_WRONLY, 0o600)
	if err != nil {
		return fmt.Errorf("localpg: open postgresql.conf: %w", err)
	}
	defer f.Close()
	if _, err := f.WriteString(appendConf); err != nil {
		return err
	}
	hba := filepath.Join(h.PGData, "pg_hba.conf")
	hbaBody := []byte(`# TYPE  DATABASE        USER            ADDRESS                 METHOD
local   all             all                                     trust
host    all             all             127.0.0.1/32            scram-sha-256
host    all             all             ::1/128                 scram-sha-256
`)
	if err := os.WriteFile(hba, hbaBody, 0o600); err != nil {
		return fmt.Errorf("localpg: write pg_hba.conf: %w", err)
	}
	return nil
}

func (h *Handle) start(ctx context.Context) error {
	// Already running?
	if h.isReady(ctx) {
		h.started = false // do not stop foreign instance
		return nil
	}
	logFile := filepath.Join(h.PGData, "pg.log")
	cmd := exec.Command(h.bin("pg_ctl"),
		"-D", h.PGData,
		"-l", logFile,
		"-o", fmt.Sprintf("-p %d -k %s", h.Port, h.Socket),
		"start",
	)
	out, err := cmd.CombinedOutput()
	if err != nil {
		return fmt.Errorf("localpg: pg_ctl start: %w\n%s", err, out)
	}
	h.started = true
	return nil
}

func (h *Handle) isReady(ctx context.Context) bool {
	args := []string{"-h", "127.0.0.1", "-p", strconv.Itoa(h.Port)}
	cmd := exec.CommandContext(ctx, h.bin("pg_isready"), args...)
	return cmd.Run() == nil
}

func (h *Handle) waitReady(ctx context.Context, timeout time.Duration) error {
	deadline := time.Now().Add(timeout)
	for time.Now().Before(deadline) {
		if h.isReady(ctx) {
			return nil
		}
		select {
		case <-ctx.Done():
			return ctx.Err()
		case <-time.After(200 * time.Millisecond):
		}
	}
	return fmt.Errorf("localpg: pg_isready timeout on 127.0.0.1:%d", h.Port)
}

func (h *Handle) psql(args ...string) error {
	cmdArgs := append([]string{"-h", h.Socket, "-p", strconv.Itoa(h.Port), "-U", "postgres"}, args...)
	cmd := exec.Command(h.bin("psql"), cmdArgs...)
	cmd.Env = append(os.Environ(), "PGHOST="+h.Socket)
	out, err := cmd.CombinedOutput()
	if err != nil {
		return fmt.Errorf("localpg: psql %v: %w\n%s", args, err, out)
	}
	return nil
}

func (h *Handle) bootstrapRoleAndDB(cfg *config.Config) error {
	user := cfg.ContextDbUser
	if user == "" {
		user = "app"
	}
	pass := cfg.ContextDbPassword
	if pass == "" {
		pass = "dev_secret_db_pass"
	}
	dbName := cfg.ContextDbName
	if dbName == "" {
		dbName = "agent_system_dev"
	}
	// Escape single quotes for SQL literals.
	esc := func(s string) string { return strings.ReplaceAll(s, "'", "''") }

	sql := fmt.Sprintf(`
DO $$
BEGIN
  IF NOT EXISTS (SELECT FROM pg_roles WHERE rolname = '%s') THEN
    CREATE ROLE %s LOGIN PASSWORD '%s';
  ELSE
    ALTER ROLE %s WITH LOGIN PASSWORD '%s';
  END IF;
END$$;
SELECT 'ok' WHERE EXISTS (SELECT FROM pg_database WHERE datname = '%s');
`, esc(user), quoteIdent(user), esc(pass), quoteIdent(user), esc(pass), esc(dbName))

	if err := h.psql("-v", "ON_ERROR_STOP=1", "-c", sql); err != nil {
		return err
	}
	// Create DB if missing (cannot run CREATE DATABASE inside DO easily on all versions).
	check := exec.Command(h.bin("psql"), "-h", h.Socket, "-p", strconv.Itoa(h.Port), "-U", "postgres", "-tAc",
		fmt.Sprintf("SELECT 1 FROM pg_database WHERE datname='%s'", esc(dbName)))
	check.Env = append(os.Environ(), "PGHOST="+h.Socket)
	out, _ := check.Output()
	if strings.TrimSpace(string(out)) != "1" {
		create := exec.Command(h.bin("createdb"),
			"-h", h.Socket, "-p", strconv.Itoa(h.Port), "-U", "postgres",
			"-O", user, dbName,
		)
		create.Env = append(os.Environ(), "PGHOST="+h.Socket)
		if o, err := create.CombinedOutput(); err != nil {
			return fmt.Errorf("localpg: createdb: %w\n%s", err, o)
		}
	}
	extSQL := fmt.Sprintf(`
CREATE EXTENSION IF NOT EXISTS vector;
CREATE EXTENSION IF NOT EXISTS "uuid-ossp";
CREATE EXTENSION IF NOT EXISTS pg_trgm;
GRANT ALL ON SCHEMA public TO %s;
`, quoteIdent(user))
	cmd := exec.Command(h.bin("psql"),
		"-h", h.Socket, "-p", strconv.Itoa(h.Port), "-U", "postgres",
		"-d", dbName, "-v", "ON_ERROR_STOP=1", "-c", extSQL,
	)
	cmd.Env = append(os.Environ(), "PGHOST="+h.Socket)
	if o, err := cmd.CombinedOutput(); err != nil {
		return fmt.Errorf("localpg: extensions: %w\n%s", err, o)
	}
	return nil
}

func quoteIdent(ident string) string {
	return `"` + strings.ReplaceAll(ident, `"`, `""`) + `"`
}

// Stop stops a postmaster we started. No-op if we attached to an already-running instance.
func (h *Handle) Stop() error {
	if h == nil || !h.started {
		return nil
	}
	cmd := exec.Command(h.bin("pg_ctl"), "-D", h.PGData, "-m", "fast", "stop")
	out, err := cmd.CombinedOutput()
	if err != nil {
		return fmt.Errorf("localpg: pg_ctl stop: %w\n%s", err, out)
	}
	h.started = false
	return nil
}