github.com/cockroachdb/cockroach@v20.2.0-alpha.1+incompatible/pkg/sql/row/kv_fetcher.go (about)

     1  // Copyright 2018 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 row
    12  
    13  import (
    14  	"context"
    15  
    16  	"github.com/cockroachdb/cockroach/pkg/kv"
    17  	"github.com/cockroachdb/cockroach/pkg/roachpb"
    18  	"github.com/cockroachdb/cockroach/pkg/sql/sqlbase"
    19  	"github.com/cockroachdb/cockroach/pkg/storage/enginepb"
    20  )
    21  
    22  // KVFetcher wraps kvBatchFetcher, providing a NextKV interface that returns the
    23  // next kv from its input.
    24  type KVFetcher struct {
    25  	kvBatchFetcher
    26  
    27  	kvs []roachpb.KeyValue
    28  
    29  	batchResponse []byte
    30  	bytesRead     int64
    31  	Span          roachpb.Span
    32  	newSpan       bool
    33  }
    34  
    35  // NewKVFetcher creates a new KVFetcher.
    36  func NewKVFetcher(
    37  	txn *kv.Txn,
    38  	spans roachpb.Spans,
    39  	reverse bool,
    40  	useBatchLimit bool,
    41  	firstBatchLimit int64,
    42  	lockStr sqlbase.ScanLockingStrength,
    43  	returnRangeInfo bool,
    44  ) (*KVFetcher, error) {
    45  	kvBatchFetcher, err := makeKVBatchFetcher(
    46  		txn, spans, reverse, useBatchLimit, firstBatchLimit, lockStr, returnRangeInfo,
    47  	)
    48  	return newKVFetcher(&kvBatchFetcher), err
    49  }
    50  
    51  func newKVFetcher(batchFetcher kvBatchFetcher) *KVFetcher {
    52  	return &KVFetcher{
    53  		kvBatchFetcher: batchFetcher,
    54  	}
    55  }
    56  
    57  // NextKV returns the next kv from this fetcher. Returns false if there are no
    58  // more kvs to fetch, the kv that was fetched, and any errors that may have
    59  // occurred.
    60  func (f *KVFetcher) NextKV(
    61  	ctx context.Context,
    62  ) (ok bool, kv roachpb.KeyValue, newSpan bool, err error) {
    63  	for {
    64  		newSpan = f.newSpan
    65  		f.newSpan = false
    66  		if len(f.kvs) != 0 {
    67  			kv = f.kvs[0]
    68  			f.kvs = f.kvs[1:]
    69  			return true, kv, newSpan, nil
    70  		}
    71  		if len(f.batchResponse) > 0 {
    72  			var key []byte
    73  			var rawBytes []byte
    74  			var err error
    75  			key, rawBytes, f.batchResponse, err = enginepb.ScanDecodeKeyValueNoTS(f.batchResponse)
    76  			if err != nil {
    77  				return false, kv, false, err
    78  			}
    79  			return true, roachpb.KeyValue{
    80  				Key: key,
    81  				Value: roachpb.Value{
    82  					RawBytes: rawBytes,
    83  				},
    84  			}, newSpan, nil
    85  		}
    86  
    87  		ok, f.kvs, f.batchResponse, f.Span, err = f.nextBatch(ctx)
    88  		if err != nil {
    89  			return ok, kv, false, err
    90  		}
    91  		if !ok {
    92  			return false, kv, false, nil
    93  		}
    94  		f.newSpan = true
    95  		f.bytesRead += int64(len(f.batchResponse))
    96  	}
    97  }