Implement DMLQueryBuilder::upsertWithReturningPks() method

CHANGELOG.md

@@ -40,6 +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)
 
 ## 1.3.0 March 21, 2024
 

src/DMLQueryBuilder.php

@@ -17,18 +17,11 @@
  */
 final class DMLQueryBuilder extends AbstractDMLQueryBuilder
 {
-    public function insertWithReturningPks(string $table, QueryInterface|array $columns, array &$params = []): string
+    public function insertWithReturningPks(string $table, array|QueryInterface $columns, array &$params = []): string
     {
         $sql = $this->insert($table, $columns, $params);
-        $returnColumns = $this->schema->getTableSchema($table)?->getPrimaryKey();
-
-        if (!empty($returnColumns)) {
-            $returnColumns = array_map($this->quoter->quoteColumnName(...), $returnColumns);
-
-            $sql .= ' RETURNING ' . implode(', ', $returnColumns);
-        }
 
-        return $sql;
+        return $this->appendReturningPksClause($sql, $table);
     }
 
     public function resetSequence(string $table, int|string|null $value = null): string
@@ -60,9 +53,9 @@ public function resetSequence(string $table, int|string|null $value = null): str
 
     public function upsert(
         string $table,
-        QueryInterface|array $insertColumns,
-        bool|array $updateColumns,
-        array &$params = []
+        array|QueryInterface $insertColumns,
+        array|bool $updateColumns = true,
+        array &$params = [],
     ): string {
         $insertSql = $this->insert($table, $insertColumns, $params);
 
@@ -93,4 +86,28 @@ public function upsert(
         return $insertSql
             . ' ON CONFLICT (' . implode(', ', $uniqueNames) . ') DO UPDATE SET ' . implode(', ', $updates);
     }
+
+    public function upsertWithReturningPks(
+        string $table,
+        array|QueryInterface $insertColumns,
+        array|bool $updateColumns = true,
+        array &$params = [],
+    ): string {
+        $sql = $this->upsert($table, $insertColumns, $updateColumns, $params);
+
+        return $this->appendReturningPksClause($sql, $table);
+    }
+
+    private function appendReturningPksClause(string $sql, string $table): string
+    {
+        $returnColumns = $this->schema->getTableSchema($table)?->getPrimaryKey();
+
+        if (!empty($returnColumns)) {
+            $returnColumns = array_map($this->quoter->quoteColumnName(...), $returnColumns);
+
+            $sql .= ' RETURNING ' . implode(', ', $returnColumns);
+        }
+
+        return $sql;
+    }
 }

tests/Provider/QueryBuilderProvider.php

@@ -317,6 +317,36 @@ public static function upsert(): array
         return $upsert;
     }
 
+    public static function upsertWithReturningPks(): array
+    {
+        $upsert = self::upsert();
+
+        foreach ($upsert as &$data) {
+            $data[3] .= ' RETURNING "id"';
+        }
+
+        $upsert['no columns to update'][3] = 'INSERT INTO "T_upsert_1" ("a") VALUES (:qp0) ON CONFLICT DO NOTHING RETURNING "a"';
+
+        return [
+            ...$upsert,
+            'composite primary key' => [
+                'notauto_pk',
+                ['id_1' => 1, 'id_2' => 2.5, 'type' => 'Test'],
+                true,
+                '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' => [
+                '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],
+            ],
+        ];
+    }
+
     public static function overlapsCondition(): array
     {
         $data = parent::overlapsCondition();