Refactor data readers

src/Command/AbstractCommand.php

@@ -7,7 +7,7 @@
 use Closure;
 use Throwable;
 use Yiisoft\Db\Exception\Exception;
-use Yiisoft\Db\Query\Data\DataReaderInterface;
+use Yiisoft\Db\Query\DataReaderInterface;
 use Yiisoft\Db\Query\QueryInterface;
 use Yiisoft\Db\QueryBuilder\DMLQueryBuilderInterface;
 use Yiisoft\Db\QueryBuilder\QueryBuilderInterface;

src/Command/CommandInterface.php

@@ -18,7 +18,7 @@
 use Yiisoft\Db\Exception\InvalidCallException;
 use Yiisoft\Db\Exception\InvalidConfigException;
 use Yiisoft\Db\Exception\NotSupportedException;
-use Yiisoft\Db\Query\Data\DataReaderInterface;
+use Yiisoft\Db\Query\DataReaderInterface;
 use Yiisoft\Db\Query\QueryInterface;
 use Yiisoft\Db\QueryBuilder\DMLQueryBuilderInterface;
 use Yiisoft\Db\Schema\Column\ColumnInterface;

src/Connection/AbstractConnection.php

@@ -38,9 +38,9 @@ public function beginTransaction(?string $isolationLevel = null): TransactionInt
         return $this->transaction;
     }
 
-    public function createBatchQueryResult(QueryInterface $query, bool $each = false): BatchQueryResultInterface
+    public function createBatchQueryResult(QueryInterface $query): BatchQueryResultInterface
     {
-        return new BatchQueryResult($query, $each);
+        return new BatchQueryResult($query);
     }
 
     public function getTablePrefix(): string

src/Connection/ConnectionInterface.php

@@ -51,11 +51,10 @@ public function beginTransaction(?string $isolationLevel = null): TransactionInt
      * Create a batch query result instance.
      *
      * @param QueryInterface $query The query to execute.
-     * @param bool $each Whether to return each row of the result set one at a time.
      *
      * @return BatchQueryResultInterface The batch query result instance.
      */
-    public function createBatchQueryResult(QueryInterface $query, bool $each = false): BatchQueryResultInterface;
+    public function createBatchQueryResult(QueryInterface $query): BatchQueryResultInterface;
 
     /**
      * Creates a command for execution.

src/Debug/CommandInterfaceProxy.php

@@ -7,7 +7,7 @@
 use Closure;
 use Throwable;
 use Yiisoft\Db\Command\CommandInterface;
-use Yiisoft\Db\Query\Data\DataReaderInterface;
+use Yiisoft\Db\Query\DataReaderInterface;
 use Yiisoft\Db\Query\QueryInterface;
 use Yiisoft\Db\Schema\Column\ColumnInterface;
 

src/Debug/ConnectionInterfaceProxy.php

@@ -38,9 +38,9 @@ public function beginTransaction(?string $isolationLevel = null): TransactionInt
         return new TransactionInterfaceDecorator($result, $this->collector);
     }
 
-    public function createBatchQueryResult(QueryInterface $query, bool $each = false): BatchQueryResultInterface
+    public function createBatchQueryResult(QueryInterface $query): BatchQueryResultInterface
     {
-        return $this->connection->createBatchQueryResult($query, $each);
+        return $this->connection->createBatchQueryResult($query);
     }
 
     public function createCommand(?string $sql = null, array $params = []): CommandInterface

src/Driver/Pdo/AbstractPdoCommand.php

@@ -20,7 +20,6 @@
 use Yiisoft\Db\Profiler\Context\CommandContext;
 use Yiisoft\Db\Profiler\ProfilerAwareInterface;
 use Yiisoft\Db\Profiler\ProfilerAwareTrait;
-use Yiisoft\Db\Query\Data\DataReader;
 use Yiisoft\Db\QueryBuilder\QueryBuilderInterface;
 
 use function restore_error_handler;
@@ -238,7 +237,7 @@ protected function internalExecute(): void
     protected function internalGetQueryResult(int $queryMode): mixed
     {
         if ($queryMode === self::QUERY_MODE_CURSOR) {
-            return new DataReader($this);
+            return new PdoDataReader($this);
         }
 
         if ($queryMode === self::QUERY_MODE_EXECUTE) {

src/Query/Data/DataReader.php → src/Driver/Pdo/PdoDataReader.php

@@ -2,15 +2,19 @@
 
 declare(strict_types=1);
 
-namespace Yiisoft\Db\Query\Data;
+namespace Yiisoft\Db\Driver\Pdo;
 
+use Closure;
 use Countable;
 use Iterator;
 use PDO;
 use PDOStatement;
-use Yiisoft\Db\Driver\Pdo\PdoCommandInterface;
 use Yiisoft\Db\Exception\InvalidCallException;
 use Yiisoft\Db\Exception\InvalidParamException;
+use Yiisoft\Db\Query\DataReaderInterface;
+use Yiisoft\Db\Query\QueryInterface;
+
+use function is_string;
 
 /**
  * Provides an abstract way to read data from a database.
@@ -21,11 +25,18 @@
  * to execute a SELECT statement and read the results.
  *
  * The class provides methods for accessing the data returned by the query.
+ *
+ * @psalm-import-type IndexBy from QueryInterface
+ * @psalm-import-type ResultCallback from QueryInterface
  */
-final class DataReader implements DataReaderInterface
+final class PdoDataReader implements DataReaderInterface
 {
-    private int $index = -1;
-    private mixed $row;
+    /** @psalm-var IndexBy|null $indexBy */
+    private Closure|string|null $indexBy = null;
+    private int $index = 0;
+    /** @psalm-var ResultCallback|null $resultCallback */
+    private Closure|null $resultCallback = null;
+    private array|false $row;
     private PDOStatement $statement;
 
     /**
@@ -40,6 +51,13 @@
         }
 
         $this->statement = $statement;
+        /** @var array|false */
+        $this->row = $statement->fetch(PDO::FETCH_ASSOC);
+    }
+
+    public function __destruct()
+    {
+        $this->statement->closeCursor();
     }
 
     /**
@@ -64,31 +82,36 @@
      */
     public function rewind(): void
     {
-        if ($this->index < 0) {
-            $this->row = $this->statement->fetch(PDO::FETCH_ASSOC);
-            $this->index = 0;
-        } else {
-            throw new InvalidCallException('DataReader cannot rewind. It is a forward-only reader.');
+        if ($this->index === 0) {
+            return;
         }
+
+        throw new InvalidCallException('DataReader cannot rewind. It is a forward-only reader.');
     }
 
-    /**
-     * Returns the index of the current row.
-     *
-     * This method is required by the interface {@see Iterator}.
-     */
-    public function key(): int
+    public function key(): int|string|null
     {
-        return $this->index;
+        if ($this->indexBy === null) {
+            return $this->index;
+        }
+
+        if ($this->row === false) {
+            return null;
+        }
+
+        if (is_string($this->indexBy)) {
+            return (string) $this->row[$this->indexBy];
+        }
+
+        return ($this->indexBy)($this->row);
     }
 
-    /**
-     * Returns the current row.
-     *
-     * This method is required by the interface {@see Iterator}.
-     */
-    public function current(): mixed
+    public function current(): array|object|false
     {
+        if ($this->resultCallback !== null && $this->row !== false) {
+            return ($this->resultCallback)([$this->row])[0];
+        }
+
         return $this->row;
     }
 
@@ -99,6 +122,7 @@
      */
     public function next(): void
     {
+        /** @var array|false */
         $this->row = $this->statement->fetch(PDO::FETCH_ASSOC);
         $this->index++;
     }
@@ -112,4 +136,16 @@
     {
         return $this->row !== false;
     }
+
+    public function indexBy(Closure|string|null $indexBy): static
+    {
+        $this->indexBy = $indexBy;
+        return $this;
+    }
+
+    public function resultCallback(Closure|null $resultCallback): static
+    {
+        $this->resultCallback = $resultCallback;
+        return $this;
+    }
 }