github.com/zuoyebang/bitalostable@v1.0.1-0.20240229032404-e3b99a834294/sstable/writer_rangekey_test.go (about)

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