refactor assets streaming

packages/core/data-transfer/src/strapi/providers/remote-destination/index.ts

@@ -333,13 +333,7 @@ class RemoteStrapiDestinationProvider implements IDestinationProvider {
   }
 
   createAssetsWriteStream(): Writable | Promise<Writable> {
-    type TransferAssetFlow = {
-      assetID: string;
-      action: 'stream';
-      data: Omit<IAsset, 'stream'>;
-      chunk: Buffer;
-    };
-    let batch: TransferAssetFlow[] = [];
+    let batch: any[] = [];
     const batchSize = 100;
     const startAssetsTransferOnce = this.#startStepOnce('assets');
 
@@ -375,23 +369,33 @@ class RemoteStrapiDestinationProvider implements IDestinationProvider {
 
         const { filename, filepath, stats, stream } = asset;
         const assetID = v4();
+        batch.push({
+          action: 'start',
+          assetID,
+          data: { filename, filepath, stats },
+        });
+        if (batch.length === batchSize) {
+          await this.#streamStep('assets', batch);
+          batch = [];
+        }
 
         for await (const chunk of stream) {
-          batch.push({
-            action: 'stream',
-            assetID,
-            data: { filename, filepath, stats },
-            chunk,
-          });
+          batch.push({ action: 'stream', assetID, data: chunk });
           if (batch.length === batchSize) {
             await this.#streamStep('assets', batch);
             batch = [];
           }
         }
+
+        batch.push({
+          action: 'end',
+          assetID,
+        });
         if (batch.length === batchSize) {
           await this.#streamStep('assets', batch);
           batch = [];
         }
+
         callback();
       },
     });

packages/core/data-transfer/src/strapi/remote/controllers/push.ts

@@ -103,54 +103,45 @@ const createPushController = (options: ILocalStrapiDestinationProviderOptions):
         });
       },
 
-      async assets(payload) {
+      async assets(payloads) {
         // TODO: close the stream upong receiving an 'end' event instead
-        if (payload === null) {
-          const assetsKeys = Object.keys(assets);
-          const previousAssetId = assetsKeys[assetsKeys.length - 1];
-          const { stream } = assets[previousAssetId];
-          stream
-            .on('close', () => {
-              delete assets[previousAssetId];
-            })
-            .on('error', (error) => {
-              throw new Error(`error while closing stream ${error}`);
-            })
-            .end();
+        if (payloads === null) {
           streams.assets?.end();
           return;
         }
-
         if (!streams.assets) {
           streams.assets = await provider.createAssetsWriteStream();
         }
 
-        payload.map(async (asset) => {
-          const { action, assetID, data, chunk } = asset;
-          const assetsKeys = Object.keys(assets);
-          if (action === 'stream') {
-            if (!assets[assetID] && assetsKeys.length !== 0) {
-              const previousAssetId = assetsKeys[assetsKeys.length - 1];
-              const { stream } = assets[previousAssetId];
-              stream
-                .on('close', () => {
-                  delete assets[previousAssetId];
-                })
-                .on('error', (error) => {
-                  throw new Error(`error while closing stream ${error}`);
-                })
-                .end();
-            }
+        payloads.forEach(async (payload) => {
+          const { action, assetID } = payload;
 
-            if (!assets[assetID] && streams.assets) {
-              assets[assetID] = { ...data, stream: new PassThrough() };
-              writeAsync(streams.assets, assets[assetID]);
-            }
+          if (action === 'start' && streams.assets) {
+            assets[assetID] = { ...payload.data, stream: new PassThrough() };
+            writeAsync(streams.assets, assets[assetID]);
+          }
+
+          if (action === 'stream') {
             // The buffer has gone through JSON operations and is now of shape { type: "Buffer"; data: UInt8Array }
             // We need to transform it back into a Buffer instance
-            const rawBuffer = chunk as unknown as { type: 'Buffer'; data: Uint8Array };
-            const buffer = Buffer.from(rawBuffer.data);
-            await writeAsync(assets[assetID].stream, buffer);
+            const rawBuffer = payload.data as unknown as { type: 'Buffer'; data: Uint8Array };
+            const chunk = Buffer.from(rawBuffer.data);
+
+            await writeAsync(assets[assetID].stream, chunk);
+          }
+
+          if (action === 'end') {
+            await new Promise<void>((resolve, reject) => {
+              const { stream } = assets[assetID];
+
+              stream
+                .on('close', () => {
+                  delete assets[assetID];
+                  resolve();
+                })
+                .on('error', reject)
+                .end();
+            });
           }
         });
       },

packages/core/data-transfer/types/remote/protocol/client/transfer/push.d.ts

@@ -24,9 +24,8 @@ type TransferStepCommands<T extends string, U> = { step: T } & TransferStepFlow<
 
 type TransferStepFlow<U> = { action: 'start' } | { action: 'stream'; data: U } | { action: 'end' };
 
-type TransferAssetFlow = {
-  assetID: string;
-  action: 'stream';
-  data: Omit<IAsset, 'stream'>;
-  chunk: Buffer;
-};
+type TransferAssetFlow = { assetID: string } & (
+  | { action: 'start'; data: Omit<IAsset, 'stream'> }
+  | { action: 'stream'; data: Buffer }
+  | { action: 'end' }
+);