yiisoft · Tigrov · Jun 8, 2025 · May 22, 2025 · May 24, 2025 · May 24, 2025
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -40,7 +40,7 @@
 - Enh #403, #404: Use `DbArrayHelper::arrange()` instead of `DbArrayHelper::index()` method (@Tigrov)
 - New #397: Realize `Schema::loadResultColumn()` method (@Tigrov)
 - New #407: Use `DateTimeColumn` class for datetime column types (@Tigrov)
-- New #408: Implement `DMLQueryBuilder::upsertWithReturningPks()` method (@Tigrov)
+- New #408, #410: Implement `DMLQueryBuilder::upsertReturning()` method (@Tigrov)
 
 ## 1.3.0 March 21, 2024
 

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

diff --git a/src/DMLQueryBuilder.php b/src/DMLQueryBuilder.php
@@ -11,6 +11,8 @@
 
 use function array_map;
 use function implode;
+use function str_ends_with;
+use function substr;
 
 /**
  * Implements a DML (Data Manipulation Language) SQL statements for PostgreSQL Server.
@@ -19,9 +21,17 @@ final class DMLQueryBuilder extends AbstractDMLQueryBuilder
 {
     public function insertWithReturningPks(string $table, array|QueryInterface $columns, array &$params = []): string
     {
-        $sql = $this->insert($table, $columns, $params);
+        $insertSql = $this->insert($table, $columns, $params);
+        $tableSchema = $this->schema->getTableSchema($table);
+        $primaryKeys = $tableSchema?->getPrimaryKey() ?? [];
+
+        if (empty($primaryKeys)) {
+            return $insertSql;
+        }
 
-        return $this->appendReturningPksClause($sql, $table);
+        $primaryKeys = array_map($this->quoter->quoteColumnName(...), $primaryKeys);
+
+        return $insertSql . ' RETURNING ' . implode(', ', $primaryKeys);
     }
 
     public function resetSequence(string $table, int|string|null $value = null): string
@@ -81,33 +91,61 @@ 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);
+        $upsertSql = $this->upsert($table, $insertColumns, $updateColumns, $params);
+
+        $returnColumns ??= $this->schema->getTableSchema($table)?->getColumnNames();
+
+        if (empty($returnColumns)) {
+            return $upsertSql;
+        }
+
+        $returnColumns = array_map($this->quoter->quoteColumnName(...), $returnColumns);
+
+        if (str_ends_with($upsertSql, ' ON CONFLICT DO NOTHING')) {
+            $tableName = $this->quoter->quoteTableName($table);
+            $dummyColumn = $this->getDummyColumn($table);
+
+            $uniqueNames = $this->prepareUpsertColumns($table, $insertColumns, $updateColumns)[0];
+            $quotedUniqueNames = array_map($this->quoter->quoteColumnName(...), $uniqueNames);
 
-        return $this->appendReturningPksClause($sql, $table);
+            $upsertSql = substr($upsertSql, 0, -10)
+                . '(' . implode(', ', $quotedUniqueNames) . ')'
+                . " DO UPDATE SET $dummyColumn = $tableName.$dummyColumn";
+        }
+
+        return $upsertSql . ' RETURNING ' . implode(', ', $returnColumns);
     }
 
-    private function appendReturningPksClause(string $sql, string $table): string
+    private function getDummyColumn(string $table): string
     {
-        $returnColumns = $this->schema->getTableSchema($table)?->getPrimaryKey();
+        /** @psalm-suppress PossiblyNullReference */
+        $columns = $this->schema->getTableSchema($table)->getColumns();
 
-        if (!empty($returnColumns)) {
-            $returnColumns = array_map($this->quoter->quoteColumnName(...), $returnColumns);
+        foreach ($columns as $column) {
+            if ($column->isPrimaryKey() || $column->isUnique()) {
+                continue;
+            }
 
-            $sql .= ' RETURNING ' . implode(', ', $returnColumns);
+            /** @psalm-suppress PossiblyNullArgument */
+            return $this->quoter->quoteColumnName($column->getName());
         }
 
-        return $sql;
+        /** @psalm-suppress PossiblyNullArgument, PossiblyFalseReference */
+        return $this->quoter->quoteColumnName(end($columns)->getName());
     }
 }
diff --git a/tests/Provider/CommandPDOProvider.php b/tests/Provider/CommandPDOProvider.php
@@ -16,7 +16,6 @@ public static function bindParam(): array
             'name' => 'user1',
             'address' => 'address1',
             'status' => 1,
-            'bool_status' => true,
             'profile_id' => 1,
         ];
 

diff --git a/tests/Provider/QueryBuilderProvider.php b/tests/Provider/QueryBuilderProvider.php
@@ -265,7 +265,7 @@ public static function upsert(): array
             ],
             'values and expressions without update part' => [
                 1 => ['{{%T_upsert}}.[[email]]' => 'dynamic@example.com', '[[ts]]' => new Expression('extract(epoch from now()) * 1000')],
-                3 => 'INSERT INTO {{%T_upsert}} ("email", "ts") VALUES (:qp0, extract(epoch from now()) * 1000) ON CONFLICT DO NOTHING',
+                3 => 'INSERT INTO "T_upsert" ("email", "ts") VALUES (:qp0, extract(epoch from now()) * 1000) ON CONFLICT DO NOTHING',
             ],
             'query, values and expressions with update part' => [
                 1 => (new Query(self::getDb()))
@@ -287,13 +287,13 @@ public static function upsert(): array
                             '[[ts]]' => new Expression('extract(epoch from now()) * 1000'),
                         ],
                     ),
-                3 => 'INSERT INTO {{%T_upsert}} ("email", [[ts]]) SELECT :phEmail AS "email", extract(epoch from now()) * 1000 AS [[ts]] ON CONFLICT DO NOTHING',
+                3 => 'INSERT INTO "T_upsert" ("email", [[ts]]) SELECT :phEmail AS "email", extract(epoch from now()) * 1000 AS [[ts]] ON CONFLICT DO NOTHING',
             ],
             'no columns to update' => [
                 3 => 'INSERT INTO "T_upsert_1" ("a") VALUES (:qp0) ON CONFLICT DO NOTHING',
             ],
             'no columns to update with unique' => [
-                3 => 'INSERT INTO {{%T_upsert}} ("email") VALUES (:qp0) ON CONFLICT DO NOTHING',
+                3 => 'INSERT INTO "T_upsert" ("email") VALUES (:qp0) ON CONFLICT DO NOTHING',
             ],
             'no unique columns in table - simple insert' => [
                 3 => 'INSERT INTO {{%animal}} ("type") VALUES (:qp0)',
@@ -317,33 +317,61 @@ 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) . '("email") DO UPDATE SET "ts" = "T_upsert"."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 ("a")'
+            . ' DO UPDATE SET "a" = "T_upsert_1"."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' => 'test@example.com', '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' => 'test@example.com', ':qp1' => 'test address', ':qp2' => 1, ':qp3' => 1],
+            ],
         ];
     }
 

diff --git a/tests/QueryBuilderTest.php b/tests/QueryBuilderTest.php
@@ -437,15 +437,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')]

diff --git a/tests/Support/Fixture/pgsql.sql b/tests/Support/Fixture/pgsql.sql
@@ -76,7 +76,6 @@ CREATE TABLE "customer" (
   name varchar(128),
   address text,
   status integer DEFAULT 0,
-  bool_status boolean DEFAULT FALSE,
   profile_id integer
 );
 
@@ -270,9 +269,9 @@ INSERT INTO "profile" (description) VALUES ('profile customer 3');
 INSERT INTO "schema1"."profile" (description) VALUES ('profile customer 1');
 INSERT INTO "schema1"."profile" (description) VALUES ('profile customer 3');
 
-INSERT INTO "customer" (email, name, address, status, bool_status, profile_id) VALUES ('user1@example.com', 'user1', 'address1', 1, true, 1);
-INSERT INTO "customer" (email, name, address, status, bool_status) VALUES ('user2@example.com', 'user2', 'address2', 1, true);
-INSERT INTO "customer" (email, name, address, status, bool_status, profile_id) VALUES ('user3@example.com', 'user3', 'address3', 2, false, 2);
+INSERT INTO "customer" (email, name, address, status, profile_id) VALUES ('user1@example.com', 'user1', 'address1', 1, 1);
+INSERT INTO "customer" (email, name, address, status) VALUES ('user2@example.com', 'user2', 'address2', 1);
+INSERT INTO "customer" (email, name, address, status, profile_id) VALUES ('user3@example.com', 'user3', 'address3', 2, 2);
 
 INSERT INTO "category" (name) VALUES ('Books');
 INSERT INTO "category" (name) VALUES ('Movies');

diff --git a/tests/Support/TestTrait.php b/tests/Support/TestTrait.php
@@ -4,8 +4,6 @@
 
 namespace Yiisoft\Db\Pgsql\Tests\Support;
 
-use Yiisoft\Db\Exception\Exception;
-use Yiisoft\Db\Exception\InvalidConfigException;
 use Yiisoft\Db\Pgsql\Connection;
 use Yiisoft\Db\Pgsql\Driver;
 use Yiisoft\Db\Pgsql\Dsn;
@@ -16,10 +14,15 @@ trait TestTrait
     private string $dsn = '';
     private string $fixture = 'pgsql.sql';
 
-    /**
-     * @throws InvalidConfigException
-     * @throws Exception
-     */
+    public static function setUpBeforeClass(): void
+    {
+        $db = self::getDb();
+
+        DbHelper::loadFixture($db, __DIR__ . '/Fixture/pgsql.sql');
+
+        $db->close();
+    }
+
     protected function getConnection(bool $fixture = false): Connection
     {
         $db = new Connection($this->getDriver(), DbHelper::getSchemaCache());
@@ -72,15 +75,6 @@ protected function setFixture(string $fixture): void
         $this->fixture = $fixture;
     }
 
-    public static function setUpBeforeClass(): void
-    {
-        $db = self::getDb();
-
-        DbHelper::loadFixture($db, __DIR__ . '/Fixture/pgsql.sql');
-
-        $db->close();
-    }
-
     protected function getDriver(): Driver
     {
         $driver = new Driver($this->getDsn(), self::getUsername(), self::getPassword());