Restore query() call for primary/replica connections

doctrine · Dec 11, 2021 · 017fe6c · 017fe6c
1 parent 92434f9
commit 017fe6c
Show file tree

Hide file tree

Showing 5 changed files with 38 additions and 48 deletions.
diff --git a/lib/Doctrine/ORM/Id/SequenceGenerator.php b/lib/Doctrine/ORM/Id/SequenceGenerator.php
@@ -4,9 +4,12 @@
 
 namespace Doctrine\ORM\Id;
 
+use Doctrine\DBAL\Connections\PrimaryReadReplicaConnection;
+use Doctrine\DBAL\Result;
 use Doctrine\ORM\EntityManager;
 use Serializable;
 
+use function assert;
 use function serialize;
 use function unserialize;
 
@@ -58,8 +61,14 @@ public function generate(EntityManager $em, $entity)
             $sql  = $conn->getDatabasePlatform()->getSequenceNextValSQL($this->_sequenceName);
 
             // Using `query` to force usage of the master server in MasterSlaveConnection
-            $this->_nextValue = (int) $conn->executeQuery($sql)->fetchOne();
-            $this->_maxValue  = $this->_nextValue + $this->_allocationSize;
+            $result = $conn instanceof PrimaryReadReplicaConnection
+                ? $conn->query($sql)
+                : $conn->executeQuery($sql);
+            assert($result instanceof Result);
+
+            $this->_nextValue = (int) $result->fetchOne();
+
+            $this->_maxValue = $this->_nextValue + $this->_allocationSize;
         }
 
         return $this->_nextValue++;

diff --git a/lib/Doctrine/ORM/Id/UuidGenerator.php b/lib/Doctrine/ORM/Id/UuidGenerator.php
@@ -5,10 +5,12 @@
 namespace Doctrine\ORM\Id;
 
 use Doctrine\DBAL\Platforms\AbstractPlatform;
+use Doctrine\DBAL\Result;
 use Doctrine\Deprecations\Deprecation;
 use Doctrine\ORM\EntityManager;
 use Doctrine\ORM\Exception\NotSupported;
 
+use function assert;
 use function method_exists;
 
 /**
@@ -39,9 +41,11 @@ public function __construct()
      */
     public function generate(EntityManager $em, $entity)
     {
-        $conn = $em->getConnection();
-        $sql  = 'SELECT ' . $conn->getDatabasePlatform()->getGuidExpression();
+        $conn   = $em->getConnection();
+        $sql    = 'SELECT ' . $conn->getDatabasePlatform()->getGuidExpression();
+        $result = $conn->query($sql);
+        assert($result instanceof Result);
 
-        return $conn->executeQuery($sql)->fetchOne();
+        return $result->fetchOne();
     }
 }
diff --git a/psalm.xml b/psalm.xml
@@ -27,7 +27,9 @@
         </DeprecatedClass>
         <DeprecatedMethod>
             <errorLevel type="suppress">
-                <!-- We're calling the deprecated method for BC here. -->
+                <!-- We're calling the deprecated methods for BC here. -->
+                <file name="lib/Doctrine/ORM/Id/SequenceGenerator.php"/>
+                <file name="lib/Doctrine/ORM/Id/UuidGenerator.php"/>
                 <file name="lib/Doctrine/ORM/Internal/SQLResultCasing.php"/>
                 <!-- We need to keep the calls for DBAL 2.13 compatibility. -->
                 <referencedMethod name="Doctrine\DBAL\Cache\QueryCacheProfile::getResultCacheDriver"/>
@@ -87,6 +89,13 @@
                 <file name="lib/Doctrine/ORM/QueryBuilder.php"/>
             </errorLevel>
         </RedundantCastGivenDocblockType>
+        <RedundantCondition>
+            <errorLevel type="suppress">
+                <!-- Fallback logic for DBAL 2 -->
+                <file name="lib/Doctrine/ORM/Id/SequenceGenerator.php"/>
+                <file name="lib/Doctrine/ORM/Id/UuidGenerator.php"/>
+            </errorLevel>
+        </RedundantCondition>
         <!-- Workaround for https://github.com/vimeo/psalm/issues/7026 -->
         <ReservedWord>
             <errorLevel type="suppress">

diff --git a/tests/Doctrine/Tests/Mocks/ConnectionMock.php b/tests/Doctrine/Tests/Mocks/ConnectionMock.php
@@ -25,9 +25,6 @@ class ConnectionMock extends Connection
     /** @var mixed */
     private $_fetchOneResult;
 
-    /** @var Exception|null */
-    private $_fetchOneException;
-
     /** @var Result|null */
     private $_queryResult;
 
@@ -116,10 +113,6 @@ public function lastInsertId($seqName = null)
      */
     public function fetchOne(string $sql, array $params = [], array $types = [])
     {
-        if ($this->_fetchOneException !== null) {
-            throw $this->_fetchOneException;
-        }
-
         return $this->_fetchOneResult;
     }
 
@@ -163,11 +156,6 @@ public function setFetchOneResult($fetchOneResult): void
         $this->_fetchOneResult = $fetchOneResult;
     }
 
-    public function setFetchOneException(?Exception $exception = null): void
-    {
-        $this->_fetchOneException = $exception;
-    }
-
     public function setDatabasePlatform(AbstractPlatform $platform): void
     {
         $this->_platformMock = $platform;

diff --git a/tests/Doctrine/Tests/ORM/Id/SequenceGeneratorTest.php b/tests/Doctrine/Tests/ORM/Id/SequenceGeneratorTest.php
@@ -4,52 +4,32 @@
 
 namespace Doctrine\Tests\ORM\Id;
 
-use BadMethodCallException;
-use Doctrine\ORM\EntityManager;
 use Doctrine\ORM\Id\SequenceGenerator;
 use Doctrine\Tests\Mocks\ArrayResultFactory;
 use Doctrine\Tests\Mocks\ConnectionMock;
 use Doctrine\Tests\OrmTestCase;
 
+use function assert;
+
 class SequenceGeneratorTest extends OrmTestCase
 {
-    /** @var EntityManager */
-    private $entityManager;
-
-    /** @var SequenceGenerator */
-    private $sequenceGenerator;
-
-    /** @var ConnectionMock */
-    private $connection;
-
-    protected function setUp(): void
-    {
-        parent::setUp();
-
-        $this->entityManager     = $this->getTestEntityManager();
-        $this->sequenceGenerator = new SequenceGenerator('seq', 10);
-        $this->connection        = $this->entityManager->getConnection();
-
-        self::assertInstanceOf(ConnectionMock::class, $this->connection);
-    }
-
     public function testGeneration(): void
     {
-        $this->connection->setFetchOneException(new BadMethodCallException(
-            'Fetch* method used. Query method should be used instead, '
-            . 'as NEXTVAL should be run on a master server in master-slave setup.'
-        ));
+        $entityManager     = $this->getTestEntityManager();
+        $sequenceGenerator = new SequenceGenerator('seq', 10);
+        $connection        = $entityManager->getConnection();
+        assert($connection instanceof ConnectionMock);
 
         for ($i = 0; $i < 42; ++$i) {
             if ($i % 10 === 0) {
-                $this->connection->setQueryResult(ArrayResultFactory::createFromArray([[(int) ($i / 10) * 10]]));
+                $connection->setQueryResult(ArrayResultFactory::createFromArray([[(int) ($i / 10) * 10]]));
             }
 
-            $id = $this->sequenceGenerator->generate($this->entityManager, null);
+            $id = $sequenceGenerator->generate($entityManager, null);
 
             self::assertSame($i, $id);
-            self::assertSame((int) ($i / 10) * 10 + 10, $this->sequenceGenerator->getCurrentMaxValue());
-            self::assertSame($i + 1, $this->sequenceGenerator->getNextValue());
+            self::assertSame((int) ($i / 10) * 10 + 10, $sequenceGenerator->getCurrentMaxValue());
+            self::assertSame($i + 1, $sequenceGenerator->getNextValue());
         }
     }
 }