From ab5a3eb6f1df0f4732f40f6e14451a8f0a547444 Mon Sep 17 00:00:00 2001
From: Daniel Bankhead <danielbankhead@google.com>
Date: Mon, 25 Sep 2023 12:21:44 -0700
Subject: [PATCH] feat: Add base `x-goog-api-client` for XML API

---
 src/transfer-manager.ts  | 63 ++++++++++++++++++++++------------------
 test/transfer-manager.ts | 22 +++++++-------
 2 files changed, 45 insertions(+), 40 deletions(-)

diff --git a/src/transfer-manager.ts b/src/transfer-manager.ts
index 18ed38798..5e3bb71c9 100644
--- a/src/transfer-manager.ts
+++ b/src/transfer-manager.ts
@@ -27,6 +27,9 @@ 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.
@@ -200,6 +203,33 @@ 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.
    *
@@ -209,38 +239,12 @@ class XMLMultiPartUploadHelper implements MultiPartUploadHelper {
     const url = `${this.baseUrl}?uploads`;
     return retry(async bail => {
       try {
-        const combinedHeaders = {
-          ...(await this.authClient.getRequestHeaders(url)),
-          ...headers,
-        };
-
-        let headerFound = false;
-
-        for (const [key, value] of Object.entries(combinedHeaders)) {
-          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)) {
-              combinedHeaders[
-                key
-              ] = `${value} gccl-gcs-cmd/${GCCL_GCS_CMD_FEATURE.UPLOAD_SHARDED}`;
-            }
-            break;
-          }
-        }
-
-        if (!headerFound) {
-          combinedHeaders[
-            'x-goog-api-client'
-          ] = `gccl-gcs-cmd/${GCCL_GCS_CMD_FEATURE.UPLOAD_SHARDED}`;
-        }
-
         const res = await this.authClient.request({
-          headers: combinedHeaders,
+          headers: this.#setGoogApiClientHeaders(headers),
           method: 'POST',
           url,
         });
+
         if (res.data && res.data.error) {
           throw res.data.error;
         }
@@ -271,7 +275,7 @@ class XMLMultiPartUploadHelper implements MultiPartUploadHelper {
     validation?: 'md5' | false
   ): Promise<void> {
     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');
@@ -318,6 +322,7 @@ class XMLMultiPartUploadHelper implements MultiPartUploadHelper {
     return retry(async bail => {
       try {
         const res = await this.authClient.request({
+          headers: this.#setGoogApiClientHeaders(),
           url,
           method: 'POST',
           body,
diff --git a/test/transfer-manager.ts b/test/transfer-manager.ts
index a443a4a15..02deccb48 100644
--- a/test/transfer-manager.ts
+++ b/test/transfer-manager.ts
@@ -474,8 +474,8 @@ describe('Transfer Manager', () => {
       );
     });
 
-    it('should set the appropriate `GCCL_GCS_CMD_KEY` - simple upload', async () => {
-      let firstCall = true;
+    it('should set the appropriate `GCCL_GCS_CMD_KEY`', async () => {
+      let called = true;
       class TestAuthClient extends AuthClient {
         async getAccessToken() {
           return {token: '', res: undefined};
@@ -486,16 +486,14 @@ describe('Transfer Manager', () => {
         }
 
         async request(opts: GaxiosOptions) {
-          if (firstCall) {
-            firstCall = false;
+          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/
-            );
-          }
+          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(
@@ -513,6 +511,8 @@ describe('Transfer Manager', () => {
       });
 
       await transferManager.uploadFileInChunks(filePath);
+
+      assert(called);
     });
   });
 });