From 8bce292019b7bacf2008e19b67b64925994a6394 Mon Sep 17 00:00:00 2001
From: Nathan VanBenschoten <nvanbenschoten@gmail.com>
Date: Mon, 30 Mar 2020 18:05:10 -0400
Subject: [PATCH] [WIP] sql: use GetRequest for point lookups

---
 pkg/sql/row/kv_batch_fetcher.go | 28 ++++++++++++++++++++++------
 pkg/workload/kv/kv.go           |  2 +-
 2 files changed, 23 insertions(+), 7 deletions(-)

diff --git a/pkg/sql/row/kv_batch_fetcher.go b/pkg/sql/row/kv_batch_fetcher.go
index 6d8ab11e7deb..cc59646409e4 100644
--- a/pkg/sql/row/kv_batch_fetcher.go
+++ b/pkg/sql/row/kv_batch_fetcher.go
@@ -283,12 +283,20 @@ func (f *txnKVFetcher) fetch(ctx context.Context) error {
 			ba.Requests[i].MustSetInner(&scans[i])
 		}
 	} else {
-		scans := make([]roachpb.ScanRequest, len(f.spans))
-		for i := range f.spans {
-			scans[i].SetSpan(f.spans[i])
-			scans[i].ScanFormat = roachpb.BATCH_RESPONSE
-			scans[i].KeyLocking = keyLocking
-			ba.Requests[i].MustSetInner(&scans[i])
+		if f.lockStr == sqlbase.ScanLockingStrength_FOR_SHARE {
+			scans := make([]roachpb.GetRequest, len(f.spans))
+			for i := range f.spans {
+				scans[i].Key = f.spans[i].Key
+				ba.Requests[i].MustSetInner(&scans[i])
+			}
+		} else {
+			scans := make([]roachpb.ScanRequest, len(f.spans))
+			for i := range f.spans {
+				scans[i].SetSpan(f.spans[i])
+				scans[i].ScanFormat = roachpb.BATCH_RESPONSE
+				scans[i].KeyLocking = keyLocking
+				ba.Requests[i].MustSetInner(&scans[i])
+			}
 		}
 	}
 	if cap(f.requestSpans) < len(f.spans) {
@@ -391,6 +399,14 @@ func (f *txnKVFetcher) nextBatch(
 				f.remainingBatches = t.BatchResponses[1:]
 			}
 			return true, t.Rows, batchResp, origSpan, nil
+		case *roachpb.GetResponse:
+			if t.Value != nil {
+				kvs = []roachpb.KeyValue{{
+					Key:   origSpan.Key,
+					Value: *t.Value,
+				}}
+			}
+			return true, kvs, batchResp, origSpan, nil
 		}
 	}
 	if f.fetchEnd {
diff --git a/pkg/workload/kv/kv.go b/pkg/workload/kv/kv.go
index 306d6542a6a0..d6762b90dd22 100644
--- a/pkg/workload/kv/kv.go
+++ b/pkg/workload/kv/kv.go
@@ -264,7 +264,7 @@ func (w *kv) Ops(urls []string, reg *histogram.Registry) (workload.QueryLoad, er
 			fmt.Fprintf(&buf, `(mod($%d, %d), $%d)`, i+1, w.shards, i+1)
 		}
 	}
-	buf.WriteString(`)`)
+	buf.WriteString(`) FOR SHARE`)
 	readStmtStr := buf.String()
 
 	// Write statement