Merge branch 'master' into v2-dev

query.go

@@ -32,6 +32,7 @@ type Query struct {
 	consistent  bool
 	limit       int
 	searchLimit int32
+	reqLimit    int
 	order       *Order
 
 	subber
@@ -178,12 +179,19 @@ func (q *Query) Limit(limit int) *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).
 // Note: limit will be capped to MaxInt32 as that is the maximum number the DynamoDB API will accept.
 func (q *Query) SearchLimit(limit int) *Query {
 	q.searchLimit = int32(min(limit, math.MaxInt32))
 	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 {
@@ -310,6 +318,7 @@ type queryIter struct {
 	err    error
 	idx    int
 	n      int
+	reqs   int
 
 	// last item evaluated
 	last Item
@@ -363,6 +372,10 @@ func (itr *queryIter) Next(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
@@ -384,8 +397,12 @@ func (itr *queryIter) Next(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.Next(ctx, out)

scan.go

@@ -25,6 +25,7 @@ type Scan struct {
 	consistent  bool
 	limit       int
 	searchLimit int32
+	reqLimit    int
 
 	segment       int32
 	totalSegments int32
@@ -127,14 +128,21 @@ func (s *Scan) Limit(limit int) *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 int) *Scan {
 	s.searchLimit = int32(min(limit, math.MaxInt32))
 	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
@@ -323,6 +331,7 @@ type scanIter struct {
 	err    error
 	idx    int
 	n      int
+	reqs   int
 
 	// last item evaluated
 	last Item
@@ -377,6 +386,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
@@ -398,8 +411,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
 		}