use doctrine registry to find the correct entity manager for class

Model/ModelManager.php

@@ -19,23 +19,23 @@ use Sonata\AdminBundle\Datagrid\DatagridInterface;
 use Sonata\AdminBundle\Datagrid\ProxyQueryInterface;
 use Sonata\AdminBundle\Exception\ModelManagerException;
 
-use Doctrine\ORM\EntityManager;
 use Doctrine\ORM\QueryBuilder;
 
 use Symfony\Component\Form\Exception\PropertyAccessDeniedException;
+use Symfony\Bundle\DoctrineBundle\Registry;
 use Sonata\DoctrineORMAdminBundle\Datagrid\ProxyQuery;
 
 class ModelManager implements ModelManagerInterface
 {
-    protected $entityManager;
+    protected $registry;
 
     /**
      *
      * @param \Doctrine\ORM\EntityManager $entityManager
      */
-    public function __construct(EntityManager $entityManager)
+    public function __construct(\Symfony\Bundle\DoctrineBundle\Registry $registry)
     {
-        $this->entityManager = $entityManager;
+        $this->registry = $registry;
     }
 
     /**
@@ -47,7 +47,7 @@ class ModelManager implements ModelManagerInterface
      */
     public function getMetadata($class)
     {
-        return $this->entityManager->getMetadataFactory()->getMetadataFor($class);
+        return $this->getEntityManager($class)->getMetadataFactory()->getMetadataFor($class);
     }
 
     /**
@@ -58,7 +58,7 @@ class ModelManager implements ModelManagerInterface
      */
     public function hasMetadata($class)
     {
-        return $this->entityManager->getMetadataFactory()->hasMetadataFor($class);
+        return $this->getEntityManager($class)->getMetadataFactory()->hasMetadataFor($class);
     }
 
     /**
@@ -96,8 +96,9 @@ class ModelManager implements ModelManagerInterface
     public function create($object)
     {
         try {
-            $this->entityManager->persist($object);
-            $this->entityManager->flush();
+            $entityManager = $this->getEntityManager($object);
+            $entityManager->persist($object);
+            $entityManager->flush();
         } catch ( \PDOException $e ) {
             throw new ModelManagerException('', 0, $e);
         }
@@ -106,8 +107,9 @@ class ModelManager implements ModelManagerInterface
     public function update($object)
     {
         try {
-            $this->entityManager->persist($object);
-            $this->entityManager->flush();
+            $entityManager = $this->getEntityManager($object);
+            $entityManager->persist($object);
+            $entityManager->flush();
         } catch ( \PDOException $e ) {
             throw new ModelManagerException('', 0, $e);
         }
@@ -116,8 +118,9 @@ class ModelManager implements ModelManagerInterface
     public function delete($object)
     {
         try {
-            $this->entityManager->remove($object);
-            $this->entityManager->flush();
+            $entityManager = $this->getEntityManager($object);
+            $entityManager->remove($object);
+            $entityManager->flush();
         } catch ( \PDOException $e ) {
             throw new ModelManagerException('', 0, $e);
         }
@@ -133,7 +136,7 @@ class ModelManager implements ModelManagerInterface
     public function find($class, $id)
     {
         $values = array_combine($this->getIdentifierFieldNames($class), explode('-', $id));
-        return $this->entityManager->getRepository($class)->find($values);
+        return $this->getEntityManager($class)->getRepository($class)->find($values);
     }
 
     /**
@@ -143,7 +146,7 @@ class ModelManager implements ModelManagerInterface
      */
     public function findBy($class, array $criteria = array())
     {
-        return $this->entityManager->getRepository($class)->findBy($criteria);
+        return $this->getEntityManager($class)->getRepository($class)->findBy($criteria);
     }
 
     /**
@@ -153,15 +156,18 @@ class ModelManager implements ModelManagerInterface
      */
     public function findOneBy($class, array $criteria = array())
     {
-        return $this->entityManager->getRepository($class)->findOneBy($criteria);
+        return $this->getEntityManager($class)->getRepository($class)->findOneBy($criteria);
     }
 
     /**
      * @return \Doctrine\ORM\EntityManager
      */
-    public function getEntityManager()
-    {
-        return $this->entityManager;
+    public function getEntityManager($class) {
+        if (is_object($class)) {
+            $class = get_class($class);
+        }
+
+        return $this->registry->getEntityManagerForClass($class);
     }
 
     /**
@@ -191,7 +197,7 @@ class ModelManager implements ModelManagerInterface
      */
     public function createQuery($class, $alias = 'o')
     {
-        $repository = $this->getEntityManager()->getRepository($class);
+        $repository = $this->getEntityManager($class)->getRepository($class);
 
         return new ProxyQuery($repository->createQueryBuilder($alias));
     }
@@ -225,11 +231,12 @@ class ModelManager implements ModelManagerInterface
      */
     public function getIdentifierValues($entity)
     {
-        if (!$this->getEntityManager()->getUnitOfWork()->isInIdentityMap($entity)) {
+        $entityManager = $this->getEntityManager($entity);
+        if (!$entityManager->getUnitOfWork()->isInIdentityMap($entity)) {
             throw new \RuntimeException('Entities passed to the choice field must be managed');
         }
 
-        return $this->getEntityManager()->getUnitOfWork()->getEntityIdentifier($entity);
+        return $entityManager->getUnitOfWork()->getEntityIdentifier($entity);
     }
 
     /**
@@ -253,7 +260,7 @@ class ModelManager implements ModelManagerInterface
         }
 
         // the entities is not managed
-        if (!$entity || !$this->getEntityManager()->getUnitOfWork()->isInIdentityMap($entity)) {
+        if (!$entity || !$this->getEntityManager($entity)->getUnitOfWork()->isInIdentityMap($entity)) {
             return null;
         }
 
@@ -301,18 +308,20 @@ class ModelManager implements ModelManagerInterface
     public function batchDelete($class, ProxyQueryInterface $queryProxy)
     {
         try {
+            $entityManager = $this->getEntityManager($class);
+
             $i = 0;
             foreach ($queryProxy->getQuery()->iterate() as $pos => $object) {
-                $this->entityManager->remove($object[0]);
+                $entityManager->remove($object[0]);
 
                 if ((++$i % 20) == 0) {
-                    $this->entityManager->flush();
-                    $this->entityManager->clear();
+                    $entityManager->flush();
+                    $entityManager->clear();
                 }
             }
 
-            $this->entityManager->flush();
-            $this->entityManager->clear();
+            $entityManager->flush();
+            $entityManager->clear();
         } catch ( \PDOException $e ) {
             throw new ModelManagerException('', 0, $e);
         }

Resources/config/doctrine_orm.xml

@@ -10,7 +10,8 @@
         </service>
 
         <service id="sonata.admin.manager.orm" class="Sonata\DoctrineORMAdminBundle\Model\ModelManager">
-            <argument type="service" id="sonata.admin.entity_manager" />
+<!--             <argument type="service" id="sonata.admin.entity_manager" /> -->
+            <argument type="service" id="doctrine" />
         </service>
 
         <!-- FormBuilder -->