Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

fix: TransferManager#downloadFileInChunks issues #2373

Merged
merged 7 commits into from
Nov 29, 2023
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>}
Comment on lines -613 to +623
Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

...and fixed JSDoc declaration

*
* @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);
rhodgkins marked this conversation as resolved.
Show resolved Hide resolved
} 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)];
}
rhodgkins marked this conversation as resolved.
Show resolved Hide resolved

/**
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