From 7b33ad39ea1fd07c0ac4483a060e93e5b488e0a3 Mon Sep 17 00:00:00 2001
From: Pauline Vos <pauline.vos@mongodb.com>
Date: Mon, 16 Jun 2025 15:12:53 +0200
Subject: [PATCH] Add annotations to query building classes

In order to help with auto-completion when users build queries from our
document models.
---
 src/Eloquent/Builder.php       |  5 ++++-
 src/Eloquent/DocumentModel.php | 23 ++++++++++++++++++++---
 tests/Models/Anniversary.php   |  4 ----
 tests/Models/HiddenAnimal.php  |  6 ------
 4 files changed, 24 insertions(+), 14 deletions(-)

diff --git a/src/Eloquent/Builder.php b/src/Eloquent/Builder.php
index 7dba7b7ab..72873f135 100644
--- a/src/Eloquent/Builder.php
+++ b/src/Eloquent/Builder.php
@@ -31,6 +31,9 @@
 /**
  * @method \MongoDB\Laravel\Query\Builder toBase()
  * @template TModel of Model
+ * @extends EloquentBuilder<TModel>
+ *
+ * @mixin \MongoDB\Laravel\Query\Builder
  */
 class Builder extends EloquentBuilder
 {
@@ -69,7 +72,7 @@ class Builder extends EloquentBuilder
     ];
 
     /**
-     * @return ($function is null ? AggregationBuilder : self)
+     * @return ($function is null ? AggregationBuilder : $this)
      *
      * @inheritdoc
      */
diff --git a/src/Eloquent/DocumentModel.php b/src/Eloquent/DocumentModel.php
index 965b1a444..200c2f455 100644
--- a/src/Eloquent/DocumentModel.php
+++ b/src/Eloquent/DocumentModel.php
@@ -50,6 +50,19 @@
 use function strlen;
 use function var_export;
 
+/**
+ * @method static QueryBuilder query()
+ * @method static self sole(array $columns)
+ * @method static bool truncate()
+ * @method static QueryBuilder where(Closure|string|array|Expression $column, mixed $operator = null, mixed $value = null, string $boolean = 'and')
+ * @method static QueryBuilder whereBetween(string $column, iterable $values, string $boolean = null, bool $not = null)
+ * @method static QueryBuilder whereIn(...$parameters, mixed $values, string $boolean = null, bool $not = null)
+ * @method static QueryBuilder whereNot(array|Closure|Expression|string $column, $operator = null, $value = null, $boolean = 'and')
+ * @method static QueryBuilder whereNotIn(...$parameters)
+ * @method static QueryBuilder whereNotNull(...$parameters)
+ * @method static QueryBuilder whereNull(string $value)
+ * @mixin Eloquent
+ */
 trait DocumentModel
 {
     use HybridRelations;
@@ -553,8 +566,12 @@ public function getParentRelation(): ?Relation
         return $this->parentRelation ?? null;
     }
 
-    /** @inheritdoc */
-    public function newEloquentBuilder($query)
+    /**
+     * @param $query
+     * @return Builder
+     */
+    #[Override]
+    public function newEloquentBuilder($query): Builder
     {
         return new Builder($query);
     }
@@ -566,7 +583,7 @@ public function qualifyColumn($column)
     }
 
     /** @inheritdoc */
-    protected function newBaseQueryBuilder()
+    protected function newBaseQueryBuilder(): QueryBuilder
     {
         $connection = $this->getConnection();
 
diff --git a/tests/Models/Anniversary.php b/tests/Models/Anniversary.php
index fb78c9a55..cb5813881 100644
--- a/tests/Models/Anniversary.php
+++ b/tests/Models/Anniversary.php
@@ -7,15 +7,11 @@
 use Illuminate\Database\Eloquent\Model;
 use MongoDB\Laravel\Eloquent\DocumentModel;
 use MongoDB\Laravel\Eloquent\Model as Eloquent;
-use MongoDB\Laravel\Query\Builder;
 
 /**
  * @property string $name
  * @property string $anniversary
  * @mixin Eloquent
- * @method static Builder create(...$values)
- * @method static Builder truncate()
- * @method static Eloquent sole(...$parameters)
  */
 class Anniversary extends Model
 {
diff --git a/tests/Models/HiddenAnimal.php b/tests/Models/HiddenAnimal.php
index 240238da0..f6217177c 100644
--- a/tests/Models/HiddenAnimal.php
+++ b/tests/Models/HiddenAnimal.php
@@ -6,17 +6,11 @@
 
 use Illuminate\Database\Eloquent\Model;
 use MongoDB\Laravel\Eloquent\DocumentModel;
-use MongoDB\Laravel\Eloquent\Model as Eloquent;
-use MongoDB\Laravel\Query\Builder;
 
 /**
  * @property string $name
  * @property string $country
  * @property bool $can_be_eaten
- * @mixin Eloquent
- * @method static Builder create(...$values)
- * @method static Builder truncate()
- * @method static Eloquent sole(...$parameters)
  */
 final class HiddenAnimal extends Model
 {