Avoid extraneous query in chunkById() methods as well

laravel · Oct 30, 2016 · 08e514c · 08e514c
1 parent 148564c
commit 08e514c
Show file tree

Hide file tree

Showing 4 changed files with 56 additions and 19 deletions.
diff --git a/src/Illuminate/Database/Eloquent/Builder.php b/src/Illuminate/Database/Eloquent/Builder.php
@@ -404,19 +404,22 @@ public function chunk($count, callable $callback)
      */
     public function chunkById($count, callable $callback, $column = 'id')
     {
-        $lastId = null;
+        $lastId = 0;
 
-        $results = $this->forPageAfterId($count, 0, $column)->get();
+        do {
+            $results = $this->forPageAfterId($count, $lastId, $column)->get();
+            $countResults = $results->count();
+
+            if ($countResults == 0) {
+                break;
+            }
 
-        while (! $results->isEmpty()) {
             if (call_user_func($callback, $results) === false) {
                 return false;
             }
 
             $lastId = $results->last()->{$column};
-
-            $results = $this->forPageAfterId($count, $lastId, $column)->get();
-        }
+        } while ($countResults == $count);
 
         return true;
     }

diff --git a/src/Illuminate/Database/Query/Builder.php b/src/Illuminate/Database/Query/Builder.php
@@ -1841,21 +1841,24 @@ public function chunk($count, callable $callback)
      */
     public function chunkById($count, callable $callback, $column = 'id', $alias = null)
     {
-        $lastId = null;
-
         $alias = $alias ?: $column;
 
-        $results = $this->forPageAfterId($count, 0, $column)->get();
+        $lastId = 0;
+
+        do {
+            $results = $this->forPageAfterId($count, $lastId, $column)->get();
+            $countResults = $results->count();
+
+            if ($countResults == 0) {
+                break;
+            }
 
-        while (! $results->isEmpty()) {
             if (call_user_func($callback, $results) === false) {
                 return false;
             }
 
             $lastId = $results->last()->{$alias};
-
-            $results = $this->forPageAfterId($count, $lastId, $column)->get();
-        }
+        } while ($countResults == $count);
 
         return true;
     }

diff --git a/tests/Database/DatabaseEloquentBuilderTest.php b/tests/Database/DatabaseEloquentBuilderTest.php
@@ -213,23 +213,38 @@ public function testChunkCanBeStoppedByReturningFalse()
         });
     }
 
-    public function testChunkPaginatesUsingId()
+    public function testChunkPaginatesUsingIdWithLastChunkComplete()
     {
         $builder = m::mock('Illuminate\Database\Eloquent\Builder[forPageAfterId,get]', [$this->getMockQueryBuilder()]);
         $builder->shouldReceive('forPageAfterId')->once()->with(2, 0, 'someIdField')->andReturn($builder);
         $builder->shouldReceive('forPageAfterId')->once()->with(2, 2, 'someIdField')->andReturn($builder);
-        $builder->shouldReceive('forPageAfterId')->once()->with(2, 10, 'someIdField')->andReturn($builder);
+        $builder->shouldReceive('forPageAfterId')->once()->with(2, 11, 'someIdField')->andReturn($builder);
 
         $builder->shouldReceive('get')->times(3)->andReturn(
             new Collection([(object) ['someIdField' => 1], (object) ['someIdField' => 2]]),
-            new Collection([(object) ['someIdField' => 10]]),
+            new Collection([(object) ['someIdField' => 10], (object) ['someIdField' => 11]]),
             new Collection([])
         );
 
         $builder->chunkById(2, function ($results) {
         }, 'someIdField');
     }
 
+    public function testChunkPaginatesUsingIdWithLastChunkPartial()
+    {
+        $builder = m::mock('Illuminate\Database\Eloquent\Builder[forPageAfterId,get]', [$this->getMockQueryBuilder()]);
+        $builder->shouldReceive('forPageAfterId')->once()->with(2, 0, 'someIdField')->andReturn($builder);
+        $builder->shouldReceive('forPageAfterId')->once()->with(2, 2, 'someIdField')->andReturn($builder);
+
+        $builder->shouldReceive('get')->times(2)->andReturn(
+            new Collection([(object) ['someIdField' => 1], (object) ['someIdField' => 2]]),
+            new Collection([(object) ['someIdField' => 10]])
+        );
+
+        $builder->chunkById(2, function ($results) {
+        }, 'someIdField');
+    }
+
     public function testPluckReturnsTheMutatedAttributesOfAModel()
     {
         $builder = $this->getBuilder();

diff --git a/tests/Database/DatabaseQueryBuilderTest.php b/tests/Database/DatabaseQueryBuilderTest.php
@@ -1734,24 +1734,40 @@ public function testTableValuedFunctionAsTableInSqlServer()
         $this->assertEquals('select * from [users](1,2)', $builder->toSql());
     }
 
-    public function testChunkPaginatesUsingId()
+    public function testChunkPaginatesUsingIdWithLastChunkComplete()
     {
         $builder = $this->getMockQueryBuilder();
 
         $builder->shouldReceive('forPageAfterId')->once()->with(2, 0, 'someIdField')->andReturn($builder);
         $builder->shouldReceive('forPageAfterId')->once()->with(2, 2, 'someIdField')->andReturn($builder);
-        $builder->shouldReceive('forPageAfterId')->once()->with(2, 10, 'someIdField')->andReturn($builder);
+        $builder->shouldReceive('forPageAfterId')->once()->with(2, 11, 'someIdField')->andReturn($builder);
 
         $builder->shouldReceive('get')->times(3)->andReturn(
             collect([(object) ['someIdField' => 1], (object) ['someIdField' => 2]]),
-            collect([(object) ['someIdField' => 10]]),
+            collect([(object) ['someIdField' => 10], (object) ['someIdField' => 11]]),
             collect([])
         );
 
         $builder->chunkById(2, function ($results) {
         }, 'someIdField');
     }
 
+    public function testChunkPaginatesUsingIdWithLastChunkPartial()
+    {
+        $builder = $this->getMockQueryBuilder();
+
+        $builder->shouldReceive('forPageAfterId')->once()->with(2, 0, 'someIdField')->andReturn($builder);
+        $builder->shouldReceive('forPageAfterId')->once()->with(2, 2, 'someIdField')->andReturn($builder);
+
+        $builder->shouldReceive('get')->times(2)->andReturn(
+            collect([(object) ['someIdField' => 1], (object) ['someIdField' => 2]]),
+            collect([(object) ['someIdField' => 10]])
+        );
+
+        $builder->chunkById(2, function ($results) {
+        }, 'someIdField');
+    }
+
     public function testChunkPaginatesUsingIdWithAlias()
     {
         $builder = $this->getMockQueryBuilder();