Add lock mode to query builder

src/Query/QueryBuilder.php

@@ -4,12 +4,14 @@
 
 use Doctrine\DBAL\Connection;
 use Doctrine\DBAL\Exception;
+use Doctrine\DBAL\LockMode;
 use Doctrine\DBAL\ParameterType;
 use Doctrine\DBAL\Query\Expression\CompositeExpression;
 use Doctrine\DBAL\Query\Expression\ExpressionBuilder;
 use Doctrine\DBAL\Result;
 use Doctrine\DBAL\Statement;
 
+use function array_filter;
 use function array_key_exists;
 use function array_keys;
 use function array_unshift;
@@ -20,6 +22,7 @@
 use function is_array;
 use function is_object;
 use function key;
+use function rtrim;
 use function strtoupper;
 use function substr;
 
@@ -135,6 +138,13 @@ class QueryBuilder
      */
     private $boundCounter = 0;
 
+    /**
+     * Lock mode to apply to SELECT queries, use one of the constants from LockMode::class
+     *
+     * @var int
+     */
+    private $lockMode = null;
+
     /**
      * Initializes a new <tt>QueryBuilder</tt>.
      *
@@ -405,6 +415,26 @@ public function getMaxResults()
         return $this->maxResults;
     }
 
+    /**
+     * Gets the current lock mode  for this query
+     */
+    public function getLockMode(): ?int
+    {
+        return $this->lockMode;
+    }
+
+    /**
+     * Set lock mode use one of the constants from LockMode::class locks are only added to SELECT queries
+     *
+     * Optimistic locking *is not* supported by DBAL.
+     */
+    public function setLockMode(int $lockMode): self
+    {
+        $this->lockMode = $lockMode;
+
+        return $this;
+    }
+
     /**
      * Either appends to or replaces a single, generic query part.
      *
@@ -1142,24 +1172,41 @@ public function resetQueryPart($queryPartName)
      */
     private function getSQLForSelect()
     {
+        $databasePlatform = $this->connection->getDatabasePlatform();
+
+        $lockSql = '';
+        switch ($this->lockMode) {
+            case LockMode::PESSIMISTIC_READ:
+                $lockSql = ' ' . $databasePlatform->getReadLockSQL();
+                break;
+
+            case LockMode::PESSIMISTIC_WRITE:
+                $lockSql = ' ' . $databasePlatform->getWriteLockSQL();
+                break;
+        }
+
+        $lockSql = rtrim($lockSql, ' ');
+
         $query = 'SELECT ' . ($this->sqlParts['distinct'] ? 'DISTINCT ' : '') .
                   implode(', ', $this->sqlParts['select']);
 
-        $query .= ($this->sqlParts['from'] ? ' FROM ' . implode(', ', $this->getFromClauses()) : '')
+        $from = $this->sqlParts['from'] ? ' FROM ' . implode(', ', $this->getFromClauses()) : '';
+
+        $query .= $databasePlatform->appendLockHint($from, $this->lockMode)
             . ($this->sqlParts['where'] !== null ? ' WHERE ' . ((string) $this->sqlParts['where']) : '')
             . ($this->sqlParts['groupBy'] ? ' GROUP BY ' . implode(', ', $this->sqlParts['groupBy']) : '')
             . ($this->sqlParts['having'] !== null ? ' HAVING ' . ((string) $this->sqlParts['having']) : '')
             . ($this->sqlParts['orderBy'] ? ' ORDER BY ' . implode(', ', $this->sqlParts['orderBy']) : '');
 
         if ($this->isLimitQuery()) {
-            return $this->connection->getDatabasePlatform()->modifyLimitQuery(
+            return $databasePlatform->modifyLimitQuery(
                 $query,
                 $this->maxResults,
                 $this->firstResult
-            );
+            ) . $lockSql;
         }
 
-        return $query;
+        return $query . $lockSql;
     }
 
     /**

tests/Functional/Query/QueryBuilderTest.php

@@ -0,0 +1,88 @@
+<?php
+
+namespace Doctrine\Tests\DBAL\Functional\Query;
+
+use Doctrine\DBAL\LockMode;
+use Doctrine\DBAL\ParameterType;
+use Doctrine\DBAL\Schema\Table;
+use Doctrine\DBAL\Tests\FunctionalTestCase;
+
+class QueryBuilderTest extends FunctionalTestCase
+{
+    /** @var bool */
+    private static $generated = false;
+
+    protected function setUp(): void
+    {
+        parent::setUp();
+
+        if (self::$generated !== false) {
+            return;
+        }
+
+        $table = new Table('lock_test_table');
+        $table->addColumn('primary_key', 'integer');
+        $table->addColumn('data', 'string');
+        $table->setPrimaryKey(['primary_key']);
+
+        $sm = $this->connection->getSchemaManager();
+        $sm->createTable($table);
+
+        $this->connection->insert('lock_test_table', [
+            'primary_key' => 1,
+            'data' => 'foo',
+        ]);
+
+        $this->connection->insert('lock_test_table', [
+            'primary_key' => 2,
+            'data' => 'bar',
+        ]);
+
+        $this->connection->insert('lock_test_table', [
+            'primary_key' => 3,
+            'data' => 'baz',
+        ]);
+
+        self::$generated = true;
+    }
+
+    public function testReadLock(): void
+    {
+        $sut   = $this->connection->createQueryBuilder();
+        $query = $sut->select('lt.*')
+            ->from('lock_test_table', 'lt')
+            ->setLockMode(LockMode::PESSIMISTIC_READ)
+            ->where($sut->expr()->eq('lt.data', '?'))
+            ->setParameter(0, 'foo', ParameterType::STRING);
+
+        $this->connection->beginTransaction();
+
+        $row = $query->execute()->fetchAssociative();
+
+        self::assertIsArray($row);
+        $row = array_change_key_case($row, CASE_LOWER);
+        self::assertEquals(['primary_key' => 1, 'data' => 'foo'], $row);
+
+        $this->connection->rollBack();
+    }
+
+    public function testWriteLock(): void
+    {
+        $sut   = $this->connection->createQueryBuilder();
+        $query = $sut->select('lt.*')
+            ->from('lock_test_table', 'lt')
+            ->setLockMode(LockMode::PESSIMISTIC_WRITE)
+            ->where($sut->expr()->eq('lt.data', '?'))
+            ->setParameter(0, 'bar', ParameterType::STRING);
+
+        $this->connection->beginTransaction();
+
+        $row = $query->execute()->fetchAssociative();
+
+        self::assertIsArray($row);
+        $row = array_change_key_case($row, CASE_LOWER);
+        self::assertEquals(['primary_key' => 2, 'data' => 'bar'], $row);
+
+        $this->connection->rollBack();
+    }
+}

tests/Query/QueryBuilderTest.php

@@ -3,7 +3,13 @@
 namespace Doctrine\DBAL\Tests\Query;
 
 use Doctrine\DBAL\Connection;
+use Doctrine\DBAL\LockMode;
 use Doctrine\DBAL\ParameterType;
+use Doctrine\DBAL\Platforms\AbstractPlatform;
+use Doctrine\DBAL\Platforms\MySQL57Platform;
+use Doctrine\DBAL\Platforms\PostgreSQL94Platform;
+use Doctrine\DBAL\Platforms\SqlitePlatform;
+use Doctrine\DBAL\Platforms\SQLServer2012Platform;
 use Doctrine\DBAL\Query\Expression\ExpressionBuilder;
 use Doctrine\DBAL\Query\QueryBuilder;
 use Doctrine\DBAL\Query\QueryException;
@@ -16,13 +22,7 @@ class QueryBuilderTest extends TestCase
 
     protected function setUp(): void
     {
-        $this->conn = $this->createMock(Connection::class);
-
-        $expressionBuilder = new ExpressionBuilder($this->conn);
-
-        $this->conn->expects(self::any())
-                   ->method('getExpressionBuilder')
-                   ->will(self::returnValue($expressionBuilder));
+        $this->conn = $this->createMockConnection(new SqlitePlatform());
     }
 
     public function testSimpleSelectWithoutFrom(): void
@@ -67,6 +67,99 @@ public function testSelectWithSimpleWhere(): void
         self::assertEquals('SELECT u.id FROM users u WHERE u.nickname = ?', (string) $qb);
     }
 
+    public function testSelectWithOptimisticLockIgnored(): void
+    {
+        $platform = new MySQL57Platform();
+
+        $qb   = new QueryBuilder($this->createMockConnection($platform));
+        $expr = $qb->expr();
+
+        $qb->select('u.id')
+            ->from('users', 'u')
+            ->setLockMode(LockMode::OPTIMISTIC)
+            ->where($expr->and($expr->eq('u.nickname', '?')));
+
+        self::assertEquals(LockMode::OPTIMISTIC, $qb->getLockMode());
+        self::assertEquals('SELECT u.id FROM users u WHERE u.nickname = ?', (string) $qb);
+    }
+
+    public function testSelectWithReadLockAppended(): void
+    {
+        $platform = new MySQL57Platform();
+
+        $qb   = new QueryBuilder($this->createMockConnection($platform));
+        $expr = $qb->expr();
+
+        $qb->select('u.id')
+            ->from('users', 'u')
+            ->setLockMode(LockMode::PESSIMISTIC_READ)
+            ->where($expr->and($expr->eq('u.nickname', '?')));
+
+        self::assertEquals('SELECT u.id FROM users u WHERE u.nickname = ? LOCK IN SHARE MODE', (string) $qb);
+    }
+
+    public function testSelectWithWriteLockAppended(): void
+    {
+        $platform = new PostgreSQL94Platform();
+
+        $qb   = new QueryBuilder($this->createMockConnection($platform));
+        $expr = $qb->expr();
+
+        $qb->select('u.id')
+            ->from('users', 'u')
+            ->setLockMode(LockMode::PESSIMISTIC_WRITE)
+            ->where($expr->and($expr->eq('u.nickname', '?')));
+
+        self::assertEquals('SELECT u.id FROM users u WHERE u.nickname = ? FOR UPDATE', (string) $qb);
+    }
+
+    public function testSelectWithWriteLockAppendedAfterLimit(): void
+    {
+        $platform = new PostgreSQL94Platform();
+
+        $qb   = new QueryBuilder($this->createMockConnection($platform));
+        $expr = $qb->expr();
+
+        $qb->select('u.id')
+            ->from('users', 'u')
+            ->setLockMode(LockMode::PESSIMISTIC_WRITE)
+            ->where($expr->and($expr->eq('u.nickname', '?')))
+            ->setFirstResult(50)
+            ->setMaxResults(10);
+
+        self::assertEquals('SELECT u.id FROM users u WHERE u.nickname = ? LIMIT 10 OFFSET 50 FOR UPDATE', (string) $qb);
+    }
+
+    public function testSelectWithReadLockTableHint(): void
+    {
+        $platform = new SQLServer2012Platform();
+
+        $qb   = new QueryBuilder($this->createMockConnection($platform));
+        $expr = $qb->expr();
+
+        $qb->select('u.id')
+            ->from('users', 'u')
+            ->setLockMode(LockMode::PESSIMISTIC_READ)
+            ->where($expr->and($expr->eq('u.nickname', '?')));
+
+        self::assertEquals('SELECT u.id FROM users u WITH (HOLDLOCK, ROWLOCK) WHERE u.nickname = ?', (string) $qb);
+    }
+
+    public function testSelectWithWriteLockTableHint(): void
+    {
+        $platform = new SQLServer2012Platform();
+
+        $qb   = new QueryBuilder($this->createMockConnection($platform));
+        $expr = $qb->expr();
+
+        $qb->select('u.id')
+            ->from('users', 'u')
+            ->setLockMode(LockMode::PESSIMISTIC_WRITE)
+            ->where($expr->and($expr->eq('u.nickname', '?')));
+
+        self::assertEquals('SELECT u.id FROM users u WITH (UPDLOCK, ROWLOCK) WHERE u.nickname = ?', (string) $qb);
+    }
+
     public function testSelectWithLeftJoin(): void
     {
         $qb   = new QueryBuilder($this->conn);
@@ -949,4 +1042,21 @@ public function testJoinWithNonUniqueAliasThrowsException(): void
 
         $qb->getSQL();
     }
+
+    private function createMockConnection(AbstractPlatform $platform): Connection
+    {
+        $conn = $this->createMock(Connection::class);
+
+        $expressionBuilder = new ExpressionBuilder($conn);
+
+        $conn->expects(self::any())
+            ->method('getExpressionBuilder')
+            ->will(self::returnValue($expressionBuilder));
+
+        $conn->expects(self::any())
+            ->method('getDatabasePlatform')
+            ->will(self::returnValue($platform));
+
+        return $conn;
+    }
 }