ceremonyclient/pebble/sstable/writer_rangekey_test.go

125 lines
2.8 KiB
Go
Raw Normal View History

2024-01-03 07:31:42 +00:00
package sstable
import (
"bytes"
"crypto/rand"
"fmt"
"strings"
"testing"
"github.com/cockroachdb/datadriven"
"github.com/cockroachdb/errors"
"github.com/cockroachdb/pebble/internal/testkeys"
"github.com/cockroachdb/pebble/objstorage/objstorageprovider"
"github.com/cockroachdb/pebble/vfs"
"github.com/stretchr/testify/require"
)
func TestWriter_RangeKeys(t *testing.T) {
var r *Reader
defer func() {
if r != nil {
require.NoError(t, r.Close())
}
}()
buildFn := func(td *datadriven.TestData) (*Reader, error) {
mem := vfs.NewMem()
f, err := mem.Create("test")
if err != nil {
return nil, err
}
// Use a "suffix-aware" Comparer, that will sort suffix-values in
// descending order of timestamp, rather than in lexical order.
cmp := testkeys.Comparer
w := NewWriter(objstorageprovider.NewFileWritable(f), WriterOptions{
Comparer: cmp,
TableFormat: TableFormatPebblev2,
})
for _, data := range strings.Split(td.Input, "\n") {
// Format. One of:
// - SET $START-$END $SUFFIX=$VALUE
// - UNSET $START-$END $SUFFIX
// - DEL $START-$END
parts := strings.Split(data, " ")
kind, startEnd := parts[0], parts[1]
startEndSplit := bytes.Split([]byte(startEnd), []byte("-"))
var start, end, suffix, value []byte
start, end = startEndSplit[0], startEndSplit[1]
switch kind {
case "SET":
sv := bytes.Split([]byte(parts[2]), []byte("="))
suffix, value = sv[0], sv[1]
err = w.RangeKeySet(start, end, suffix, value)
case "UNSET":
suffix = []byte(parts[2])
err = w.RangeKeyUnset(start, end, suffix)
case "DEL":
err = w.RangeKeyDelete(start, end)
default:
return nil, errors.Newf("unexpected key kind: %s", kind)
}
if err != nil {
return nil, err
}
// Scramble the bytes in each of the input arrays. This helps with
// flushing out subtle bugs due to byte slice re-use.
for _, slice := range [][]byte{start, end, suffix, value} {
_, _ = rand.Read(slice)
}
}
if err = w.Close(); err != nil {
return nil, err
}
f, err = mem.Open("test")
if err != nil {
return nil, err
}
r, err = newReader(f, ReaderOptions{Comparer: cmp})
if err != nil {
return nil, err
}
return r, nil
}
datadriven.RunTest(t, "testdata/writer_range_keys", func(t *testing.T, td *datadriven.TestData) string {
switch td.Cmd {
case "build":
if r != nil {
_ = r.Close()
r = nil
}
var err error
r, err = buildFn(td)
if err != nil {
return err.Error()
}
iter, err := r.NewRawRangeKeyIter()
if err != nil {
return err.Error()
}
defer iter.Close()
var buf bytes.Buffer
for s := iter.First(); s != nil; s = iter.Next() {
_, _ = fmt.Fprintf(&buf, "%s\n", s)
}
return buf.String()
default:
return fmt.Sprintf("unknown command: %s", td.Cmd)
}
})
}