SQL AST: narrow result types with WHERE condition (#596)

Author: hemberger

src/SqlAst/ParserInference.php

@@ -19,7 +19,9 @@
 use SqlFtw\Sql\Dml\TableReference\Join;
 use SqlFtw\Sql\Dml\TableReference\TableReferenceSubquery;
 use SqlFtw\Sql\Dml\TableReference\TableReferenceTable;
+use SqlFtw\Sql\Expression\Asterisk;
 use SqlFtw\Sql\Expression\Identifier;
+use SqlFtw\Sql\Expression\SimpleName;
 use SqlFtw\Sql\SqlSerializable;
 use staabm\PHPStanDba\QueryReflection\QueryReflection;
 use staabm\PHPStanDba\SchemaReflection\Join as SchemaJoin;
@@ -50,16 +52,18 @@ public function narrowResultType(string $queryString, ConstantArrayType $resultT
         // returns a Generator. will not parse anything if you don't iterate over it
         $commands = $parser->parse($queryString);
 
-        $fromColumns = null;
+        $selectColumns = null;
         $fromTable = null;
+        $where = null;
         $joins = [];
         foreach ($commands as [$command]) {
             // Parser does not throw exceptions. this allows to parse partially invalid code and not fail on first error
             if ($command instanceof SelectCommand) {
-                if (null === $fromColumns) {
-                    $fromColumns = $command->getColumns();
+                if (null === $selectColumns) {
+                    $selectColumns = $command->getColumns();
                 }
                 $from = $command->getFrom();
+                $where = $command->getWhere();
 
                 if (null === $from) {
                     // no FROM clause, use an empty Table to signify this
@@ -128,13 +132,21 @@ public function narrowResultType(string $queryString, ConstantArrayType $resultT
             // not parsable atm, return un-narrowed type
             return $resultType;
         }
-        if (null === $fromColumns) {
+        if (null === $selectColumns) {
             throw new ShouldNotHappenException();
         }
 
-        $queryScope = new QueryScope($fromTable, $joins);
+        $queryScope = new QueryScope($fromTable, $joins, $where);
 
-        foreach ($fromColumns as $i => $column) {
+        // If we're selecting '*', get the selected columns from the table
+        if (\count($selectColumns) === 1 && $selectColumns[0]->getExpression() instanceof Asterisk) {
+            $selectColumns = [];
+            foreach ($fromTable->getColumns() as $column) {
+                $selectColumns[] = new SelectExpression(new SimpleName($column->getName()));
+            }
+        }
+
+        foreach ($selectColumns as $i => $column) {
             $expression = $column->getExpression();
 
             $offsetType = new ConstantIntegerType($i);

src/SqlAst/QueryScope.php

@@ -12,6 +12,7 @@
 use PHPStan\Type\NullType;
 use PHPStan\Type\Type;
 use PHPStan\Type\TypeCombinator;
+use SqlFtw\Sql\Expression\BinaryOperator;
 use SqlFtw\Sql\Expression\BoolValue;
 use SqlFtw\Sql\Expression\CaseExpression;
 use SqlFtw\Sql\Expression\ComparisonOperator;
@@ -26,6 +27,7 @@
 use SqlFtw\Sql\Expression\Parentheses;
 use SqlFtw\Sql\Expression\SimpleName;
 use SqlFtw\Sql\Expression\StringValue;
+use SqlFtw\Sql\SqlSerializable;
 use staabm\PHPStanDba\SchemaReflection\Column;
 use staabm\PHPStanDba\SchemaReflection\Join;
 use staabm\PHPStanDba\SchemaReflection\Table;
@@ -47,13 +49,19 @@ final class QueryScope
      */
     private $joinedTables;
 
+    /**
+     * @var ?SqlSerializable
+     */
+    private $whereCondition;
+
     /**
      * @param list<Join> $joinedTables
      */
-    public function __construct(Table $fromTable, array $joinedTables)
+    public function __construct(Table $fromTable, array $joinedTables, ?SqlSerializable $whereCondition)
     {
         $this->fromTable = $fromTable;
         $this->joinedTables = $joinedTables;
+        $this->whereCondition = $whereCondition;
 
         $this->extensions = [
             new PositiveIntReturnTypeExtension(),
@@ -78,6 +86,24 @@ public function __construct(Table $fromTable, array $joinedTables)
      * @param Identifier|Literal|ExpressionNode $expression
      */
     public function getType($expression): Type
+    {
+        $resultType = $this->resolveExpression($expression);
+
+        if ($this->whereCondition !== null && $expression instanceof SimpleName) {
+            $resultType = $this->narrowWhereCondition(
+                $this->whereCondition,
+                $expression->getName(),
+                $resultType
+            );
+        }
+
+        return $resultType;
+    }
+
+    /**
+     * @param Identifier|Literal|ExpressionNode $expression
+     */
+    private function resolveExpression($expression): Type
     {
         if ($expression instanceof NullLiteral) {
             return new NullType();
@@ -219,4 +245,53 @@ private function narrowJoinCondition(Column $column, Join $join): ?Type
 
         return null;
     }
+
+    private function narrowWhereCondition(SqlSerializable $op, string $name, Type $valueType): Type
+    {
+        // If condition is in parentheses, try again with its contents
+        if ($op instanceof Parentheses) {
+            return $this->narrowWhereCondition($op->getContents(), $name, $valueType);
+        }
+
+        // Only binary ops are currently supported
+        if (! ($op instanceof BinaryOperator)) {
+            return $valueType;
+        }
+
+        $left = $op->getLeft();
+        $right = $op->getRight();
+        $operator = $op->getOperator()->getValue();
+
+        // Handle compound conditions
+        if ($operator === Operator::AND) {
+            if ($left instanceof BinaryOperator) {
+                $valueType = $this->narrowWhereCondition($left, $name, $valueType);
+            }
+            if ($right instanceof BinaryOperator) {
+                $valueType = $this->narrowWhereCondition($right, $name, $valueType);
+            }
+            return $valueType;
+        }
+
+        // Only simple names are currently supported
+        if (! ($left instanceof SimpleName)) {
+            return $valueType;
+        }
+        if ($left->getName() !== $name) {
+            return $valueType;
+        }
+
+        // Handle NULL comparisons
+        if ($right instanceof NullLiteral) {
+            if ($operator === Operator::IS_NOT) {
+                return TypeCombinator::removeNull($valueType);
+            }
+            if ($operator === Operator::IS) {
+                return TypeCombinator::intersect($valueType, new NullType());
+            }
+        }
+
+        // Unsupported operator
+        return $valueType;
+    }
 }