Append to the end of existing SPARK_CONNECT_USER_AGENT instead of o…

…verwritting it (#870) ## Changes fixes #854 ## Tests * unit
databricks · Sep 27, 2023 · ced4413 · ced4413
1 parent 16615b1
commit ced4413
Show file tree

Hide file tree

Showing 3 changed files with 34 additions and 2 deletions.
diff --git a/packages/databricks-vscode/src/cluster/ClusterLoader.ts b/packages/databricks-vscode/src/cluster/ClusterLoader.ts
@@ -194,7 +194,13 @@ export class ClusterLoader implements Disposable {
                     this.onStopRequested(() => {
                         clearInterval(timer);
                         resolve();
-                    })
+                    }),
+                    {
+                        dispose: () => {
+                            clearInterval(timer);
+                            resolve();
+                        },
+                    }
                 );
             });
         }

diff --git a/packages/databricks-vscode/src/utils/envVarGenerators.test.ts b/packages/databricks-vscode/src/utils/envVarGenerators.test.ts
@@ -88,13 +88,20 @@ describe(__filename, () => {
     describe("getDbConnectEnvVars", () => {
         const mockWorkspacePath = Uri.file("example");
         let mockAuthProvider: AuthProvider;
+        let existingEnv: any;
+
         beforeEach(() => {
             when(mockApiClient.product).thenReturn("test");
             when(mockApiClient.productVersion).thenReturn("0.0.1");
             mockAuthProvider = mock(AuthProvider);
             when(mockDatabricksWorkspace.authProvider).thenReturn(
                 instance(mockAuthProvider)
             );
+            existingEnv = Object.assign({}, process.env);
+        });
+
+        afterEach(() => {
+            process.env = existingEnv;
         });
 
         it("should generate correct dbconnect env vars when auth type is not profile", async () => {
@@ -111,6 +118,21 @@ describe(__filename, () => {
             });
         });
 
+        it("should append our user agent any existing SPARK_CONNECT_USER_AGENT in VS Code parent env", async () => {
+            process.env.SPARK_CONNECT_USER_AGENT = "existing";
+            when(mockAuthProvider.authType).thenReturn("azure-cli");
+
+            const actual = await getDbConnectEnvVars(
+                instance(mockConnectionManager),
+                mockWorkspacePath
+            );
+
+            assert.deepEqual(actual, {
+                SPARK_CONNECT_USER_AGENT: "existing test/0.0.1",
+                DATABRICKS_PROJECT_ROOT: mockWorkspacePath.fsPath,
+            });
+        });
+
         it("should generate correct dbconnect env vars when auth type is profile", async () => {
             when(mockAuthProvider.authType).thenReturn("profile");
             const mockConfig = mock(Config);

diff --git a/packages/databricks-vscode/src/utils/envVarGenerators.ts b/packages/databricks-vscode/src/utils/envVarGenerators.ts
@@ -42,6 +42,7 @@ function getUserAgent(connectionManager: ConnectionManager) {
     if (!client) {
         return;
     }
+
     return `${client.product}/${client.productVersion}`;
 }
 
@@ -104,9 +105,12 @@ export async function getDbConnectEnvVars(
     workspacePath: Uri
 ) {
     const userAgent = getUserAgent(connectionManager);
+    const existingSparkUa = process.env.SPARK_CONNECT_USER_AGENT ?? "";
     /* eslint-disable @typescript-eslint/naming-convention */
     return {
-        SPARK_CONNECT_USER_AGENT: userAgent,
+        //We append our user agent to any existing SPARK_CONNECT_USER_AGENT defined in the
+        //environment of the parent process of VS Code.
+        SPARK_CONNECT_USER_AGENT: [existingSparkUa, userAgent].join(" ").trim(),
         DATABRICKS_PROJECT_ROOT: workspacePath.fsPath,
         ...((await getSparkRemoteEnvVar(connectionManager)) || {}),
     };