From aa05c14a5b284557be17fcba2fc7cc734e97f7df Mon Sep 17 00:00:00 2001
From: Feek <fiachra@squareup.com>
Date: Sat, 6 Jun 2020 23:42:09 -0700
Subject: [PATCH] fix: relations should also be able to call query builder
 methods

---
 .../RelationReturnTypeProvider.php            |  6 +++++-
 tests/acceptance/EloquentBuilderTypes.feature | 19 +++++++++++++++++++
 .../acceptance/EloquentRelationTypes.feature  | 14 ++++++++++++++
 3 files changed, 38 insertions(+), 1 deletion(-)

diff --git a/src/ReturnTypeProvider/RelationReturnTypeProvider.php b/src/ReturnTypeProvider/RelationReturnTypeProvider.php
index f84baee..9ceeb5d 100644
--- a/src/ReturnTypeProvider/RelationReturnTypeProvider.php
+++ b/src/ReturnTypeProvider/RelationReturnTypeProvider.php
@@ -11,6 +11,7 @@ use Illuminate\Database\Eloquent\Relations\HasOne;
 use Illuminate\Database\Eloquent\Relations\HasOneOrMany;
 use Illuminate\Database\Eloquent\Relations\HasOneThrough;
 use Illuminate\Database\Eloquent\Relations\Relation;
+use Illuminate\Database\Query\Builder as QueryBuilder;
 use PhpParser\Node\Expr\MethodCall;
 use Psalm\CodeLocation;
 use Psalm\Context;
@@ -51,7 +52,10 @@ final class RelationReturnTypeProvider implements MethodReturnTypeProviderInterf
         // returns an instance of ITSELF, rather than the instance of the builder. That explains this nonsense
 
         // If this method name is on the builder object, proxy it over there
-        if ($source->getCodebase()->methods->methodExists(new MethodIdentifier(Builder::class, $method_name_lowercase))) {
+
+        if ($source->getCodebase()->methods->methodExists(new MethodIdentifier(Builder::class, $method_name_lowercase)) ||
+            $source->getCodebase()->methods->methodExists(new MethodIdentifier(QueryBuilder::class, $method_name_lowercase))
+        ) {
             if (!$template_type_parameters) {
                 return null;
             }
diff --git a/tests/acceptance/EloquentBuilderTypes.feature b/tests/acceptance/EloquentBuilderTypes.feature
index 97fc9c4..7fe4ca2 100644
--- a/tests/acceptance/EloquentBuilderTypes.feature
+++ b/tests/acceptance/EloquentBuilderTypes.feature
@@ -145,3 +145,22 @@ Feature: Eloquent Builder Types
     Then I see these errors
       | MixedInferredReturnType | Could not verify return type 'Illuminate\Database\Eloquent\Builder<User>' for UserRepository::test_failure |
       | MixedReturnStatement    | Could not infer a return type                                                                              |
+
+  Scenario: can call methods on underlying query builder
+    Given I have the following code
+    """
+    <?php declare(strict_types=1);
+
+    use Tests\Psalm\LaravelPlugin\Models\User;
+    use \Illuminate\Database\Eloquent\Builder;
+
+    /**
+    * @psalm-param Builder<User> $builder
+    * @psalm-return Builder<User>
+    */
+    function test(Builder $builder): Builder {
+      return $builder->orderBy('id', 'ASC');
+    }
+    """
+    When I run Psalm
+    Then I see no errors
diff --git a/tests/acceptance/EloquentRelationTypes.feature b/tests/acceptance/EloquentRelationTypes.feature
index 5c6e6af..0b7972f 100644
--- a/tests/acceptance/EloquentRelationTypes.feature
+++ b/tests/acceptance/EloquentRelationTypes.feature
@@ -214,3 +214,17 @@ Feature: Eloquent Relation Types
     """
     When I run Psalm
     Then I see no errors
+
+  Scenario: Relationships return themselves when the proxied method is a query builder method
+    Given I have the following code
+    """
+    /**
+    * @param HasOne<Phone> $relationship
+    * @psalm-return HasOne<Phone>
+    */
+    function test(HasOne $relationship): HasOne {
+      return $relationship->orderBy('id', 'ASC');
+    }
+    """
+    When I run Psalm
+    Then I see no errors