wal_writer_test.go

package litestream_test

import (
	"bytes"
	"database/sql"
	"encoding/binary"
	"io"
	"os"
	"path/filepath"
	"testing"

	"github.com/benbjohnson/litestream"
	"github.com/benbjohnson/litestream/internal/testingutil"
	_ "github.com/mattn/go-sqlite3"
)

func TestWALWriter_Static(t *testing.T) {
	testDir := filepath.Join("testdata", "wal-writer", "static")
	tempDir := t.TempDir()

	// Read in WAL file generated by sqlite3
	buf, err := os.ReadFile(filepath.Join(testDir, "db-wal"))
	if err != nil {
		t.Fatal(err)
	}

	// Create new WAL file.
	if err := os.WriteFile(filepath.Join(tempDir, "db-wal"), nil, 0666); err != nil {
		t.Fatal(err)
	}

	w := litestream.NewWALWriter(filepath.Join(tempDir, "db-wal"), 0666, 4096)
	w.Salt0 = binary.BigEndian.Uint32(buf[16:])
	w.Salt1 = binary.BigEndian.Uint32(buf[20:])

	if err := w.Open(); err != nil {
		t.Fatal(err)
	} else if err := w.WriteHeader(); err != nil {
		t.Fatal(err)
	}

	for b := buf[litestream.WALHeaderSize:]; len(b) > 0; b = b[litestream.WALFrameHeaderSize+4096:] {
		pgno := binary.BigEndian.Uint32(b[0:])
		commit := binary.BigEndian.Uint32(b[4:])
		if err := w.WriteFrame(pgno, commit, b[litestream.WALFrameHeaderSize:][:4096]); err != nil {
			t.Fatal(err)
		}
	}

	if err := w.Close(); err != nil {
		t.Fatal(err)
	}

	// Read generated WAL and compare with original.
	if buf2, err := os.ReadFile(filepath.Join(tempDir, "db-wal")); err != nil {
		t.Fatal(err)
	} else if !bytes.Equal(buf, buf2) {
		t.Fatal("wal file mismatch")
	}
}

func TestWALWriter_Live(t *testing.T) {
	testDir := filepath.Join("testdata", "wal-writer", "live")
	tempDir := t.TempDir()

	// Copy DB file into temporary dir.
	testingutil.CopyFile(t, filepath.Join(testDir, "db"), filepath.Join(tempDir, "db"))

	// Open database.
	db, err := sql.Open("sqlite3", filepath.Join(tempDir, "db"))
	if err != nil {
		t.Fatal(err)
	}
	defer db.Close()

	// Verify that table is empty.
	var n int
	if err := db.QueryRow(`SELECT COUNT(*) FROM t`).Scan(&n); err != nil {
		t.Fatal(err)
	} else if got, want := n, 0; got != want {
		t.Fatalf("init: n=%d, want %d", got, want)
	}

	// Copy WAL file into place.
	testingutil.CopyFile(t, filepath.Join(testDir, "db-wal"), filepath.Join(tempDir, "db-wal"))

	// Invalidate both copies of the WAL index headers.
	f, err := os.OpenFile(filepath.Join(tempDir, "db-shm"), os.O_RDWR, 0666)
	if err != nil {
		t.Fatal(err)
	}
	defer f.Close()

	// Read index header.
	idx := make([]byte, 136)
	if _, err := io.ReadFull(f, idx); err != nil {
		t.Fatal(err)
	}

	// Invalidate "isInit" flags
	idx[12], idx[48+12] = 0, 0

	// Write header back into index.
	if _, err := f.Seek(0, io.SeekStart); err != nil {
		t.Fatal(err)
	} else if _, err := f.Write(idx); err != nil {
		t.Fatal(err)
	}

	// Verify that table now has one row.
	if err := db.QueryRow(`SELECT COUNT(*) FROM t`).Scan(&n); err != nil {
		t.Fatal(err)
	} else if got, want := n, 1; got != want {
		t.Fatalf("post-wal: n=%d, want %d", got, want)
	}
}