Added "update and return" to the query builder

CHANGELOG.md

@@ -12,6 +12,7 @@
 	- `Statistics::populationVariance()`
 	- `Statistics::sampleStandardDeviation()`
 	- `Statistics::populationStandardDeviation()`
+* Added "update and return" functionality to the query builder. The feature is currently supported by the `Firebird`, `PostgreSQL`, `SQLite` and `SQL Server` compilers.
 
 #### Changes
 

src/mako/database/midgard/Query.php

@@ -201,6 +201,42 @@ public function update(array $values): int
 		return $updated;
 	}
 
+	/**
+	 * Updates data from the chosen table and returns a result set.
+	 *
+	 * @return ResultSet<int, TClass>
+	 */
+	public function updateAndReturn(array $values, array $return = ['*']): ResultSet
+	{
+		// Execute "beforeUpdate" hooks
+
+		foreach ($this->model->getHooks('beforeUpdate') as $hook) {
+			$values = $hook($values, $this);
+		}
+
+		// Update record(s)
+
+		if ($return !== ['*']) {
+			$return = array_unique([$this->model->getPrimaryKey(), ...$return]);
+		}
+
+		$updated = $this->updateAndReturnAll($values, $return, false, PDO::FETCH_ASSOC);
+
+		// Execute "afterUpdate" hooks
+
+		foreach ($this->model->getHooks('afterUpdate') as $hook) {
+			$hook($updated);
+		}
+
+		// Return updated records
+
+		if (!empty($updated)) {
+			$updated = $this->hydrateModelsAndLoadIncludes($updated);
+		}
+
+		return $this->createResultSet($updated);
+	}
+
 	/**
 	 * {@inheritDoc}
 	 */

src/mako/database/query/Query.php

@@ -1216,7 +1216,7 @@ protected function createResultSet(array $results): ResultSet
 	}
 
 	/**
-	 * Executes a SELECT query and returns an array containing all of the result set rows.
+	 * Executes a SELECT query and returns an array or result set containing all of the result set rows.
 	 */
 	protected function fetchAll(bool $returnResultSet, mixed ...$fetchMode): array|ResultSet
 	{
@@ -1228,7 +1228,7 @@ protected function fetchAll(bool $returnResultSet, mixed ...$fetchMode): array|R
 	}
 
 	/**
-	 * Executes a SELECT query and returns an array containing all of the result set rows.
+	 * Executes a SELECT query and returns a result set containing all of the result set rows.
 	 *
 	 * @return ResultSet<int, Result>
 	 */
@@ -1490,6 +1490,28 @@ public function update(array $values): int
 		return $this->connection->queryAndCount($query['sql'], $query['params']);
 	}
 
+	/**
+	 * Updates data from the chosen table and returns an array or result set.
+	 */
+	protected function updateAndReturnAll(array $values, array $return, bool $returnResultSet, mixed ...$fetchMode): array|ResultSet
+	{
+		$query = $this->compiler->updateAndReturn($values, $return);
+
+		$results = $this->connection->all($query['sql'], $query['params'], ...$fetchMode);
+
+		return $returnResultSet ? $this->createResultSet($results) : $results;
+	}
+
+	/**
+	 * Updates data from the chosen table and returns a result set.
+	 *
+	 * @return ResultSet<int, Result>
+	 */
+	public function updateAndReturn(array $values, array $return = ['*']): ResultSet
+	{
+		return $this->updateAndReturnAll($values, $return, true, PDO::FETCH_CLASS, Result::class);
+	}
+
 	/**
 	 * Increments column value.
 	 */

src/mako/database/query/compilers/Compiler.php

@@ -828,6 +828,16 @@ public function update(array $values): array
 		return ['sql' => $sql, 'params' => $this->params];
 	}
 
+	/**
+	 * Compiles an UPDATE query with a RETURNING clause.
+	 *
+	 * @return array{sql: string, params: array}
+	 */
+	public function updateAndReturn(array $values, array $return): array
+	{
+		throw new DatabaseException(sprintf('The [ %s ] query compiler does not support update and return queries.', static::class));
+	}
+
 	/**
 	 * Compiles a DELETE query.
 	 *

src/mako/database/query/compilers/Firebird.php

@@ -75,4 +75,16 @@ public function lock(null|bool|string $lock): string
 
 		return $lock === true ? ' FOR UPDATE WITH LOCK' : ($lock === false ? ' WITH LOCK' : " {$lock}");
 	}
+
+	/**
+	 * {@inheritDoc}
+	 */
+	public function updateAndReturn(array $values, array $return): array
+	{
+		$query = $this->update($values);
+
+		$query['sql'] .= ' RETURNING ' . $this->columns($return);
+
+		return $query;
+	}
 }

src/mako/database/query/compilers/Postgres.php

@@ -121,4 +121,16 @@ public function insertOrUpdate(array $insertValues, array $updateValues, array $
 
 		return ['sql' => $sql, 'params' => $this->params];
 	}
+
+	/**
+	 * {@inheritDoc}
+	 */
+	public function updateAndReturn(array $values, array $return): array
+	{
+		$query = $this->update($values);
+
+		$query['sql'] .= ' RETURNING ' . $this->columns($return);
+
+		return $query;
+	}
 }

src/mako/database/query/compilers/SQLServer.php

@@ -11,6 +11,9 @@
 use mako\database\query\Raw;
 use mako\database\query\Subquery;
 
+use function array_map;
+use function explode;
+use function implode;
 use function str_replace;
 
 /**
@@ -114,4 +117,20 @@ protected function offset(?int $offset): string
 
 		return '';
 	}
+
+	/**
+	 * {@inheritDoc}
+	 */
+	public function updateAndReturn(array $values, array $return): array
+	{
+		$sql = $this->query->getPrefix()
+		. 'UPDATE '
+		. $this->escapeTableName($this->query->getTable())
+		. ' SET '
+		. $this->updateColumns($values)
+		. ' OUTPUT ' . implode(', ', array_map(fn ($column) => "inserted.{$column}", explode(', ', $this->columns($return))))
+		. $this->wheres($this->query->getWheres());
+
+		return ['sql' => $sql, 'params' => $this->params];
+	}
 }

src/mako/database/query/compilers/SQLite.php

@@ -108,4 +108,16 @@ public function insertOrUpdate(array $insertValues, array $updateValues, array $
 
 		return ['sql' => $sql, 'params' => $this->params];
 	}
+
+	/**
+	 * {@inheritDoc}
+	 */
+	public function updateAndReturn(array $values, array $return): array
+	{
+		$query = $this->update($values);
+
+		$query['sql'] .= ' RETURNING ' . $this->columns($return);
+
+		return $query;
+	}
 }

tests/integration/database/midgard/ORMTest.php

@@ -676,4 +676,19 @@ public function testSubquery(): void
 
 		$this->assertSame(3, count($counters));
 	}
+
+	/**
+	 *
+	 */
+	public function testUpdateAndReturn(): void
+	{
+		$updated = (new TestUser)->where('id', '=', 1)->updateAndReturn(['username' => 'bax'], ['username']);
+
+		$this->assertInstanceOf(ResultSet::class, $updated);
+		$this->assertInstanceOf(TestUser::class, $updated[0]);
+
+		$this->assertEquals(1, $updated[0]->id);
+		$this->assertEquals('bax', $updated[0]->username);
+		$this->assertSame(['id' => 1, 'username' => 'bax'], $updated[0]->toArray());
+	}
 }

tests/integration/database/query/compilers/BaseCompilerTest.php

@@ -10,6 +10,8 @@
 use LogicException;
 use mako\database\exceptions\NotFoundException;
 use mako\database\query\Query;
+use mako\database\query\Result;
+use mako\database\query\ResultSet;
 use mako\database\query\Subquery;
 use mako\pagination\PaginationFactoryInterface;
 use mako\pagination\PaginationInterface;
@@ -455,4 +457,22 @@ public function testEnum(): void
 
 		$this->assertSame(UsernameEnum::foo->name, $user->username);
 	}
+
+	/**
+	 *
+	 */
+	public function testUpdateAndReturn(): void
+	{
+		$query = new Query($this->connectionManager->getConnection());
+
+		$updated = $query->table('users')->where('id', '=', 1)->updateAndReturn(['username' => 'bax'], ['id', 'username']);
+
+		$this->assertInstanceOf(ResultSet::class, $updated);
+		$this->assertInstanceOf(Result::class, $updated[0]);
+
+		$this->assertSame(1, $updated[0]->id);
+		$this->assertSame('bax', $updated[0]->username);
+
+		$this->assertEquals('UPDATE "users" SET "username" = \'bax\' WHERE "id" = 1 RETURNING "id", "username"', $this->connectionManager->getConnection()->getLog()[0]['query']);
+	}
 }

tests/unit/database/query/compilers/FirebirdCompilerTest.php

@@ -315,4 +315,19 @@ public function testOrWhereDate(): void
 		$this->assertEquals('SELECT * FROM "foobar" WHERE "foo" = ? OR CAST("date" AS DATE) = ?', $query['sql']);
 		$this->assertEquals(['bar', '2019-07-05'], $query['params']);
 	}
+
+	/**
+	 *
+	 */
+	public function testUpdateAndReturn(): void
+	{
+		$query = $this->getBuilder();
+
+		$query->where('id', '=', 1);
+
+		$query = $query->getCompiler()->updateAndReturn(['foo' => 'bar'], ['id', 'foo']);
+
+		$this->assertEquals('UPDATE "foobar" SET "foo" = ? WHERE "id" = ? RETURNING "id", "foo"', $query['sql']);
+		$this->assertEquals(['bar', 1], $query['params']);
+	}
 }

tests/unit/database/query/compilers/PostgresCompilerTest.php

@@ -400,4 +400,19 @@ public function testInsertOrUpdateWithMultipleConstraints(): void
 		$this->assertEquals('INSERT INTO "foobar" ("foo") VALUES (?) ON CONFLICT ("foo", "bar") DO UPDATE SET "foo" = ?', $query['sql']);
 		$this->assertEquals(['bar', 'dupe'], $query['params']);
 	}
+
+	/**
+	 *
+	 */
+	public function testUpdateAndReturn(): void
+	{
+		$query = $this->getBuilder();
+
+		$query->where('id', '=', 1);
+
+		$query = $query->getCompiler()->updateAndReturn(['foo' => 'bar'], ['id', 'foo']);
+
+		$this->assertEquals('UPDATE "foobar" SET "foo" = ? WHERE "id" = ? RETURNING "id", "foo"', $query['sql']);
+		$this->assertEquals(['bar', 1], $query['params']);
+	}
 }

tests/unit/database/query/compilers/SQLServerCompilerTest.php

@@ -454,4 +454,19 @@ public function testOrWhereDate(): void
 		$this->assertEquals('SELECT * FROM [foobar] WHERE [foo] = ? OR CAST([date] AS DATE) = ?', $query['sql']);
 		$this->assertEquals(['bar', '2019-07-05'], $query['params']);
 	}
+
+	/**
+	 *
+	 */
+	public function testUpdateAndReturn(): void
+	{
+		$query = $this->getBuilder();
+
+		$query->where('id', '=', 1);
+
+		$query = $query->getCompiler()->updateAndReturn(['foo' => 'bar'], ['id', 'foo']);
+
+		$this->assertEquals('UPDATE [foobar] SET [foo] = ? OUTPUT inserted.[id], inserted.[foo] WHERE [id] = ?', $query['sql']);
+		$this->assertEquals(['bar', 1], $query['params']);
+	}
 }

tests/unit/database/query/compilers/SQLiteCompilerTest.php

@@ -329,4 +329,19 @@ public function testInsertOrUpdateWithMultipleConstraints(): void
 		$this->assertEquals('INSERT INTO "foobar" ("foo") VALUES (?) ON CONFLICT ("foo", "bar") DO UPDATE SET "foo" = ?', $query['sql']);
 		$this->assertEquals(['bar', 'dupe'], $query['params']);
 	}
+
+	/**
+	 *
+	 */
+	public function testUpdateAndReturn(): void
+	{
+		$query = $this->getBuilder();
+
+		$query->where('id', '=', 1);
+
+		$query = $query->getCompiler()->updateAndReturn(['foo' => 'bar'], ['id', 'foo']);
+
+		$this->assertEquals('UPDATE "foobar" SET "foo" = ? WHERE "id" = ? RETURNING "id", "foo"', $query['sql']);
+		$this->assertEquals(['bar', 1], $query['params']);
+	}
 }