Support query inference with multiple entity managers

README.md

@@ -107,6 +107,25 @@ $kernel->boot();
 return $kernel->getContainer()->get('doctrine')->getManager();
 ```
 
+If your application uses multiple entity managers, return the Doctrine manager
+registry from the loader. PHPStan Doctrine will use it to pick the object manager
+that owns the entity being analysed:
+
+```php
+// tests/object-manager.php
+
+use App\Kernel;
+use Symfony\Component\Dotenv\Dotenv;
+
+require __DIR__ . '/../vendor/autoload.php';
+
+(new Dotenv())->bootEnv(__DIR__ . '/../.env');
+
+$kernel = new Kernel($_SERVER['APP_ENV'], (bool) $_SERVER['APP_DEBUG']);
+$kernel->boot();
+return $kernel->getContainer()->get('doctrine');
+```
+
 ## Query type inference
 
 This extension can infer the result type of DQL queries when an `objectManagerLoader` is provided.

src/Rules/Doctrine/ORM/DqlRule.php

@@ -58,19 +58,16 @@ public function processNode(Node $node, Scope $scope): array
 			return [];
 		}
 
-		$objectManager = $this->objectMetadataResolver->getObjectManager();
-		if ($objectManager === null) {
-			return [];
-		}
-		if (!$objectManager instanceof $entityManagerInterface) {
-			return [];
-		}
-
-		/** @var EntityManagerInterface $objectManager */
-		$objectManager = $objectManager;
-
 		$messages = [];
 		foreach ($dqls as $dql) {
+			$objectManager = $this->objectMetadataResolver->getObjectManagerForDql($dql->getValue());
+			if (!$objectManager instanceof $entityManagerInterface) {
+				continue;
+			}
+
+			/** @var EntityManagerInterface $objectManager */
+			$objectManager = $objectManager;
+
 			$query = $objectManager->createQuery($dql->getValue());
 			try {
 				$query->getAST();

src/Rules/Doctrine/ORM/QueryBuilderDqlRule.php

@@ -92,21 +92,18 @@ public function processNode(Node $node, Scope $scope): array
 			return [];
 		}
 
-		$objectManager = $this->objectMetadataResolver->getObjectManager();
-		if ($objectManager === null) {
-			return [];
-		}
-
 		$entityManagerInterface = 'Doctrine\ORM\EntityManagerInterface';
-		if (!$objectManager instanceof $entityManagerInterface) {
-			return [];
-		}
-
-		/** @var EntityManagerInterface $objectManager */
-		$objectManager = $objectManager;
 
 		$messages = [];
 		foreach ($dqls as $dql) {
+			$objectManager = $this->objectMetadataResolver->getObjectManagerForDql($dql->getValue());
+			if (!$objectManager instanceof $entityManagerInterface) {
+				continue;
+			}
+
+			/** @var EntityManagerInterface $objectManager */
+			$objectManager = $objectManager;
+
 			try {
 				$objectManager->createQuery($dql->getValue())->getAST();
 			} catch (QueryException $e) {

src/Type/Doctrine/CreateQueryDynamicReturnTypeExtension.php

@@ -89,7 +89,7 @@ public function getTypeFromMethodCall(
 			if ($type instanceof ConstantStringType) {
 				$queryString = $type->getValue();
 
-				$em = $this->objectMetadataResolver->getObjectManager();
+				$em = $this->objectMetadataResolver->getObjectManagerForDql($queryString);
 				if (!$em instanceof EntityManagerInterface) {
 					return new QueryType($queryString, null, null);
 				}

src/Type/Doctrine/GetRepositoryDynamicReturnTypeExtension.php

@@ -104,9 +104,10 @@ public function getTypeFromMethodCall(
 		}
 
 		$repositoryTypes = [];
+		$managerName = $this->getManagerName($scope, $methodCall->getArgs());
 		foreach ($objectNames as $objectName) {
 			try {
-				$repositoryClass = $this->getRepositoryClass($objectName, $defaultRepositoryClass);
+				$repositoryClass = $this->getRepositoryClass($objectName, $defaultRepositoryClass, $managerName);
 			} catch (\Doctrine\Persistence\Mapping\MappingException | MappingException | AnnotationException $e) {
 				return $this->getDefaultReturnType($scope, $methodCall->getArgs(), $methodReflection, $defaultRepositoryClass);
 			}
@@ -138,7 +139,24 @@ private function getDefaultReturnType(Scope $scope, array $args, MethodReflectio
 		return $defaultType;
 	}
 
-	private function getRepositoryClass(string $className, string $defaultRepositoryClass): string
+	/**
+	 * @param Arg[] $args
+	 */
+	private function getManagerName(Scope $scope, array $args): ?string
+	{
+		if (count($args) < 2) {
+			return null;
+		}
+
+		$managerNames = $scope->getType($args[1]->value)->getConstantStrings();
+		if (count($managerNames) !== 1) {
+			return null;
+		}
+
+		return $managerNames[0]->getValue();
+	}
+
+	private function getRepositoryClass(string $className, string $defaultRepositoryClass, ?string $managerName): string
 	{
 		if (!$this->reflectionProvider->hasClass($className)) {
 			return $defaultRepositoryClass;
@@ -149,12 +167,29 @@ private function getRepositoryClass(string $className, string $defaultRepository
 			return $defaultRepositoryClass;
 		}
 
+		if ($managerName !== null) {
+			$objectManager = $this->metadataResolver->getObjectManagerByName($managerName);
+			if ($objectManager !== null) {
+				$metadata = $objectManager->getClassMetadata($classReflection->getName());
+				$odmMetadataClass = 'Doctrine\ODM\MongoDB\Mapping\ClassMetadata';
+				if ($metadata instanceof $odmMetadataClass) {
+					/** @var ClassMetadata<object> $odmMetadata */
+					$odmMetadata = $metadata;
+					return $odmMetadata->customRepositoryClassName ?? $defaultRepositoryClass;
+				}
+
+				if ($metadata instanceof \Doctrine\ORM\Mapping\ClassMetadata) {
+					return $metadata->customRepositoryClassName ?? $defaultRepositoryClass;
+				}
+			}
+		}
+
 		$metadata = $this->metadataResolver->getClassMetadata($classReflection->getName());
 		if ($metadata !== null) {
 			return $metadata->customRepositoryClassName ?? $defaultRepositoryClass;
 		}
 
-		$objectManager = $this->metadataResolver->getObjectManager();
+		$objectManager = $this->metadataResolver->getObjectManagerForClass($classReflection->getName());
 		if ($objectManager === null) {
 			return $defaultRepositoryClass;
 		}

src/Type/Doctrine/ObjectMetadataResolver.php

@@ -7,14 +7,20 @@
 use Doctrine\ORM\EntityManagerInterface;
 use Doctrine\ORM\Mapping\ClassMetadata;
 use Doctrine\ORM\Mapping\MappingException;
+use Doctrine\Persistence\ManagerRegistry;
 use Doctrine\Persistence\ObjectManager;
 use PHPStan\Doctrine\Mapping\ClassMetadataFactory;
 use PHPStan\ShouldNotHappenException;
 use ReflectionException;
+use Throwable;
+use function array_merge;
 use function class_exists;
+use function count;
 use function is_file;
 use function is_readable;
 use function method_exists;
+use function preg_match_all;
+use function reset;
 use function sprintf;
 use const PHP_VERSION_ID;
 
@@ -23,8 +29,8 @@ final class ObjectMetadataResolver
 
 	private ?string $objectManagerLoader = null;
 
-	/** @var ObjectManager|false|null */
-	private $objectManager;
+	/** @var ObjectManager|ManagerRegistry|false|null */
+	private $objectManagerLoaderResult;
 
 	private ?ClassMetadataFactory $metadataFactory = null;
 
@@ -47,23 +53,99 @@ public function hasObjectManagerLoader(): bool
 	/** @api */
 	public function getObjectManager(): ?ObjectManager
 	{
-		if ($this->objectManager === false) {
+		$objectManagerLoaderResult = $this->getObjectManagerLoaderResult();
+		if (!$objectManagerLoaderResult instanceof ManagerRegistry) {
+			return $objectManagerLoaderResult;
+		}
+
+		return $objectManagerLoaderResult->getManager();
+	}
+
+	/**
+	 * @param class-string $className
+	 */
+	public function getObjectManagerForClass(string $className): ?ObjectManager
+	{
+		$objectManagerLoaderResult = $this->getObjectManagerLoaderResult();
+		if (!$objectManagerLoaderResult instanceof ManagerRegistry) {
+			return $objectManagerLoaderResult;
+		}
+
+		$objectManager = $objectManagerLoaderResult->getManagerForClass($className);
+		if ($objectManager instanceof ObjectManager) {
+			return $objectManager;
+		}
+
+		return $this->getObjectManager();
+	}
+
+	public function getObjectManagerByName(string $name): ?ObjectManager
+	{
+		$objectManagerLoaderResult = $this->getObjectManagerLoaderResult();
+		if (!$objectManagerLoaderResult instanceof ManagerRegistry) {
 			return null;
 		}
 
-		if ($this->objectManager !== null) {
-			return $this->objectManager;
+		try {
+			return $objectManagerLoaderResult->getManager($name);
+		} catch (Throwable $e) {
+			return null;
+		}
+	}
+
+	public function getObjectManagerForDql(string $dql): ?ObjectManager
+	{
+		$objectManagerLoaderResult = $this->getObjectManagerLoaderResult();
+		if (!$objectManagerLoaderResult instanceof ManagerRegistry) {
+			return $objectManagerLoaderResult;
+		}
+
+		preg_match_all('~\b(?:FROM|UPDATE)\s+([\\\\A-Za-z_][\\\\A-Za-z0-9_]*)(?:\s+|$)~i', $dql, $matches);
+		preg_match_all('~\bDELETE\s+(?:FROM\s+)?([\\\\A-Za-z_][\\\\A-Za-z0-9_]*)(?:\s+|$)~i', $dql, $deleteMatches);
+		foreach (array_merge($matches[1], $deleteMatches[1]) as $className) {
+			if (!class_exists($className)) {
+				continue;
+			}
+
+			$objectManager = $objectManagerLoaderResult->getManagerForClass($className);
+			if ($objectManager !== null) {
+				return $objectManager;
+			}
+		}
+
+		return $this->getObjectManager();
+	}
+
+	/**
+	 * @return ObjectManager|ManagerRegistry|null
+	 */
+	private function getObjectManagerLoaderResult()
+	{
+		if ($this->objectManagerLoaderResult === false) {
+			return null;
+		}
+
+		if ($this->objectManagerLoaderResult !== null) {
+			return $this->objectManagerLoaderResult;
 		}
 
 		if ($this->objectManagerLoader === null) {
-			$this->objectManager = false;
+			$this->objectManagerLoaderResult = false;
 
 			return null;
 		}
 
-		$this->objectManager = $this->loadObjectManager($this->objectManagerLoader);
+		$objectManagerLoaderResult = $this->loadObjectManager($this->objectManagerLoader);
+		if ($objectManagerLoaderResult instanceof ManagerRegistry) {
+			$objectManagers = $objectManagerLoaderResult->getManagers();
+			if (count($objectManagers) === 1) {
+				$objectManagerLoaderResult = reset($objectManagers);
+			}
+		}
+
+		$this->objectManagerLoaderResult = $objectManagerLoaderResult;
 
-		return $this->objectManager;
+		return $this->objectManagerLoaderResult;
 	}
 
 	public function isNativeLazyObjectsEnabled(): bool
@@ -99,7 +181,7 @@ public function isTransient(string $className): bool
 			return true;
 		}
 
-		$objectManager = $this->getObjectManager();
+		$objectManager = $this->getObjectManagerForClass($className);
 
 		try {
 			if ($objectManager === null) {
@@ -143,7 +225,7 @@ public function getClassMetadata(string $className): ?ClassMetadata
 			return null;
 		}
 
-		$objectManager = $this->getObjectManager();
+		$objectManager = $this->getObjectManagerForClass($className);
 
 		try {
 			if ($objectManager === null) {
@@ -172,7 +254,10 @@ public function getClassMetadata(string $className): ?ClassMetadata
 		return $ormMetadata;
 	}
 
-	private function loadObjectManager(string $objectManagerLoader): ?ObjectManager
+	/**
+	 * @return ObjectManager|ManagerRegistry|null
+	 */
+	private function loadObjectManager(string $objectManagerLoader)
 	{
 		if (!is_file($objectManagerLoader)) {
 			throw new ShouldNotHappenException(sprintf(

src/Type/Doctrine/QueryBuilder/QueryBuilderGetQueryDynamicReturnTypeExtension.php

@@ -191,7 +191,7 @@ public function getTypeFromMethodCall(
 
 	private function getQueryType(string $dql): Type
 	{
-		$em = $this->objectMetadataResolver->getObjectManager();
+		$em = $this->objectMetadataResolver->getObjectManagerForDql($dql);
 		if (!$em instanceof EntityManagerInterface) {
 			return new QueryType($dql, null);
 		}

tests/Type/Doctrine/MultipleEntityManagersQueryTypeInferenceTest.php

@@ -0,0 +1,35 @@
+<?php declare(strict_types = 1);
+
+namespace PHPStan\Type\Doctrine;
+
+use PHPStan\Testing\TypeInferenceTestCase;
+
+class MultipleEntityManagersQueryTypeInferenceTest extends TypeInferenceTestCase
+{
+
+	/** @return iterable<mixed> */
+	public function dataFileAsserts(): iterable
+	{
+		yield from $this->gatherAssertTypes(__DIR__ . '/data/QueryResult/multipleEntityManagers.php');
+	}
+
+	/**
+	 * @dataProvider dataFileAsserts
+	 * @param mixed ...$args
+	 */
+	public function testFileAsserts(
+		string $assertType,
+		string $file,
+		...$args
+	): void
+	{
+		$this->assertFileAsserts($assertType, $file, ...$args);
+	}
+
+	/** @return string[] */
+	public static function getAdditionalConfigFiles(): array
+	{
+		return [__DIR__ . '/data/QueryResult/config-multiple-entity-managers.neon'];
+	}
+
+}