mesh-catalog/modules/postgres/cmd/postgres-provider/client_test.go

package main

import (
	"context"
	"errors"
	"os"
	"path/filepath"
	"strings"
	"testing"
	"time"

	"github.com/jackc/pgx/v5/pgconn"
)

var ctx = context.Background()

func TestQuoting(t *testing.T) {
	if got := Ident(`we"ird`); got != `"we""ird"` {
		t.Fatal(got)
	}
	if got := Literal(`it's`); got != `'it''s'` {
		t.Fatal(got)
	}
	if got := RoleStatement(false, "mesh_ace_letta", "p'w"); got != `CREATE ROLE "mesh_ace_letta" WITH LOGIN PASSWORD 'p''w' VALID UNTIL 'infinity'` {
		t.Fatal(got)
	}
	if got := RoleStatement(true, "r", "p"); !strings.HasPrefix(got, `ALTER ROLE "r" WITH LOGIN`) {
		t.Fatal(got)
	}
}

func TestExtensionStatementIsIdempotentAndQuoted(t *testing.T) {
	if got := ExtensionStatement("vector"); got != `CREATE EXTENSION IF NOT EXISTS "vector"` {
		t.Fatal(got)
	}
	if got := ExtensionStatement("uuid-ossp"); got != `CREATE EXTENSION IF NOT EXISTS "uuid-ossp"` {
		t.Fatal(got)
	}
	// A name cannot leave its quotes.
	if got := ExtensionStatement(`x"; DROP DATABASE y; --`); got != `CREATE EXTENSION IF NOT EXISTS "x""; DROP DATABASE y; --"` {
		t.Fatal(got)
	}
}

func TestExtensionsFromAContribution(t *testing.T) {
	got, err := Extensions(map[string]any{"name": "letta"})
	if err != nil || got != nil {
		t.Fatal(got, err)
	}
	got, err = Extensions(map[string]any{"extensions": []any{"vector", " vector ", "pg_trgm"}})
	if err != nil || strings.Join(got, ",") != "vector,pg_trgm" {
		t.Fatal(got, err)
	}
	for _, bad := range []any{"vector", []any{"vector", 3}, []any{""}, map[string]any{}} {
		if _, err := Extensions(map[string]any{"extensions": bad}); err == nil {
			t.Fatalf("%v accepted", bad)
		}
	}
}

func TestAnExtensionTheServerDoesNotOfferIsRefusedByName(t *testing.T) {
	f, c := newFake(admin)
	f.answer(`FROM pg_available_extensions`, []string{"name"}, []string{"vector"}, []string{"pg_trgm"})
	err := c.EnsureExtensions(ctx, "mesh_ace_letta", []string{"vector", "made_up", "plv8"})
	if err == nil || !strings.Contains(err.Error(), `"made_up"`) || !strings.Contains(err.Error(), `"plv8"`) {
		t.Fatalf("refusal does not name the extensions: %v", err)
	}
	if has(f.statements(), `CREATE EXTENSION`) {
		t.Fatalf("installed something although one was refused: %v", f.statements())
	}
}

func TestExtensionsAreInstalledAsTheAdminInTheConsumersDatabase(t *testing.T) {
	f, c := newFake(admin)
	f.answer(`FROM pg_available_extensions`, []string{"name"}, []string{"vector"})
	if err := c.EnsureExtensions(ctx, "mesh_ace_letta", []string{"vector"}); err != nil {
		t.Fatal(err)
	}
	last := f.calls[len(f.calls)-1]
	if last.SQL != `CREATE EXTENSION IF NOT EXISTS "vector"` || last.Database != "mesh_ace_letta" || last.User != "postgres" || last.Password != "admin-secret" {
		t.Fatalf("%+v", last)
	}
	noDrop(t, f.statements())

	f.reset()
	if err := c.EnsureExtensions(ctx, "mesh_ace_letta", nil); err != nil || len(f.calls) != 0 {
		t.Fatalf("asked the server for nothing to do: %v %v", f.calls, err)
	}
}

func TestCreateIsIdempotent(t *testing.T) {
	f, c := newFake(admin)
	if err := c.CreateDatabaseAndRole(ctx, "mesh_ace_letta", "mesh_ace_letta", "pw"); err != nil {
		t.Fatal(err)
	}
	sql := f.statements()
	if !has(sql, `^CREATE ROLE "mesh_ace_letta"`) || !has(sql, `^CREATE DATABASE "mesh_ace_letta" OWNER "mesh_ace_letta"$`) {
		t.Fatal(strings.Join(sql, "\n"))
	}

	// Both exist: the password is set again, the database is not made again.
	f, c = newFake(admin)
	f.answer(`FROM pg_roles`, []string{"?column?"}, []string{"1"})
	f.answer(`FROM pg_database`, []string{"?column?"}, []string{"1"})
	if err := c.CreateDatabaseAndRole(ctx, "mesh_ace_letta", "mesh_ace_letta", "pw"); err != nil {
		t.Fatal(err)
	}
	sql = f.statements()
	if !has(sql, `^ALTER ROLE "mesh_ace_letta" WITH LOGIN PASSWORD 'pw' VALID UNTIL 'infinity'`) || has(sql, `CREATE DATABASE`) {
		t.Fatal(strings.Join(sql, "\n"))
	}
	if !has(sql, `^GRANT ALL PRIVILEGES ON DATABASE "mesh_ace_letta" TO "mesh_ace_letta"$`) {
		t.Fatal(strings.Join(sql, "\n"))
	}
	noDrop(t, sql)
}

func TestWithdrawingLocksTheLoginAndKeepsTheDatabase(t *testing.T) {
	f, c := newFake(admin)
	f.answer(`FROM pg_roles`, []string{"?column?"}, []string{"1"})
	if err := c.LockRole(ctx, "mesh_anchor_mail"); err != nil {
		t.Fatal(err)
	}
	sql := f.statements()
	if !has(sql, `ALTER ROLE "mesh_anchor_mail" NOLOGIN`) || !has(sql, `pg_terminate_backend.*usename = 'mesh_anchor_mail'`) {
		t.Fatal(strings.Join(sql, "\n"))
	}
	noDrop(t, sql)

	// No such role: nothing to lock, nothing done.
	f, c = newFake(admin)
	if err := c.LockRole(ctx, "gone"); err != nil || has(f.statements(), `ALTER`) {
		t.Fatal(f.statements(), err)
	}
}

func TestRetiringRenamesAsideLocksTheOwnerAndDropsNothing(t *testing.T) {
	f, c := newFake(admin)
	f.answer(`pg_get_userbyid`, []string{"owner"}, []string{"mesh_anchor_mail"})
	aside, err := c.RetireDatabase(ctx, "mesh_anchor_mail", time.Date(2026, 10, 5, 9, 0, 0, 0, time.UTC))
	if err != nil || aside != "mesh_anchor_mail_deleted_20261005" {
		t.Fatal(aside, err)
	}
	sql := f.statements()
	if !has(sql, `ALTER DATABASE "mesh_anchor_mail" RENAME TO "mesh_anchor_mail_deleted_20261005"`) || !has(sql, `ALTER ROLE "mesh_anchor_mail" NOLOGIN`) {
		t.Fatal(strings.Join(sql, "\n"))
	}
	noDrop(t, sql)

	f, c = newFake(admin)
	f.answer(`pg_get_userbyid`, []string{"owner"}, []string{"x"})
	f.answer(`_deleted_`, []string{"?column?"}, []string{"1"})
	if _, err := c.RetireDatabase(ctx, "x", time.Now()); err == nil || has(f.statements(), `RENAME`) {
		t.Fatal("renamed over a database already set aside")
	}
}

func TestARetiredNameFitsPostgres(t *testing.T) {
	name := RetiredName(strings.Repeat("x", 70), time.Date(2026, 10, 5, 0, 0, 0, 0, time.UTC))
	if len(name) > 63 || !strings.HasSuffix(name, "_deleted_20261005") {
		t.Fatal(name)
	}
}

func TestACallersStatementRunsAsTheReaderAsGivenReadOnly(t *testing.T) {
	conn := admin
	conn.ReaderPassword = "reader-secret"
	f, c := newFake(conn)
	f.answer(`^COMMIT; DROP`, []string{"name", "n"}, []string{"alpha", "1"}, []string{"b,eta", "2"})
	statement := "COMMIT; DROP TABLE everything"
	r, err := c.ReadOnlyQuery(ctx, "inventory", statement)
	if err != nil {
		t.Fatal(err)
	}
	asked := f.calls[len(f.calls)-1]
	if asked.User != Reader || asked.Password != "reader-secret" || asked.Database != "inventory" {
		t.Fatalf("not as the reader: %+v", asked)
	}
	if asked.SQL != statement {
		t.Fatalf("not sent as given: %q", asked.SQL)
	}
	if asked.Options["default_transaction_read_only"] != "on" || asked.Options["statement_timeout"] != "60s" {
		t.Fatalf("session not read-only from its first statement: %v", asked.Options)
	}
	rows := r.Maps()
	if r.Command != "COMMIT" || len(rows) != 2 || rows[1]["name"] != "b,eta" || rows[0]["n"] != "1" {
		t.Fatalf("%+v", r)
	}
}

func TestTheReaderIsMadeAsTheAdminWithEveryAttributeOnce(t *testing.T) {
	conn := admin
	conn.ReaderPassword = "reader-secret"
	f, c := newFake(conn)
	for _, q := range []string{"SELECT 1", "SELECT 2"} {
		if _, err := c.ReadOnlyQuery(ctx, "inventory", q); err != nil {
			t.Fatal(err)
		}
	}
	var ddl []string
	readerCalls := 0
	for _, k := range f.calls {
		switch k.User {
		case "postgres":
			if k.Password != "admin-secret" {
				t.Fatal("admin call without the admin's password")
			}
			ddl = append(ddl, k.SQL)
		case Reader:
			readerCalls++
		}
	}
	creates := 0
	for _, s := range ddl {
		if strings.HasPrefix(s, `CREATE ROLE "`+Reader+`"`) {
			creates++
			for _, a := range []string{"LOGIN", "NOSUPERUSER", "NOCREATEDB", "NOCREATEROLE", "NOREPLICATION", "NOBYPASSRLS"} {
				if !strings.Contains(s, " "+a+" ") {
					t.Fatalf("%s missing from %s", a, s)
				}
			}
		}
	}
	// Each query is two statements as the reader: what it can read that is secret, then the caller's.
	if creates != 1 || readerCalls != 4 {
		t.Fatalf("made %d times, %d reader calls", creates, readerCalls)
	}
	if !has(ddl, `default_transaction_read_only = on`) {
		t.Fatal(strings.Join(ddl, "\n"))
	}
	// novox/hq issue 337: the reader reads only what is declared, so it is never given the role that
	// reads everything — and an install that gave it that role has it taken back.
	if has(ddl, `GRANT pg_read_all_data`) || !has(ddl, `^REVOKE pg_read_all_data FROM "`+Reader+`"$`) {
		t.Fatal("the reader is given, or keeps, pg_read_all_data:\n" + strings.Join(ddl, "\n"))
	}
}

func TestWithoutTheReadersPasswordNothingRunsAsTheAdmin(t *testing.T) {
	f, c := newFake(admin)
	_, err := c.ReadOnlyQuery(ctx, "inventory", "SELECT 1")
	if err == nil || !strings.Contains(err.Error(), "refused rather than run as the admin") {
		t.Fatal(err)
	}
	if len(f.calls) != 0 {
		t.Fatalf("ran %v", f.calls)
	}
}

func TestAFailedReaderSetupIsAskedAgain(t *testing.T) {
	conn := admin
	conn.ReaderPassword = "r"
	f, c := newFake(conn)
	f.fail(`^CREATE ROLE`, errors.New("server busy"))
	if _, err := c.ReadOnlyQuery(ctx, "db", "SELECT 1"); err == nil {
		t.Fatal("no error")
	}
	f.rules = nil
	if _, err := c.ReadOnlyQuery(ctx, "db", "SELECT 1"); err != nil {
		t.Fatal(err)
	}
}

func TestHolds(t *testing.T) {
	f, c := newFake(admin)
	f.answer(`FROM pg_extension`, []string{"extname"}, []string{"plpgsql"})
	ok, err := c.Holds(ctx, "db", "db", "pw", nil)
	if err != nil || !ok {
		t.Fatal(ok, err)
	}
	if k := f.calls[0]; k.User != "db" || k.Password != "pw" || k.Database != "db" {
		t.Fatalf("not checked as the consumer: %+v", k)
	}
	// A wanted extension missing: not held, so it is applied again.
	if ok, err := c.Holds(ctx, "db", "db", "pw", []string{"vector"}); err != nil || ok {
		t.Fatal(ok, err)
	}
	// The server saying the login is wrong or gone: not held.
	for _, code := range []string{"28P01", "28000", "3D000", "42501"} {
		f.dialErr = func(Login) error { return &pgconn.PgError{Code: code} }
		if ok, err := c.Holds(ctx, "db", "db", "pw", nil); err != nil || ok {
			t.Fatal(code, ok, err)
		}
	}
	// Unable to ask is not evidence of loss.
	f.dialErr = func(Login) error { return errors.New("connection refused") }
	if _, err := c.Holds(ctx, "db", "db", "pw", nil); err == nil {
		t.Fatal("an unreachable server reported as a lost login")
	}
}

func TestClientFromEnv(t *testing.T) {
	dir := t.TempDir()
	write := func(name, content string) string {
		p := filepath.Join(dir, name)
		if err := os.WriteFile(p, []byte(content), 0o600); err != nil {
			t.Fatal(err)
		}
		return p
	}
	env := map[string]string{
		"MESH_PROVISION_POSTGRES":            "postgres://postgres@127.0.0.1:6852/postgres?sslmode=disable",
		"MESH_PROVISION_PASSWORD_FILE":       write("superuser.secret", "admin\n"),
		"MESH_POSTGRES_READER_PASSWORD_FILE": write("reader.secret", "from-the-file\n"),
		"MESH_PROVISION_POSTGRES_PORT":       "${port:mesh-store}",
	}
	c, err := ClientFromEnv(func(k string) string { return env[k] })
	if err != nil {
		t.Fatal(err)
	}
	if c.conn.Host != "127.0.0.1" || c.conn.Port != 6852 || c.conn.User != "postgres" || c.conn.Password != "admin" ||
		c.conn.ReaderPassword != "from-the-file" || c.conn.SSLMode != "disable" {
		t.Fatalf("%+v", c.conn)
	}
	env["MESH_PROVISION_POSTGRES_PORT"] = "7000"
	c, _ = ClientFromEnv(func(k string) string { return env[k] })
	if c.conn.Port != 7000 {
		t.Fatal(c.conn.Port)
	}
	if _, err := ClientFromEnv(func(string) string { return "" }); err == nil {
		t.Fatal("no host and no password accepted")
	}
}

func TestNullsStayNull(t *testing.T) {
	v := "x"
	r := Result{Fields: []string{"a", "b"}, Rows: [][]*string{{&v, nil}}}
	m := r.Maps()[0]
	if m["a"] != "x" || m["b"] != nil {
		t.Fatal(m)
	}
}