Allow service repositories to be used with more than one manager #1044

Repository/ContainerRepositoryFactory.php

@@ -54,6 +54,11 @@ public function getRepository(EntityManagerInterface $entityManager, $entityName
                     throw new RuntimeException(sprintf('The service "%s" must implement ObjectRepository (or extend a base class, like ServiceEntityRepository).', $repositoryServiceId));
                 }
 
+                // Use the correct manager when supported
+                if (method_exists($repository, 'withManager')) {
+                    return $repository->withManager($entityManager);
+                }
+
                 return $repository;
             }
 

Repository/ServiceEntityRepository.php

@@ -2,8 +2,12 @@
 
 namespace Doctrine\Bundle\DoctrineBundle\Repository;
 
+use Doctrine\Bundle\DoctrineBundle\Repository\ServiceEntityRepositoryInterface;
 use Doctrine\Common\Persistence\ManagerRegistry;
+use Doctrine\ORM\EntityManager;
 use Doctrine\ORM\EntityRepository;
+use Doctrine\Bundle\DoctrineBundle\Registry;
+use InvalidArgumentException;
 use LogicException;
 
 /**
@@ -14,28 +18,103 @@
  *
  * class YourEntityRepository extends ServiceEntityRepository
  * {
- *     public function __construct(RegistryInterface $registry)
+ *     public function __construct(RegistryInterface $registry, ...$args)
  *     {
- *         parent::__construct($registry, YourEntity::class);
+ *         parent::__construct($registry, YourEntity::class, func_get_args());
  *     }
  * }
  */
 class ServiceEntityRepository extends EntityRepository implements ServiceEntityRepositoryInterface
 {
+    /**
+     * All the arguments used by the constructor of the user-defined service repository
+     * 
+     * @var array
+     */
+    private $args;
+
+    /**
+     * @var Registry
+     */
+    private $registry;
+
+    /**
+     * @var EntityManager
+     */
+    private $manager;
+
+    /**
+     * @var array
+     */
+    private static $nonDefaultUserRepositoriesIds = [];
+
     /**
      * @param string $entityClass The class name of the entity this repository manages
      */
-    public function __construct(ManagerRegistry $registry, $entityClass)
+    public function __construct(ManagerRegistry $registry, string $entityClass, ...$args)
     {
-        $manager = $registry->getManagerForClass($entityClass);
+        if (!isset($args[0])) {
+            throw new InvalidArgumentException('You repository constructor uses an incorrect definition');
+        }
+
+        $this->args = $args[0];
 
-        if ($manager === null) {
+        // Is there a specific entity manager to use here? To be searched in the extra args we internally manage.
+        $userSpecifiedManager = null;
+
+        foreach ($this->args as $arg) {
+            if ($arg instanceof EntityManager) {
+                $userSpecifiedManager = $arg;
+                break;
+            }
+        }
+
+        // Default manager: "first one defined for the entity"
+        $manager = $userSpecifiedManager ? : $registry->getManagerForClass($entityClass);
+
+        if (null === $manager) {
             throw new LogicException(sprintf(
                 'Could not find the entity manager for class "%s". Check your Doctrine configuration to make sure it is configured to load this entity’s metadata.',
                 $entityClass
             ));
         }
 
+        $this->registry = $registry;
+        $this->manager = $manager;
+
         parent::__construct($manager, $manager->getClassMetadata($entityClass));
     }
+
+    /**
+     *
+     * @param string|EntityManager $entityManagerRef
+     *
+     * @return self
+     */
+    public function withManager($entityManagerRef)
+    {
+        // getFrom: instance
+        if ($entityManagerRef instanceof EntityManager) {
+            $userSpecifiedManager = $entityManagerRef;
+        } else { // getFrom: name
+            $userSpecifiedManager = $this->registry->getManager($entityManagerRef);
+        }
+
+        // If a different manager than the autowired-one is required, instantiate a new user's service-repository with the right one.
+        if ($userSpecifiedManager !== $this->manager) {
+            $managerInstanceId = spl_object_hash($userSpecifiedManager);
+
+            if (!isset(self::$nonDefaultUserRepositoriesIds[$managerInstanceId])) {
+                // Use the very-same arguments for the new instance but the manager name is added to be used as the default manager.
+                $args = $this->args;
+                $args[] = $userSpecifiedManager;
+
+                self::$nonDefaultUserRepositoriesIds[$managerInstanceId] = new static(...$args);
+            }
+
+            return self::$nonDefaultUserRepositoriesIds[$managerInstanceId];
+        }
+
+        return $this;
+    }
 }

Tests/DependencyInjection/Fixtures/Bundles/RepositoryServiceBundle/Repository/TestCustomServiceRepoRepository.php

@@ -8,8 +8,8 @@
 
 class TestCustomServiceRepoRepository extends ServiceEntityRepository
 {
-    public function __construct(RegistryInterface $registry)
+    public function __construct(RegistryInterface $registry, ...$args)
     {
-        parent::__construct($registry, TestCustomServiceRepoEntity::class);
+        parent::__construct($registry, TestCustomServiceRepoEntity::class, func_get_args());
     }
 }

Tests/Repository/ServiceEntityRepositoryTest.php

@@ -15,6 +15,6 @@ class ServiceEntityRepositoryTest extends TestCase
     public function testConstructorThrowsExceptionWhenNoManagerFound()
     {
         $registry = $this->getMockBuilder(ManagerRegistry::class)->getMock();
-        new ServiceEntityRepository($registry, TestEntity::class);
+        new ServiceEntityRepository($registry, TestEntity::class, [[]]);
     }
 }