From 5e5b79329ebe2e25ccc15fb99c5d7e7d3effd147 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Milan=20Mat=C4=9Bj=C4=8Dek?= <milan.matejcek@gmail.com>
Date: Thu, 17 Oct 2024 09:18:12 +0200
Subject: [PATCH] feat(ChunkWriteService): support iterable directly

---
 .../Contracts/ChunkWriteServiceContract.php     |  8 +++++---
 src/Database/Services/ChunkWriteService.php     |  5 +++--
 .../Database/Services/ChunkWriteServiceTest.php | 17 ++++++++++++++++-
 3 files changed, 24 insertions(+), 6 deletions(-)

diff --git a/src/Database/Contracts/ChunkWriteServiceContract.php b/src/Database/Contracts/ChunkWriteServiceContract.php
index 48106823..0846afb6 100644
--- a/src/Database/Contracts/ChunkWriteServiceContract.php
+++ b/src/Database/Contracts/ChunkWriteServiceContract.php
@@ -5,14 +5,16 @@
 namespace LaraStrict\Database\Contracts;
 
 use Closure;
-use Generator;
 use Illuminate\Database\Eloquent\Model;
 use LaraStrict\Database\Entities\ChunkWriteStateEntity;
 
+/**
+ * @phpstan-type TData iterable<int|string, Model>
+ */
 interface ChunkWriteServiceContract
 {
     /**
-     * @param Closure(): Generator<int, Model> $closure
+     * @param Closure(): TData|TData $closure
      */
-    public function write(Closure $closure): ChunkWriteStateEntity;
+    public function write(Closure|iterable $closure): ChunkWriteStateEntity;
 }
diff --git a/src/Database/Services/ChunkWriteService.php b/src/Database/Services/ChunkWriteService.php
index ec2aa3ca..24c24d5a 100644
--- a/src/Database/Services/ChunkWriteService.php
+++ b/src/Database/Services/ChunkWriteService.php
@@ -12,11 +12,12 @@
 
 final class ChunkWriteService implements ChunkWriteServiceContract
 {
-    public function write(Closure $closure, int $batchSize = 0): ChunkWriteStateEntity
+    public function write(Closure|iterable $closure, int $batchSize = 0): ChunkWriteStateEntity
     {
         $writeState = new ChunkWriteStateEntity(batchSize: $batchSize);
 
-        foreach ($closure() as $model) {
+        $source = $closure instanceof Closure ? $closure() : $closure;
+        foreach ($source as $model) {
             $this->add($model, $writeState);
         }
 
diff --git a/tests/Unit/Database/Services/ChunkWriteServiceTest.php b/tests/Unit/Database/Services/ChunkWriteServiceTest.php
index c8e4442a..e6e0df96 100644
--- a/tests/Unit/Database/Services/ChunkWriteServiceTest.php
+++ b/tests/Unit/Database/Services/ChunkWriteServiceTest.php
@@ -5,12 +5,16 @@
 namespace Tests\LaraStrict\Unit\Database\Services;
 
 use Closure;
+use LaraStrict\Database\Contracts\ChunkWriteServiceContract;
 use LaraStrict\Database\Entities\ChunkWriteStateEntity;
 use LaraStrict\Database\Services\ChunkWriteService;
 use LaraStrict\Tests\Traits\SqlTestEnable;
 use PHPUnit\Framework\Assert;
 use PHPUnit\Framework\TestCase;
 
+/**
+ * @phpstan-import-type TData from ChunkWriteServiceContract
+ */
 final class ChunkWriteServiceTest extends TestCase
 {
     use SqlTestEnable;
@@ -38,6 +42,14 @@ static function () {
                     );
                 },
             ],
+            [
+                static function (self $self) {
+                    $self->assert(
+                        new ChunkWriteStateEntity(32768, TestModel::class, 3, 2),
+                        [new TestModel(), new TestModel(), new TestModel()],
+                    );
+                },
+            ],
         ];
     }
 
@@ -50,7 +62,10 @@ public function test(Closure $assert): void
         $assert($this);
     }
 
-    public function assert(ChunkWriteStateEntity $expected, Closure $data): void
+    /**
+     * @param Closure(): TData|TData $data
+     */
+    public function assert(ChunkWriteStateEntity $expected, Closure|iterable $data): void
     {
         $state = (new ChunkWriteService())->write($data);
         Assert::assertEquals($expected, $state);