Replace upsertWithReturningPks() with upsertWithReturning() in DMLQueryBuilder class

CHANGELOG.md

@@ -38,8 +38,8 @@
 - New #348: Realize `Schema::loadResultColumn()` method (@Tigrov)
 - New #354: Add `FOR` clause to query (@vjik)
 - New #355: Use `DateTimeColumn` class for datetime column types (@Tigrov)
-- New #356: Implement `DMLQueryBuilder::upsertWithReturningPks()` method (@Tigrov)
-- Enh #356: Refactor `Command::insertWithReturningPks()` and `DMLQueryBuilder::upsert()` methods (@Tigrov)
+- New #356, #357: Implement `DMLQueryBuilder::upsertReturning()` method (@Tigrov)
+- Enh #356, #357: Refactor `Command::insertWithReturningPks()` and `DMLQueryBuilder::upsert()` methods (@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

@@ -27,35 +27,49 @@ 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 Sqlite for tables without auto increment when inserting sub-query.'
+                );
+            }
+
+            break;
+        }
+
         $params = [];
-        $sql = $this->db->getQueryBuilder()->insert($table, $columns, $params);
-        $this->setSql($sql)->bindValues($params);
+        $insertSql = $this->db->getQueryBuilder()->insert($table, $columns, $params);
+        $this->setSql($insertSql)->bindValues($params);
 
         if ($this->execute() === 0) {
             return false;
         }
 
-        $tableSchema = $this->db->getSchema()->getTableSchema($table);
-        $tablePrimaryKeys = $tableSchema?->getPrimaryKey() ?? [];
-
-        if (empty($tablePrimaryKeys)) {
+        if (empty($primaryKeys)) {
             return [];
         }
 
-        if ($columns instanceof QueryInterface) {
-            throw new NotSupportedException(__METHOD__ . '() not supported for QueryInterface by SQLite.');
-        }
-
         $result = [];
 
-        /** @var TableSchema $tableSchema */
-        foreach ($tablePrimaryKeys as $name) {
+        foreach ($primaryKeys as $name) {
             /** @var ColumnInterface $column */
-            $column = $tableSchema->getColumn($name);
+            $column = $tableColumns[$name];
 
             if ($column->isAutoIncrement()) {
                 $value = $this->db->getLastInsertId();
             } else {
+                /** @var array $columns */
                 $value = $columns[$name] ?? $column->getDefaultValue();
             }
 

src/DMLQueryBuilder.php

@@ -80,27 +80,55 @@ public function upsert(
             }
         }
 
-        [$updates, $params] = $this->prepareUpdateSets($table, $updateColumns, $params);
+        $quotedUniqueNames = array_map($this->quoter->quoteColumnName(...), $uniqueNames);
+        $updates = $this->prepareUpdateSets($table, $updateColumns, $params);
 
         return $insertSql
-            . ' ON CONFLICT (' . implode(', ', $uniqueNames) . ') DO UPDATE SET ' . implode(', ', $updates);
+            . ' ON CONFLICT (' . implode(', ', $quotedUniqueNames) . ')'
+            . ' DO UPDATE SET ' . implode(', ', $updates);
     }
 
-    public function upsertWithReturningPks(
+    public function upsertReturning(
         string $table,
         array|QueryInterface $insertColumns,
         array|bool $updateColumns = true,
+        array|null $returnColumns = null,
         array &$params = [],
     ): string {
-        $sql = $this->upsert($table, $insertColumns, $updateColumns, $params);
-        $returnColumns = $this->schema->getTableSchema($table)?->getPrimaryKey();
+        $upsertSql = $this->upsert($table, $insertColumns, $updateColumns, $params);
 
-        if (!empty($returnColumns)) {
-            $returnColumns = array_map($this->quoter->quoteColumnName(...), $returnColumns);
+        $returnColumns ??= $this->schema->getTableSchema($table)?->getColumnNames();
 
-            $sql .= ' RETURNING ' . implode(', ', $returnColumns);
+        if (empty($returnColumns)) {
+            return $upsertSql;
         }
 
-        return $sql;
+        $returnColumns = array_map($this->quoter->quoteColumnName(...), $returnColumns);
+
+        if (str_ends_with($upsertSql, ' ON CONFLICT DO NOTHING')) {
+            $dummyColumn = $this->getDummyColumn($table);
+
+            $upsertSql = substr($upsertSql, 0, -10) . "DO UPDATE SET $dummyColumn = $dummyColumn";
+        }
+
+        return $upsertSql . ' RETURNING ' . implode(', ', $returnColumns);
+    }
+
+    private function getDummyColumn(string $table): string
+    {
+        /** @psalm-suppress PossiblyNullReference */
+        $columns = $this->schema->getTableSchema($table)->getColumns();
+
+        foreach ($columns as $column) {
+            if ($column->isPrimaryKey() || $column->isUnique()) {
+                continue;
+            }
+
+            /** @psalm-suppress PossiblyNullArgument */
+            return $this->quoter->quoteColumnName($column->getName());
+        }
+
+        /** @psalm-suppress PossiblyNullArgument, PossiblyFalseReference */
+        return $this->quoter->quoteColumnName(end($columns)->getName());
     }
 }

tests/CommandTest.php

@@ -10,6 +10,7 @@
 use Yiisoft\Db\Exception\Exception;
 use Yiisoft\Db\Exception\InvalidConfigException;
 use Yiisoft\Db\Exception\NotSupportedException;
+use Yiisoft\Db\Query\Query;
 use Yiisoft\Db\Sqlite\Tests\Support\TestTrait;
 use Yiisoft\Db\Tests\Common\CommonCommandTest;
 
@@ -326,6 +327,21 @@ public function testGetRawSql(string $sql, array $params, string $expectedRawSql
         parent::testGetRawSql($sql, $params, $expectedRawSql);
     }
 
+    public function testInsertWithReturningPksWithSubqueryAndNoAutoincrement(): void
+    {
+        $db = $this->getConnection(true);
+        $command = $db->createCommand();
+
+        $query = (new Query($db))->select(['order_id' => 1, 'item_id' => 2, 'quantity' => 3, 'subtotal' => 4]);
+
+        $this->expectException(NotSupportedException::class);
+        $this->expectExceptionMessage(
+            'Yiisoft\Db\Sqlite\Command::insertWithReturningPks() is not supported by Sqlite for tables without auto increment when inserting sub-query.'
+        );
+
+        $command->insertWithReturningPks('order_item', $query);
+    }
+
     /**
      * @throws Throwable
      * @throws InvalidConfigException

tests/Provider/QueryBuilderProvider.php

@@ -13,6 +13,7 @@
 use Yiisoft\Db\Tests\Support\TraversableObject;
 
 use function array_replace;
+use function array_splice;
 
 final class QueryBuilderProvider extends \Yiisoft\Db\Tests\Provider\QueryBuilderProvider
 {
@@ -160,7 +161,7 @@ public static function upsert(): array
             ],
             'values and expressions without update part' => [
                 3 => <<<SQL
-                INSERT INTO {{%T_upsert}} (`email`, `ts`) VALUES (:qp0, CURRENT_TIMESTAMP) ON CONFLICT DO NOTHING
+                INSERT INTO `T_upsert` (`email`, `ts`) VALUES (:qp0, CURRENT_TIMESTAMP) ON CONFLICT DO NOTHING
                 SQL,
             ],
             'query, values and expressions with update part' => [
@@ -170,7 +171,7 @@ public static function upsert(): array
             ],
             'query, values and expressions without update part' => [
                 3 => <<<SQL
-                INSERT INTO {{%T_upsert}} (`email`, [[ts]]) SELECT :phEmail AS `email`, CURRENT_TIMESTAMP AS [[ts]] ON CONFLICT DO NOTHING
+                INSERT INTO `T_upsert` (`email`, [[ts]]) SELECT :phEmail AS `email`, CURRENT_TIMESTAMP AS [[ts]] ON CONFLICT DO NOTHING
                 SQL,
             ],
             'no columns to update' => [
@@ -180,7 +181,7 @@ public static function upsert(): array
             ],
             'no columns to update with unique' => [
                 3 => <<<SQL
-                INSERT INTO {{%T_upsert}} (`email`) VALUES (:qp0) ON CONFLICT DO NOTHING
+                INSERT INTO `T_upsert` (`email`) VALUES (:qp0) ON CONFLICT DO NOTHING
                 SQL,
             ],
             'no unique columns in table - simple insert' => [
@@ -199,33 +200,60 @@ public static function upsert(): array
         return $upsert;
     }
 
-    public static function upsertWithReturningPks(): array
+    public static function upsertReturning(): array
     {
         $upsert = self::upsert();
 
-        foreach ($upsert as &$data) {
-            $data[3] .= ' RETURNING `id`';
+        $withoutUpdate = [
+            'regular values without update part',
+            'query without update part',
+            'values and expressions without update part',
+            'query, values and expressions without update part',
+            'no columns to update with unique',
+        ];
+
+        foreach ($upsert as $name => &$data) {
+            array_splice($data, 3, 0, [['id']]);
+            if (in_array($name, $withoutUpdate, true)) {
+                $data[4] = substr($data[4], 0, -10) . 'DO UPDATE SET `ts` = `ts`';
+            }
+
+            $data[4] .= ' RETURNING `id`';
         }
 
-        $upsert['no columns to update'][3] = 'INSERT INTO `T_upsert_1` (`a`) VALUES (:qp0) ON CONFLICT DO NOTHING RETURNING `a`';
+        $upsert['no columns to update'][3] = ['a'];
+        $upsert['no columns to update'][4] = 'INSERT INTO `T_upsert_1` (`a`) VALUES (:qp0) ON CONFLICT DO UPDATE SET `a` = `a` RETURNING `a`';
 
         return [
             ...$upsert,
             'composite primary key' => [
                 'notauto_pk',
                 ['id_1' => 1, 'id_2' => 2.5, 'type' => 'Test'],
                 true,
+                ['id_1', 'id_2'],
                 'INSERT INTO `notauto_pk` (`id_1`, `id_2`, `type`) VALUES (:qp0, :qp1, :qp2)'
                 . ' ON CONFLICT (`id_1`, `id_2`) DO UPDATE SET `type`=EXCLUDED.`type` RETURNING `id_1`, `id_2`',
                 [':qp0' => 1, ':qp1' => 2.5, ':qp2' => 'Test'],
             ],
-            'no primary key' => [
+            'no return columns' => [
                 'type',
                 ['int_col' => 3, 'char_col' => 'a', 'float_col' => 1.2, 'bool_col' => true],
                 true,
+                [],
                 'INSERT INTO `type` (`int_col`, `char_col`, `float_col`, `bool_col`) VALUES (:qp0, :qp1, :qp2, :qp3)',
                 [':qp0' => 3, ':qp1' => 'a', ':qp2' => 1.2, ':qp3' => true],
             ],
+            'return all columns' => [
+                'T_upsert',
+                ['email' => '[email protected]', 'address' => 'test address', 'status' => 1, 'profile_id' => 1],
+                true,
+                null,
+                'INSERT INTO `T_upsert` (`email`, `address`, `status`, `profile_id`) VALUES (:qp0, :qp1, :qp2, :qp3)'
+                . ' ON CONFLICT (`email`) DO UPDATE SET'
+                . ' `address`=EXCLUDED.`address`, `status`=EXCLUDED.`status`, `profile_id`=EXCLUDED.`profile_id`'
+                . ' RETURNING `id`, `ts`, `email`, `recovery_email`, `address`, `status`, `orders`, `profile_id`',
+                [':qp0' => '[email protected]', ':qp1' => 'test address', ':qp2' => 1, ':qp3' => 1],
+            ],
         ];
     }
 

tests/QueryBuilderTest.php

@@ -623,15 +623,16 @@ public function testUpsert(
         parent::testUpsert($table, $insertColumns, $updateColumns, $expectedSql, $expectedParams);
     }
 
-    #[DataProviderExternal(QueryBuilderProvider::class, 'upsertWithReturningPks')]
-    public function testUpsertWithReturningPks(
+    #[DataProviderExternal(QueryBuilderProvider::class, 'upsertReturning')]
+    public function testUpsertReturning(
         string $table,
         array|QueryInterface $insertColumns,
         array|bool $updateColumns,
+        array|null $returnColumns,
         string $expectedSql,
         array $expectedParams
     ): void {
-        parent::testUpsertWithReturningPks($table, $insertColumns, $updateColumns, $expectedSql, $expectedParams);
+        parent::testUpsertReturning($table, $insertColumns, $updateColumns, $returnColumns, $expectedSql, $expectedParams);
     }
 
     #[DataProviderExternal(QueryBuilderProvider::class, 'selectScalar')]