github.com/cockroachdb/pebble@v1.1.1-0.20240513155919-3622ade60459/sstable/writer_rangekey_test.go (about)

     1  package sstable
     2  
     3  import (
     4  	"bytes"
     5  	"crypto/rand"
     6  	"fmt"
     7  	"strings"
     8  	"testing"
     9  
    10  	"github.com/cockroachdb/datadriven"
    11  	"github.com/cockroachdb/errors"
    12  	"github.com/cockroachdb/pebble/internal/testkeys"
    13  	"github.com/cockroachdb/pebble/objstorage/objstorageprovider"
    14  	"github.com/cockroachdb/pebble/vfs"
    15  	"github.com/stretchr/testify/require"
    16  )
    17  
    18  func TestWriter_RangeKeys(t *testing.T) {
    19  	var r *Reader
    20  	defer func() {
    21  		if r != nil {
    22  			require.NoError(t, r.Close())
    23  		}
    24  	}()
    25  
    26  	buildFn := func(td *datadriven.TestData) (*Reader, error) {
    27  		mem := vfs.NewMem()
    28  		f, err := mem.Create("test")
    29  		if err != nil {
    30  			return nil, err
    31  		}
    32  
    33  		// Use a "suffix-aware" Comparer, that will sort suffix-values in
    34  		// descending order of timestamp, rather than in lexical order.
    35  		cmp := testkeys.Comparer
    36  		w := NewWriter(objstorageprovider.NewFileWritable(f), WriterOptions{
    37  			Comparer:    cmp,
    38  			TableFormat: TableFormatPebblev2,
    39  		})
    40  		for _, data := range strings.Split(td.Input, "\n") {
    41  			// Format. One of:
    42  			// - SET $START-$END $SUFFIX=$VALUE
    43  			// - UNSET $START-$END $SUFFIX
    44  			// - DEL $START-$END
    45  			parts := strings.Split(data, " ")
    46  			kind, startEnd := parts[0], parts[1]
    47  
    48  			startEndSplit := bytes.Split([]byte(startEnd), []byte("-"))
    49  
    50  			var start, end, suffix, value []byte
    51  			start, end = startEndSplit[0], startEndSplit[1]
    52  
    53  			switch kind {
    54  			case "SET":
    55  				sv := bytes.Split([]byte(parts[2]), []byte("="))
    56  				suffix, value = sv[0], sv[1]
    57  				err = w.RangeKeySet(start, end, suffix, value)
    58  			case "UNSET":
    59  				suffix = []byte(parts[2])
    60  				err = w.RangeKeyUnset(start, end, suffix)
    61  			case "DEL":
    62  				err = w.RangeKeyDelete(start, end)
    63  			default:
    64  				return nil, errors.Newf("unexpected key kind: %s", kind)
    65  			}
    66  			if err != nil {
    67  				return nil, err
    68  			}
    69  
    70  			// Scramble the bytes in each of the input arrays. This helps with
    71  			// flushing out subtle bugs due to byte slice re-use.
    72  			for _, slice := range [][]byte{start, end, suffix, value} {
    73  				_, _ = rand.Read(slice)
    74  			}
    75  		}
    76  
    77  		if err = w.Close(); err != nil {
    78  			return nil, err
    79  		}
    80  
    81  		f, err = mem.Open("test")
    82  		if err != nil {
    83  			return nil, err
    84  		}
    85  
    86  		r, err = newReader(f, ReaderOptions{Comparer: cmp})
    87  		if err != nil {
    88  			return nil, err
    89  		}
    90  
    91  		return r, nil
    92  	}
    93  
    94  	datadriven.RunTest(t, "testdata/writer_range_keys", func(t *testing.T, td *datadriven.TestData) string {
    95  		switch td.Cmd {
    96  		case "build":
    97  			if r != nil {
    98  				_ = r.Close()
    99  				r = nil
   100  			}
   101  
   102  			var err error
   103  			r, err = buildFn(td)
   104  			if err != nil {
   105  				return err.Error()
   106  			}
   107  
   108  			iter, err := r.NewRawRangeKeyIter()
   109  			if err != nil {
   110  				return err.Error()
   111  			}
   112  			defer iter.Close()
   113  
   114  			var buf bytes.Buffer
   115  			for s := iter.First(); s != nil; s = iter.Next() {
   116  				_, _ = fmt.Fprintf(&buf, "%s\n", s)
   117  			}
   118  			return buf.String()
   119  
   120  		default:
   121  			return fmt.Sprintf("unknown command: %s", td.Cmd)
   122  		}
   123  	})
   124  }