Add locking methods to the query builder

src/Connection.php

@@ -430,7 +430,7 @@ private function detectDatabasePlatform(): AbstractPlatform
      *
      * @throws Throwable
      */
-    private function getDatabasePlatformVersion()
+    public function getDatabasePlatformVersion()
     {
         // Driver does not support version specific platforms.
         if (! $this->_driver instanceof VersionAwarePlatformDriver) {

src/Driver/AbstractMySQLDriver.php

@@ -11,6 +11,7 @@
 use Doctrine\DBAL\Platforms\MariaDb1027Platform;
 use Doctrine\DBAL\Platforms\MariaDb1043Platform;
 use Doctrine\DBAL\Platforms\MariaDb1052Platform;
+use Doctrine\DBAL\Platforms\MariaDb1060Platform;
 use Doctrine\DBAL\Platforms\MySQL57Platform;
 use Doctrine\DBAL\Platforms\MySQL80Platform;
 use Doctrine\DBAL\Platforms\MySQLPlatform;
@@ -39,6 +40,10 @@ public function createDatabasePlatformForVersion($version)
 
         if ($mariadb) {
             $mariaDbVersion = $this->getMariaDbMysqlVersionNumber($version);
+            if (version_compare($mariaDbVersion, '10.6.0', '>=')) {
+                return new MariaDb1060Platform();
+            }
+
             if (version_compare($mariaDbVersion, '10.5.2', '>=')) {
                 return new MariaDb1052Platform();
             }

src/Id/TableGenerator.php

@@ -6,7 +6,6 @@
 use Doctrine\DBAL\Driver;
 use Doctrine\DBAL\DriverManager;
 use Doctrine\DBAL\Exception;
-use Doctrine\DBAL\LockMode;
 use Doctrine\Deprecations\Deprecation;
 use Throwable;
 
@@ -115,11 +114,13 @@ public function nextValue($sequence)
         $this->conn->beginTransaction();
 
         try {
-            $platform = $this->conn->getDatabasePlatform();
-            $sql      = 'SELECT sequence_value, sequence_increment_by'
-                . ' FROM ' . $platform->appendLockHint($this->generatorTableName, LockMode::PESSIMISTIC_WRITE)
-                . ' WHERE sequence_name = ? ' . $platform->getWriteLockSQL();
-            $row      = $this->conn->fetchAssociative($sql, [$sequence]);
+            $qb  = $this->conn->createQueryBuilder();
+            $sql = $qb->select('t.sequence_value', 't.sequence_increment_by')
+                ->from($this->generatorTableName, 't')
+                ->where('t.sequence_name = ?')
+                ->lockForUpdate()
+                ->getSQL();
+            $row = $this->conn->fetchAssociative($sql, [$sequence]);
 
             if ($row !== false) {
                 $row = array_change_key_case($row, CASE_LOWER);

src/Platforms/AbstractPlatform.php

@@ -17,6 +17,7 @@
 use Doctrine\DBAL\Exception\InvalidLockMode;
 use Doctrine\DBAL\LockMode;
 use Doctrine\DBAL\Platforms\Keywords\KeywordList;
+use Doctrine\DBAL\Query\QueryLock;
 use Doctrine\DBAL\Schema\AbstractSchemaManager;
 use Doctrine\DBAL\Schema\Column;
 use Doctrine\DBAL\Schema\ColumnDiff;
@@ -54,6 +55,7 @@
 use function is_bool;
 use function is_int;
 use function is_string;
+use function ksort;
 use function preg_quote;
 use function preg_replace;
 use function sprintf;
@@ -62,6 +64,7 @@
 use function strpos;
 use function strtolower;
 use function strtoupper;
+use function trim;
 
 /**
  * Base class for all DatabasePlatforms. The DatabasePlatforms are the central
@@ -1765,6 +1768,49 @@ public function getForUpdateSQL()
         return 'FOR UPDATE';
     }
 
+    /**
+     * Returns the SKIP LOCKED expression.
+     */
+    public function getSkipLockedSQL(): string
+    {
+        return 'SKIP LOCKED';
+    }
+
+    public function getLocksSql(QueryLock ...$locks): string
+    {
+        return trim(implode(' ', $this->getLocksSqlList(...$locks)));
+    }
+
+    /** @return string[] */
+    protected function getLocksSqlList(QueryLock ...$locks): array
+    {
+        $locksSqlList = [];
+        foreach ($locks as $lock) {
+            switch ($lock->value()) {
+                case QueryLock::forUpdate()->value():
+                    $locksSqlList[0] = $this->getForUpdateSQL();
+                    break;
+                case QueryLock::skipLocked()->value():
+                    $locksSqlList[1] = $this->getSkipLockedSQL();
+                    break;
+            }
+        }
+
+        ksort($locksSqlList);
+
+        return $locksSqlList;
+    }
+
+    public function isLockLocatedAfterFrom(): bool
+    {
+        return false;
+    }
+
+    public function isLockLocatedAtTheEnd(): bool
+    {
+        return true;
+    }
+
     /**
      * Honors that some SQL vendors such as MsSql use table hints for locking instead of the
      * ANSI SQL FOR UPDATE specification.

src/Platforms/DB2111Platform.php

@@ -37,4 +37,9 @@ protected function doModifyLimitQuery($query, $limit, $offset)
 
         return $query;
     }
+
+    public function getSkipLockedSQL(): string
+    {
+        return 'SKIP LOCKED DATA';
+    }
 }

src/Platforms/DB2Platform.php

@@ -1041,4 +1041,9 @@ public function createSchemaManager(Connection $connection): DB2SchemaManager
     {
         return new DB2SchemaManager($connection, $this);
     }
+
+    public function getSkipLockedSQL(): string
+    {
+        return '';
+    }
 }

src/Platforms/MariaDBPlatform.php

@@ -10,6 +10,11 @@
  */
 class MariaDBPlatform extends MySQLPlatform
 {
+    public function getSkipLockedSQL(): string
+    {
+        return '';
+    }
+
     /**
      * {@inheritDoc}
      *

src/Platforms/MariaDb1060Platform.php

@@ -0,0 +1,21 @@
+<?php
+
+namespace Doctrine\DBAL\Platforms;
+
+/**
+ * Provides the behavior, features and SQL dialect of the MariaDB 10.6 (10.6.0 GA) database platform.
+ *
+ * Note: Should not be used with versions prior to 10.6.0.
+ */
+class MariaDb1060Platform extends MariaDb1052Platform
+{
+    /**
+     * Returns the FOR UPDATE SKIP LOCKED expression.
+     * This method will become obsolete once the minimum MariaDb version is at least 10.6.0,
+     * as this method already exists in the base AbstractPlatform class.
+     */
+    public function getSkipLockedSQL(): string
+    {
+        return 'SKIP LOCKED';
+    }
+}

src/Platforms/MySQL57Platform.php

@@ -96,4 +96,12 @@ protected function initializeDoctrineTypeMappings()
 
         $this->doctrineTypeMapping['json'] = Types::JSON;
     }
+
+    /**
+     * Returns '', as SKIP LOCKED is only available since MySQL 8.
+     */
+    public function getSkipLockedSQL(): string
+    {
+        return '';
+    }
 }

src/Platforms/MySQL80Platform.php

@@ -25,4 +25,14 @@ protected function getReservedKeywordsClass()
 
         return Keywords\MySQL80Keywords::class;
     }
+
+    /**
+     * Returns the SKIP LOCKED expression.
+     * When support for MySQL 5.7 is removed, this method can be removed from this
+     * class as it is already implemented in the base class.
+     */
+    public function getSkipLockedSQL(): string
+    {
+        return 'SKIP LOCKED';
+    }
 }

src/Platforms/PostgreSQL100Platform.php

@@ -27,4 +27,12 @@ protected function getReservedKeywordsClass(): string
 
         return PostgreSQL100Keywords::class;
     }
+
+    /**
+     * Returns the SKIP LOCKED expression.
+     */
+    public function getSkipLockedSQL(): string
+    {
+        return 'SKIP LOCKED';
+    }
 }

src/Platforms/PostgreSQL94Platform.php

@@ -9,4 +9,11 @@
  */
 class PostgreSQL94Platform extends PostgreSQLPlatform
 {
+    /**
+     * Returns the SKIP LOCKED expression.
+     */
+    public function getSkipLockedSQL(): string
+    {
+        return '';
+    }
 }

src/Platforms/SQLServerPlatform.php

@@ -5,6 +5,7 @@
 use Doctrine\DBAL\Connection;
 use Doctrine\DBAL\Exception\InvalidLockMode;
 use Doctrine\DBAL\LockMode;
+use Doctrine\DBAL\Query\QueryLock;
 use Doctrine\DBAL\Schema\Column;
 use Doctrine\DBAL\Schema\ColumnDiff;
 use Doctrine\DBAL\Schema\ForeignKeyConstraint;
@@ -33,8 +34,10 @@
 use function is_bool;
 use function is_numeric;
 use function is_string;
+use function ltrim;
 use function preg_match;
 use function preg_match_all;
+use function rtrim;
 use function sprintf;
 use function str_replace;
 use function strpos;
@@ -1613,7 +1616,33 @@ public function appendLockHint(string $fromClause, int $lockMode): string
      */
     public function getForUpdateSQL()
     {
-        return ' ';
+        return 'WITH (UPDLOCK, ROWLOCK)';
+    }
+
+    public function getSkipLockedSQL(): string
+    {
+        return 'WITH (READPAST)';
+    }
+
+    public function getLocksSql(QueryLock ...$locks): string
+    {
+        $locksSqlList = $this->getLocksSqlList(...$locks);
+
+        foreach ($locksSqlList as $key => $lockSql) {
+            $locksSqlList[$key] = rtrim(ltrim($lockSql, 'WITH ('), ')');
+        }
+
+        return 'WITH (' . implode(', ', $locksSqlList) . ')';
+    }
+
+    public function isLockLocatedAfterFrom(): bool
+    {
+        return true;
+    }
+
+    public function isLockLocatedAtTheEnd(): bool
+    {
+        return false;
     }
 
     /**

src/Platforms/SqlitePlatform.php

@@ -740,6 +740,11 @@ public function getForUpdateSQL()
         return '';
     }
 
+    public function getSkipLockedSQL(): string
+    {
+        return '';
+    }
+
     /**
      * {@inheritDoc}
      *

src/Query/QueryBuilder.php

@@ -75,6 +75,7 @@ class QueryBuilder
         'having'   => null,
         'orderBy'  => [],
         'values'   => [],
+        'locks'    => [],
     ];
 
     /**
@@ -1356,18 +1357,23 @@ private function getSQLForSelect(): string
         $query = 'SELECT ' . ($this->sqlParts['distinct'] ? 'DISTINCT ' : '') .
                   implode(', ', $this->sqlParts['select']);
 
+        $platform = $this->connection->getDatabasePlatform();
+
+        $locksSql = $this->hasLocks() ? ' ' . $platform->getLocksSql(...$this->sqlParts['locks']) : '';
+
         $query .= ($this->sqlParts['from'] ? ' FROM ' . implode(', ', $this->getFromClauses()) : '')
+            . ($platform->isLockLocatedAfterFrom() ? $locksSql : '')
             . ($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(
-                $query,
-                $this->maxResults,
-                $this->firstResult,
-            );
+            $query = $platform->modifyLimitQuery($query, $this->maxResults, $this->firstResult);
+        }
+
+        if ($platform->isLockLocatedAtTheEnd()) {
+            $query .= $locksSql;
         }
 
         return $query;
@@ -1623,4 +1629,33 @@ public function disableResultCache(): self
 
         return $this;
     }
+
+    private function hasLocks(): bool
+    {
+        return $this->sqlParts['locks'] !== [];
+    }
+
+    /**
+     * Sets a lock on the queried rows, until the end of the transaction
+     *
+     * @return $this
+     */
+    public function lockForUpdate(): self
+    {
+        $this->sqlParts['locks'][] = QueryLock::forUpdate();
+
+        return $this;
+    }
+
+    /**
+     * Sets a lock on the queried rows, until the end of the transaction
+     *
+     * @return $this
+     */
+    public function skipLocked(): self
+    {
+        $this->sqlParts['locks'][] = QueryLock::skipLocked();
+
+        return $this;
+    }
 }