llmgw/internal/store/copy_test.go

package store

import (
	"context"
	"os"
	"path/filepath"
	"regexp"
	"slices"
	"strings"
	"testing"
	"time"
)

func TestEveryTableIsInTheCopyOrder(t *testing.T) {
	// A table added by a migration and not added here is copied by nothing,
	// and the destination comes up quietly missing it — discovered when
	// somebody looks for a setting that was there before the move. Read out of
	// the migrations rather than listed twice, so the two cannot disagree.
	var declared []string
	entries, err := migrationFS.ReadDir("migrations/sqlite")
	if err != nil {
		t.Fatal(err)
	}
	for _, e := range entries {
		body, err := migrationFS.ReadFile("migrations/sqlite/" + e.Name())
		if err != nil {
			t.Fatal(err)
		}
		for _, line := range strings.Split(string(body), "\n") {
			line = strings.TrimSpace(line)
			rest, ok := strings.CutPrefix(line, "CREATE TABLE ")
			if !ok {
				continue
			}
			name, _, _ := strings.Cut(rest, " ")
			name = strings.TrimSuffix(strings.TrimSpace(name), "(")
			declared = append(declared, name)
		}
	}
	if len(declared) == 0 {
		t.Fatal("found no tables in the migrations, so this test proves nothing")
	}
	renamed := renamedTables(t)
	for _, name := range declared {
		if to, ok := renamed[name]; ok {
			// A table rebuilt under a temporary name and renamed back is the
			// table it became, not another one.
			name = to
		}
		if slices.Contains(copyOrder, name) {
			continue
		}
		// Left behind on purpose, with the reason written down beside it.
		if why, ok := notCopied[name]; ok {
			if why == "" {
				t.Errorf("%q is excluded from the copy with no reason given", name)
			}
			continue
		}
		t.Errorf("%q is created by a migration but nothing copies it", name)
	}
}

func TestCopyRefusesADestinationThatAlreadyHasRows(t *testing.T) {
	// Merging two stores means deciding what happens when both hold a project
	// with the same id and different contents, and every answer is a guess
	// about which one somebody meant.
	ctx := context.Background()
	src, dst := openDB(t), openDB(t)
	tree(t, src)
	tree(t, dst)

	_, err := src.CopyTo(ctx, dst)
	if err == nil || !strings.Contains(err.Error(), "empty store") {
		t.Errorf("copying into a populated store gave %v", err)
	}
}

func TestACopiedStoreStillAuthenticatesTheSameKey(t *testing.T) {
	// The reason this exists. Keys are stored hashed, so a migration that
	// loses them cannot recreate them: every person and every script would
	// need a new one issued.
	ctx := context.Background()
	src := openDB(t)
	proj := tree(t, src)
	secret, k := mustKey(t, src, "k1", proj)

	if err := src.WriteUsage(ctx, []Usage{{
		At: time.Now(), KeyID: k.ID, ProjectID: proj, Model: "m", ServedModel: "m",
		PromptTokens: 3, CompletionTokens: 4, CostMicros: 99, Status: 200,
	}}); err != nil {
		t.Fatal(err)
	}

	dst, err := Open(ctx, filepath.Join(t.TempDir(), "copy.db"))
	if err != nil {
		t.Fatal(err)
	}
	defer dst.Close()

	res, err := src.CopyTo(ctx, dst)
	if err != nil {
		t.Fatalf("copying: %v", err)
	}
	if len(res) != len(copyOrder) {
		t.Errorf("copied %d tables of %d", len(res), len(copyOrder))
	}

	// The same secret the caller already holds still resolves, in the copy.
	got, err := dst.KeyByHash(ctx, HashKey(secret))
	if err != nil {
		t.Fatalf("the copied store does not know the key: %v", err)
	}
	if got.ID != k.ID || got.ProjectID != proj {
		t.Errorf("key came back as %+v", got)
	}

	// And the spend it had, so a budget is not quietly reset by moving.
	spend, err := dst.SpendSince(ctx, time.Now().Add(-time.Hour))
	if err != nil {
		t.Fatal(err)
	}
	if spend[k.ID] != 99 {
		t.Errorf("spend after the copy = %d, want 99", spend[k.ID])
	}
}

func TestCopyCarriesConversationsAndNotes(t *testing.T) {
	ctx := context.Background()
	src := openDB(t)
	proj := tree(t, src)
	if err := src.SetThreadRetention(ctx, proj, 24); err != nil {
		t.Fatal(err)
	}
	if err := src.AppendThread(ctx, proj, "th1", []Message{
		{Role: "user", Content: "remembered"},
	}); err != nil {
		t.Fatal(err)
	}

	dst, err := Open(ctx, filepath.Join(t.TempDir(), "copy.db"))
	if err != nil {
		t.Fatal(err)
	}
	defer dst.Close()
	if _, err := src.CopyTo(ctx, dst); err != nil {
		t.Fatal(err)
	}

	msgs, err := dst.Thread(ctx, proj, "th1")
	if err != nil {
		t.Fatalf("the conversation did not survive: %v", err)
	}
	if len(msgs) != 1 || msgs[0].Content != "remembered" {
		t.Errorf("thread came back as %+v", msgs)
	}
	// The retention setting is a column on projects, so it rides along — a
	// destination that kept the conversations and forgot the promise about
	// how long would be worse than one that kept neither.
	if d, err := dst.ThreadRetention(ctx, proj); err != nil || d == 0 {
		t.Errorf("retention after the copy = %v (err %v)", d, err)
	}
}

func TestCopyingAnEmptyStoreIsFine(t *testing.T) {
	ctx := context.Background()
	src := openDB(t)
	dst, err := Open(ctx, filepath.Join(t.TempDir(), "copy.db"))
	if err != nil {
		t.Fatal(err)
	}
	defer dst.Close()
	if _, err := src.CopyTo(ctx, dst); err != nil {
		t.Errorf("copying an empty store failed: %v", err)
	}
	_ = os.Remove("")
}

func TestTheCopyOrderSatisfiesTheForeignKeys(t *testing.T) {
	// Completeness is not enough: a table copied before the one it references
	// fails on a constraint partway through, leaving some tables copied and
	// some not. The first version of copyOrder followed the order the
	// migrations create tables in, and users is created after api_keys even
	// though api_keys points at it — so the copy died three tables in, on a
	// real database, after the gateway had already been stopped.
	//
	// Derived from the migrations rather than restated, so a new reference
	// fails here rather than during somebody's migration.
	deps := map[string]map[string]bool{}
	note := func(table, ref string) {
		if table == ref {
			return
		}
		if deps[table] == nil {
			deps[table] = map[string]bool{}
		}
		deps[table][ref] = true
	}

	entries, err := migrationFS.ReadDir("migrations/sqlite")
	if err != nil {
		t.Fatal(err)
	}
	body := ""
	for _, e := range entries {
		b, err := migrationFS.ReadFile("migrations/sqlite/" + e.Name())
		if err != nil {
			t.Fatal(err)
		}
		body += string(b)
	}

	// Table bodies, so a REFERENCES is attributed to the table holding it,
	// under the name it ends with for one rebuilt and renamed.
	renamed := renamedTables(t)
	for _, m := range regexpCreateTable.FindAllStringSubmatch(body, -1) {
		table := m[1]
		if to, ok := renamed[table]; ok {
			table = to
		}
		for _, r := range regexpReferences.FindAllStringSubmatch(m[2], -1) {
			note(table, r[1])
		}
	}
	// And columns added later, which is exactly how api_keys got its user_id.
	for _, m := range regexpAlterRef.FindAllStringSubmatch(body, -1) {
		note(m[1], m[2])
	}
	if len(deps) == 0 {
		t.Fatal("found no foreign keys, so this test proves nothing")
	}

	position := map[string]int{}
	for i, t := range copyOrder {
		position[t] = i
	}
	for table, refs := range deps {
		for ref := range refs {
			if position[ref] > position[table] {
				t.Errorf("%s is copied before %s, which it references", table, ref)
			}
		}
	}
}

var (
	regexpCreateTable = regexp.MustCompile(`(?s)CREATE TABLE (\w+)\s*\((.*?)\n\);`)
	regexpReferences  = regexp.MustCompile(`REFERENCES\s+(\w+)`)
	regexpAlterRef    = regexp.MustCompile(`ALTER TABLE (\w+) ADD COLUMN[^;]*REFERENCES\s+(\w+)`)
	regexpRename      = regexp.MustCompile(`ALTER TABLE (\w+) RENAME TO (\w+)`)
)

// renamedTables maps a table a migration renames to the name it gets.
func renamedTables(t *testing.T) map[string]string {
	t.Helper()
	out := map[string]string{}
	entries, err := migrationFS.ReadDir("migrations/sqlite")
	if err != nil {
		t.Fatal(err)
	}
	for _, e := range entries {
		b, err := migrationFS.ReadFile("migrations/sqlite/" + e.Name())
		if err != nil {
			t.Fatal(err)
		}
		for _, m := range regexpRename.FindAllStringSubmatch(string(b), -1) {
			out[m[1]] = m[2]
		}
	}
	return out
}

func TestAnEmptyBlobDoesNotBecomeNullInTheCopy(t *testing.T) {
	// key_claims.sealed is emptied the moment a claim is used, so a claimed
	// row holds an empty blob in a NOT NULL column. An empty blob scans as a
	// typed nil slice and a real NULL as an untyped one; pgx cannot tell them
	// apart, so the distinction is restored before the insert.
	if got := keepEmptyBytes([]byte(nil)); got == nil {
		t.Error("an empty byte column would be inserted as NULL")
	} else if b, ok := got.([]byte); !ok || len(b) != 0 {
		t.Errorf("empty bytes came out as %#v", got)
	}
	// A genuine NULL arrives as an untyped nil and must stay one.
	if got := keepEmptyBytes(nil); got != nil {
		t.Errorf("a NULL became %#v", got)
	}
	// Everything else passes through.
	if got := keepEmptyBytes(int64(7)); got != int64(7) {
		t.Errorf("a number was rewritten to %#v", got)
	}
	if got := keepEmptyBytes([]byte("x")); string(got.([]byte)) != "x" {
		t.Errorf("a non-empty blob was rewritten to %#v", got)
	}
}