Query: Determine omitted & target columns correctly

Icinga · Sep 23, 2024 · d141e1f · d141e1f
1 parent e9d6117
commit d141e1f
Show file tree

Hide file tree

Showing 2 changed files with 58 additions and 1 deletion.
diff --git 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 ($value, $column) use ($toExclude) {
+                    if (is_string($column) && isset($toExclude[$column])) {
+                        return false;
+                    }
+
+                    if (is_string($column)) {
+                        return ! in_array($column, $toExclude, true);
+                    }
+
+                    return ! in_array($value, $toExclude, true);
+                }, ARRAY_FILTER_USE_BOTH);
             }
 
             if (! empty($customAliases)) {

diff --git 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,29 @@ 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)
+            ->withColumns(['aliased_expr', 'username']);
+
+        $this->assertEquals(
+            [
+                'user.id',
+                '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 +510,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');