Merge pull request #26 from agorapulse/feature/limit-dynamodb-queries

limit DynamoDB queries and scans

gradle.properties

@@ -1,4 +1,4 @@
-version = 1.0.4.3
+version = 1.0.4.4
 micronautVersion = 1.0.4
 gruVersion = 0.7.1
 druVersion = 0.5.0

micronaut-aws-sdk/src/main/groovy/com/agorapulse/micronaut/aws/dynamodb/builder/DefaultQueryBuilder.groovy

@@ -81,6 +81,12 @@ class DefaultQueryBuilder<T> implements QueryBuilder<T> {
         return this
     }
 
+    @Override
+    QueryBuilder<T> limit(int max) {
+        this.max = max
+        return this
+    }
+
     DefaultQueryBuilder<T> offset(Object exclusiveStartHashKeyValue, Object exclusiveRangeStartKey) {
         this.exclusiveHashStartKey = exclusiveStartHashKeyValue
         this.exclusiveRangeStartKey = exclusiveRangeStartKey
@@ -94,7 +100,11 @@ class DefaultQueryBuilder<T> implements QueryBuilder<T> {
 
     @Override
     Flowable<T> query(IDynamoDBMapper mapper) {
-        return FlowableQueryResultHelper.generate(metadata.itemClass, mapper, resolveExpression(mapper))
+        Flowable<T> results = FlowableQueryResultHelper.generate(metadata.itemClass, mapper, resolveExpression(mapper))
+        if (max < Integer.MAX_VALUE) {
+            return results.take(max)
+        }
+        return results
     }
 
     @Override
@@ -164,5 +174,6 @@ class DefaultQueryBuilder<T> implements QueryBuilder<T> {
     private Object hashKey
     private Object exclusiveHashStartKey
     private Object exclusiveRangeStartKey
+    private int max = Integer.MAX_VALUE
     private Consumer<DynamoDBQueryExpression<T>> configurer = { } as Consumer<DynamoDBQueryExpression<T>>
 }

micronaut-aws-sdk/src/main/groovy/com/agorapulse/micronaut/aws/dynamodb/builder/DefaultScanBuilder.groovy

@@ -62,6 +62,12 @@ class DefaultScanBuilder<T> implements ScanBuilder<T> {
         return this
     }
 
+    @Override
+    ScanBuilder<T> limit(int max) {
+        this.max = max
+        return this
+    }
+
     DefaultScanBuilder<T> offset(Object exclusiveHashStartKeyValue, Object exclusiveRangeStartKeyValue) {
         this.exclusiveHashStartKey = exclusiveHashStartKeyValue
         this.exclusiveRangeStartKey = exclusiveRangeStartKeyValue
@@ -75,7 +81,11 @@ class DefaultScanBuilder<T> implements ScanBuilder<T> {
 
     @Override
     Flowable<T> scan(IDynamoDBMapper mapper) {
-        return FlowableQueryResultHelper.generate(metadata.itemClass, mapper, resolveExpression(mapper))
+        Flowable<T> results = FlowableQueryResultHelper.generate(metadata.itemClass, mapper, resolveExpression(mapper))
+        if (max < Integer.MAX_VALUE) {
+            return results.take(max)
+        }
+        return results
     }
 
     @Override
@@ -135,4 +145,5 @@ class DefaultScanBuilder<T> implements ScanBuilder<T> {
     private Consumer<DynamoDBScanExpression> configurer = { } as Consumer<DynamoDBScanExpression>
     private Object exclusiveHashStartKey
     private Object exclusiveRangeStartKey
+    private int max = Integer.MAX_VALUE
 }

micronaut-aws-sdk/src/main/groovy/com/agorapulse/micronaut/aws/dynamodb/builder/QueryBuilder.java

@@ -187,6 +187,16 @@ default QueryBuilder<T> and(
      */
     QueryBuilder<T> page(int page);
 
+    /**
+     * Sets the maximum number of items to be returned from the queries.
+     *
+     * This is a shortcut for calling <code>{@link io.reactivex.Flowable#take(long)}</code> on the result Flowable.
+     *
+     * @param max the maximum number of items returned
+     * @return self
+     */
+    QueryBuilder<T> limit(int max);
+
 
     /**
      * Sets the query offset by defining the exclusive start hash and range key (hash and range key of the last entity returned).

micronaut-aws-sdk/src/main/groovy/com/agorapulse/micronaut/aws/dynamodb/builder/ScanBuilder.java

@@ -132,6 +132,16 @@ default ScanBuilder<T> and(
      */
     ScanBuilder<T> page(int page);
 
+    /**
+     * Sets the maximum number of items to be returned from the queries.
+     *
+     * This is a shortcut for calling <code>{@link io.reactivex.Flowable#take(long)}</code> on the result Flowable.
+     *
+     * @param max the maximum number of items returned
+     * @return self
+     */
+    ScanBuilder<T> limit(int max);
+
     /**
      * Sets the scan offset by defining the exclusive start hash and range key (hash and range key of the last entity returned).
      * @param exclusiveStartKeyValue exclusive start key hash value

micronaut-aws-sdk/src/test/groovy/com/agorapulse/micronaut/aws/dynamodb/DefaultDynamoDBServiceSpec.groovy

@@ -248,9 +248,11 @@ class DefaultDynamoDBServiceSpec extends Specification {
             s.queryByRangeIndex('1', 'bar').blockingSingle().parentId == null // projection
             s.queryByRangeIndex('1', 'bar').blockingSingle().rangeIndex == 'bar' // projection
             s.queryByDates('1', REFERENCE_DATE.minusDays(1).toDate(), REFERENCE_DATE.plusDays(2).toDate()).count().blockingGet() == 2
+            s.queryByDatesWithLimit('1', REFERENCE_DATE.minusDays(1).toDate(), REFERENCE_DATE.plusDays(2).toDate(), 1).count().blockingGet() == 1
             s.queryByDates('3', REFERENCE_DATE.plusDays(9).toDate(), REFERENCE_DATE.plusDays(20).toDate()).count().blockingGet() == 1
 
             s.scanAllByRangeIndex('bar').count().blockingGet() == 4
+            s.scanAllByRangeIndexWithLimit('bar', 2).count().blockingGet() == 2
 
             s.increment('1001', '1')
             s.increment('1001', '1')
@@ -376,6 +378,15 @@ interface DynamoDBItemDBService {
     })
     Flowable<DynamoDBEntity> queryByDates(String hashKey, Date after, Date before)
 
+    @Query({
+        query(DynamoDBEntity) {
+            hash hashKey
+            range { between DynamoDBEntity.DATE_INDEX, after, before }
+            limit max
+        }
+    })
+    Flowable<DynamoDBEntity> queryByDatesWithLimit(String hashKey, Date after, Date before, int max)
+
     void delete(DynamoDBEntity entity)
     void delete(String hashKey, String rangeKey)
 
@@ -430,6 +441,16 @@ interface DynamoDBItemDBService {
     Flowable<DynamoDBEntity> scanAllByRangeIndex(String foo)                            // <5>
     // end::sample-scan[]
 
+    @Scan({
+        scan(DynamoDBEntity) {
+            filter {
+                eq DynamoDBEntity.RANGE_INDEX, foo
+            }
+            limit max
+        }
+    })
+    Flowable<DynamoDBEntity> scanAllByRangeIndexWithLimit(String foo, int max)
+
 // tag::service-footer[]
 }
 // end::service-footer[]