Skip to content

Commit

Permalink
fix: TransferManager#downloadFileInChunks issues (#2373)
Browse files Browse the repository at this point in the history
* 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>
  • Loading branch information
rhodgkins and danielbankhead committed Nov 29, 2023
1 parent 58bee1a commit 65950f3
Show file tree
Hide file tree
Showing 2 changed files with 74 additions and 13 deletions.
50 changes: 37 additions & 13 deletions src/transfer-manager.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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';
Expand Down Expand Up @@ -105,6 +111,7 @@ export interface DownloadFileInChunksOptions {
chunkSizeBytes?: number;
destination?: string;
validation?: 'crc32c' | false;
noReturnData?: boolean;
}

export interface UploadFileInChunksOptions {
Expand Down Expand Up @@ -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
* ```
Expand All @@ -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)
Expand Down Expand Up @@ -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)];
}

/**
Expand Down
37 changes: 37 additions & 0 deletions test/transfer-manager.ts
Original file line number Diff line number Diff line change
Expand Up @@ -249,6 +249,7 @@ describe('Transfer Manager', () => {
{
metadata: {
size: 1024,
crc32c: 'AAAAAA==',
},
},
]);
Expand All @@ -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 = () => {
Expand All @@ -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(
Expand Down

0 comments on commit 65950f3

Please sign in to comment.