fix: TransferManager#downloadFileInChunks issues (#2373)

* First pass at fixes for #2372 * Added JSDoc for DownloadFileInChunksOptions#noReturnData * Changed type of DownloadFileInChunksOptions#noReturnData to be clearer on the default * Update src/transfer-manager.ts Fixed typo Co-authored-by: Daniel Bankhead <dan@danielbankhead.com> * TransferManager#downloadFileInChunks test case for failed CRC32C validation * TransferManager#downloadFileInChunks test case for return values and noReturnData option * Fixed JSDocs for TransferManager#downloadFileInChunks --------- Co-authored-by: Daniel Bankhead <dan@danielbankhead.com>
googleapis · Nov 29, 2023 · 65950f3 · 65950f3
1 parent 58bee1a
commit 65950f3
Show file tree

Hide file tree

Showing 2 changed files with 74 additions and 13 deletions.
diff --git a/src/transfer-manager.ts b/src/transfer-manager.ts
@@ -15,7 +15,13 @@
  */
 
 import {Bucket, UploadOptions, UploadResponse} from './bucket.js';
-import {DownloadOptions, DownloadResponse, File} from './file.js';
+import {
+  DownloadOptions,
+  DownloadResponse,
+  File,
+  FileExceptionMessages,
+  RequestError,
+} from './file.js';
 import pLimit from 'p-limit';
 import * as path from 'path';
 import {createReadStream, promises as fsp} from 'fs';
@@ -105,6 +111,7 @@ export interface DownloadFileInChunksOptions {
   chunkSizeBytes?: number;
   destination?: string;
   validation?: 'crc32c' | false;
+  noReturnData?: boolean;
 }
 
 export interface UploadFileInChunksOptions {
@@ -604,15 +611,16 @@ export class TransferManager {
    * to use when downloading the file.
    * @property {number} [chunkSizeBytes] The size in bytes of each chunk to be downloaded.
    * @property {string | boolean} [validation] Whether or not to perform a CRC32C validation check when download is complete.
+   * @property {boolean} [noReturnData] Whether or not to return the downloaded data. A `true` value here would be useful for files with a size that will not fit into memory.
    *
    */
   /**
    * Download a large file in chunks utilizing parallel download operations. This is a convenience method
    * that utilizes {@link File#download} to perform the download.
    *
-   * @param {object} [file | string] {@link File} to download.
+   * @param {File | string} fileOrName {@link File} to download.
    * @param {DownloadFileInChunksOptions} [options] Configuration options.
-   * @returns {Promise<DownloadResponse>}
+   * @returns {Promise<void | DownloadResponse>}
    *
    * @example
    * ```
@@ -639,7 +647,8 @@ export class TransferManager {
     let limit = pLimit(
       options.concurrencyLimit || DEFAULT_PARALLEL_CHUNKED_DOWNLOAD_LIMIT
     );
-    const promises: Promise<{bytesWritten: number; buffer: Buffer}>[] = [];
+    const noReturnData = Boolean(options.noReturnData);
+    const promises: Promise<Buffer | void>[] = [];
     const file: File =
       typeof fileOrName === 'string'
         ? this.bucket.file(fileOrName)
@@ -667,24 +676,39 @@ export class TransferManager {
             end: chunkEnd,
             [GCCL_GCS_CMD_KEY]: GCCL_GCS_CMD_FEATURE.DOWNLOAD_SHARDED,
           });
-          return fileToWrite.write(resp[0], 0, resp[0].length, chunkStart);
+          const result = await fileToWrite.write(
+            resp[0],
+            0,
+            resp[0].length,
+            chunkStart
+          );
+          if (noReturnData) return;
+          return result.buffer;
         })
       );
 
       start += chunkSize;
     }
 
-    let results: DownloadResponse;
+    let chunks: Array<Buffer | void>;
     try {
-      const data = await Promise.all(promises);
-      results = data.map(result => result.buffer) as DownloadResponse;
-      if (options.validation === 'crc32c') {
-        await CRC32C.fromFile(filePath);
-      }
-      return results;
+      chunks = await Promise.all(promises);
     } finally {
-      fileToWrite.close();
+      await fileToWrite.close();
+    }
+
+    if (options.validation === 'crc32c' && fileInfo[0].metadata.crc32c) {
+      const downloadedCrc32C = await CRC32C.fromFile(filePath);
+      if (!downloadedCrc32C.validate(fileInfo[0].metadata.crc32c)) {
+        const mismatchError = new RequestError(
+          FileExceptionMessages.DOWNLOAD_MISMATCH
+        );
+        mismatchError.code = 'CONTENT_DOWNLOAD_MISMATCH';
+        throw mismatchError;
+      }
     }
+    if (noReturnData) return;
+    return [Buffer.concat(chunks as Buffer[], size)];
   }
 
   /**

diff --git a/test/transfer-manager.ts b/test/transfer-manager.ts
@@ -249,6 +249,7 @@ describe('Transfer Manager', () => {
         {
           metadata: {
             size: 1024,
+            crc32c: 'AAAAAA==',
           },
         },
       ]);
@@ -265,6 +266,26 @@ describe('Transfer Manager', () => {
       assert.strictEqual(downloadCallCount, 1);
     });
 
+    it('should return downloaded data', async () => {
+      sandbox.stub(file, 'download').callsFake(() => {
+        return Promise.resolve([Buffer.alloc(100)]);
+      });
+
+      const data = await transferManager.downloadFileInChunks(file);
+      assert.deepStrictEqual(data, [Buffer.alloc(1024)]);
+    });
+
+    it('should not return downloaded data when noReturnData flag is set', async () => {
+      sandbox.stub(file, 'download').callsFake(() => {
+        return Promise.resolve([Buffer.alloc(100)]);
+      });
+
+      const data = await transferManager.downloadFileInChunks(file, {
+        noReturnData: true,
+      });
+      assert.strictEqual(data, undefined);
+    });
+
     it('should call fromFile when validation is set to crc32c', async () => {
       let callCount = 0;
       file.download = () => {
@@ -279,6 +300,22 @@ describe('Transfer Manager', () => {
       assert.strictEqual(callCount, 1);
     });
 
+    it('should throw an error if crc32c validation fails', async () => {
+      file.download = () => {
+        return Promise.resolve([Buffer.alloc(0)]) as Promise<DownloadResponse>;
+      };
+      CRC32C.fromFile = () => {
+        return Promise.resolve(new CRC32C(1)); // Set non-expected initial value
+      };
+
+      await assert.rejects(
+        transferManager.downloadFileInChunks(file, {validation: 'crc32c'}),
+        {
+          code: 'CONTENT_DOWNLOAD_MISMATCH',
+        }
+      );
+    });
+
     it('should set the appropriate `GCCL_GCS_CMD_KEY`', async () => {
       sandbox.stub(file, 'download').callsFake(async options => {
         assert.strictEqual(