From b513a7fbd3aa0b7e9a959c2b20cce4dc507b98ab Mon Sep 17 00:00:00 2001
From: Yonas Habteab <yonas.habteab@icinga.com>
Date: Fri, 20 Sep 2024 09:36:29 +0200
Subject: [PATCH] Query: Determine omitted & target columns correctly

---
 src/Query.php       | 13 +++++++++++-
 tests/QueryTest.php | 49 +++++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 61 insertions(+), 1 deletion(-)

diff --git a/src/Query.php b/src/Query.php
index 0e19dd1..c82e2e7 100644
--- a/src/Query.php
+++ b/src/Query.php
@@ -454,7 +454,18 @@ public function assembleSelect()
         foreach ($resolved as $target) {
             $targetColumns = $resolved[$target]->getArrayCopy();
             if (isset($omitted[$target])) {
-                $targetColumns = array_diff($targetColumns, $omitted[$target]->getArrayCopy());
+                $toExclude = $omitted[$target]->getArrayCopy();
+                $targetColumns = array_filter($targetColumns, function ($column, $alias) use ($toExclude) {
+                    if (is_string($alias) && isset($toExclude[$alias])) {
+                        return false;
+                    }
+
+                    if (is_string($alias)) {
+                        return ! in_array($alias, $toExclude, true);
+                    }
+
+                    return ! in_array($column, $toExclude, true);
+                }, ARRAY_FILTER_USE_BOTH);
             }
 
             if (! empty($customAliases)) {
diff --git a/tests/QueryTest.php b/tests/QueryTest.php
index 4034a9e..898df7b 100644
--- a/tests/QueryTest.php
+++ b/tests/QueryTest.php
@@ -4,6 +4,7 @@
 
 use ipl\Orm\Exception\InvalidRelationException;
 use ipl\Orm\Query;
+use ipl\Orm\ResolvedExpression;
 use ipl\Sql\Expression;
 use ipl\Tests\Sql\TestCase;
 
@@ -401,6 +402,32 @@ public function testWithColumnsSupportsExpressions()
         );
     }
 
+    public function testWithColumnsSupportsAliasedExpressions()
+    {
+        $model = new class () extends User {
+            public function getColumns(): array
+            {
+                return ['username', 'aliased_expr' => new Expression('1')];
+            }
+        };
+
+        $query = (new Query())
+            ->setModel($model)
+            ->columns('username')
+            ->withColumns('aliased_expr');
+
+        $this->assertEquals(
+            [
+                'user.username',
+                'aliased_expr' => new ResolvedExpression(
+                    new Expression('1'),
+                    $query->getResolver()->requireAndResolveColumns(['aliased_expr'])
+                ),
+            ],
+            $query->assembleSelect()->getColumns()
+        );
+    }
+
     public function testWithColumnsHandlesCustomAliasesCorrectly()
     {
         $query = (new Query())
@@ -486,6 +513,28 @@ public function testWithoutColumnsOverridesColumnsAndWithColumns()
         );
     }
 
+    public function testWithoutColumnsExcludesAliasedExpressions()
+    {
+        $model = new class () extends User {
+            public function getColumns(): array
+            {
+                return [
+                    'username',
+                    'aliased_expr' => new Expression('1')
+                ];
+            }
+        };
+
+        $query = (new Query())
+            ->setModel($model)
+            ->withoutColumns(['aliased_expr', 'id']);
+
+        $this->assertSame(
+            ['user.username'],
+            $query->assembleSelect()->getColumns()
+        );
+    }
+
     public function testWithoutColumnsDoesNotWorkWithExpressions()
     {
         $expression = new Expression('1');