PHPORM-238 Add support for withCount and other aggregations

src/Eloquent/Builder.php

@@ -7,23 +7,32 @@
 use Illuminate\Database\Eloquent\Builder as EloquentBuilder;
 use Illuminate\Database\Eloquent\Collection;
 use Illuminate\Database\Eloquent\Model;
+use Illuminate\Database\Eloquent\Relations\Relation;
+use Illuminate\Support\Str;
+use InvalidArgumentException;
 use MongoDB\BSON\Document;
 use MongoDB\Builder\Type\QueryInterface;
 use MongoDB\Builder\Type\SearchOperatorInterface;
 use MongoDB\Driver\CursorInterface;
 use MongoDB\Driver\Exception\WriteException;
 use MongoDB\Laravel\Connection;
+use MongoDB\Laravel\Eloquent\Model as DocumentModel;
 use MongoDB\Laravel\Helpers\QueriesRelationships;
 use MongoDB\Laravel\Query\AggregationBuilder;
+use MongoDB\Laravel\Relations\EmbedsOneOrMany;
+use MongoDB\Laravel\Relations\HasMany;
 use MongoDB\Model\BSONDocument;
 
 use function array_key_exists;
 use function array_merge;
 use function collect;
+use function count;
+use function explode;
 use function is_array;
 use function is_object;
 use function iterator_to_array;
 use function property_exists;
+use function sprintf;
 
 /**
  * @method \MongoDB\Laravel\Query\Builder toBase()
@@ -34,6 +43,9 @@ class Builder extends EloquentBuilder
     private const DUPLICATE_KEY_ERROR = 11000;
     use QueriesRelationships;
 
+    /** @var array{relation: Relation, function: string, constraints: array, column: string, alias: string}[] */
+    private array $withAggregate = [];
+
     /**
      * The methods that should be returned from query builder.
      *
@@ -294,6 +306,91 @@ public function createOrFirst(array $attributes = [], array $values = [])
         }
     }
 
+    public function withAggregate($relations, $column, $function = null)
+    {
+        if (empty($relations)) {
+            return $this;
+        }
+
+        $relations = is_array($relations) ? $relations : [$relations];
+
+        foreach ($this->parseWithRelations($relations) as $name => $constraints) {
+            $segments = explode(' ', $name);
+
+            $name = $segments[0];
+            $alias = (count($segments) === 3 && Str::lower($segments[1]) === 'as' ? $segments[2] : Str::snake($name) . '_' . $function);
+
+            $relation = $this->getRelationWithoutConstraints($name);
+
+            if (! DocumentModel::isDocumentModel($relation->getRelated())) {
+                throw new InvalidArgumentException('WithAggregate does not support hybrid relations');
+            }
+
+            if ($relation instanceof EmbedsOneOrMany) {
+                $subQuery = $this->newQuery();
+                $constraints($subQuery);
+                if ($subQuery->getQuery()->wheres) {
+                    // @see https://jira.mongodb.org/browse/PHPORM-292
+                    throw new InvalidArgumentException('Constraints are not supported for embedded relations');
+                }
+
+                switch ($function) {
+                    case 'count':
+                        $this->getQuery()->project([$alias => ['$size' => ['$ifNull' => ['$' . $name, []]]]]);
+                        break;
+                    case 'min':
+                    case 'max':
+                    case 'avg':
+                        $this->getQuery()->project([$alias => ['$' . $function => '$' . $name . '.' . $column]]);
+                        break;
+                    default:
+                        throw new InvalidArgumentException(sprintf('Invalid aggregate function "%s"', $function));
+                }
+            } else {
+                $this->withAggregate[$alias] = [
+                    'relation' => $relation,
+                    'function' => $function,
+                    'constraints' => $constraints,
+                    'column' => $column,
+                    'alias' => $alias,
+                ];
+            }
+        }
+
+        return $this;
+    }
+
+    public function eagerLoadRelations(array $models)
+    {
+        if ($this->withAggregate) {
+            $modelIds = collect($models)->pluck($this->model->getKeyName())->all();
+
+            foreach ($this->withAggregate as $withAggregate) {
+                if ($withAggregate['relation'] instanceof HasMany) {
+                    $results = $withAggregate['relation']->newQuery()
+                        ->where($withAggregate['constraints'])
+                        ->whereIn($withAggregate['relation']->getForeignKeyName(), $modelIds)
+                        ->groupBy($withAggregate['relation']->getForeignKeyName())
+                        ->aggregate($withAggregate['function'], [$withAggregate['column']]);
+
+                    foreach ($models as $model) {
+                        $value = $withAggregate['function'] === 'count' ? 0 : null;
+                        foreach ($results as $result) {
+                            if ($model->getKey() === $result->{$withAggregate['relation']->getForeignKeyName()}) {
+                                $value = $result->aggregate;
+                                break;
+                            }
+                        }
+
+                        $model->setAttribute($withAggregate['alias'], $value);
+                    }
+                }
+            }
+        }
+
+        return parent::eagerLoadRelations($models);
+    }
+
     /**
      * Add the "updated at" column to an array of values.
      * TODO Remove if https://github.com/laravel/framework/commit/6484744326531829341e1ff886cc9b628b20d73e

src/Query/Builder.php

@@ -357,7 +357,7 @@ public function toMql(): array
 
                     $aggregations = blank($this->aggregate['columns']) ? [] : $this->aggregate['columns'];
 
-                    if ($column === '*' && $function === 'count' && ! $this->groups) {
+                    if (in_array('*', $aggregations) && $function === 'count' && empty($group['_id'])) {
                         $options = $this->inheritConnectionOptions($this->options);
 
                         return ['countDocuments' => [$wheres, $options]];
@@ -611,7 +611,7 @@ public function aggregate($function = null, $columns = ['*'])
 
         $this->bindings['select'] = [];
 
-        $results = $this->get($columns);
+        $results = $this->get();
 
         // Once we have executed the query, we will reset the aggregate property so
         // that more select queries can be executed against the database without
@@ -650,6 +650,14 @@ public function aggregateByGroup(string $function, array $columns = ['*'])
         return $this->aggregate($function, $columns);
     }
 
+    public function count($columns = '*')
+    {
+        // Can be removed when available in Laravel: https://github.com/laravel/framework/pull/53209
+        $results = $this->aggregate(__FUNCTION__, Arr::wrap($columns));
+
+        return $results instanceof Collection ? $results : (int) $results;
+    }
+
     /** @inheritdoc */
     public function exists()
     {