feat: add Query Builder lockForUpdate()

system/Database/BaseBuilder.php

@@ -110,6 +110,11 @@ class BaseBuilder
      */
     protected $QBOffset = false;
 
+    /**
+     * QB FOR UPDATE flag
+     */
+    protected bool $QBLockForUpdate = false;
+
     /**
      * QB ORDER BY data
      *
@@ -1620,6 +1625,16 @@ public function limit(?int $value = null, ?int $offset = 0)
         return $this;
     }
 
+    /**
+     * Locks the selected rows for update.
+     */
+    public function lockForUpdate(): static
+    {
+        $this->QBLockForUpdate = true;
+
+        return $this;
+    }
+
     /**
      * Sets the OFFSET value
      *
@@ -1801,20 +1816,26 @@ public function countAllResults(bool $reset = true)
         }
 
         // We cannot use a LIMIT when getting the single row COUNT(*) result
-        $limit = $this->QBLimit;
+        $limit         = $this->QBLimit;
+        $lockForUpdate = $this->QBLockForUpdate;
 
-        $this->QBLimit = false;
+        $this->QBLimit         = false;
+        $this->QBLockForUpdate = false;
 
-        if ($this->QBDistinct === true || ! empty($this->QBGroupBy)) {
-            // We need to backup the original SELECT in case DBPrefix is used
-            $select = $this->QBSelect;
-            $sql    = $this->countString . $this->db->protectIdentifiers('numrows') . "\nFROM (\n" . $this->compileSelect() . "\n) CI_count_all_results";
+        try {
+            if ($this->QBDistinct === true || ! empty($this->QBGroupBy)) {
+                // We need to backup the original SELECT in case DBPrefix is used
+                $select = $this->QBSelect;
+                $sql    = $this->countString . $this->db->protectIdentifiers('numrows') . "\nFROM (\n" . $this->compileSelect() . "\n) CI_count_all_results";
 
-            // Restore SELECT part
-            $this->QBSelect = $select;
-            unset($select);
-        } else {
-            $sql = $this->compileSelect($this->countString . $this->db->protectIdentifiers('numrows'));
+                // Restore SELECT part
+                $this->QBSelect = $select;
+                unset($select);
+            } else {
+                $sql = $this->compileSelect($this->countString . $this->db->protectIdentifiers('numrows'));
+            }
+        } finally {
+            $this->QBLockForUpdate = $lockForUpdate;
         }
 
         if ($this->testMode) {
@@ -3223,9 +3244,19 @@ protected function compileSelect($selectOverride = false): string
             $sql = $this->_limit($sql . "\n");
         }
 
+        $sql .= $this->compileLockForUpdate();
+
         return $this->unionInjection($sql);
     }
 
+    /**
+     * Compile the SELECT lock clause.
+     */
+    protected function compileLockForUpdate(): string
+    {
+        return $this->QBLockForUpdate ? "\nFOR UPDATE" : '';
+    }
+
     /**
      * Checks if the ignore option is supported by
      * the Database Driver for the specific statement.
@@ -3533,17 +3564,18 @@ protected function resetRun(array $qbResetItems)
     protected function resetSelect()
     {
         $this->resetRun([
-            'QBSelect'   => [],
-            'QBJoin'     => [],
-            'QBWhere'    => [],
-            'QBGroupBy'  => [],
-            'QBHaving'   => [],
-            'QBOrderBy'  => [],
-            'QBNoEscape' => [],
-            'QBDistinct' => false,
-            'QBLimit'    => false,
-            'QBOffset'   => false,
-            'QBUnion'    => [],
+            'QBSelect'        => [],
+            'QBJoin'          => [],
+            'QBWhere'         => [],
+            'QBGroupBy'       => [],
+            'QBHaving'        => [],
+            'QBOrderBy'       => [],
+            'QBNoEscape'      => [],
+            'QBDistinct'      => false,
+            'QBLimit'         => false,
+            'QBOffset'        => false,
+            'QBLockForUpdate' => false,
+            'QBUnion'         => [],
         ]);
 
         if ($this->db instanceof BaseConnection) {

system/Database/OCI8/Builder.php

@@ -213,6 +213,18 @@ protected function _limit(string $sql, bool $offsetIgnore = false): string
         return $sql . ' OFFSET ' . $offset . ' ROWS FETCH NEXT ' . $this->QBLimit . ' ROWS ONLY';
     }
 
+    /**
+     * Compile the SELECT lock clause.
+     */
+    protected function compileLockForUpdate(): string
+    {
+        if ($this->QBLockForUpdate && ($this->QBLimit !== false || $this->QBOffset)) {
+            throw new DatabaseException('OCI8 does not support lockForUpdate() with limit() or offset().');
+        }
+
+        return parent::compileLockForUpdate();
+    }
+
     /**
      * Generates a platform-specific batch update string from the supplied data
      */

system/Database/SQLSRV/Builder.php

@@ -31,6 +31,8 @@
  */
 class Builder extends BaseBuilder
 {
+    private const LOCK_FOR_UPDATE_HINT = ' WITH (UPDLOCK, ROWLOCK)';
+
     /**
      * ORDER BY random keyword
      *
@@ -74,7 +76,17 @@ protected function _fromTables(): string
         $from = [];
 
         foreach ($this->QBFrom as $value) {
-            $from[] = str_starts_with($value, '(SELECT') ? $value : $this->getFullName($value);
+            if (str_starts_with($value, '(SELECT')) {
+                if ($this->QBLockForUpdate) {
+                    throw new DatabaseException('SQLSRV does not support lockForUpdate() on subqueries.');
+                }
+
+                $from[] = $value;
+
+                continue;
+            }
+
+            $from[] = $this->getFullName($value) . ($this->QBLockForUpdate ? self::LOCK_FOR_UPDATE_HINT : '');
         }
 
         return implode(', ', $from);
@@ -675,9 +687,23 @@ protected function compileSelect($selectOverride = false): string
             $sql = $this->_limit($sql . "\n");
         }
 
+        $sql .= $this->compileLockForUpdate();
+
         return $this->unionInjection($sql);
     }
 
+    /**
+     * Compile the SELECT lock clause.
+     */
+    protected function compileLockForUpdate(): string
+    {
+        if ($this->QBLockForUpdate && $this->QBFrom === []) {
+            throw new DatabaseException('SQLSRV does not support lockForUpdate() without a FROM table.');
+        }
+
+        return '';
+    }
+
     /**
      * Compiles the select statement based on the other functions called
      * and runs the query

system/Database/SQLite3/Builder.php

@@ -55,6 +55,18 @@ class Builder extends BaseBuilder
         'insert' => 'OR IGNORE',
     ];
 
+    /**
+     * Compile the SELECT lock clause.
+     */
+    protected function compileLockForUpdate(): string
+    {
+        if ($this->QBLockForUpdate) {
+            throw new DatabaseException('SQLite3 does not support lockForUpdate().');
+        }
+
+        return '';
+    }
+
     /**
      * Replace statement
      *

tests/system/Database/Builder/CountTest.php

@@ -14,6 +14,7 @@
 namespace CodeIgniter\Database\Builder;
 
 use CodeIgniter\Database\BaseBuilder;
+use CodeIgniter\Database\SQLSRV\Builder as SQLSRVBuilder;
 use CodeIgniter\Test\CIUnitTestCase;
 use CodeIgniter\Test\Mock\MockConnection;
 use PHPUnit\Framework\Attributes\Group;
@@ -55,6 +56,34 @@ public function testCountAllResults(): void
         $this->assertSame($expectedSQL, str_replace("\n", ' ', $answer));
     }
 
+    public function testCountAllResultsDoesNotUseLockForUpdate(): void
+    {
+        $builder = new BaseBuilder('jobs', $this->db);
+        $builder->testMode();
+
+        $answer = $builder->where('id >', 3)->lockForUpdate()->countAllResults(false);
+
+        $expectedSQL = 'SELECT COUNT(*) AS "numrows" FROM "jobs" WHERE "id" > :id:';
+
+        $this->assertSame($expectedSQL, str_replace("\n", ' ', $answer));
+        $this->assertSame('SELECT * FROM "jobs" WHERE "id" > 3 FOR UPDATE', str_replace("\n", ' ', $builder->getCompiledSelect(false)));
+    }
+
+    public function testCountAllResultsWithSQLSRVDoesNotUseLockForUpdate(): void
+    {
+        $this->db = new MockConnection(['DBDriver' => 'SQLSRV', 'database' => 'test', 'schema' => 'dbo']);
+
+        $builder = new SQLSRVBuilder('jobs', $this->db);
+        $builder->testMode();
+
+        $answer = $builder->where('id >', 3)->lockForUpdate()->countAllResults(false);
+
+        $expectedSQL = 'SELECT COUNT(*) AS "numrows" FROM "test"."dbo"."jobs" WHERE "id" > :id:';
+
+        $this->assertSame($expectedSQL, str_replace("\n", ' ', $answer));
+        $this->assertSame('SELECT * FROM "test"."dbo"."jobs" WITH (UPDLOCK, ROWLOCK) WHERE "id" > 3', str_replace("\n", ' ', $builder->getCompiledSelect(false)));
+    }
+
     public function testCountAllResultsWithGroupBy(): void
     {
         $builder = new BaseBuilder('jobs', $this->db);