Implement DMLQueryBuilder::upsertWithReturning() method

CHANGELOG.md

@@ -39,8 +39,8 @@
 - Enh #389, #390: Use `DbArrayHelper::arrange()` instead of `DbArrayHelper::index()` method (@Tigrov)
 - New #387: Realize `Schema::loadResultColumn()` method (@Tigrov)
 - New #393: Use `DateTimeColumn` class for datetime column types (@Tigrov)
-- New #394: Implement `Command::upsertWithReturningPks()` method (@Tigrov)
-- Enh #394: Refactor `Command::insertWithReturningPks()` method (@Tigrov)
+- New #394, #395: Implement `Command::upsertReturning()` method (@Tigrov)
+- Enh #394, #395: Refactor `Command::insertWithReturningPks()` method (@Tigrov)
 
 ## 1.2.0 March 21, 2024
 

phpunit.xml.dist

@@ -1,5 +1,12 @@
 <?xml version="1.0" encoding="UTF-8"?>
-<phpunit xmlns:xsi="http://www.w3.org/2001/XMLSchema-instance" backupGlobals="false" colors="true" bootstrap="vendor/autoload.php" failOnRisky="true" failOnWarning="true" executionOrder="default" resolveDependencies="true" xsi:noNamespaceSchemaLocation="https://schema.phpunit.de/10.1/phpunit.xsd">
+<phpunit
+    xmlns:xsi="http://www.w3.org/2001/XMLSchema-instance"
+    colors="true"
+    bootstrap="vendor/autoload.php"
+    failOnRisky="true"
+    failOnWarning="true"
+    xsi:noNamespaceSchemaLocation="vendor/phpunit/phpunit/schema/10.4.xsd"
+>
   <coverage/>
   <php>
     <ini name="error_reporting" value="-1"/>

src/Command.php

@@ -4,6 +4,8 @@
 
 namespace Yiisoft\Db\Mysql;
 
+use PDO;
+use PDOStatement;
 use Yiisoft\Db\Driver\Pdo\AbstractPdoCommand;
 use Yiisoft\Db\Exception\IntegrityException;
 use Yiisoft\Db\Exception\NotSupportedException;
@@ -22,21 +24,103 @@ final class Command extends AbstractPdoCommand
 {
     public function insertWithReturningPks(string $table, array|QueryInterface $columns): array|false
     {
+        $tableSchema = $this->db->getSchema()->getTableSchema($table);
+        $primaryKeys = $tableSchema?->getPrimaryKey() ?? [];
+        $tableColumns = $tableSchema?->getColumns() ?? [];
+
+        foreach ($primaryKeys as $name) {
+            /** @var ColumnInterface $column */
+            $column = $tableColumns[$name];
+
+            if ($column->isAutoIncrement()) {
+                continue;
+            }
+
+            if ($columns instanceof QueryInterface) {
+                throw new NotSupportedException(
+                    __METHOD__ . '() is not supported by MySQL for tables without auto increment when inserting sub-query.'
+                );
+            }
+
+            break;
+        }
+
         $params = [];
-        $sql = $this->db->getQueryBuilder()->insert($table, $columns, $params);
+        $insertSql = $this->db->getQueryBuilder()->insert($table, $columns, $params);
+        $this->setSql($insertSql)->bindValues($params);
+
+        if ($this->execute() === 0) {
+            return false;
+        }
+
+        if (empty($primaryKeys)) {
+            return [];
+        }
 
-        return $this->executeWithReturningPks($sql, $params, $table, $columns, __METHOD__);
+        $result = [];
+
+        foreach ($primaryKeys as $name) {
+            /** @var ColumnInterface $column */
+            $column = $tableColumns[$name];
+
+            if ($column->isAutoIncrement()) {
+                $value = $this->db->getLastInsertId();
+            } else {
+                /** @var array $columns */
+                $value = $columns[$name] ?? $column->getDefaultValue();
+            }
+
+            if ($this->phpTypecasting) {
+                $value = $column->phpTypecast($value);
+            }
+
+            $result[$name] = $value;
+        }
+
+        return $result;
     }
 
-    public function upsertWithReturningPks(
+    public function upsertReturning(
         string $table,
         array|QueryInterface $insertColumns,
-        array|bool $updateColumns = true
+        array|bool $updateColumns = true,
+        array|null $returnColumns = null,
     ): array|false {
+        $returnColumns ??= $this->db->getTableSchema($table)?->getColumnNames();
+
+        if (empty($returnColumns)) {
+            $this->upsert($table, $insertColumns, $updateColumns)->execute();
+            return [];
+        }
+
         $params = [];
-        $sql = $this->db->getQueryBuilder()->upsert($table, $insertColumns, $updateColumns, $params);
+        $sql = $this->getQueryBuilder()
+            ->upsertReturning($table, $insertColumns, $updateColumns, $returnColumns, $params);
+
+        $this->setSql($sql)->bindValues($params);
+        $this->queryInternal(self::QUERY_MODE_EXECUTE);
+
+        /** @psalm-var PDOStatement $this->pdoStatement */
+        $this->pdoStatement->nextRowset();
+        /** @psalm-var array<string,mixed>|false $result */
+        $result = $this->pdoStatement->fetch(PDO::FETCH_ASSOC);
+        $this->pdoStatement->closeCursor();
 
-        return $this->executeWithReturningPks($sql, $params, $table, $insertColumns, __METHOD__);
+        if (!$this->phpTypecasting || $result === false) {
+            return $result;
+        }
+
+        $columns = $this->db->getTableSchema($table)?->getColumns();
+
+        if (empty($columns)) {
+            return $result;
+        }
+
+        foreach ($result as $name => &$value) {
+            $value = $columns[$name]->phpTypecast($value);
+        }
+
+        return $result;
     }
 
     protected function queryInternal(int $queryMode): mixed
@@ -67,51 +151,4 @@ public function showDatabases(): array
 
         return $this->setSql($sql)->queryColumn();
     }
-
-    private function executeWithReturningPks(
-        string $sql,
-        array $params,
-        string $table,
-        array|QueryInterface $columns,
-        string $method,
-    ): array|false {
-        $tableSchema = $this->db->getSchema()->getTableSchema($table);
-        $primaryKeys = $tableSchema?->getPrimaryKey() ?? [];
-
-        if ($columns instanceof QueryInterface && !empty($primaryKeys)) {
-            throw new NotSupportedException($method . '() not supported for QueryInterface by MySQL.');
-        }
-
-        $this->setSql($sql)->bindValues($params);
-
-        if ($this->execute() === 0) {
-            return false;
-        }
-
-        if (empty($primaryKeys)) {
-            return [];
-        }
-
-        $result = [];
-
-        /** @var TableSchema $tableSchema */
-        foreach ($primaryKeys as $name) {
-            /** @var ColumnInterface $column */
-            $column = $tableSchema->getColumn($name);
-
-            if ($column->isAutoIncrement()) {
-                $value = $this->db->getLastInsertId();
-            } else {
-                $value = $columns[$name] ?? $column->getDefaultValue();
-            }
-
-            if ($this->phpTypecasting) {
-                $value = $column->phpTypecast($value);
-            }
-
-            $result[$name] = $value;
-        }
-
-        return $result;
-    }
 }

src/DMLQueryBuilder.php

@@ -7,10 +7,17 @@
 use Yiisoft\Db\Exception\InvalidArgumentException;
 use Yiisoft\Db\Exception\NotSupportedException;
 use Yiisoft\Db\Expression\Expression;
+use Yiisoft\Db\Expression\ExpressionInterface;
 use Yiisoft\Db\Query\QueryInterface;
 use Yiisoft\Db\QueryBuilder\AbstractDMLQueryBuilder;
 
+use function array_diff;
+use function array_intersect;
+use function array_key_exists;
+use function array_keys;
+use function array_map;
 use function implode;
+use function is_array;
 use function str_replace;
 
 /**
@@ -68,28 +75,128 @@ public function upsert(
         if ($updateColumns === true) {
             $updateColumns = [];
             /** @psalm-var string[] $updateNames */
-            foreach ($updateNames as $quotedName) {
-                $updateColumns[$quotedName] = new Expression('VALUES(' . $quotedName . ')');
+            foreach ($updateNames as $name) {
+                $updateColumns[$name] = new Expression('VALUES(' . $this->quoter->quoteColumnName($name) . ')');
             }
         }
 
         if (empty($updateColumns)) {
             return str_replace('INSERT INTO', 'INSERT IGNORE INTO', $insertSql);
         }
 
-        [$updates, $params] = $this->prepareUpdateSets($table, $updateColumns, $params);
+        $updates = $this->prepareUpdateSets($table, $updateColumns, $params);
 
         return $insertSql . ' ON DUPLICATE KEY UPDATE ' . implode(', ', $updates);
     }
 
-    /** @throws NotSupportedException */
-    public function upsertWithReturningPks(
+    public function upsertReturning(
         string $table,
         array|QueryInterface $insertColumns,
         array|bool $updateColumns = true,
+        array|null $returnColumns = null,
         array &$params = [],
     ): string {
-        throw new NotSupportedException(__METHOD__ . '() is not supported by MySQL.');
+        $tableSchema = $this->schema->getTableSchema($table);
+        $returnColumns ??= $tableSchema?->getColumnNames();
+
+        $upsertSql = $this->upsert($table, $insertColumns, $updateColumns, $params);
+
+        if (empty($returnColumns)) {
+            return $upsertSql;
+        }
+
+        $quoter = $this->quoter;
+        /** @var TableSchema $tableSchema */
+        $uniqueColumns = $tableSchema->getPrimaryKey()
+            ?: $this->prepareUpsertColumns($table, $insertColumns, $updateColumns)[0];
+
+        if (empty($uniqueColumns)) {
+            $returnValues = $this->prepareColumnValues($tableSchema, $returnColumns, $insertColumns, $params);
+            $selectValues = [];
+
+            foreach ($returnValues as $name => $value) {
+                $selectValues[] = $value . ' ' . $quoter->quoteColumnName($name);
+            }
+
+            return $upsertSql . ';' . 'SELECT ' . implode(', ', $selectValues);
+        }
+
+        if (is_array($updateColumns) && !empty(array_intersect($uniqueColumns, array_keys($updateColumns)))) {
+            throw new NotSupportedException(
+                __METHOD__ . '() is not supported by MySQL when updating primary key or unique values.'
+            );
+        }
+
+        $uniqueValues = $this->prepareColumnValues($tableSchema, $uniqueColumns, $insertColumns, $params);
+
+        if (empty(array_diff($returnColumns, array_keys($uniqueValues)))) {
+            $selectValues = [];
+
+            foreach ($returnColumns as $name) {
+                $selectValues[] = $uniqueValues[$name] . ' ' . $quoter->quoteColumnName($name);
+            }
+
+            return $upsertSql . ';' . 'SELECT ' . implode(', ', $selectValues);
+        }
+
+        $conditions = [];
+
+        foreach ($uniqueValues as $name => $value) {
+            if (array_key_exists($value, $params) && $params[$value] === null) {
+                throw new NotSupportedException(
+                    __METHOD__ . '() is not supported by MySQL when inserting `null` primary key or unique values.'
+                );
+            }
+
+            $conditions[] = $quoter->quoteColumnName($name) . ' = ' . $value;
+        }
+
+        $quotedReturnColumns = array_map($quoter->quoteColumnName(...), $returnColumns);
+
+        return $upsertSql . ';'
+            . 'SELECT ' . implode(', ', $quotedReturnColumns)
+            . ' FROM ' . $this->quoter->quoteTableName($table)
+            . ' WHERE ' . implode(' AND ', $conditions);
+    }
+
+    /**
+     * @param string[] $columnNames
+     *
+     * @return string[] Prepared column values for using in a SQL statement.
+     * @psalm-return array<string, string>
+     */
+    private function prepareColumnValues(
+        TableSchema $tableSchema,
+        array $columnNames,
+        array|QueryInterface $insertColumns,
+        array &$params,
+    ): array {
+        $columnValues = [];
+
+        $tableColumns = $tableSchema->getColumns();
+
+        foreach ($columnNames as $name) {
+            $column = $tableColumns[$name];
+
+            if ($column->isAutoIncrement()) {
+                $columnValues[$name] = 'LAST_INSERT_ID()';
+            } elseif ($insertColumns instanceof QueryInterface) {
+                throw new NotSupportedException(
+                    self::class . '::upsertReturning() is not supported by MySQL'
+                    . ' for tables without auto increment when inserting sub-query.'
+                );
+            } else {
+                $value = $insertColumns[$name] ?? $column->getDefaultValue();
+
+                if ($value instanceof ExpressionInterface) {
+                    $columnValues[$name] = $this->queryBuilder->buildExpression($value, $params);
+                } else {
+                    $columnValues[$name] = $this->queryBuilder->bindParam($value, $params);
+                }
+            }
+        }
+
+        return $columnValues;
     }
 
     protected function prepareInsertValues(string $table, array|QueryInterface $columns, array $params = []): array

tests/CommandTest.php

@@ -6,7 +6,6 @@
 
 use PHPUnit\Framework\Attributes\DataProviderExternal;
 use Yiisoft\Db\Exception\NotSupportedException;
-use Yiisoft\Db\Expression\Expression;
 use Yiisoft\Db\Mysql\Tests\Provider\CommandProvider;
 use Yiisoft\Db\Mysql\Tests\Support\TestTrait;
 use Yiisoft\Db\Query\Query;
@@ -100,16 +99,19 @@ public function testGetRawSql(string $sql, array $params, string $expectedRawSql
         parent::testGetRawSql($sql, $params, $expectedRawSql);
     }
 
-    public function testInsertWithReturningPksWithQuery(): void
+    public function testInsertWithReturningPksWithSubqueryAndNoAutoincrement(): void
     {
         $db = $this->getConnection(true);
         $command = $db->createCommand();
-        $query = (new Query($db))->select(new Expression("'new category'"));
+
+        $query = (new Query($db))->select(['order_id' => 1, 'item_id' => 2, 'quantity' => 3, 'subtotal' => 4]);
 
         $this->expectException(NotSupportedException::class);
-        $this->expectExceptionMessage('Yiisoft\Db\Mysql\Command::insertWithReturningPks() not supported for QueryInterface by MySQL.');
+        $this->expectExceptionMessage(
+            'Yiisoft\Db\Mysql\Command::insertWithReturningPks() is not supported by MySQL for tables without auto increment when inserting sub-query.'
+        );
 
-        $command->insertWithReturningPks('category', $query);
+        $command->insertWithReturningPks('order_item', $query);
     }
 
     #[DataProviderExternal(CommandProvider::class, 'update')]
@@ -130,18 +132,6 @@ public function testUpsert(array $firstData, array $secondData): void
         parent::testUpsert($firstData, $secondData);
     }
 
-    public function testUpsertWithReturningPksWithQuery(): void
-    {
-        $db = $this->getConnection(true);
-        $command = $db->createCommand();
-        $query = (new Query($db))->select(new Expression("'new category'"));
-
-        $this->expectException(NotSupportedException::class);
-        $this->expectExceptionMessage('Yiisoft\Db\Mysql\Command::upsertWithReturningPks() not supported for QueryInterface by MySQL.');
-
-        $command->upsertWithReturningPks('category', $query);
-    }
-
     public function testShowDatabases(): void
     {
         $this->assertSame([self::getDatabaseName()], self::getDb()->createCommand()->showDatabases());