feat: enhance pluck method to support key-value pairs and handle tabl…

src/QueryBuilder.php

@@ -797,14 +797,103 @@ public function first(string ...$columns): mixed
     }
 
     /**
-     * Retrieve a single column from the result
+     * Get a collection of values from a given column.
+     *
+     * Like Laravel's pluck, this method temporarily sets the select columns
+     * only if no explicit select() was previously called (via onceWithColumns).
+     * It then strips table prefixes / aliases from the column name to correctly
+     * extract values from the result set.
+     *
+     * @param Expression|string $column The column to pluck values from.
+     * @param Expression|string|null $key Optional column to use as keys in the returned collection.
+     * @return array
      * @throws QueryBuilderException
      */
-    public function pluck(string $column): array
+    public function pluck(Expression|string $column, Expression|string|null $key = null): array
     {
-        $this->select($column);
-        $result = $this->get();
-        return $result->pluck($column);
+        $queryResult = $this->onceWithColumns(
+            is_null($key) || $key === $column ? [$column] : [$column, $key],
+            function () {
+                return $this->get();
+            }
+        );
+
+        if ($queryResult->isEmpty()) {
+            return [];
+        }
+
+        // Strip table prefix or alias so we can look up the column in the result rows.
+        $column = $this->stripTableForPluck($column);
+        $key = $this->stripTableForPluck($key);
+
+        $results = [];
+        foreach ($queryResult as $row) {
+            $row = (array)$row;
+            $value = $row[$column] ?? null;
+            if (is_null($key)) {
+                $results[] = $value;
+            } else {
+                $results[$row[$key] ?? null] = $value;
+            }
+        }
+
+        return $results;
+    }
+
+    /**
+     * Execute the given callback while selecting the given columns.
+     *
+     * If columns have already been explicitly set via select(), they are preserved.
+     * After running the callback, the columns are restored to their original value.
+     *
+     * @param array $columns Columns to use if none were set.
+     * @param callable $callback The callback that executes the query.
+     * @return mixed The result of the callback.
+     */
+    private function onceWithColumns(array $columns, callable $callback): mixed
+    {
+        $original = $this->columns;
+
+        if (empty($original)) {
+            $this->columns = $columns;
+        }
+
+        $result = $callback();
+
+        $this->columns = $original;
+
+        return $result;
+    }
+
+    /**
+     * Strip off the table name or alias from a column identifier.
+     *
+     * Handles both Expression objects and plain strings.
+     * Examples:
+     *   "users.id"              → "id"
+     *   "users.id AS user_id"   → "user_id"
+     *   "id AS user_id"         → "user_id"
+     *
+     * @param Expression|string|null $column
+     * @return string|null
+     */
+    private function stripTableForPluck(Expression|string|null $column): ?string
+    {
+        if (is_null($column)) {
+            return null;
+        }
+
+        $columnString = $column instanceof Expression
+            ? $column->getValue()
+            : $column;
+
+        // If the column contains " AS " (case-insensitive), split on that.
+        // Otherwise, split on "." to strip the table prefix.
+        $separator = str_contains(strtolower($columnString), ' as ') ? ' as ' : '\.';
+
+        $parts = preg_split('~' . $separator . '~i', $columnString);
+
+        return end($parts);
     }
 
     /**

tests/MysqlDialectTest.php

@@ -5,6 +5,8 @@
 use Abdulelahragih\QueryBuilder\Builders\JoinClauseBuilder;
 use Abdulelahragih\QueryBuilder\Builders\WhereQueryBuilder;
 use Abdulelahragih\QueryBuilder\Data\QueryBuilderException;
+use Abdulelahragih\QueryBuilder\Data\Collection;
+use Abdulelahragih\QueryBuilder\Grammar\Expression;
 use Abdulelahragih\QueryBuilder\QueryBuilder;
 use Abdulelahragih\QueryBuilder\Tests\Traits\TestTrait;
 use Error;
@@ -417,9 +419,70 @@ public function testPluck()
             ->table('users')
             ->limit(3)
             ->pluck('id');
+        $this->assertIsArray($result);
         $this->assertEquals([1, 2, 3], $result);
     }
 
+    public function testPluckWithTablePrefix()
+    {
+        $builder = new QueryBuilder($this->pdo);
+        $result = $builder
+            ->table('users')
+            ->pluck('users.id');
+        $this->assertEquals([1, 2, 3], $result);
+    }
+
+    public function testPluckPreservesExistingSelect()
+    {
+        $builder = new QueryBuilder($this->pdo);
+        $result = $builder
+            ->table('users')
+            ->select($builder->raw('id AS user_id'))
+            ->pluck('user_id');
+        $this->assertEquals([1, 2, 3], $result);
+    }
+
+    public function testPluckWithKey()
+    {
+        $builder = new QueryBuilder($this->pdo);
+        $result = $builder
+            ->table('users')
+            ->pluck('name', 'id');
+        $this->assertEquals([1 => 'Sam', 2 => 'John', 3 => 'Jane'], $result);
+    }
+
+    public function testPluckEmptyResult()
+    {
+        $builder = new QueryBuilder($this->pdo);
+        $result = $builder
+            ->table('users')
+            ->where('id', '=', 999)
+            ->pluck('id');
+        $this->assertIsArray($result);
+        $this->assertEmpty($result);
+    }
+
+    public function testPluckWithExpression()
+    {
+        $builder = new QueryBuilder($this->pdo);
+        $result = $builder
+            ->table('users')
+            ->limit(3)
+            ->pluck(Expression::make('id'));
+        $this->assertEquals([1, 2, 3], $result);
+    }
+
+    public function testPluckWithJoinAndAlias()
+    {
+        $builder = new QueryBuilder($this->pdo);
+        $result = $builder
+            ->table('posts')
+            ->join('users', 'users.id', '=', 'posts.user_id')
+            ->select($builder->raw('posts.user_id AS uid'))
+            ->pluck('uid');
+        $this->assertEquals([1, 1, 2], $result);
+    }
+
     public function testSingleInsert()
     {
         $builder = new QueryBuilder($this->pdo);

tests/PostgresQueryBuilderTest.php

@@ -5,6 +5,7 @@
 use Abdulelahragih\QueryBuilder\Builders\JoinClauseBuilder;
 use Abdulelahragih\QueryBuilder\Builders\WhereQueryBuilder;
 use Abdulelahragih\QueryBuilder\Data\QueryBuilderException;
+use Abdulelahragih\QueryBuilder\Data\Collection;
 use Abdulelahragih\QueryBuilder\DB;
 use Abdulelahragih\QueryBuilder\Grammar\Dialects\PostgresDialect;
 use Abdulelahragih\QueryBuilder\Grammar\Expression;
@@ -439,9 +440,70 @@ public function testPluck()
             ->table('users')
             ->limit(3)
             ->pluck('id');
+        $this->assertIsArray($result);
         $this->assertEquals([1, 2, 3], $result);
     }
 
+    public function testPluckWithTablePrefix()
+    {
+        $builder = new QueryBuilder($this->pdo, new PostgresDialect());
+        $result = $builder
+            ->table('users')
+            ->pluck('users.id');
+        $this->assertEquals([1, 2, 3], $result);
+    }
+
+    public function testPluckPreservesExistingSelect()
+    {
+        $builder = new QueryBuilder($this->pdo, new PostgresDialect());
+        $result = $builder
+            ->table('users')
+            ->select($builder->raw('id AS user_id'))
+            ->pluck('user_id');
+        $this->assertEquals([1, 2, 3], $result);
+    }
+
+    public function testPluckWithKey()
+    {
+        $builder = new QueryBuilder($this->pdo, new PostgresDialect());
+        $result = $builder
+            ->table('users')
+            ->pluck('name', 'id');
+        $this->assertEquals([1 => 'Sam', 2 => 'John', 3 => 'Jane'], $result);
+    }
+
+    public function testPluckEmptyResult()
+    {
+        $builder = new QueryBuilder($this->pdo, new PostgresDialect());
+        $result = $builder
+            ->table('users')
+            ->where('id', '=', 999)
+            ->pluck('id');
+        $this->assertIsArray($result);
+        $this->assertEmpty($result);
+    }
+
+    public function testPluckWithExpression()
+    {
+        $builder = new QueryBuilder($this->pdo, new PostgresDialect());
+        $result = $builder
+            ->table('users')
+            ->limit(3)
+            ->pluck(Expression::make('id'));
+        $this->assertEquals([1, 2, 3], $result);
+    }
+
+    public function testPluckWithJoinAndAlias()
+    {
+        $builder = new QueryBuilder($this->pdo, new PostgresDialect());
+        $result = $builder
+            ->table('posts')
+            ->join('users', 'users.id', '=', 'posts.user_id')
+            ->select($builder->raw('posts.user_id AS uid'))
+            ->pluck('uid');
+        $this->assertEquals([1, 1, 2], $result);
+    }
+
     public function testSingleInsert()
     {
         $builder = new QueryBuilder($this->pdo, new PostgresDialect());

tests/Traits/TestTrait.php

@@ -30,5 +30,18 @@ private function seedFakeData()
         $this->pdo->exec("INSERT INTO users (id, name) VALUES (1, 'Sam');");
         $this->pdo->exec("INSERT INTO users (id, name) VALUES (2, 'John');");
         $this->pdo->exec("INSERT INTO users (id, name) VALUES (3, 'Jane');");
+
+        // create table posts for join tests
+        $this->pdo->exec('
+        CREATE TABLE IF NOT EXISTS posts (
+            id INT PRIMARY KEY,
+            user_id INT,
+            title VARCHAR(255)
+        );
+        ');
+        $this->pdo->exec('DELETE FROM posts');
+        $this->pdo->exec("INSERT INTO posts (id, user_id, title) VALUES (1, 1, 'Post 1');");
+        $this->pdo->exec("INSERT INTO posts (id, user_id, title) VALUES (2, 1, 'Post 2');");
+        $this->pdo->exec("INSERT INTO posts (id, user_id, title) VALUES (3, 2, 'Post 3');");
     }
 }