Refactor Query::column() (#815)

yiisoft · Apr 15, 2024 · e380ae2 · e380ae2
1 parent 3f977a5
commit e380ae2
Show file tree

Hide file tree

Showing 6 changed files with 44 additions and 32 deletions.
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -2,6 +2,7 @@
 
 ## 2.0.0 under development
 
+- Enh #815: Refactor `Query::column()` method (@Tigrov) 
 - Enh #816: Allow scalar values for `$columns` parameter of `Query::select()` and `Query::addSelect()` methods (@Tigrov)
 - Enh #806: Non-unique placeholder names inside `Expression::$params` will be replaced with unique names (@Tigrov)
 - Enh #806: Build `Expression` instances inside `Expression::$params` when build a query using `QueryBuilder` (@Tigrov)

diff --git a/src/Query/Query.php b/src/Query/Query.php
@@ -16,11 +16,15 @@
 use Yiisoft\Db\Helper\DbArrayHelper;
 use Yiisoft\Db\QueryBuilder\QueryBuilderInterface;
 
+use function array_column;
+use function array_combine;
 use function array_key_exists;
+use function array_map;
 use function array_merge;
 use function array_shift;
 use function array_unshift;
 use function count;
+use function current;
 use function gettype;
 use function is_array;
 use function is_int;
@@ -29,7 +33,6 @@
 use function key;
 use function preg_match;
 use function preg_split;
-use function reset;
 use function str_contains;
 use function strcasecmp;
 use function strlen;
@@ -84,6 +87,7 @@ class Query implements QueryInterface
     protected array $params = [];
     protected array $union = [];
     protected array $withQueries = [];
+    /** @psalm-var Closure(array):array-key|string|null $indexBy */
     protected Closure|string|null $indexBy = null;
     protected ExpressionInterface|int|null $limit = null;
     protected ExpressionInterface|int|null $offset = null;
@@ -256,39 +260,37 @@ public function column(): array
             return $this->createCommand()->queryColumn();
         }
 
-        if (is_string($this->indexBy) && count($this->select) === 1) {
-            if (!str_contains($this->indexBy, '.') && count($tables = $this->getTablesUsedInFrom()) > 0) {
-                $this->select[] = key($tables) . '.' . $this->indexBy;
-            } else {
-                $this->select[] = $this->indexBy;
+        if (is_string($this->indexBy)) {
+            if (count($this->select) === 1) {
+                if (!str_contains($this->indexBy, '.') && count($tables = $this->getTablesUsedInFrom()) > 0) {
+                    $this->select[] = key($tables) . '.' . $this->indexBy;
+                } else {
+                    $this->select[] = $this->indexBy;
+                }
             }
-        }
 
-        $rows = $this->createCommand()->queryAll();
-        $results = [];
-        $column = null;
+            $rows = $this->createCommand()->queryAll();
+
+            if (empty($rows)) {
+                return [];
+            }
 
-        if (is_string($this->indexBy)) {
             if (($dotPos = strpos($this->indexBy, '.')) === false) {
                 $column = $this->indexBy;
             } else {
                 $column = substr($this->indexBy, $dotPos + 1);
             }
+
+            return array_column($rows, key(current($rows)), $column);
         }
 
-        /** @psalm-var array<array-key, array<string, string>> $rows */
-        foreach ($rows as $row) {
-            $value = reset($row);
+        $rows = $this->createCommand()->queryAll();
 
-            if ($this->indexBy instanceof Closure) {
-                /** @psalm-suppress MixedArrayOffset */
-                $results[($this->indexBy)($row)] = $value;
-            } else {
-                $results[$row[$column] ?? $row[$this->indexBy]] = $value;
-            }
+        if (empty($rows)) {
+            return [];
         }
 
-        return $results;
+        return array_combine(array_map($this->indexBy, $rows), array_column($rows, key(current($rows))));
     }
 
     public function count(string $sql = '*'): int|string

diff --git a/src/Query/QueryInterface.php b/src/Query/QueryInterface.php
@@ -169,6 +169,8 @@ public function getHaving(): string|array|ExpressionInterface|null;
 
     /**
      * @return Closure|string|null The "index by" value.
+     *
+     * @psalm-return Closure(array):array-key|string|null
      */
     public function getIndexBy(): Closure|string|null;
 

diff --git a/src/Query/QueryPartsInterface.php b/src/Query/QueryPartsInterface.php
@@ -326,6 +326,8 @@ public function having(array|ExpressionInterface|string|null $condition, array $
      *     // return the index value corresponding to $row
      * }
      * ```
+     *
+     * @psalm-param Closure(array):array-key|string|null $column
      */
     public function indexBy(string|Closure|null $column): static;
 

diff --git a/tests/AbstractQueryTest.php b/tests/AbstractQueryTest.php
@@ -757,6 +757,22 @@ public function testColumnWithIndexBy(): void
             2 => 'user2',
             3 => 'user3',
         ], $query->column());
+
+        $query = (new Query($db))
+            ->select('name')
+            ->from('customer')
+            ->indexBy('id')
+            ->where(['id' => null]);
+
+        $this->assertSame([], $query->column());
+
+        $query = (new Query($db))
+            ->select(['name', 'id'])
+            ->from('customer')
+            ->indexBy(fn (array $row) => $row['id'] * 2)
+            ->where(['id' => null]);
+
+        $this->assertSame([], $query->column());
     }
 
     /**

diff --git a/tests/Common/CommonQueryTest.php b/tests/Common/CommonQueryTest.php
@@ -10,17 +10,6 @@
 
 abstract class CommonQueryTest extends AbstractQueryTest
 {
-    public function testColumnWithIndexBy(): void
-    {
-        $db = $this->getConnection(true);
-
-        $query = (new Query($db))->select('customer.name')->from('customer')->indexBy('customer.id');
-
-        $this->assertSame([1 => 'user1', 2 => 'user2', 3 => 'user3'], $query->column());
-
-        $db->close();
-    }
-
     public function testColumnIndexByWithClosure()
     {
         $db = $this->getConnection(true);