Introduce a DynamicMethodReturnTypeExtension for the DBAL QueryBuilder execute method

fixes #117

Author: mitelg

composer.json

@@ -20,6 +20,7 @@
 		"doctrine/annotations": "^1.11.0",
 		"doctrine/collections": "^1.0",
 		"doctrine/common": "^2.7 || ^3.0",
+		"doctrine/dbal": "^2.11.0",
 		"doctrine/mongodb-odm": "^1.3 || ^2.1",
 		"doctrine/orm": "^2.5",
 		"doctrine/persistence": "^1.1 || ^2.0",

extension.neon

@@ -212,6 +212,10 @@ services:
 		arguments:
 			class: Doctrine\ORM\Query\Expr
 			argumentsProcessor: @doctrineQueryBuilderArgumentsProcessor
+	-
+		class: PHPStan\Type\Doctrine\DBAL\QueryBuilder\QueryBuilderExecuteMethodExtension
+		tags:
+			- phpstan.broker.dynamicMethodReturnTypeExtension
 
 	# Type descriptors
 	-

src/Type/Doctrine/DBAL/QueryBuilder/QueryBuilderExecuteMethodExtension.php

@@ -0,0 +1,57 @@
+<?php declare(strict_types = 1);
+
+namespace PHPStan\Type\Doctrine\DBAL\QueryBuilder;
+
+use Doctrine\DBAL\Driver\ResultStatement;
+use Doctrine\DBAL\Query\QueryBuilder;
+use PhpParser\Node\Expr\MethodCall;
+use PhpParser\Node\Identifier;
+use PHPStan\Analyser\Scope;
+use PHPStan\Reflection\MethodReflection;
+use PHPStan\Reflection\ParametersAcceptorSelector;
+use PHPStan\Type\DynamicMethodReturnTypeExtension;
+use PHPStan\Type\ObjectType;
+use PHPStan\Type\Type;
+
+class QueryBuilderExecuteMethodExtension implements DynamicMethodReturnTypeExtension
+{
+
+	public function getClass(): string
+	{
+		return QueryBuilder::class;
+	}
+
+	public function isMethodSupported(MethodReflection $methodReflection): bool
+	{
+		return $methodReflection->getName() === 'execute';
+	}
+
+	public function getTypeFromMethodCall(MethodReflection $methodReflection, MethodCall $methodCall, Scope $scope): Type
+	{
+		$defaultReturnType = ParametersAcceptorSelector::selectSingle($methodReflection->getVariants())->getReturnType();
+
+		$queryBuilderType = new ObjectType(QueryBuilder::class);
+		$var = $methodCall->var;
+		while ($var instanceof MethodCall) {
+			$varType = $scope->getType($var->var);
+			if (!$queryBuilderType->isSuperTypeOf($varType)->yes()) {
+				return $defaultReturnType;
+			}
+
+			$nameObject = $var->name;
+			if (!($nameObject instanceof Identifier)) {
+				return $defaultReturnType;
+			}
+
+			$name = $nameObject->toString();
+			if ($name === 'select' || $name === 'addSelect') {
+				return new ObjectType(ResultStatement::class);
+			}
+
+			$var = $var->var;
+		}
+
+		return $defaultReturnType;
+	}
+
+}

tests/DoctrineIntegration/ORM/EntityManagerIntegrationTest.php

@@ -18,6 +18,7 @@ public function dataTopics(): array
 			['entityManagerMergeReturn'],
 			['customRepositoryUsage'],
 			['queryBuilder'],
+			['dbalQueryBuilderExecuteDynamicReturn'],
 		];
 	}
 

tests/DoctrineIntegration/ORM/data/dbalQueryBuilderExecuteDynamicReturn-7.json

@@ -0,0 +1,22 @@
+[
+    {
+        "message": "Cannot call method fetchAll() on Doctrine\\DBAL\\Driver\\ResultStatement|int.",
+        "line": 43,
+        "ignorable": true
+    },
+    {
+        "message": "Cannot call method fetchAll() on Doctrine\\DBAL\\Driver\\ResultStatement|int.",
+        "line": 60,
+        "ignorable": true
+    },
+    {
+        "message": "Cannot call method fetchAll() on Doctrine\\DBAL\\Driver\\ResultStatement|int.",
+        "line": 76,
+        "ignorable": true
+    },
+    {
+        "message": "Cannot call method fetchAll() on Doctrine\\DBAL\\Driver\\ResultStatement|int.",
+        "line": 87,
+        "ignorable": true
+    }
+]

tests/DoctrineIntegration/ORM/data/dbalQueryBuilderExecuteDynamicReturn.php

@@ -0,0 +1,89 @@
+<?php declare(strict_types=1);
+
+namespace PHPStan\DoctrineIntegration\ORM\DbalQueryBuilderExecuteDynamicReturn;
+
+use Doctrine\DBAL\Connection;
+use Doctrine\DBAL\Query\QueryBuilder;
+
+class Example
+{
+	/** @var Connection */
+	private $connection;
+
+	public function __construct(Connection $connection)
+	{
+		$this->connection = $connection;
+	}
+
+	/**
+	 * @return mixed[]
+	 */
+	public function testCaseOne(int $userId): array
+	{
+		return $this->connection->createQueryBuilder()
+			->select('*')
+			->from('user')
+			->where('user.id = :id')
+			->setParameter('id', $userId)
+			->execute()
+			->fetchAll();
+	}
+
+	/**
+	 * @return mixed[]
+	 */
+	public function testCaseTwo(int $userId): array
+	{
+		$qb = $this->connection->createQueryBuilder();
+		$qb->select('*');
+		$qb->from('user');
+		$qb->where('user.id = :id');
+		$qb->setParameter('id', $userId);
+
+		return $qb->execute()->fetchAll();
+	}
+
+	/**
+	 * @return mixed[]
+	 */
+	public function testCaseThree(?int $userId = null): array
+	{
+		$qb = $this->connection->createQueryBuilder();
+		$qb->select('*');
+		$qb->from('user');
+
+		if ($userId !== null) {
+			$qb->where('user.id = :id');
+			$qb->setParameter('id', $userId);
+		}
+
+		return $qb->execute()->fetchAll();
+	}
+
+	/**
+	 * @return mixed[]
+	 */
+	public function testCaseFourA(?int $userId = null): array
+	{
+		$qb = $this->connection->createQueryBuilder();
+		$qb->select('*');
+		$qb->from('user');
+
+		if ($userId !== null) {
+			return $this->testCaseFourB($qb, $userId);
+		}
+
+		return $qb->execute()->fetchAll();
+	}
+
+	/**
+	 * @return mixed[]
+	 */
+	private function testCaseFourB(QueryBuilder $qb, int $userId): array
+	{
+		$qb->where('user.id = :id');
+		$qb->setParameter('id', $userId);
+
+		return $qb->execute()->fetchAll();
+	}
+}