SQL AST: fix aggregate functions with GROUP BY (#599)

staabm · Apr 13, 2023 · e01951c · e01951c
1 parent b4e5054
commit e01951c
Show file tree

Hide file tree

Showing 8 changed files with 4,428 additions and 1,325 deletions.
diff --git a/src/SqlAst/AvgReturnTypeExtension.php b/src/SqlAst/AvgReturnTypeExtension.php
@@ -15,6 +15,16 @@
 
 final class AvgReturnTypeExtension implements QueryFunctionReturnTypeExtension
 {
+    /**
+     * @var bool
+     */
+    private $hasGroupBy;
+
+    public function __construct(bool $hasGroupBy)
+    {
+        $this->hasGroupBy = $hasGroupBy;
+    }
+
     public function isFunctionSupported(FunctionCall $expression): bool
     {
         return \in_array($expression->getFunction()->getName(), [BuiltInFunction::AVG], true);
@@ -33,6 +43,7 @@ public function getReturnType(FunctionCall $expression, QueryScope $scope): ?Typ
         if ($argType->isNull()->yes()) {
             return $argType;
         }
+        $containsNull = TypeCombinator::containsNull($argType);
         $argType = TypeCombinator::removeNull($argType);
 
         if ($argType instanceof UnionType) {
@@ -45,7 +56,10 @@ public function getReturnType(FunctionCall $expression, QueryScope $scope): ?Typ
             $newType = $this->convertToAvgType($argType);
         }
 
-        return TypeCombinator::addNull($newType);
+        if ($containsNull || ! $this->hasGroupBy) {
+            $newType = TypeCombinator::addNull($newType);
+        }
+        return $newType;
     }
 
     private function convertToAvgType(Type $argType): Type

diff --git a/src/SqlAst/MinMaxReturnTypeExtension.php b/src/SqlAst/MinMaxReturnTypeExtension.php
@@ -11,6 +11,16 @@
 
 final class MinMaxReturnTypeExtension implements QueryFunctionReturnTypeExtension
 {
+    /**
+     * @var bool
+     */
+    private $hasGroupBy;
+
+    public function __construct(bool $hasGroupBy)
+    {
+        $this->hasGroupBy = $hasGroupBy;
+    }
+
     public function isFunctionSupported(FunctionCall $expression): bool
     {
         return \in_array($expression->getFunction()->getName(), [BuiltInFunction::MIN, BuiltInFunction::MAX], true);
@@ -26,6 +36,9 @@ public function getReturnType(FunctionCall $expression, QueryScope $scope): ?Typ
 
         $argType = $scope->getType($args[0]);
 
-        return TypeCombinator::addNull($argType);
+        if (! $this->hasGroupBy) {
+            $argType = TypeCombinator::addNull($argType);
+        }
+        return $argType;
     }
 }
diff --git a/src/SqlAst/ParserInference.php b/src/SqlAst/ParserInference.php
@@ -55,6 +55,7 @@ public function narrowResultType(string $queryString, ConstantArrayType $resultT
         $selectColumns = null;
         $fromTable = null;
         $where = null;
+        $groupBy = null;
         $joins = [];
         foreach ($commands as [$command]) {
             // Parser does not throw exceptions. this allows to parse partially invalid code and not fail on first error
@@ -64,6 +65,7 @@ public function narrowResultType(string $queryString, ConstantArrayType $resultT
                 }
                 $from = $command->getFrom();
                 $where = $command->getWhere();
+                $groupBy = $command->getGroupBy();
 
                 if (null === $from) {
                     // no FROM clause, use an empty Table to signify this
@@ -136,7 +138,7 @@ public function narrowResultType(string $queryString, ConstantArrayType $resultT
             throw new ShouldNotHappenException();
         }
 
-        $queryScope = new QueryScope($fromTable, $joins, $where);
+        $queryScope = new QueryScope($fromTable, $joins, $where, $groupBy !== null);
 
         // If we're selecting '*', get the selected columns from the table
         if (\count($selectColumns) === 1 && $selectColumns[0]->getExpression() instanceof Asterisk) {

diff --git a/src/SqlAst/QueryScope.php b/src/SqlAst/QueryScope.php
@@ -57,7 +57,7 @@ final class QueryScope
     /**
      * @param list<Join> $joinedTables
      */
-    public function __construct(Table $fromTable, array $joinedTables, ?SqlSerializable $whereCondition)
+    public function __construct(Table $fromTable, array $joinedTables, ?SqlSerializable $whereCondition, bool $hasGroupBy)
     {
         $this->fromTable = $fromTable;
         $this->joinedTables = $joinedTables;
@@ -73,12 +73,12 @@ public function __construct(Table $fromTable, array $joinedTables, ?SqlSerializa
             new InstrReturnTypeExtension(),
             new StrCaseReturnTypeExtension(),
             new ReplaceReturnTypeExtension(),
-            new AvgReturnTypeExtension(),
-            new SumReturnTypeExtension(),
+            new AvgReturnTypeExtension($hasGroupBy),
+            new SumReturnTypeExtension($hasGroupBy),
             new IsNullReturnTypeExtension(),
             new AbsReturnTypeExtension(),
             new RoundReturnTypeExtension(),
-            new MinMaxReturnTypeExtension(),
+            new MinMaxReturnTypeExtension($hasGroupBy),
         ];
     }
 

diff --git a/src/SqlAst/SumReturnTypeExtension.php b/src/SqlAst/SumReturnTypeExtension.php
@@ -14,6 +14,16 @@
 
 final class SumReturnTypeExtension implements QueryFunctionReturnTypeExtension
 {
+    /**
+     * @var bool
+     */
+    private $hasGroupBy;
+
+    public function __construct(bool $hasGroupBy)
+    {
+        $this->hasGroupBy = $hasGroupBy;
+    }
+
     public function isFunctionSupported(FunctionCall $expression): bool
     {
         return \in_array($expression->getFunction()->getName(), [BuiltInFunction::SUM], true);
@@ -43,10 +53,9 @@ public function getReturnType(FunctionCall $expression, QueryScope $scope): ?Typ
             $result = $argType;
         }
 
-        if ($containsNull) {
+        if ($containsNull || ! $this->hasGroupBy) {
             $result = TypeCombinator::addNull($result);
         }
-
         return $result;
     }
 }