Add limit support (#76)

Close #73

Author: baopham

README.md

@@ -112,6 +112,11 @@ $model->chunk(10, function ($records) {
 
     }
 });
+
+// limit
+// Use this with caution unless your limit is small.
+// DynamoDB has a limit of 1MB so if your limit is very big, the results will not be expected.
+$model->where('name', 'foo')->take(3)->get();
 ```
 
 * Or if you want to sync your DB table with a DynamoDb table, use trait `BaoPham\DynamoDb\ModelTrait`, it will call a `PutItem` after the model is saved.

src/DynamoDbQueryBuilder.php

@@ -9,6 +9,13 @@
 
 class DynamoDbQueryBuilder
 {
+    /**
+     * The maximum number of records to return.
+     *
+     * @var int
+     */
+    public $limit;
+
     /**
      * @var array
      */
@@ -33,6 +40,30 @@ class DynamoDbQueryBuilder
      */
     protected $lastEvaluatedKey;
 
+    /**
+     * Alias to set the "limit" value of the query.
+     *
+     * @param  int  $value
+     * @return DynamoDbQueryBuilder\static
+     */
+    public function take($value)
+    {
+        return $this->limit($value);
+    }
+
+    /**
+     * Set the "limit" value of the query.
+     *
+     * @param  int  $value
+     * @return $this
+     */
+    public function limit($value)
+    {
+        $this->limit = $value;
+
+        return $this;
+    }
+
     public function where($column, $operator = null, $value = null, $boolean = 'and')
     {
         if ($boolean != 'and') {
@@ -209,6 +240,10 @@ public function count()
 
     protected function getAll($columns = [], $limit = -1, $use_iterator = true)
     {
+        if ($limit === -1 && isset($this->limit)) {
+            $limit = $this->limit;
+        }
+
         if ($conditionValue = $this->conditionsContainKey()) {
             if ($this->conditionsAreExactSearch()) {
                 $item = $this->find($conditionValue, $columns);
@@ -237,14 +272,14 @@ protected function getAll($columns = [], $limit = -1, $use_iterator = true)
 
         // If the $where is not empty, we run getIterator.
         if (!empty($this->where)) {
-
             // Index key condition exists, then use Query instead of Scan.
             // However, Query only supports a few conditions.
             if ($index = $this->conditionsContainIndexKey()) {
                 $keysInfo = $index['keysInfo'];
                 $isCompositeKey = isset($keysInfo['range']);
                 $hashKeyOperator = array_get($this->where, $keysInfo['hash'] . '.ComparisonOperator');
                 $isValidQueryDynamoDbOperator = ComparisonOperator::isValidQueryDynamoDbOperator($hashKeyOperator);
+
                 if ($isValidQueryDynamoDbOperator && $isCompositeKey) {
                     $rangeKeyOperator = array_get($this->where, $keysInfo['range'] . '.ComparisonOperator');
                     $isValidQueryDynamoDbOperator = ComparisonOperator::isValidQueryDynamoDbOperator($rangeKeyOperator, true);
@@ -254,10 +289,13 @@ protected function getAll($columns = [], $limit = -1, $use_iterator = true)
                     $op = 'Query';
                     $query['IndexName'] = $index['name'];
                     $query['KeyConditions'][$keysInfo['hash']] = array_get($this->where, $keysInfo['hash']);
+
                     if ($isCompositeKey) {
                         $query['KeyConditions'][$keysInfo['range']] = array_get($this->where, $keysInfo['range']);
                     }
+
                     $nonKeyConditions = array_except($this->where, array_values($keysInfo));
+
                     if (!empty($nonKeyConditions)) {
                         $query['QueryFilter'] = $nonKeyConditions;
                     }
@@ -274,6 +312,10 @@ protected function getAll($columns = [], $limit = -1, $use_iterator = true)
 
         if ($use_iterator) {
             $iterator = $this->client->getIterator($op, $query);
+
+            if (isset($query['Limit'])) {
+                $iterator = new \LimitIterator($iterator, 0, $query['Limit']);
+            }
         } else {
             if ($op === 'Scan') {
                 $res = $this->client->scan($query);

tests/DynamoDbCompositeModelTest.php

@@ -185,7 +185,7 @@ public function testStaticMethods()
 
     public function testConditionContainingCompositeIndexKey()
     {
-        $fooItem = $this->seed([
+        $this->seed([
             'id' => ['S' => 'id1'],
             'id2' => ['S' => '2'],
             'name' => ['S' => 'Foo'],
@@ -234,7 +234,7 @@ public function testConditionContainingCompositeIndexKey()
 
     public function testConditionsDoNotContainAllCompositeKeys()
     {
-        $fooItem = $this->seed([
+        $this->seed([
             'id' => ['S' => 'id1'],
             'id2' => ['S' => '2'],
             'name' => ['S' => 'Foo'],

tests/DynamoDbModelTest.php

@@ -318,7 +318,7 @@ public function testDifferentQueries()
         $barQuery = $this->testModel->where('name', 'Bar');
 
         $this->assertEquals(1, $fooQuery->count());
-        $this->assertEquals(1, $fooQuery->count());
+        $this->assertEquals(1, $barQuery->count());
         $this->assertEquals($this->testModel->unmarshalItem($expectedFoo), $fooQuery->first()->toArray());
         $this->assertEquals($this->testModel->unmarshalItem($expectedBar), $barQuery->first()->toArray());
     }
@@ -332,9 +332,9 @@ public function testChunkScan()
         $iteration = 1;
 
         $this->testModel->chunk(2, function ($results) use (&$iteration) {
-            if ($iteration == 1) {
+            if ($iteration === 1) {
                 $this->assertEquals(2, count($results));
-            } else if ($iteration == 2) {
+            } else if ($iteration === 2) {
                 $this->assertEquals(1, count($results));
             }
 
@@ -415,6 +415,27 @@ public function testSerialize()
         $this->assertInstanceOf(get_class($this->testModel), $unserializedItems->first());
     }
 
+    public function testLimit()
+    {
+        $count = 10;
+        for ($i = 0; $i < $count; $i++) {
+            $this->seed(['name' => ['S' => 'foo']]);
+        }
+
+        $items = $this->testModel->take(3)->get();
+        $this->assertCount(3, $items->toArray());
+
+        $items = $this->testModel->limit(3)->get();
+        $this->assertCount(3, $items->toArray());
+
+        // Ensure "limit" is reset in new queries
+        $items = $this->testModel->get();
+        $this->assertCount($count, $items->toArray());
+
+        $items = $this->testModel->where('name', 'foo')->take(4)->get();
+        $this->assertCount(4, $items->toArray());
+    }
+
     protected function seed($attributes = [])
     {
         $item = [