Merge pull request #118 from jasonlfunk/main

Stream large snapshots
spatie · Apr 20, 2022 · 69fa5c7 · 69fa5c7
2 parents ca25ad0 + 1bc9632
commit 69fa5c7
Show file tree

Hide file tree

Showing 4 changed files with 120 additions and 5 deletions.
diff --git a/README.md b/README.md
@@ -168,6 +168,12 @@ By default, `snapshot:load` will drop all existing tables in the database. If yo
 php artisan snapshot:load my-first-dump --drop-tables=0
 ```
 
+By default, `snapshot:load` will load the entire snapshot into memory which may cause problems when using large files. To avoid this, you can pass the `--stream` option to stream the snapshot to the database one statement at a time:
+
+```bash
+php artisan snapshot:load my-first-dump --stream
+```
+
 To list all the dumps run:
 
 ```bash

diff --git a/src/Commands/Load.php b/src/Commands/Load.php
@@ -12,7 +12,7 @@ class Load extends Command
     use AsksForSnapshotName;
     use ConfirmableTrait;
 
-    protected $signature = 'snapshot:load {name?} {--connection=} {--force} --disk {--latest} {--drop-tables=1}';
+    protected $signature = 'snapshot:load {name?} {--connection=} {--force} {--stream} --disk {--latest} {--drop-tables=1}';
 
     protected $description = 'Load up a snapshot.';
 
@@ -45,6 +45,10 @@ public function handle()
             return;
         }
 
+        if ($this->option('stream') ?: false) {
+            $snapshot->useStream();
+        }
+
         $snapshot->load($this->option('connection'), (bool) $this->option('drop-tables'));
 
         $this->info("Snapshot `{$name}` loaded!");

diff --git a/src/Snapshot.php b/src/Snapshot.php
@@ -3,12 +3,15 @@
 namespace Spatie\DbSnapshots;
 
 use Carbon\Carbon;
-use Illuminate\Filesystem\FilesystemAdapter as Disk;
 use Illuminate\Support\Facades\DB;
-use Spatie\DbSnapshots\Events\DeletedSnapshot;
-use Spatie\DbSnapshots\Events\DeletingSnapshot;
+use Illuminate\Support\LazyCollection;
+use Illuminate\Support\Facades\Storage;
 use Spatie\DbSnapshots\Events\LoadedSnapshot;
+use Spatie\DbSnapshots\Events\DeletedSnapshot;
 use Spatie\DbSnapshots\Events\LoadingSnapshot;
+use Spatie\DbSnapshots\Events\DeletingSnapshot;
+use Spatie\TemporaryDirectory\TemporaryDirectory;
+use Illuminate\Filesystem\FilesystemAdapter as Disk;
 
 class Snapshot
 {
@@ -20,6 +23,10 @@ class Snapshot
 
     public ?string $compressionExtension = null;
 
+    private bool $useStream = false;
+
+    const STREAM_BUFFER_SIZE = 16384;
+
     public function __construct(Disk $disk, string $fileName)
     {
         $this->disk = $disk;
@@ -36,6 +43,13 @@ public function __construct(Disk $disk, string $fileName)
         $this->name = pathinfo($fileName, PATHINFO_FILENAME);
     }
 
+    public function useStream()
+    {
+        $this->useStream = true;
+
+        return $this;
+    }
+
     public function load(string $connectionName = null, bool $dropTables = true): void
     {
         event(new LoadingSnapshot($this));
@@ -48,15 +62,72 @@ public function load(string $connectionName = null, bool $dropTables = true): vo
             $this->dropAllCurrentTables();
         }
 
+        $this->useStream ? $this->loadStream($connectionName) : $this->loadAsync($connectionName);
+
+        event(new LoadedSnapshot($this));
+    }
+
+    protected function loadAsync(string $connectionName = null)
+    {
         $dbDumpContents = $this->disk->get($this->fileName);
 
         if ($this->compressionExtension === 'gz') {
             $dbDumpContents = gzdecode($dbDumpContents);
         }
 
         DB::connection($connectionName)->unprepared($dbDumpContents);
+    }
 
-        event(new LoadedSnapshot($this));
+    protected function isASqlComment(string $line): bool
+    {
+        return substr($line, 0, 2) === '--';
+    }
+
+    protected function shouldIgnoreLine(string $line): bool
+    {
+        $line = trim($line);
+        return empty($line) || $this->isASqlComment($line);
+    }
+
+    protected function loadStream(string $connectionName = null)
+    {
+        LazyCollection::make(function() {
+            $stream = $this->disk->readStream($this->fileName);
+
+            $statement = '';
+            while(!feof($stream)) {
+                $chunk = $this->compressionExtension === 'gz'
+                        ? gzdecode(gzread($stream, self::STREAM_BUFFER_SIZE))
+                        : fread($stream, self::STREAM_BUFFER_SIZE);
+
+                $lines = explode("\n", $chunk);
+                foreach($lines as $idx => $line) {
+                    if ($this->shouldIgnoreLine($line)) {
+                        continue;
+                    }
+
+                    $statement .= $line;
+
+                    // Carry-over the last line to the next chunk since it
+                    // is possible that this chunk finished mid-line right on
+                    // a semi-colon.
+                    if (count($lines) == $idx + 1) {
+                        break;
+                    }
+
+                    if (substr(trim($statement), -1, 1) === ';') {
+                        yield $statement;
+                        $statement = '';
+                    }
+                }
+            }
+
+            if (substr(trim($statement), -1, 1) === ';') {
+                yield $statement;
+            }
+        })->each(function (string $statement) use($connectionName) {
+            DB::connection($connectionName)->unprepared($statement);
+        });
     }
 
     public function delete(): void

diff --git a/tests/Commands/LoadTest.php b/tests/Commands/LoadTest.php
@@ -38,6 +38,40 @@ public function it_can_load_a_snapshot()
         $this->assertSnapshotLoaded('snapshot2');
     }
 
+    /** @test */
+    public function it_can_load_a_snapshot_via_streaming()
+    {
+        $this->assertSnapshotNotLoaded('snapshot2');
+
+        $this->command
+            ->shouldReceive('choice')
+            ->once()
+            ->andReturn('snapshot2');
+
+        Artisan::call('snapshot:load', [
+            '--stream' => true
+        ]);
+
+        $this->assertSnapshotLoaded('snapshot2');
+    }
+
+    /** @test */
+    public function it_can_load_a_compressed_snapshot_via_streaming()
+    {
+        $this->assertSnapshotNotLoaded('snapshot4');
+
+        $this->command
+            ->shouldReceive('choice')
+            ->once()
+            ->andReturn('snapshot4');
+
+        Artisan::call('snapshot:load', [
+            '--stream' => true
+        ]);
+
+        $this->assertSnapshotLoaded('snapshot4');
+    }
+
     /** @test */
     public function it_drops_tables_when_loading_a_snapshot()
     {