support PDO::FETCH_COLUMN in fetchAll() (#301)

Co-authored-by: Markus Staab <[email protected]>

Author: staabm

src/Extensions/PdoQueryDynamicReturnTypeExtension.php

@@ -18,6 +18,7 @@
 use PHPStan\Type\TypeCombinator;
 use staabm\PHPStanDba\PdoReflection\PdoStatementReflection;
 use staabm\PHPStanDba\QueryReflection\QueryReflection;
+use staabm\PHPStanDba\QueryReflection\QueryReflector;
 
 final class PdoQueryDynamicReturnTypeExtension implements DynamicMethodReturnTypeExtension
 {
@@ -83,6 +84,10 @@ private function inferType(MethodCall $methodCall, Expr $queryExpr, Scope $scope
             if (null === $reflectionFetchType) {
                 return null;
             }
+            // not yet implemented on query()
+            if (QueryReflector::FETCH_TYPE_COLUMN === $reflectionFetchType) {
+                return null;
+            }
         }
 
         $queryReflection = new QueryReflection();

src/Extensions/PdoStatementFetchDynamicReturnTypeExtension.php

@@ -12,13 +12,15 @@
 use PHPStan\Reflection\ParametersAcceptorSelector;
 use PHPStan\Type\ArrayType;
 use PHPStan\Type\Constant\ConstantBooleanType;
+use PHPStan\Type\Constant\ConstantIntegerType;
 use PHPStan\Type\DynamicMethodReturnTypeExtension;
 use PHPStan\Type\IntegerType;
 use PHPStan\Type\MixedType;
 use PHPStan\Type\Type;
 use PHPStan\Type\TypeCombinator;
 use staabm\PHPStanDba\PdoReflection\PdoStatementReflection;
 use staabm\PHPStanDba\QueryReflection\QueryReflection;
+use staabm\PHPStanDba\QueryReflection\QueryReflector;
 
 final class PdoStatementFetchDynamicReturnTypeExtension implements DynamicMethodReturnTypeExtension
 {
@@ -76,7 +78,23 @@ private function inferType(MethodReflection $methodReflection, MethodCall $metho
             }
         }
 
-        $rowType = $pdoStatementReflection->getRowType($statementType, $fetchType);
+        if ('fetchAll' === $methodReflection->getName() && QueryReflector::FETCH_TYPE_COLUMN === $fetchType) {
+            $columnIndex = 0;
+
+            if (\count($args) > 1) {
+                $columnIndexType = $scope->getType($args[1]->value);
+                if ($columnIndexType instanceof ConstantIntegerType) {
+                    $columnIndex = $columnIndexType->getValue();
+                } else {
+                    return null;
+                }
+            }
+
+            $rowType = $pdoStatementReflection->getColumnRowType($statementType, $columnIndex);
+        } else {
+            $rowType = $pdoStatementReflection->getRowType($statementType, $fetchType);
+        }
+
         if (null === $rowType) {
             return null;
         }

src/PdoReflection/PdoStatementReflection.php

@@ -64,6 +64,8 @@ public function getFetchType(Type $fetchModeType): ?int
             return QueryReflector::FETCH_TYPE_NUMERIC;
         } elseif (PDO::FETCH_BOTH === $fetchModeType->getValue()) {
             return QueryReflector::FETCH_TYPE_BOTH;
+        } elseif (PDO::FETCH_COLUMN === $fetchModeType->getValue()) {
+            return QueryReflector::FETCH_TYPE_COLUMN;
         }
 
         return null;
@@ -132,6 +134,20 @@ public function getRowType(Type $statementType, int $fetchType): ?Type
         return null;
     }
 
+    public function getColumnRowType(Type $statementType, int $columnIndex): ?Type
+    {
+        $statementType = $this->getRowType($statementType, QueryReflector::FETCH_TYPE_NUMERIC);
+
+        if ($statementType instanceof ConstantArrayType) {
+            $valueTypes = $statementType->getValueTypes();
+            if (\array_key_exists($columnIndex, $valueTypes)) {
+                return $valueTypes[$columnIndex];
+            }
+        }
+
+        return null;
+    }
+
     /**
      * @param QueryReflector::FETCH_TYPE* $fetchType
      */

src/QueryReflection/QueryReflector.php

@@ -14,6 +14,8 @@ interface QueryReflector
     public const FETCH_TYPE_BOTH = 5;
     public const FETCH_TYPE_KEY_VALUE = 6;
 
+    public const FETCH_TYPE_COLUMN = 50;
+
     public function validateQueryString(string $queryString): ?Error;
 
     /**

tests/default/data/pdo-stmt-fetch.php

@@ -28,6 +28,15 @@ public function fetchAll(PDO $pdo)
         $all = $stmt->fetchAll(PDO::FETCH_ASSOC);
         assertType('array<int, array{email: string, adaid: int<0, 4294967295>}>', $all);
 
+        $all = $stmt->fetchAll(PDO::FETCH_COLUMN);
+        assertType('array<int, string>', $all);
+
+        $all = $stmt->fetchAll(PDO::FETCH_COLUMN, 0);
+        assertType('array<int, string>', $all);
+
+        $all = $stmt->fetchAll(PDO::FETCH_COLUMN, 1);
+        assertType('array<int, int<0, 4294967295>>', $all);
+
         // not yet supported fetch types
         $all = $stmt->fetchAll(PDO::FETCH_OBJ);
         assertType('array', $all); // XXX since php8 this cannot return false