diff --git a/src/file.ts b/src/file.ts index a236723bc..65fc20b58 100644 --- a/src/file.ts +++ b/src/file.ts @@ -61,6 +61,7 @@ import { ApiError, Duplexify, DuplexifyConstructor, + GCCL_GCS_CMD_KEY, } from './nodejs-common/util'; // eslint-disable-next-line @typescript-eslint/no-var-requires const duplexify: DuplexifyConstructor = require('duplexify'); @@ -221,6 +222,7 @@ type PublicResumableUploadOptions = export interface CreateResumableUploadOptions extends Pick { preconditionOpts?: PreconditionOptions; + [GCCL_GCS_CMD_KEY]?: resumableUpload.UploadConfig[typeof GCCL_GCS_CMD_KEY]; } export type CreateResumableUploadResponse = [string]; @@ -371,6 +373,7 @@ export interface CreateReadStreamOptions { start?: number; end?: number; decompress?: boolean; + [GCCL_GCS_CMD_KEY]?: string; } export interface SaveOptions extends CreateWriteStreamOptions { @@ -1580,12 +1583,16 @@ class File extends ServiceObject { headers.Range = `bytes=${tailRequest ? end : `${start}-${end}`}`; } - const reqOpts = { + const reqOpts: DecorateRequestOptions = { uri: '', headers, qs: query, }; + if (options[GCCL_GCS_CMD_KEY]) { + reqOpts[GCCL_GCS_CMD_KEY] = options[GCCL_GCS_CMD_KEY]; + } + this.requestStream(reqOpts) .on('error', err => { throughStream.destroy(err); @@ -1738,6 +1745,7 @@ class File extends ServiceObject { userProject: options.userProject || this.userProject, retryOptions: retryOptions, params: options?.preconditionOpts || this.instancePreconditionOpts, + [GCCL_GCS_CMD_KEY]: options[GCCL_GCS_CMD_KEY], }, callback! ); @@ -3907,6 +3915,7 @@ class File extends ServiceObject { params: options?.preconditionOpts || this.instancePreconditionOpts, chunkSize: options?.chunkSize, highWaterMark: options?.highWaterMark, + [GCCL_GCS_CMD_KEY]: options[GCCL_GCS_CMD_KEY], }); uploadStream @@ -3951,6 +3960,7 @@ class File extends ServiceObject { name: this.name, }, uri: uri, + [GCCL_GCS_CMD_KEY]: options[GCCL_GCS_CMD_KEY], }; if (this.generation !== undefined) { diff --git a/src/nodejs-common/service.ts b/src/nodejs-common/service.ts index d20d19735..69fc9dfb6 100644 --- a/src/nodejs-common/service.ts +++ b/src/nodejs-common/service.ts @@ -21,6 +21,7 @@ import {Interceptor} from './service-object'; import { BodyResponseCallback, DecorateRequestOptions, + GCCL_GCS_CMD_KEY, MakeAuthenticatedRequest, PackageJson, util, @@ -253,6 +254,12 @@ export class Service { } gccl-invocation-id/${uuid.v4()}`, }; + if (reqOpts[GCCL_GCS_CMD_KEY]) { + reqOpts.headers[ + 'x-goog-api-client' + ] += ` gccl-gcs-cmd/${reqOpts[GCCL_GCS_CMD_KEY]}`; + } + if (reqOpts.shouldReturnStream) { return this.makeAuthenticatedRequest(reqOpts) as {} as r.Request; } else { diff --git a/src/nodejs-common/util.ts b/src/nodejs-common/util.ts index 0d4fc6374..c119fd71b 100644 --- a/src/nodejs-common/util.ts +++ b/src/nodejs-common/util.ts @@ -36,6 +36,14 @@ import {getRuntimeTrackingString} from '../util'; const packageJson = require('../../../package.json'); +/** + * A unique symbol for providing a `gccl-gcs-cmd` value + * for the `X-Goog-API-Client` header. + * + * E.g. the `V` in `X-Goog-API-Client: gccl-gcs-cmd/V` + **/ +export const GCCL_GCS_CMD_KEY = Symbol.for('GCCL_GCS_CMD'); + // eslint-disable-next-line @typescript-eslint/no-var-requires const duplexify: DuplexifyConstructor = require('duplexify'); @@ -214,7 +222,9 @@ export interface MakeWritableStreamOptions { request?: r.Options; makeAuthenticatedRequest( - reqOpts: r.OptionsWithUri, + reqOpts: r.OptionsWithUri & { + [GCCL_GCS_CMD_KEY]?: string; + }, fnobj: { onAuthenticated( err: Error | null, @@ -233,6 +243,7 @@ export interface DecorateRequestOptions extends r.CoreOptions { interceptors_?: Interceptor[]; shouldReturnStream?: boolean; projectId?: string; + [GCCL_GCS_CMD_KEY]?: string; } export interface ParsedHttpResponseBody { @@ -530,7 +541,9 @@ export class Util { body: writeStream, }, ], - } as {} as r.OptionsWithUri; + } as {} as r.OptionsWithUri & { + [GCCL_GCS_CMD_KEY]?: string; + }; options.makeAuthenticatedRequest(reqOpts, { onAuthenticated(err, authenticatedReqOpts) { @@ -539,7 +552,9 @@ export class Util { return; } - requestDefaults.headers = util._getDefaultHeaders(); + requestDefaults.headers = util._getDefaultHeaders( + reqOpts[GCCL_GCS_CMD_KEY] + ); const request = teenyRequest.defaults(requestDefaults); request(authenticatedReqOpts!, (err, resp, body) => { util.handleResp(err, resp, body, (err, data) => { @@ -863,7 +878,9 @@ export class Util { maxRetryValue = config.retryOptions.maxRetries; } - requestDefaults.headers = this._getDefaultHeaders(); + requestDefaults.headers = this._getDefaultHeaders( + reqOpts[GCCL_GCS_CMD_KEY] + ); const options = { request: teenyRequest.defaults(requestDefaults), retries: autoRetryValue !== false ? maxRetryValue : 0, @@ -1014,13 +1031,19 @@ export class Util { : [optionsOrCallback as T, cb as C]; } - _getDefaultHeaders() { - return { + _getDefaultHeaders(gcclGcsCmd?: string) { + const headers = { 'User-Agent': util.getUserAgentFromPackageJson(packageJson), 'x-goog-api-client': `${getRuntimeTrackingString()} gccl/${ packageJson.version } gccl-invocation-id/${uuid.v4()}`, }; + + if (gcclGcsCmd) { + headers['x-goog-api-client'] += ` gccl-gcs-cmd/${gcclGcsCmd}`; + } + + return headers; } } diff --git a/src/resumable-upload.ts b/src/resumable-upload.ts index 521b6586f..f0194869f 100644 --- a/src/resumable-upload.ts +++ b/src/resumable-upload.ts @@ -27,6 +27,7 @@ import retry = require('async-retry'); import {RetryOptions, PreconditionOptions} from './storage'; import * as uuid from 'uuid'; import {getRuntimeTrackingString} from './util'; +import {GCCL_GCS_CMD_KEY} from './nodejs-common/util'; const NOT_FOUND_STATUS_CODE = 404; const RESUMABLE_INCOMPLETE_STATUS_CODE = 308; @@ -193,6 +194,8 @@ export interface UploadConfig extends Pick { * Configuration options for retrying retryable errors. */ retryOptions: RetryOptions; + + [GCCL_GCS_CMD_KEY]?: string; } export interface ConfigMetadata { @@ -274,6 +277,7 @@ export class Upload extends Writable { private localWriteCache: Buffer[] = []; private localWriteCacheByteLength = 0; private upstreamEnded = false; + #gcclGcsCmd?: string; constructor(cfg: UploadConfig) { super(cfg); @@ -347,6 +351,8 @@ export class Upload extends Writable { : NaN; this.contentLength = isNaN(contentLength) ? '*' : contentLength; + this.#gcclGcsCmd = cfg[GCCL_GCS_CMD_KEY]; + this.once('writing', () => { if (this.uri) { this.continueUploading(); @@ -585,6 +591,14 @@ export class Upload extends Writable { delete metadata.contentType; } + let googAPIClient = `${getRuntimeTrackingString()} gccl/${ + packageJson.version + } gccl-invocation-id/${this.currentInvocationId.uri}`; + + if (this.#gcclGcsCmd) { + googAPIClient += ` gccl-gcs-cmd/${this.#gcclGcsCmd}`; + } + // Check if headers already exist before creating new ones const reqOpts: GaxiosOptions = { method: 'POST', @@ -598,9 +612,7 @@ export class Upload extends Writable { ), data: metadata, headers: { - 'x-goog-api-client': `${getRuntimeTrackingString()} gccl/${ - packageJson.version - } gccl-invocation-id/${this.currentInvocationId.uri}`, + 'x-goog-api-client': googAPIClient, ...headers, }, }; @@ -766,10 +778,16 @@ export class Upload extends Writable { }, }); + let googAPIClient = `${getRuntimeTrackingString()} gccl/${ + packageJson.version + } gccl-invocation-id/${this.currentInvocationId.chunk}`; + + if (this.#gcclGcsCmd) { + googAPIClient += ` gccl-gcs-cmd/${this.#gcclGcsCmd}`; + } + const headers: GaxiosOptions['headers'] = { - 'x-goog-api-client': `${getRuntimeTrackingString()} gccl/${ - packageJson.version - } gccl-invocation-id/${this.currentInvocationId.chunk}`, + 'x-goog-api-client': googAPIClient, }; // If using multiple chunk upload, set appropriate header @@ -904,15 +922,21 @@ export class Upload extends Writable { } private async getAndSetOffset() { + let googAPIClient = `${getRuntimeTrackingString()} gccl/${ + packageJson.version + } gccl-invocation-id/${this.currentInvocationId.offset}`; + + if (this.#gcclGcsCmd) { + googAPIClient += ` gccl-gcs-cmd/${this.#gcclGcsCmd}`; + } + const opts: GaxiosOptions = { method: 'PUT', url: this.uri!, headers: { 'Content-Length': 0, 'Content-Range': 'bytes */*', - 'x-goog-api-client': `${getRuntimeTrackingString()} gccl/${ - packageJson.version - } gccl-invocation-id/${this.currentInvocationId.offset}`, + 'x-goog-api-client': googAPIClient, }, }; try { diff --git a/src/transfer-manager.ts b/src/transfer-manager.ts index af0a2d3ae..344b4438a 100644 --- a/src/transfer-manager.ts +++ b/src/transfer-manager.ts @@ -26,6 +26,10 @@ import * as retry from 'async-retry'; import {ApiError} from './nodejs-common'; import {GaxiosResponse, Headers} from 'gaxios'; import {createHash} from 'crypto'; +import {GCCL_GCS_CMD_KEY} from './nodejs-common/util'; +import {getRuntimeTrackingString} from './util'; + +const packageJson = require('../../package.json'); /** * Default number of concurrently executing promises to use when calling uploadManyFiles. @@ -64,6 +68,21 @@ const UPLOAD_IN_CHUNKS_DEFAULT_CHUNK_SIZE = 32 * 1024 * 1024; const DEFAULT_PARALLEL_CHUNKED_UPLOAD_LIMIT = 2; const EMPTY_REGEX = '(?:)'; + +/** + * The `gccl-gcs-cmd` value for the `X-Goog-API-Client` header. + * Example: `gccl-gcs-cmd/tm.upload_many` + * + * @see {@link GCCL_GCS_CMD}. + * @see {@link GCCL_GCS_CMD_KEY}. + */ +const GCCL_GCS_CMD_FEATURE = { + UPLOAD_MANY: 'tm.upload_many', + DOWNLOAD_MANY: 'tm.download_many', + UPLOAD_SHARDED: 'tm.upload_sharded', + DOWNLOAD_SHARDED: 'tm.download_sharded', +}; + export interface UploadManyFilesOptions { concurrencyLimit?: number; skipIfExists?: boolean; @@ -170,8 +189,9 @@ class XMLMultiPartUploadHelper implements MultiPartUploadHelper { this.uploadId = uploadId || ''; this.bucket = bucket; this.fileName = fileName; - // eslint-disable-next-line prettier/prettier - this.baseUrl = `https://${bucket.name}.${new URL(this.bucket.storage.apiEndpoint).hostname}/${fileName}`; + this.baseUrl = `https://${bucket.name}.${ + new URL(this.bucket.storage.apiEndpoint).hostname + }/${fileName}`; this.xmlBuilder = new XMLBuilder({arrayNodeName: 'Part'}); this.xmlParser = new XMLParser(); this.partsMap = partsMap || new Map(); @@ -183,28 +203,52 @@ class XMLMultiPartUploadHelper implements MultiPartUploadHelper { }; } + #setGoogApiClientHeaders(headers: Headers = {}): Headers { + let headerFound = false; + + for (const [key, value] of Object.entries(headers)) { + if (key.toLocaleLowerCase().trim() === 'x-goog-api-client') { + headerFound = true; + + // Prepend command feature to value, if not already there + if (!value.includes(GCCL_GCS_CMD_FEATURE.UPLOAD_SHARDED)) { + headers[ + key + ] = `${value} gccl-gcs-cmd/${GCCL_GCS_CMD_FEATURE.UPLOAD_SHARDED}`; + } + break; + } + } + + // If the header isn't present, add it + if (!headerFound) { + headers['x-goog-api-client'] = `${getRuntimeTrackingString()} gccl/${ + packageJson.version + } gccl-gcs-cmd/${GCCL_GCS_CMD_FEATURE.UPLOAD_SHARDED}`; + } + + return headers; + } + /** * Initiates a multipart upload (MPU) to the XML API and stores the resultant upload id. * * @returns {Promise} */ - async initiateUpload(headers: {[key: string]: string} = {}): Promise { + async initiateUpload(headers: Headers = {}): Promise { const url = `${this.baseUrl}?uploads`; return retry(async bail => { try { const res = await this.authClient.request({ + headers: this.#setGoogApiClientHeaders(headers), method: 'POST', url, - headers, }); + if (res.data && res.data.error) { throw res.data.error; } - const parsedXML = this.xmlParser.parse<{ - InitiateMultipartUploadResult: { - UploadId: string; - }; - }>(res.data); + const parsedXML = this.xmlParser.parse(res.data); this.uploadId = parsedXML.InitiateMultipartUploadResult.UploadId; } catch (e) { this.#handleErrorResponse(e as Error, bail); @@ -227,7 +271,7 @@ class XMLMultiPartUploadHelper implements MultiPartUploadHelper { validation?: 'md5' | false ): Promise { const url = `${this.baseUrl}?partNumber=${partNumber}&uploadId=${this.uploadId}`; - let headers: Headers = {}; + let headers: Headers = this.#setGoogApiClientHeaders(); if (validation === 'md5') { const hash = createHash('md5').update(chunk).digest('base64'); @@ -274,6 +318,7 @@ class XMLMultiPartUploadHelper implements MultiPartUploadHelper { return retry(async bail => { try { const res = await this.authClient.request({ + headers: this.#setGoogApiClientHeaders(), url, method: 'POST', body, @@ -426,6 +471,7 @@ export class TransferManager { const passThroughOptionsCopy: UploadOptions = { ...options.passthroughOptions, + [GCCL_GCS_CMD_KEY]: GCCL_GCS_CMD_FEATURE.UPLOAD_MANY, }; passThroughOptionsCopy.destination = filePath; @@ -435,6 +481,7 @@ export class TransferManager { passThroughOptionsCopy.destination ); } + promises.push( limit(() => this.bucket.upload(filePath, passThroughOptionsCopy as UploadOptions) @@ -519,6 +566,7 @@ export class TransferManager { for (const file of files) { const passThroughOptionsCopy = { ...options.passthroughOptions, + [GCCL_GCS_CMD_KEY]: GCCL_GCS_CMD_FEATURE.DOWNLOAD_MANY, }; if (options.prefix) { @@ -531,6 +579,7 @@ export class TransferManager { if (options.stripPrefix) { passThroughOptionsCopy.destination = file.name.replace(regex, ''); } + promises.push(limit(() => file.download(passThroughOptionsCopy))); } @@ -601,9 +650,15 @@ export class TransferManager { chunkEnd = chunkEnd > size ? size : chunkEnd; promises.push( limit(() => - file.download({start: chunkStart, end: chunkEnd}).then(resp => { - return fileToWrite.write(resp[0], 0, resp[0].length, chunkStart); - }) + file + .download({ + start: chunkStart, + end: chunkEnd, + [GCCL_GCS_CMD_KEY]: GCCL_GCS_CMD_FEATURE.DOWNLOAD_SHARDED, + }) + .then(resp => { + return fileToWrite.write(resp[0], 0, resp[0].length, chunkStart); + }) ) ); diff --git a/test/file.ts b/test/file.ts index 4aa19968d..e399db904 100644 --- a/test/file.ts +++ b/test/file.ts @@ -59,6 +59,7 @@ import { BaseMetadata, SetMetadataOptions, } from '../src/nodejs-common/service-object'; +import {GCCL_GCS_CMD_KEY} from '../src/nodejs-common/util'; class HTTPError extends Error { code: number; @@ -1037,6 +1038,24 @@ describe('File', () => { file.createReadStream(options).resume(); }); + it('should pass the `GCCL_GCS_CMD_KEY` to `requestStream`', done => { + const expected = 'expected/value'; + + file.requestStream = (opts: DecorateRequestOptions) => { + assert.equal(opts[GCCL_GCS_CMD_KEY], expected); + + process.nextTick(() => done()); + + return duplexify(); + }; + + file + .createReadStream({ + [GCCL_GCS_CMD_KEY]: expected, + }) + .resume(); + }); + describe('authenticating', () => { it('should create an authenticated request', done => { file.requestStream = (opts: DecorateRequestOptions) => { @@ -4932,6 +4951,7 @@ describe('File', () => { makeWritableStreamOverride = (stream: {}, options_: any) => { assert.deepStrictEqual(options_.metadata, options.metadata); assert.deepStrictEqual(options_.request, { + [GCCL_GCS_CMD_KEY]: undefined, qs: { name: file.name, predefinedAcl: options.predefinedAcl, diff --git a/test/nodejs-common/service.ts b/test/nodejs-common/service.ts index a131d6489..c92b24df7 100644 --- a/test/nodejs-common/service.ts +++ b/test/nodejs-common/service.ts @@ -29,6 +29,7 @@ import { import { BodyResponseCallback, DecorateRequestOptions, + GCCL_GCS_CMD_KEY, MakeAuthenticatedRequest, MakeAuthenticatedRequestFactoryConfig, util, @@ -526,6 +527,23 @@ describe('Service', () => { service.request_(reqOpts, assert.ifError); }); + it('should add the `gccl-gcs-cmd` to the api-client header when provided', done => { + const expected = 'example.expected/value'; + service.makeAuthenticatedRequest = (reqOpts: DecorateRequestOptions) => { + const pkg = service.packageJson; + const r = new RegExp( + `^gl-node/${process.versions.node} gccl/${pkg.version} gccl-invocation-id/(?[^W]+) gccl-gcs-cmd/${expected}$` + ); + assert.ok(r.test(reqOpts.headers!['x-goog-api-client'])); + done(); + }; + + service.request_( + {...reqOpts, [GCCL_GCS_CMD_KEY]: expected}, + assert.ifError + ); + }); + describe('projectIdRequired', () => { describe('false', () => { it('should include the projectId', done => { diff --git a/test/nodejs-common/util.ts b/test/nodejs-common/util.ts index f8269bcf4..1f59b1c40 100644 --- a/test/nodejs-common/util.ts +++ b/test/nodejs-common/util.ts @@ -40,6 +40,7 @@ import { DecorateRequestOptions, Duplexify, DuplexifyConstructor, + GCCL_GCS_CMD_KEY, GoogleErrorBody, GoogleInnerError, MakeAuthenticatedRequestFactoryConfig, @@ -546,6 +547,7 @@ describe('common/util', () => { qs: { uploadType: 'media', }, + [GCCL_GCS_CMD_KEY]: 'some.value', } as DecorateRequestOptions; util.makeWritableStream(dup, { @@ -556,6 +558,7 @@ describe('common/util', () => { assert.strictEqual(request.method, req.method); assert.deepStrictEqual(request.qs, req.qs); assert.strictEqual(request.uri, req.uri); + assert.strictEqual(request[GCCL_GCS_CMD_KEY], req[GCCL_GCS_CMD_KEY]); // eslint-disable-next-line @typescript-eslint/no-explicit-any const mp = request.multipart as any[]; diff --git a/test/resumable-upload.ts b/test/resumable-upload.ts index 943140fe6..326a63fc3 100644 --- a/test/resumable-upload.ts +++ b/test/resumable-upload.ts @@ -35,6 +35,7 @@ import { PROTOCOL_REGEX, } from '../src/resumable-upload'; import {GaxiosOptions, GaxiosError, GaxiosResponse} from 'gaxios'; +import {GCCL_GCS_CMD_KEY} from '../src/nodejs-common/util'; nock.disableNetConnect(); @@ -51,7 +52,7 @@ const RESUMABLE_INCOMPLETE_STATUS_CODE = 308; const CHUNK_SIZE_MULTIPLE = 2 ** 18; const queryPath = '/?userProject=user-project-id'; const X_GOOG_API_HEADER_REGEX = - /^gl-node\/(?[^W]+) gccl\/(?[^W]+) gccl-invocation-id\/(?[^W]+)$/; + /^gl-node\/(?[^W]+) gccl\/(?[^W]+) gccl-invocation-id\/(?[^W]+) gccl-gcs-cmd\/(?[^W]+)$/; function mockAuthorizeRequest( code = 200, @@ -113,6 +114,7 @@ describe('resumable-upload', () => { authConfig: {keyFile}, apiEndpoint: API_ENDPOINT, retryOptions: {...RETRY_OPTIONS}, + [GCCL_GCS_CMD_KEY]: 'sample.command', }); }); @@ -1494,7 +1496,7 @@ describe('resumable-upload', () => { describe('#makeRequest', () => { it('should set encryption headers', async () => { const key = crypto.randomBytes(32); - const up = upload({ + up = upload({ bucket: 'BUCKET', file: FILE, key, diff --git a/test/transfer-manager.ts b/test/transfer-manager.ts index 9eb2d7877..2c76d1e84 100644 --- a/test/transfer-manager.ts +++ b/test/transfer-manager.ts @@ -31,26 +31,31 @@ import { } from '../src'; import * as assert from 'assert'; import * as path from 'path'; -import * as stream from 'stream'; import * as fs from 'fs'; import * as fsp from 'fs/promises'; import * as sinon from 'sinon'; -import {GaxiosResponse} from 'gaxios'; +import {GaxiosOptions, GaxiosResponse} from 'gaxios'; +import {GCCL_GCS_CMD_KEY} from '../src/nodejs-common/util'; +import {AuthClient, GoogleAuth} from 'google-auth-library'; +import {tmpdir} from 'os'; describe('Transfer Manager', () => { const BUCKET_NAME = 'test-bucket'; - const STORAGE = sinon.createStubInstance(Storage); - STORAGE.retryOptions = { - autoRetry: true, - maxRetries: 3, - retryDelayMultiplier: 2, - totalTimeout: 600, - maxRetryDelay: 60, - retryableErrorFn: (err: ApiError) => { - return err.code === 500; - }, - idempotencyStrategy: IdempotencyStrategy.RetryConditional, - }; + const STORAGE = sinon.stub( + new Storage({ + retryOptions: { + autoRetry: true, + maxRetries: 3, + retryDelayMultiplier: 2, + totalTimeout: 600, + maxRetryDelay: 60, + retryableErrorFn: (err: ApiError) => { + return err.code === 500; + }, + idempotencyStrategy: IdempotencyStrategy.RetryConditional, + }, + }) + ); let sandbox: sinon.SinonSandbox; let transferManager: TransferManager; let bucket: Bucket; @@ -130,6 +135,19 @@ describe('Transfer Manager', () => { const result = await transferManager.uploadManyFiles(paths); assert.strictEqual(result[0][0].name, paths[0]); }); + + it('should set the appropriate `GCCL_GCS_CMD_KEY`', async () => { + const paths = ['/a/b/foo/bar.txt']; + + sandbox.stub(bucket, 'upload').callsFake(async (_path, options) => { + assert.strictEqual( + (options as UploadOptions)[GCCL_GCS_CMD_KEY], + 'tm.upload_many' + ); + }); + + await transferManager.uploadManyFiles(paths, {prefix: 'hello/world'}); + }); }); describe('downloadManyFiles', () => { @@ -178,6 +196,19 @@ describe('Transfer Manager', () => { }); await transferManager.downloadManyFiles([file], {stripPrefix}); }); + + it('should set the appropriate `GCCL_GCS_CMD_KEY`', async () => { + const file = new File(bucket, 'first.txt'); + + sandbox.stub(file, 'download').callsFake(async options => { + assert.strictEqual( + (options as DownloadOptions)[GCCL_GCS_CMD_KEY], + 'tm.download_many' + ); + }); + + await transferManager.downloadManyFiles([file]); + }); }); describe('downloadFileInChunks', () => { @@ -223,14 +254,26 @@ describe('Transfer Manager', () => { await transferManager.downloadFileInChunks(file, {validation: 'crc32c'}); assert.strictEqual(callCount, 1); }); + + it('should set the appropriate `GCCL_GCS_CMD_KEY`', async () => { + sandbox.stub(file, 'download').callsFake(async options => { + assert.strictEqual( + (options as DownloadOptions)[GCCL_GCS_CMD_KEY], + 'tm.download_sharded' + ); + return [Buffer.alloc(100)]; + }); + + await transferManager.downloadFileInChunks(file); + }); }); describe('uploadFileInChunks', () => { let mockGeneratorFunction: MultiPartHelperGenerator; let fakeHelper: sinon.SinonStubbedInstance; - let readStreamStub: sinon.SinonStub; - const path = '/a/b/c.txt'; - const pThrough = new stream.PassThrough(); + let readStreamSpy: sinon.SinonSpy; + let directory: string; + let filePath: string; class FakeXMLHelper implements MultiPartUploadHelper { bucket: Bucket; fileName: string; @@ -255,10 +298,18 @@ describe('Transfer Manager', () => { } } - beforeEach(() => { - readStreamStub = sandbox - .stub(fs, 'createReadStream') - .returns(pThrough as unknown as fs.ReadStream); + before(async () => { + directory = await fsp.mkdtemp( + path.join(tmpdir(), 'tm-uploadFileInChunks-') + ); + + filePath = path.join(directory, 't.txt'); + + await fsp.writeFile(filePath, 'hello'); + }); + + beforeEach(async () => { + readStreamSpy = sandbox.spy(fs, 'createReadStream'); mockGeneratorFunction = (bucket, fileName, uploadId, partsMap) => { fakeHelper = sandbox.createStubInstance(FakeXMLHelper); fakeHelper.uploadId = uploadId || ''; @@ -271,12 +322,16 @@ describe('Transfer Manager', () => { }; }); + after(async () => { + await fsp.rm(directory, {force: true, recursive: true}); + }); + it('should call initiateUpload, uploadPart, and completeUpload', async () => { - process.nextTick(() => { - pThrough.push('hello world'); - pThrough.end(); - }); - await transferManager.uploadFileInChunks(path, {}, mockGeneratorFunction); + await transferManager.uploadFileInChunks( + filePath, + {}, + mockGeneratorFunction + ); assert.strictEqual(fakeHelper.initiateUpload.calledOnce, true); assert.strictEqual(fakeHelper.uploadPart.calledOnce, true); assert.strictEqual(fakeHelper.completeUpload.calledOnce, true); @@ -286,14 +341,14 @@ describe('Transfer Manager', () => { const options = {highWaterMark: 32 * 1024 * 1024, start: 0}; await transferManager.uploadFileInChunks( - path, + filePath, { chunkSizeBytes: 32 * 1024 * 1024, }, mockGeneratorFunction ); - assert.strictEqual(readStreamStub.calledOnceWith(path, options), true); + assert.strictEqual(readStreamSpy.calledOnceWith(filePath, options), true); }); it('should set the correct start offset when called with an existing parts map', async () => { @@ -303,7 +358,7 @@ describe('Transfer Manager', () => { }; await transferManager.uploadFileInChunks( - path, + filePath, { uploadId: '123', partsMap: new Map([ @@ -315,12 +370,12 @@ describe('Transfer Manager', () => { mockGeneratorFunction ); - assert.strictEqual(readStreamStub.calledOnceWith(path, options), true); + assert.strictEqual(readStreamSpy.calledOnceWith(filePath, options), true); }); it('should not call initiateUpload if an uploadId is provided', async () => { await transferManager.uploadFileInChunks( - path, + filePath, { uploadId: '123', partsMap: new Map([ @@ -353,7 +408,7 @@ describe('Transfer Manager', () => { }; assert.rejects( transferManager.uploadFileInChunks( - path, + filePath, {autoAbortFailure: false}, mockGeneratorFunction ), @@ -382,7 +437,7 @@ describe('Transfer Manager', () => { }; await transferManager.uploadFileInChunks( - path, + filePath, {headers: headersToAdd}, mockGeneratorFunction ); @@ -413,14 +468,50 @@ describe('Transfer Manager', () => { return fakeHelper; }; - process.nextTick(() => { - pThrough.push('hello world'); - pThrough.end(); - }); - assert.doesNotThrow(() => - transferManager.uploadFileInChunks(path, {}, mockGeneratorFunction) + transferManager.uploadFileInChunks(filePath, {}, mockGeneratorFunction) ); }); + + it('should set the appropriate `GCCL_GCS_CMD_KEY`', async () => { + let called = true; + class TestAuthClient extends AuthClient { + async getAccessToken() { + return {token: '', res: undefined}; + } + + async getRequestHeaders() { + return {}; + } + + async request(opts: GaxiosOptions) { + called = true; + + assert(opts.headers); + assert('x-goog-api-client' in opts.headers); + assert.match( + opts.headers['x-goog-api-client'], + /gccl-gcs-cmd\/tm.upload_sharded/ + ); + + return { + data: Buffer.from( + ` + 1 + ` + ), + headers: {}, + } as GaxiosResponse; + } + } + + transferManager.bucket.storage.authClient = new GoogleAuth({ + authClient: new TestAuthClient(), + }); + + await transferManager.uploadFileInChunks(filePath); + + assert(called); + }); }); });