Remove pipeline creation from query builder and add execution methods

Author: GromNaN

src/Query/AggregationBuilder.php

@@ -5,24 +5,28 @@
 namespace MongoDB\Laravel\Query;
 
 use Illuminate\Support\Collection as LaravelCollection;
+use Illuminate\Support\LazyCollection;
+use InvalidArgumentException;
+use Iterator;
 use MongoDB\Builder\BuilderEncoder;
 use MongoDB\Builder\Stage\FluentFactoryTrait;
 use MongoDB\Collection as MongoDBCollection;
+use MongoDB\Driver\CursorInterface;
 use MongoDB\Laravel\Collection as LaravelMongoDBCollection;
 
 use function array_replace;
 use function collect;
+use function sprintf;
+use function str_starts_with;
 
-final class AggregationBuilder
+class AggregationBuilder
 {
     use FluentFactoryTrait;
 
     public function __construct(
         private MongoDBCollection|LaravelMongoDBCollection $collection,
-        array $pipeline = [],
-        private array $options = [],
+        private readonly array $options = [],
     ) {
-        $this->pipeline = $pipeline;
     }
 
     /**
@@ -31,6 +35,10 @@ public function __construct(
      */
     public function addRawStage(string $operator, mixed $value): static
     {
+        if (! str_starts_with($operator, '$')) {
+            throw new InvalidArgumentException(sprintf('The stage name "%s" is invalid. It must start with a "$" sign.', $operator));
+        }
+
         $this->pipeline[] = [$operator => $value];
 
         return $this;
@@ -39,7 +47,42 @@ public function addRawStage(string $operator, mixed $value): static
     /**
      * Execute the aggregation pipeline and return the results.
      */
-    public function get(array $options = []): LaravelCollection
+    public function get(array $options = []): LaravelCollection|LazyCollection
+    {
+        $cursor = $this->execute($options);
+
+        return collect($cursor->toArray());
+    }
+
+    /**
+     * Execute the aggregation pipeline and return the results in a lazy collection.
+     */
+    public function cursor($options = []): LazyCollection
+    {
+        $cursor = $this->execute($options);
+
+        return LazyCollection::make(function () use ($cursor) {
+            foreach ($cursor as $item) {
+                yield $item;
+            }
+        });
+    }
+
+    /**
+     * Execute the aggregation pipeline and return the first result.
+     */
+    public function first(array $options = []): mixed
+    {
+        return (clone $this)
+            ->limit(1)
+            ->cursor($options)
+            ->first();
+    }
+
+    /**
+     * Execute the aggregation pipeline and return MongoDB cursor.
+     */
+    private function execute(array $options): CursorInterface&Iterator
     {
         $encoder = new BuilderEncoder();
         $pipeline = $encoder->encode($this->getPipeline());
@@ -50,6 +93,6 @@ public function get(array $options = []): LaravelCollection
             $options,
         );
 
-        return collect($this->collection->aggregate($pipeline, $options)->toArray());
+        return $this->collection->aggregate($pipeline, $options);
     }
 }

src/Query/Builder.php

@@ -21,10 +21,7 @@
 use MongoDB\BSON\ObjectID;
 use MongoDB\BSON\Regex;
 use MongoDB\BSON\UTCDateTime;
-use MongoDB\Builder\Accumulator;
-use MongoDB\Builder\Expression\FieldPath;
 use MongoDB\Builder\Stage\FluentFactoryTrait;
-use MongoDB\Builder\Variable;
 use MongoDB\Driver\Cursor;
 use Override;
 use RuntimeException;
@@ -284,59 +281,6 @@ public function dump(mixed ...$args)
         return $this;
     }
 
-    private function getAggregationBuilder(): AggregationBuilder
-    {
-        if (! trait_exists(FluentFactoryTrait::class)) {
-            throw new BadMethodCallException('Aggregation builder requires package mongodb/builder 0.2+');
-        }
-
-        $agg = new AggregationBuilder($this->collection, [], $this->options);
-
-        $wheres = $this->compileWheres();
-
-        if (count($wheres)) {
-            $agg->match(...$wheres);
-        }
-
-        // Distinct query
-        if ($this->distinct) {
-            // Return distinct results directly
-            $column = $this->columns[0] ?? '_id';
-
-            $agg->group(
-                _id: \MongoDB\Builder\Expression::fieldPath($column),
-                _document: Accumulator::first(Variable::root()),
-            );
-            $agg->replaceRoot(
-                newRoot: new FieldPath('_document'),
-            );
-        }
-
-        if ($this->orders) {
-            $agg->sort(...$this->orders);
-        }
-
-        if ($this->offset) {
-            $agg->skip($this->offset);
-        }
-
-        if ($this->limit) {
-            $agg->limit($this->limit);
-        }
-
-        // Normal query
-        // Add custom projections.
-        if ($this->projections || $this->columns) {
-            $columns = in_array('*', $this->columns) ? [] : $this->columns;
-            $projection = array_fill_keys($columns, true) + $this->projections;
-            if ($projection) {
-                $agg->project(...$projection);
-            }
-        }
-
-        return $agg;
-    }
-
     /**
      * Return the MongoDB query to be run in the form of an element array like ['method' => [arguments]].
      *
@@ -597,14 +541,23 @@ public function generateCacheKey()
     }
 
     /** @return ($function is null ? AggregationBuilder : mixed) */
-    public function aggregate($function = null, $columns = [])
+    public function aggregate($function = null, $columns = ['*'])
     {
         if ($function === null) {
+            if (! trait_exists(FluentFactoryTrait::class)) {
+                // This error will be unreachable when the mongodb/builder package will be merged into mongodb/mongodb
+                throw new BadMethodCallException('Aggregation builder requires package mongodb/builder 0.2+');
+            }
+
             if ($columns !== [] && $columns !== ['*']) {
                 throw new InvalidArgumentException('Columns cannot be specified to create an aggregation builder. Add a $project stage instead.');
             }
 
-            return $this->getAggregationBuilder();
+            if ($this->wheres) {
+                throw new BadMethodCallException('Aggregation builder does not support previous query-builder instructions. Use a $match stage instead.');
+            }
+
+            return new AggregationBuilder($this->collection, $this->options);
         }
 
         $this->aggregate = [

tests/Query/AggregationBuilderTest.php

@@ -6,6 +6,8 @@
 
 use DateTimeImmutable;
 use Illuminate\Support\Collection;
+use Illuminate\Support\LazyCollection;
+use InvalidArgumentException;
 use MongoDB\BSON\Document;
 use MongoDB\BSON\ObjectId;
 use MongoDB\BSON\UTCDateTime;
@@ -23,24 +25,25 @@ class AggregationBuilderTest extends TestCase
     public function tearDown(): void
     {
         User::truncate();
+
+        parent::tearDown();
     }
 
-    public function testCreateFromQueryBuilder(): void
+    public function testCreateAggregationBuilder(): void
     {
         User::insert([
             ['name' => 'John Doe', 'birthday' => new UTCDateTime(new DateTimeImmutable('1989-01-01'))],
             ['name' => 'Jane Doe', 'birthday' => new UTCDateTime(new DateTimeImmutable('1990-01-01'))],
         ]);
 
         // Create the aggregation pipeline from the query builder
-        $pipeline = User::where('name', 'John Doe')
-            ->limit(10)
-            ->offset(0)
-            ->aggregate();
+        $pipeline = User::aggregate();
 
         $this->assertInstanceOf(AggregationBuilder::class, $pipeline);
 
         $pipeline
+            ->match(name: 'John Doe')
+            ->limit(10)
             ->addFields(
                 // Requires MongoDB 5.0+
                 year: Expression::year(
@@ -67,13 +70,22 @@ public function testCreateFromQueryBuilder(): void
 
         // Execute the pipeline and validate the results
         $results = $pipeline->get();
-
         $this->assertInstanceOf(Collection::class, $results);
         $this->assertCount(1, $results);
         $this->assertInstanceOf(ObjectId::class, $results->first()['_id']);
         $this->assertSame('John Doe', $results->first()['name']);
         $this->assertIsInt($results->first()['year']);
         $this->assertArrayNotHasKey('birthday', $results->first());
+
+        // Execute the pipeline and validate the results in a lazy collection
+        $results = $pipeline->cursor();
+        $this->assertInstanceOf(LazyCollection::class, $results);
+
+        // Execute the pipeline and return the first result
+        $result = $pipeline->first();
+        $this->assertIsArray($result);
+        $this->assertInstanceOf(ObjectId::class, $result['_id']);
+        $this->assertSame('John Doe', $result['name']);
     }
 
     public function testAddRawStage(): void
@@ -95,47 +107,15 @@ public function testAddRawStage(): void
         $this->assertSamePipeline($expected, $pipeline->getPipeline());
     }
 
-    public function testDistinct(): void
+    public function testAddRawStageInvalid(): void
     {
-        User::insert([
-            ['name' => 'Jane Doe', 'birthday' => new UTCDateTime(new DateTimeImmutable('1991-01-01'))],
-            ['name' => 'John Doe', 'birthday' => new UTCDateTime(new DateTimeImmutable('1989-01-01'))],
-            ['name' => 'John Doe', 'birthday' => new UTCDateTime(new DateTimeImmutable('1990-01-01'))],
-        ]);
-
-        // Create the aggregation pipeline from the query builder
-        $pipeline = User::orderBy('name')
-            ->distinct('name')
-            ->select('name', 'birthday')
-            ->aggregate();
-
-        $expected = [
-            [
-                '$group' => [
-                    '_id' => '$name',
-                    '_document' => ['$first' => '$$ROOT'],
-                ],
-            ],
-            [
-                '$replaceRoot' => ['newRoot' => '$_document'],
-            ],
-            [
-                '$sort' => ['name' => 1],
-            ],
-            [
-                '$project' => ['birthday' => true, 'name' => true],
-            ],
-        ];
-
-        $this->assertSamePipeline($expected, $pipeline->getPipeline());
+        $collection = $this->createMock(MongoDBCollection::class);
 
-        $results = $pipeline->get();
+        $pipeline = new AggregationBuilder($collection);
 
-        $this->assertCount(2, $results);
-        $this->assertSame('Jane Doe', $results[0]['name']);
-        $this->assertSame('1991-01-01', $results[0]['birthday']->toDateTime()->format('Y-m-d'));
-        $this->assertSame('John Doe', $results[1]['name']);
-        $this->assertSame('1989-01-01', $results[1]['birthday']->toDateTime()->format('Y-m-d'));
+        $this->expectException(InvalidArgumentException::class);
+        $this->expectExceptionMessage('The stage name "match" is invalid. It must start with a "$" sign.');
+        $pipeline->addRawStage('match', ['name' => 'John Doe']);
     }
 
     private static function assertSamePipeline(array $expected, Pipeline $pipeline): void