use generator to optimize memory

Author: zajca

src/Backend/BackendImportAdapterInterface.php

@@ -4,6 +4,7 @@
 
 namespace Keboola\Db\ImportExport\Backend;
 
+use Generator;
 use Keboola\Db\ImportExport\ImportOptions;
 use Keboola\Db\ImportExport\Storage\DestinationInterface;
 use Keboola\Db\ImportExport\Storage\SourceInterface;
@@ -12,12 +13,9 @@ interface BackendImportAdapterInterface
 {
     public function __construct(SourceInterface $source);
 
-    /**
-     * @return string[]
-     */
     public function getCopyCommands(
         DestinationInterface $destination,
         ImportOptions $importOptions,
         string $stagingTableName
-    ): array;
+    ): Generator;
 }

src/Backend/Snowflake/SnowflakeImportAdapterInterface.php

@@ -4,6 +4,7 @@
 
 namespace Keboola\Db\ImportExport\Backend\Snowflake;
 
+use Generator;
 use Keboola\Db\ImportExport\Backend\ImportState;
 use Keboola\Db\ImportExport\ImportOptions;
 use Keboola\Db\ImportExport\Backend\BackendImportAdapterInterface;
@@ -15,11 +16,10 @@ interface SnowflakeImportAdapterInterface extends BackendImportAdapterInterface
     /**
      * Snowflake import is handled differently for copy table2table and file2table
      *
-     * @param string[] $commands - sql commands array
      * @return int - number of imported rows
      */
     public function executeCopyCommands(
-        array $commands,
+        Generator $commands,
         Connection $connection,
         DestinationInterface $destination,
         ImportOptions $importOptions,

src/Storage/ABS/SnowflakeImportAdapter.php

@@ -4,6 +4,7 @@
 
 namespace Keboola\Db\ImportExport\Storage\ABS;
 
+use Generator;
 use Keboola\Csv\CsvOptions;
 use Keboola\Db\ImportExport\Backend\ImporterInterface;
 use Keboola\Db\ImportExport\Backend\ImportState;
@@ -35,7 +36,7 @@ public function __construct(SourceInterface $source)
      * @param Table $destination
      */
     public function executeCopyCommands(
-        array $commands,
+        Generator $commands,
         Connection $connection,
         DestinationInterface $destination,
         ImportOptions $importOptions,
@@ -62,11 +63,10 @@ public function getCopyCommands(
         DestinationInterface $destination,
         ImportOptions $importOptions,
         string $stagingTableName
-    ): array {
+    ): Generator {
         $filesToImport = $this->source->getManifestEntries();
-        $commands = [];
         foreach (array_chunk($filesToImport, ImporterInterface::SLICED_FILES_CHUNK_SIZE) as $entries) {
-            $commands[] = sprintf(
+            yield sprintf(
                 'COPY INTO %s.%s 
 FROM %s
 CREDENTIALS=(AZURE_SAS_TOKEN=\'%s\')
@@ -88,7 +88,6 @@ function ($entry) {
                 )
             );
         }
-        return $commands;
     }
 
     private function getCsvCopyCommandOptions(

src/Storage/Snowflake/SnowflakeImportAdapter.php

@@ -4,9 +4,9 @@
 
 namespace Keboola\Db\ImportExport\Storage\Snowflake;
 
+use Generator;
 use Keboola\Db\ImportExport\Backend\ImportState;
 use Keboola\Db\ImportExport\ImportOptions;
-use Keboola\Db\ImportExport\Backend\Snowflake\Helper\QuoteHelper;
 use Keboola\Db\ImportExport\Backend\Snowflake\SnowflakeImportAdapterInterface;
 use Keboola\Db\ImportExport\Storage\DestinationInterface;
 use Keboola\Db\ImportExport\Storage\SourceInterface;
@@ -32,14 +32,14 @@ public function __construct(SourceInterface $source)
      * @param Table $destination
      */
     public function executeCopyCommands(
-        array $commands,
+        Generator $commands,
         Connection $connection,
         DestinationInterface $destination,
         ImportOptions $importOptions,
         ImportState $importState
     ): int {
         $importState->startTimer('copyToStaging');
-        $connection->query($commands[0]);
+        $connection->query($commands->current());
         $rows = $connection->fetchAll(sprintf(
             'SELECT COUNT(*) AS "count" FROM %s.%s',
             QueryBuilder::quoteIdentifier($destination->getSchema()),
@@ -56,7 +56,7 @@ public function getCopyCommands(
         DestinationInterface $destination,
         ImportOptions $importOptions,
         string $stagingTableName
-    ): array {
+    ): Generator {
         $quotedColumns = array_map(function ($column) {
             return QueryBuilder::quoteIdentifier($column);
         }, $importOptions->getColumns());
@@ -75,6 +75,6 @@ public function getCopyCommands(
             QueryBuilder::quoteIdentifier($this->source->getTableName()),
         );
 
-        return [$sql];
+        yield $sql;
     }
 }

src/Storage/Snowflake/Table.php

@@ -8,7 +8,6 @@
 use Keboola\Db\ImportExport\Backend\BackendImportAdapterInterface;
 use Keboola\Db\ImportExport\Backend\ExporterInterface;
 use Keboola\Db\ImportExport\Backend\ImporterInterface;
-use Keboola\Db\ImportExport\Backend\Snowflake\Helper\QuoteHelper;
 use Keboola\Db\ImportExport\Backend\Snowflake\Importer as SnowflakeImporter;
 use Keboola\Db\ImportExport\Storage\DestinationInterface;
 use Keboola\Db\ImportExport\Storage\NoBackendAdapterException;

tests/unit/SourceStorage/ABS/BaseFileTest.php renamed to tests/unit/Storage/ABS/BaseFileTest.php

@@ -4,6 +4,7 @@
 
 namespace Tests\Keboola\Db\ImportExportUnit\Storage\ABS;
 
+use Generator;
 use Keboola\Csv\CsvOptions;
 use Keboola\Db\ImportExport\Backend\ImportState;
 use Keboola\Db\ImportExport\ImportOptions;
@@ -30,9 +31,14 @@ public function testExecuteCopyCommands(): void
         /** @var ImportOptions|MockObject $options */
         $options = self::createMock(ImportOptions::class);
 
+        $generator = function (): Generator {
+            yield 'cmd1';
+            yield 'cmd2';
+        };
+
         $adapter = new SnowflakeImportAdapter($source);
         $rows = $adapter->executeCopyCommands(
-            ['cmd1', 'cmd2'],
+            $generator(),
             $connection,
             new Storage\Snowflake\Table('', ''),
             $options,
@@ -60,16 +66,16 @@ public function testGetCopyCommands(): void
             'stagingTable'
         );
 
-        self::assertSame([
+        self::assertSame(
             <<<EOT
 COPY INTO "schema"."stagingTable" 
 FROM 'containerUrl'
 CREDENTIALS=(AZURE_SAS_TOKEN='sasToken')
 FILE_FORMAT = (TYPE=CSV FIELD_DELIMITER = ',' FIELD_OPTIONALLY_ENCLOSED_BY = '\"' ESCAPE_UNENCLOSED_FIELD = NONE)
 FILES = ('azure://url')
 EOT,
-
-        ], $commands);
+            $commands->current()
+        );
     }
 
     public function testGetCopyCommandsChunk(): void
@@ -105,21 +111,28 @@ public function testGetCopyCommandsChunk(): void
             return sprintf("'%s'", $file);
         }, $cmd2Files));
 
-        self::assertSame([
+        self::assertSame(
             <<<EOT
 COPY INTO "schema"."stagingTable" 
 FROM 'containerUrl'
 CREDENTIALS=(AZURE_SAS_TOKEN='sasToken')
 FILE_FORMAT = (TYPE=CSV FIELD_DELIMITER = ',' FIELD_OPTIONALLY_ENCLOSED_BY = '\"' ESCAPE_UNENCLOSED_FIELD = NONE)
 FILES = ($cmd1Files)
 EOT,
+            $commands->current()
+        );
+
+        $commands->next();
+
+        self::assertSame(
             <<<EOT
 COPY INTO "schema"."stagingTable" 
 FROM 'containerUrl'
 CREDENTIALS=(AZURE_SAS_TOKEN='sasToken')
 FILE_FORMAT = (TYPE=CSV FIELD_DELIMITER = ',' FIELD_OPTIONALLY_ENCLOSED_BY = '\"' ESCAPE_UNENCLOSED_FIELD = NONE)
 FILES = ($cmd2Files)
 EOT,
-        ], $commands);
+            $commands->current()
+        );
     }
 }