github.com/cockroachdb/cockroach@v20.2.0-alpha.1+incompatible/pkg/kv/kvserver/batcheval/cmd_scan_test.go (about)

     1  // Copyright 2020 The Cockroach Authors.
     2  //
     3  // Use of this software is governed by the Business Source License
     4  // included in the file licenses/BSL.txt.
     5  //
     6  // As of the Change Date specified in that file, in accordance with
     7  // the Business Source License, use of this software will be governed
     8  // by the Apache License, Version 2.0, included in the file
     9  // licenses/APL.txt.
    10  
    11  package batcheval
    12  
    13  import (
    14  	"context"
    15  	"fmt"
    16  	"testing"
    17  
    18  	"github.com/cockroachdb/cockroach/pkg/roachpb"
    19  	"github.com/cockroachdb/cockroach/pkg/storage"
    20  	"github.com/cockroachdb/cockroach/pkg/testutils"
    21  	"github.com/cockroachdb/cockroach/pkg/util/hlc"
    22  	"github.com/cockroachdb/cockroach/pkg/util/leaktest"
    23  	"github.com/stretchr/testify/require"
    24  )
    25  
    26  func TestScanReverseScanTargetBytes(t *testing.T) {
    27  	defer leaktest.AfterTest(t)()
    28  
    29  	// Sanity checks for the TargetBytes scan option. We're not checking the specifics here, just
    30  	// that the plumbing works. TargetBytes is tested in-depth via TestMVCCHistories.
    31  
    32  	const (
    33  		tbNone = 0      // no limit, i.e. should return all kv pairs
    34  		tbOne  = 1      // one byte = return first key only
    35  		tbLots = 100000 // de facto ditto tbNone
    36  	)
    37  	testutils.RunTrueAndFalse(t, "reverse", func(t *testing.T, reverse bool) {
    38  		for _, tb := range []int64{tbNone, tbOne, tbLots} {
    39  			t.Run(fmt.Sprintf("targetBytes=%d", tb), func(t *testing.T) {
    40  				for _, sf := range []roachpb.ScanFormat{roachpb.KEY_VALUES, roachpb.BATCH_RESPONSE} {
    41  					t.Run(fmt.Sprintf("format=%s", sf), func(t *testing.T) {
    42  						testScanReverseScanInner(t, tb, sf, reverse, tb != tbOne)
    43  					})
    44  				}
    45  			})
    46  		}
    47  	})
    48  }
    49  
    50  func testScanReverseScanInner(
    51  	t *testing.T, tb int64, sf roachpb.ScanFormat, reverse bool, expBoth bool,
    52  ) {
    53  	defer leaktest.AfterTest(t)()
    54  
    55  	ctx := context.Background()
    56  	k1, k2 := roachpb.Key("a"), roachpb.Key("b")
    57  	ts := hlc.Timestamp{WallTime: 1}
    58  
    59  	eng := storage.NewDefaultInMem()
    60  	defer eng.Close()
    61  
    62  	// Write to k1 and k2.
    63  	for _, k := range []roachpb.Key{k1, k2} {
    64  		err := storage.MVCCPut(ctx, eng, nil, k, ts, roachpb.MakeValueFromString("value-"+string(k)), nil)
    65  		require.NoError(t, err)
    66  	}
    67  
    68  	var req roachpb.Request
    69  	var resp roachpb.Response
    70  	if !reverse {
    71  		req = &roachpb.ScanRequest{ScanFormat: sf}
    72  		resp = &roachpb.ScanResponse{}
    73  	} else {
    74  		req = &roachpb.ReverseScanRequest{ScanFormat: sf}
    75  		resp = &roachpb.ReverseScanResponse{}
    76  	}
    77  	req.SetHeader(roachpb.RequestHeader{Key: k1, EndKey: roachpb.KeyMax})
    78  
    79  	cArgs := CommandArgs{
    80  		Args: req,
    81  		Header: roachpb.Header{
    82  			Timestamp:   ts,
    83  			TargetBytes: tb,
    84  		},
    85  	}
    86  
    87  	if !reverse {
    88  		_, err := Scan(ctx, eng, cArgs, resp)
    89  		require.NoError(t, err)
    90  	} else {
    91  		_, err := ReverseScan(ctx, eng, cArgs, resp)
    92  		require.NoError(t, err)
    93  	}
    94  	expN := 1
    95  	if expBoth {
    96  		expN = 2
    97  	}
    98  
    99  	require.EqualValues(t, expN, resp.Header().NumKeys)
   100  	require.NotZero(t, resp.Header().NumBytes)
   101  
   102  	var rows []roachpb.KeyValue
   103  	if !reverse {
   104  		rows = resp.(*roachpb.ScanResponse).Rows
   105  	} else {
   106  		rows = resp.(*roachpb.ReverseScanResponse).Rows
   107  	}
   108  
   109  	if rows != nil {
   110  		require.Len(t, rows, expN)
   111  	}
   112  }