llmgw/internal/store/migrate0010_test.go

package store

import (
	"context"
	"database/sql"
	"path/filepath"
	"strings"
	"testing"
)

// Migration 0010 rebuilds api_keys on SQLite, a kind of migration this store
// had never run. Checked against a store that already has what a running
// cluster has: a person in a project, their key, a pending claim pointing at
// the key, and usage — including the chat page's, which has no key row.
func TestMigration0010KeepsAnExistingStoreWhole(t *testing.T) {
	ctx := context.Background()
	path := filepath.Join(t.TempDir(), "old.db")
	db, err := sql.Open("sqlite", path+"?_pragma=foreign_keys(1)")
	if err != nil {
		t.Fatal(err)
	}
	defer db.Close()
	db.SetMaxOpenConns(1)

	if _, err := db.ExecContext(ctx,
		`CREATE TABLE schema_migrations (version TEXT PRIMARY KEY, applied_at INTEGER NOT NULL)`); err != nil {
		t.Fatal(err)
	}
	names, err := migrationNames("sqlite")
	if err != nil {
		t.Fatal(err)
	}
	for _, name := range names {
		if name >= "0010" {
			break
		}
		body, _ := migrationFS.ReadFile("migrations/sqlite/" + name)
		if err := applyOne(ctx, db, "sqlite", name, string(body)); err != nil {
			t.Fatalf("%s: %v", name, err)
		}
	}

	for _, q := range []string{
		`INSERT INTO organizations (id, name) VALUES ('o', 'Acme')`,
		`INSERT INTO teams (id, org_id, name) VALUES ('t', 'o', 'ml')`,
		`INSERT INTO projects (id, team_id, name) VALUES ('p', 't', 'training')`,
		`INSERT INTO users (id, email, name, password_hash, state, project_id, created_at, approved_at, approved_by)
		 VALUES ('u', 'a@x', 'a', 'h', 'active', 'p', 1, 2, 'admin@x')`,
		`INSERT INTO users (id, email, name, password_hash, state, created_at)
		 VALUES ('w', 'w@x', 'w', 'h', 'pending', 1)`,
		`INSERT INTO api_keys (id, key_hash, key_prefix, name, project_id, created_at, user_id)
		 VALUES ('k', X'01', 'sk-1', 'laptop', 'p', 1, 'u')`,
		`INSERT INTO api_keys (id, key_hash, key_prefix, name, project_id, created_at)
		 VALUES ('ops', X'02', 'sk-2', 'ci', 'p', 1)`,
		`INSERT INTO key_claims (token_hash, key_id, sealed, created_at, expires_at) VALUES (X'aa', 'k', X'00', 1, 9)`,
		`INSERT INTO usage (ts, key_id, project_id, model, served_model, prompt_tokens, completion_tokens,
		   cost_micros, status, latency_ms) VALUES (5, 'k', 'p', 'm', 'm', 1, 1, 3, 200, 1)`,
		`INSERT INTO usage (ts, key_id, project_id, model, served_model, prompt_tokens, completion_tokens,
		   cost_micros, status, latency_ms) VALUES (6, 'session:u', 'p', 'm', 'm', 1, 1, 4, 200, 1)`,
		`INSERT INTO usage (ts, key_id, project_id, model, served_model, prompt_tokens, completion_tokens,
		   cost_micros, status, latency_ms) VALUES (7, 'ops', 'p', 'm', 'm', 1, 1, 5, 200, 1)`,
	} {
		if _, err := db.ExecContext(ctx, q); err != nil {
			t.Fatalf("%s: %v", q, err)
		}
	}

	if err := migrate(ctx, db, "sqlite"); err != nil {
		t.Fatalf("migrating: %v", err)
	}

	one := func(q string, args ...any) string {
		t.Helper()
		var v sql.NullString
		if err := db.QueryRowContext(ctx, q, args...).Scan(&v); err != nil {
			t.Fatalf("%s: %v", q, err)
		}
		return v.String
	}
	if got := one(`SELECT COUNT(*) FROM api_keys`); got != "2" {
		t.Errorf("keys after the rebuild: %s", got)
	}
	if got := one(`SELECT user_id FROM api_keys WHERE id = 'k'`); got != "u" {
		t.Errorf("the key's owner: %q", got)
	}
	if got := one(`SELECT key_id FROM key_claims`); got != "k" {
		t.Errorf("the claim lost its key: %q", got)
	}
	if got := one(`SELECT scope_type || ':' || scope_id || ':' || role || ':' || added_by FROM memberships WHERE user_id = 'u'`); got != "project:p:member:admin@x" {
		t.Errorf("the person's project did not become a membership: %q", got)
	}
	if got := one(`SELECT COUNT(*) FROM memberships WHERE user_id = 'w'`); got != "0" {
		t.Error("a pending signup was given a membership")
	}
	var users []string
	rows, _ := db.QueryContext(ctx, `SELECT key_id || '=' || user_id FROM usage ORDER BY ts`)
	for rows.Next() {
		var s string
		_ = rows.Scan(&s)
		users = append(users, s)
	}
	rows.Close()
	if strings.Join(users, ",") != "k=u,session:u=u,ops=" {
		t.Errorf("usage attributed as %v", users)
	}
	// The rebuilt table takes a personal key, and still refuses a project
	// that does not exist.
	if _, err := db.ExecContext(ctx, `INSERT INTO api_keys (id, key_hash, key_prefix, name, created_at, user_id)
		VALUES ('pk', X'03', 'sk-3', 'mine', 1, 'u')`); err != nil {
		t.Errorf("a personal key was refused: %v", err)
	}
	if _, err := db.ExecContext(ctx, `INSERT INTO api_keys (id, key_hash, key_prefix, name, project_id, created_at)
		VALUES ('bad', X'04', 'sk-4', 'x', 'nope', 1)`); err == nil {
		t.Error("the rebuilt table lost its foreign key to projects")
	}
	fk, err := db.QueryContext(ctx, `PRAGMA foreign_key_check`)
	if err != nil {
		t.Fatal(err)
	}
	if fk.Next() {
		t.Error("foreign keys are violated after the migration")
	}
	fk.Close()
	// And they are back on for everything after it.
	if one(`PRAGMA foreign_keys`) != "1" {
		t.Error("the migration left foreign keys off")
	}
}