Refactor DMLQueryBuilder (#313)

* Refactor DMLQueryBuilder * Add line to CHANGELOG.md * Improve performance of quoting column names up to 10% using `array_map()` * Remove `Generator` from tests * Add type `bool|array` to parameter `$updateColumns` of `upsert()` method
yiisoft · Nov 1, 2023 · cf03e81 · cf03e81
1 parent c467070
commit cf03e81
Show file tree

Hide file tree

Showing 4 changed files with 34 additions and 48 deletions.
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -7,6 +7,7 @@
 - Enh #302: Refactor `ColumnSchema` (@Tigrov)
 - Bug #302: Fix incorrect convert string value for BIT type (@Tigrov)
 - Bug #309: Fix retrieving sequence name from default value (@Tigrov)
+- Bug #313: Refactor `DMLQueryBuilder`, related with yiisoft/db#746 (@Tigrov)
 - Chg #319: Remove use of abstract type `SchemaInterface::TYPE_JSONB` (@Tigrov)
 
 ## 1.1.0 July 24, 2023

diff --git a/src/DMLQueryBuilder.php b/src/DMLQueryBuilder.php
@@ -9,8 +9,8 @@
 use Yiisoft\Db\Query\QueryInterface;
 use Yiisoft\Db\QueryBuilder\AbstractDMLQueryBuilder;
 
+use function array_map;
 use function implode;
-use function reset;
 
 /**
  * Implements a DML (Data Manipulation Language) SQL statements for PostgreSQL Server.
@@ -20,20 +20,15 @@ final class DMLQueryBuilder extends AbstractDMLQueryBuilder
     public function insertWithReturningPks(string $table, QueryInterface|array $columns, array &$params = []): string
     {
         $sql = $this->insert($table, $columns, $params);
-
-        $tableSchema = $this->schema->getTableSchema($table);
-
-        $returnColumns = [];
-        if ($tableSchema !== null) {
-            $returnColumns = $tableSchema->getPrimaryKey();
-        }
+        $returnColumns = $this->schema->getTableSchema($table)?->getPrimaryKey();
 
         if (!empty($returnColumns)) {
-            $returning = [];
-            foreach ($returnColumns as $name) {
-                $returning[] = $this->quoter->quoteColumnName($name);
-            }
-            $sql .= ' RETURNING ' . implode(', ', $returning);
+            $returnColumns = array_map(
+                [$this->quoter, 'quoteColumnName'],
+                $returnColumns,
+            );
+
+            $sql .= ' RETURNING ' . implode(', ', $returnColumns);
         }
 
         return $sql;
@@ -43,73 +38,59 @@ public function resetSequence(string $table, int|string $value = null): string
     {
         $tableSchema = $this->schema->getTableSchema($table);
 
-        if ($tableSchema !== null && ($sequence = $tableSchema->getSequenceName()) !== null) {
-            /**
-             * @link https://www.postgresql.org/docs/8.1/static/functions-sequence.html
-             */
-            $sequence = $this->quoter->quoteTableName($sequence);
-            $table = $this->quoter->quoteTableName($table);
+        if ($tableSchema === null) {
+            throw new InvalidArgumentException("Table not found: '$table'.");
+        }
 
-            if ($value === null) {
-                $pk = $tableSchema->getPrimaryKey();
-                $key = $this->quoter->quoteColumnName(reset($pk));
-                $value = "(SELECT COALESCE(MAX($key),0) FROM $table)+1";
-            }
+        $sequence = $tableSchema->getSequenceName();
 
-            return "SELECT SETVAL('$sequence',$value,false)";
+        if ($sequence === null) {
+            throw new InvalidArgumentException("There is not sequence associated with table '$table'.");
         }
 
-        if ($tableSchema === null) {
-            throw new InvalidArgumentException("Table not found: '$table'.");
+        /** @link https://www.postgresql.org/docs/8.1/static/functions-sequence.html */
+        $sequence = $this->quoter->quoteTableName($sequence);
+
+        if ($value === null) {
+            $table = $this->quoter->quoteTableName($table);
+            $key = $tableSchema->getPrimaryKey()[0];
+            $key = $this->quoter->quoteColumnName($key);
+            $value = "(SELECT COALESCE(MAX($key),0) FROM $table)+1";
         }
 
-        throw new InvalidArgumentException("There is not sequence associated with table '$table'.");
+        return "SELECT SETVAL('$sequence',$value,false)";
     }
 
     public function upsert(
         string $table,
         QueryInterface|array $insertColumns,
-        $updateColumns,
+        bool|array $updateColumns,
         array &$params = []
     ): string {
         $insertSql = $this->insert($table, $insertColumns, $params);
 
-        /** @psalm-var array $uniqueNames */
-        [$uniqueNames, , $updateNames] = $this->prepareUpsertColumns(
-            $table,
-            $insertColumns,
-            $updateColumns,
-        );
+        [$uniqueNames, , $updateNames] = $this->prepareUpsertColumns($table, $insertColumns, $updateColumns);
 
         if (empty($uniqueNames)) {
             return $insertSql;
         }
 
-        if ($updateNames === []) {
+        if ($updateColumns === false || $updateNames === []) {
             /** there are no columns to update */
-            $updateColumns = false;
-        }
-
-        if ($updateColumns === false) {
             return "$insertSql ON CONFLICT DO NOTHING";
         }
 
         if ($updateColumns === true) {
             $updateColumns = [];
 
-            /** @psalm-var string $name */
+            /** @psalm-var string[] $updateNames */
             foreach ($updateNames as $name) {
                 $updateColumns[$name] = new Expression(
                     'EXCLUDED.' . $this->quoter->quoteColumnName($name)
                 );
             }
         }
 
-        /**
-         * @psalm-var array $updateColumns
-         * @psalm-var string[] $uniqueNames
-         * @psalm-var string[] $updates
-         */
         [$updates, $params] = $this->prepareUpdateSets($table, $updateColumns, $params);
 
         return $insertSql

diff --git a/tests/Provider/QueryBuilderProvider.php b/tests/Provider/QueryBuilderProvider.php
@@ -365,6 +365,11 @@ public static function upsert(): array
                     'VALUES (:qp0, :qp1, :qp2, :qp3) ON CONFLICT ("email") ' .
                     'DO UPDATE SET "address"=EXCLUDED."address", "status"=EXCLUDED."status", "profile_id"=EXCLUDED."profile_id"',
             ],
+            'regular values with unique at not the first position' => [
+                3 => 'INSERT INTO "T_upsert" ("address", "email", "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"',
+            ],
             'regular values with update part' => [
                 2 => ['address' => 'foo {{city}}', 'status' => 2, 'orders' => new Expression('"T_upsert"."orders" + 1')],
                 3 => 'INSERT INTO "T_upsert" ("email", "address", "status", "profile_id") ' .

diff --git a/tests/QueryBuilderTest.php b/tests/QueryBuilderTest.php
@@ -4,7 +4,6 @@
 
 namespace Yiisoft\Db\Pgsql\Tests;
 
-use Generator;
 use Throwable;
 use Yiisoft\Db\Driver\Pdo\PdoConnectionInterface;
 use Yiisoft\Db\Exception\Exception;
@@ -253,7 +252,7 @@ public function testAddUnique(string $name, string $table, array|string $columns
     /**
      * @dataProvider \Yiisoft\Db\Pgsql\Tests\Provider\QueryBuilderProvider::batchInsert
      */
-    public function testBatchInsert(string $table, array $columns, iterable|Generator $rows, string $expected): void
+    public function testBatchInsert(string $table, array $columns, iterable $rows, string $expected): void
     {
         parent::testBatchInsert($table, $columns, $rows, $expected);
     }