Port postgres to Go and install the extensions a consumer asks for
letta crash-loops on 'type "vector" does not exist': pgvector is not a trusted extension, so only the provider's superuser can create it, and the provisioner never did. A contribution may now name extensions; the provider creates each (IF NOT EXISTS, available ones only) in the consumer's database on every pass. Go per the standing rule for a TypeScript module that changes. letta asks for vector.
This commit is contained in:
jochen
@@ -0,0 +1,325 @@
|
||||
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)
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
if creates != 1 || readerCalls != 2 {
|
||||
t.Fatalf("made %d times, %d reader calls", creates, readerCalls)
|
||||
}
|
||||
if !has(ddl, `^GRANT pg_read_all_data TO "`+Reader+`"$`) || !has(ddl, `default_transaction_read_only = on`) {
|
||||
t.Fatal(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)
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user