Query: Determine omitted & target columns correctly

src/Query.php

@@ -1,4 +1,4 @@
 <?php
 
 namespace ipl\Orm;
 
@@ -135,7 +135,7 @@
             throw new InvalidArgumentException('Argument $class must be a string');
         }
 
         if (! (new ReflectionClass($class))->newInstanceWithoutConstructor() instanceof ResultSet) {
             throw new InvalidArgumentException(
                 $class . ' must be an instance of ' . ResultSet::class
             );
@@ -454,7 +454,18 @@
         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)) {
@@ -733,9 +744,9 @@
         $columnStorage = new SplObjectStorage();
 
         foreach ($columns as list($target, $alias, $column)) {
             if (! $columnStorage->contains($target)) {
                 $resolved = new ArrayObject();
                 $columnStorage->attach($target, $resolved);
             } else {
                 $resolved = $columnStorage[$target];
             }

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');