pizzasql/pkg/storage/pizzakv_integration_test.go

package storage

import (
	"bytes"
	"errors"
	"fmt"
	"os"
	"os/exec"
	"path/filepath"
	"strings"
	"sync"
	"sync/atomic"
	"testing"
	"time"

	"github.com/goccy/go-json"
)

func startPizzaKVTest(t testing.TB, binary, socket, database string) func() {
	t.Helper()
	cmd := exec.Command(binary, "-unix="+socket, "-path="+database)
	cmd.Stdout = os.Stderr
	cmd.Stderr = os.Stderr
	if err := cmd.Start(); err != nil {
		t.Fatalf("start PizzaKV: %v", err)
	}
	var once sync.Once
	stop := func() {
		once.Do(func() {
			_ = cmd.Process.Kill()
			_ = cmd.Wait()
		})
	}
	t.Cleanup(stop)
	return stop
}

func waitPizzaKVPool(t testing.TB, socket string) *KVPool {
	t.Helper()
	addr := "unix:" + socket
	deadline := time.Now().Add(10 * time.Second)
	for time.Now().Before(deadline) {
		pool, err := NewKVPool(addr, 2, 5*time.Second)
		if err == nil {
			return pool
		}
		time.Sleep(20 * time.Millisecond)
	}
	t.Fatal("PizzaKV did not become ready")
	return nil
}

func shortPizzaKVSocket(t testing.TB) string {
	t.Helper()
	dir, err := os.MkdirTemp("/tmp", "pkv-")
	if err != nil {
		t.Fatal(err)
	}
	t.Cleanup(func() { _ = os.RemoveAll(dir) })
	return filepath.Join(dir, "s")
}

func TestPizzaKVIntegration(t *testing.T) {
	binary := os.Getenv("PIZZAKV_BIN")
	if binary == "" {
		t.Skip("PIZZAKV_BIN is not set")
	}

	dir := t.TempDir()
	socket := shortPizzaKVSocket(t)
	database := filepath.Join(dir, "integration.pkvdb")
	stop := startPizzaKVTest(t, binary, socket, database)
	pool := waitPizzaKVPool(t, socket)

	schemas := NewSchemaManager(pool, "integration")
	tables := NewTableManager(pool, schemas, "integration")
	if err := schemas.CreateTable(&Schema{
		Name: "events",
		Columns: []Column{
			{Name: "id", Type: "INTEGER", Nullable: false, PrimaryKey: true},
			{Name: "payload", Type: "BLOB", Nullable: true},
		},
	}); err != nil {
		t.Fatalf("create table: %v", err)
	}

	for i := int64(1); i <= 20; i++ {
		payload := []byte{byte(i), 0, '|', '\r', '\n'}
		if err := tables.Insert("events", Row{"id": i, "payload": payload}); err != nil {
			t.Fatalf("insert %d: %v", i, err)
		}
	}
	largePayload := bytes.Repeat([]byte{0xab}, 2*1024*1024)
	if err := tables.Insert("events", Row{"id": int64(21), "payload": largePayload}); err != nil {
		t.Fatalf("insert large row: %v", err)
	}
	largeRows, err := tables.Select("events", func(row Row) bool { return row["id"] == int64(21) })
	if err != nil || len(largeRows) != 1 {
		t.Fatalf("scan large row: len=%d err=%v", len(largeRows), err)
	}
	if got, ok := largeRows[0]["payload"].([]byte); !ok || !bytes.Equal(got, largePayload) {
		t.Fatalf("large binary payload mismatch: len=%d type=%T", len(got), largeRows[0]["payload"])
	}
	rows, err := tables.SelectWithLimit("events", nil, 3, 2)
	if err != nil || len(rows) != 3 {
		t.Fatalf("limited select: len=%d err=%v", len(rows), err)
	}
	row, err := tables.GetByPK("events", "1")
	if err != nil {
		t.Fatalf("point read: %v", err)
	}
	if got, ok := row["payload"].([]byte); !ok || !bytes.Equal(got, []byte{1, 0, '|', '\r', '\n'}) {
		t.Fatalf("binary payload = %v (%T)", row["payload"], row["payload"])
	}
	if count, err := tables.CountFast("events"); err != nil || count != 21 {
		t.Fatalf("count = %d, err=%v", count, err)
	}

	if err := pool.Close(); err != nil {
		t.Fatalf("close pool: %v", err)
	}
	stop()
	startPizzaKVTest(t, binary, socket, database)
	pool = waitPizzaKVPool(t, socket)
	defer pool.Close()
	schemas = NewSchemaManager(pool, "integration")
	tables = NewTableManager(pool, schemas, "integration")
	row, err = tables.GetByPK("events", "1")
	if err != nil {
		t.Fatalf("point read after restart: %v", err)
	}
	if got, ok := row["payload"].([]byte); !ok || !bytes.Equal(got, []byte{1, 0, '|', '\r', '\n'}) {
		t.Fatalf("binary payload after restart = %v (%T)", row["payload"], row["payload"])
	}
	if count, err := tables.CountFast("events"); err != nil || count != 21 {
		t.Fatalf("count after restart = %d, err=%v", count, err)
	}
	if err := schemas.DropTable("events"); err != nil {
		t.Fatalf("drop table: %v", err)
	}
	if schemas.TableExists("events") {
		t.Fatal("table still exists")
	}
}

// A composite key must survive a real engine restart: the encoded key is the
// durable row key, so both reads and duplicate detection keep working after the
// schema and rows are reloaded from PizzaKV.
func TestPizzaKVCompositePrimaryKeyIntegration(t *testing.T) {
	binary := os.Getenv("PIZZAKV_BIN")
	if binary == "" {
		t.Skip("PIZZAKV_BIN is not set")
	}

	dir := t.TempDir()
	socket := shortPizzaKVSocket(t)
	database := filepath.Join(dir, "composite.pkvdb")
	stop := startPizzaKVTest(t, binary, socket, database)
	pool := waitPizzaKVPool(t, socket)

	schemas := NewSchemaManager(pool, "composite")
	tables := NewTableManager(pool, schemas, "composite")
	if err := schemas.CreateTable(&Schema{
		Name:              "pairs",
		PrimaryKeyColumns: []string{"a", "b"},
		Columns: []Column{
			{Name: "a", Type: "TEXT", Nullable: false},
			{Name: "b", Type: "INTEGER", Nullable: false},
			{Name: "v", Type: "TEXT", Nullable: true},
		},
	}); err != nil {
		t.Fatalf("create table: %v", err)
	}
	for _, row := range []Row{
		{"a": "x", "b": int64(1), "v": "one"},
		{"a": "x", "b": int64(2), "v": "two"},
	} {
		if err := tables.Insert("pairs", row); err != nil {
			t.Fatalf("insert %v: %v", row, err)
		}
	}

	if err := pool.Close(); err != nil {
		t.Fatalf("close pool: %v", err)
	}
	stop()
	startPizzaKVTest(t, binary, socket, database)
	pool = waitPizzaKVPool(t, socket)
	defer pool.Close()

	schemas = NewSchemaManager(pool, "composite")
	tables = NewTableManager(pool, schemas, "composite")
	schema, err := schemas.GetSchema("pairs")
	if err != nil {
		t.Fatalf("reload schema: %v", err)
	}
	if strings.Join(schema.PrimaryKeyList(), ",") != "a,b" {
		t.Fatalf("reloaded key = %v, want [a b]", schema.PrimaryKeyList())
	}
	row, err := tables.GetByPK("pairs", "1:x1:1")
	if err != nil || row["v"] != "one" {
		t.Fatalf("point read after restart = %v, err %v", row, err)
	}
	rows, err := tables.Select("pairs", nil)
	if err != nil || len(rows) != 2 {
		t.Fatalf("scan after restart = %v, err %v", rows, err)
	}
	var duplicate *PrimaryKeyViolationError
	if err := tables.Insert("pairs", Row{"a": "x", "b": int64(1), "v": "again"}); !errors.As(err, &duplicate) {
		t.Fatalf("duplicate after restart = %v, want PrimaryKeyViolationError", err)
	}
	if err := tables.Insert("pairs", Row{"a": "x", "b": int64(3), "v": "three"}); err != nil {
		t.Fatalf("insert after restart: %v", err)
	}
	if count, err := tables.CountFast("pairs"); err != nil || count != 3 {
		t.Fatalf("count after restart = %d, err %v", count, err)
	}
}

func TestPizzaKVCompareBatchIntegration(t *testing.T) {
	binary := os.Getenv("PIZZAKV_BIN")
	if binary == "" {
		t.Skip("PIZZAKV_BIN is not set")
	}

	dir := t.TempDir()
	socket := shortPizzaKVSocket(t)
	database := filepath.Join(dir, "compare.pkvdb")
	startPizzaKVTest(t, binary, socket, database)
	pool := waitPizzaKVPool(t, socket)
	defer pool.Close()

	var seededLSN uint64
	if err := pool.WithClient(func(c *KVClient) error {
		lsn, committed, err := c.CompareBatchWrite(
			[]CompareCheck{{Key: []byte("k"), LSN: 0}},
			[]BatchOp{{Op: batchPut, Key: []byte("k"), Value: []byte("v")}},
			nil,
		)
		if err != nil {
			return err
		}
		if !committed || lsn == 0 {
			return fmt.Errorf("absent check expected commit, committed=%v lsn=%d", committed, lsn)
		}
		seededLSN = lsn
		return nil
	}); err != nil {
		t.Fatalf("seed: %v", err)
	}

	if err := pool.WithClient(func(c *KVClient) error {
		lsn, committed, err := c.CompareBatchWrite(
			[]CompareCheck{{Key: []byte("k"), LSN: 0}},
			[]BatchOp{{Op: batchPut, Key: []byte("k"), Value: []byte("x")}},
			nil,
		)
		if err != nil {
			return err
		}
		if committed || lsn != 0 {
			return fmt.Errorf("expected conflict, committed=%v lsn=%d", committed, lsn)
		}
		return nil
	}); err != nil {
		t.Fatalf("stale conflict: %v", err)
	}

	const workers = 8
	var wins atomic.Int32
	var wg sync.WaitGroup
	errCh := make(chan error, workers)
	for i := 0; i < workers; i++ {
		wg.Add(1)
		go func() {
			defer wg.Done()
			if err := pool.WithClient(func(c *KVClient) error {
				lsn, committed, err := c.CompareBatchWrite(
					[]CompareCheck{{Key: []byte("k"), LSN: seededLSN}},
					[]BatchOp{{Op: batchPut, Key: []byte("k"), Value: []byte("winner")}},
					nil,
				)
				if err != nil {
					return err
				}
				if committed {
					if lsn <= seededLSN {
						return fmt.Errorf("winner lsn %d did not advance past %d", lsn, seededLSN)
					}
					wins.Add(1)
				}
				return nil
			}); err != nil {
				errCh <- err
			}
		}()
	}
	wg.Wait()
	close(errCh)
	for err := range errCh {
		if err != nil {
			t.Fatalf("concurrent: %v", err)
		}
	}
	if got := wins.Load(); got != 1 {
		t.Fatalf("exactly one transaction should win, got %d", got)
	}
}

func TestPizzaKVLegacyMigrationIntegration(t *testing.T) {
	binary := os.Getenv("PIZZAKV_BIN")
	if binary == "" {
		t.Skip("PIZZAKV_BIN is not set")
	}

	dir := t.TempDir()
	source := filepath.Join(dir, "legacy.db")
	destination := filepath.Join(dir, "legacy.pkvdb")
	socket := shortPizzaKVSocket(t)
	schema, err := json.Marshal(&Schema{
		Name:       "events",
		Columns:    []Column{{Name: "id", Type: "INTEGER", PrimaryKey: true}, {Name: "name", Type: "TEXT", Nullable: true}},
		PrimaryKey: "id",
		NextRowID:  2,
	})
	if err != nil {
		t.Fatal(err)
	}
	legacy := []byte(fmt.Sprintf(
		"W|legacy:_schema:events|%s\rW|legacy:_sys:tables|[\"events\"]\rW|legacy:_data:events:1|{\"id\":1,\"name\":\"old\",\"_rowid_\":1}\r",
		schema,
	))
	if err := os.WriteFile(source, legacy, 0o600); err != nil {
		t.Fatal(err)
	}
	if output, err := exec.Command(binary, "-migrate="+source, "-path="+destination).CombinedOutput(); err != nil {
		t.Fatalf("migrate legacy database: %v\n%s", err, output)
	}
	unchanged, err := os.ReadFile(source)
	if err != nil {
		t.Fatal(err)
	}
	if !bytes.Equal(unchanged, legacy) {
		t.Fatal("migration modified the legacy source")
	}

	startPizzaKVTest(t, binary, socket, destination)
	pool := waitPizzaKVPool(t, socket)
	defer pool.Close()
	schemas := NewSchemaManager(pool, "legacy")
	tables := NewTableManager(pool, schemas, "legacy")
	row, err := tables.GetByPK("events", "1")
	if err != nil || row["name"] != "old" {
		t.Fatalf("read migrated row: row=%v err=%v", row, err)
	}
	if err := tables.Insert("events", Row{"id": int64(2), "name": "new"}); err != nil {
		t.Fatalf("insert binary row after migration: %v", err)
	}
	rows, err := tables.Select("events", nil)
	if err != nil || len(rows) != 2 {
		t.Fatalf("mixed legacy/binary scan: len=%d err=%v", len(rows), err)
	}
}

func BenchmarkPizzaKVStorage(b *testing.B) {
	binary := os.Getenv("PIZZAKV_BIN")
	if binary == "" {
		b.Skip("PIZZAKV_BIN is not set")
	}

	dir := b.TempDir()
	socket := shortPizzaKVSocket(b)
	startPizzaKVTest(b, binary, socket, filepath.Join(dir, "benchmark.pkvdb"))
	pool := waitPizzaKVPool(b, socket)
	defer pool.Close()
	schemas := NewSchemaManager(pool, "benchmark")
	tables := NewTableManager(pool, schemas, "benchmark")
	if err := schemas.CreateTable(&Schema{
		Name: "events",
		Columns: []Column{
			{Name: "id", Type: "INTEGER", Nullable: false, PrimaryKey: true},
			{Name: "symbol", Type: "TEXT", Nullable: false},
			{Name: "price", Type: "REAL", Nullable: false},
		},
	}); err != nil {
		b.Fatal(err)
	}
	rows := make([]Row, 10_000)
	for i := range rows {
		rows[i] = Row{"id": int64(i + 1), "symbol": fmt.Sprintf("PIZZA-%03d", i%100), "price": float64(i) / 100}
	}
	if n, err := tables.InsertBulk("events", rows); err != nil || n != len(rows) {
		b.Fatalf("seed: n=%d err=%v", n, err)
	}

	b.Run("point_read", func(b *testing.B) {
		b.ReportAllocs()
		for b.Loop() {
			if _, err := tables.GetByPK("events", "5000"); err != nil {
				b.Fatal(err)
			}
		}
	})
	b.Run("limit_10", func(b *testing.B) {
		b.ReportAllocs()
		for b.Loop() {
			if _, err := tables.SelectWithLimit("events", nil, 10, 0); err != nil {
				b.Fatal(err)
			}
		}
	})
	b.Run("scan_10000", func(b *testing.B) {
		b.ReportAllocs()
		for b.Loop() {
			if rows, err := tables.Select("events", nil); err != nil || len(rows) != 10_000 {
				b.Fatalf("scan: len=%d err=%v", len(rows), err)
			}
		}
	})
}