From d5fb45284a7660b407ce466c029efe1e9748f967 Mon Sep 17 00:00:00 2001
From: Greg <greg@toki.waseda.jp>
Date: Sun, 5 May 2024 06:28:40 +0900
Subject: [PATCH] add RequestLimit to Query and Scan (#231)

---
 query.go | 17 +++++++++++++++++
 scan.go  | 19 ++++++++++++++++++-
 2 files changed, 35 insertions(+), 1 deletion(-)

diff --git a/query.go b/query.go
index 2cfd6fc..a4593dc 100644
--- a/query.go
+++ b/query.go
@@ -31,6 +31,7 @@ type Query struct {
 	consistent  bool
 	limit       int64
 	searchLimit int64
+	reqLimit    int
 	order       *Order
 
 	subber
@@ -177,11 +178,18 @@ func (q *Query) Limit(limit int64) *Query {
 // SearchLimit specifies the maximum amount of results to examine.
 // If a filter is not specified, the number of results will be limited.
 // If a filter is specified, the number of results to consider for filtering will be limited.
+// SearchLimit > 0 implies RequestLimit(1).
 func (q *Query) SearchLimit(limit int64) *Query {
 	q.searchLimit = limit
 	return q
 }
 
+// RequestLimit specifies the maximum amount of requests to make against DynamoDB's API.
+func (q *Query) RequestLimit(limit int) *Query {
+	q.reqLimit = limit
+	return q
+}
+
 // Order specifies the desired result order.
 // Requires a range key (a.k.a. partition key) to be specified.
 func (q *Query) Order(order Order) *Query {
@@ -320,6 +328,7 @@ type queryIter struct {
 	err    error
 	idx    int
 	n      int64
+	reqs   int
 
 	// last item evaluated
 	last map[string]*dynamodb.AttributeValue
@@ -379,6 +388,10 @@ func (itr *queryIter) NextWithContext(ctx context.Context, out interface{}) bool
 		if itr.output.LastEvaluatedKey == nil || itr.query.searchLimit > 0 {
 			return false
 		}
+		// have we hit the request limit?
+		if itr.query.reqLimit > 0 && itr.reqs == itr.query.reqLimit {
+			return false
+		}
 
 		// no, prepare next request and reset index
 		itr.input.ExclusiveStartKey = itr.output.LastEvaluatedKey
@@ -400,8 +413,12 @@ func (itr *queryIter) NextWithContext(ctx context.Context, out interface{}) bool
 	if len(itr.output.LastEvaluatedKey) > len(itr.exLEK) {
 		itr.exLEK = itr.output.LastEvaluatedKey
 	}
+	itr.reqs++
 
 	if len(itr.output.Items) == 0 {
+		if itr.query.reqLimit > 0 && itr.reqs == itr.query.reqLimit {
+			return false
+		}
 		if itr.output.LastEvaluatedKey != nil {
 			// we need to retry until we get some data
 			return itr.NextWithContext(ctx, out)
diff --git a/scan.go b/scan.go
index e98cd6e..449e7be 100644
--- a/scan.go
+++ b/scan.go
@@ -22,6 +22,7 @@ type Scan struct {
 	consistent  bool
 	limit       int64
 	searchLimit int64
+	reqLimit    int
 
 	segment       int64
 	totalSegments int64
@@ -120,14 +121,21 @@ func (s *Scan) Limit(limit int64) *Scan {
 	return s
 }
 
-// SearchLimit specifies a maximum amount of results to evaluate.
+// SearchLimit specifies the maximum amount of results to evaluate.
 // Use this along with StartFrom and Iter's LastEvaluatedKey to split up results.
 // Note that DynamoDB limits result sets to 1MB.
+// SearchLimit > 0 implies RequestLimit(1).
 func (s *Scan) SearchLimit(limit int64) *Scan {
 	s.searchLimit = limit
 	return s
 }
 
+// RequestLimit specifies the maximum amount of requests to make against DynamoDB's API.
+func (s *Scan) RequestLimit(limit int) *Scan {
+	s.reqLimit = limit
+	return s
+}
+
 // ConsumedCapacity will measure the throughput capacity consumed by this operation and add it to cc.
 func (s *Scan) ConsumedCapacity(cc *ConsumedCapacity) *Scan {
 	s.cc = cc
@@ -335,6 +343,7 @@ type scanIter struct {
 	err    error
 	idx    int
 	n      int64
+	reqs   int
 
 	// last item evaluated
 	last map[string]*dynamodb.AttributeValue
@@ -395,6 +404,10 @@ redo:
 		if itr.output.LastEvaluatedKey == nil || itr.scan.searchLimit > 0 {
 			return false
 		}
+		// have we hit the request limit?
+		if itr.scan.reqLimit > 0 && itr.reqs == itr.scan.reqLimit {
+			return false
+		}
 
 		// no, prepare next request and reset index
 		itr.input.ExclusiveStartKey = itr.output.LastEvaluatedKey
@@ -416,8 +429,12 @@ redo:
 	if len(itr.output.LastEvaluatedKey) > len(itr.exLEK) {
 		itr.exLEK = itr.output.LastEvaluatedKey
 	}
+	itr.reqs++
 
 	if len(itr.output.Items) == 0 {
+		if itr.scan.reqLimit > 0 && itr.reqs == itr.scan.reqLimit {
+			return false
+		}
 		if itr.output.LastEvaluatedKey != nil {
 			goto redo
 		}