Export DATABRICKS_HOST env var with correct protocol

packages/databricks-vscode/src/cli/SyncTasks.test.ts

@@ -103,7 +103,8 @@ describe(__filename, () => {
                     /* eslint-disable @typescript-eslint/naming-convention */
                     DATABRICKS_CLI_UPSTREAM: "databricks-vscode",
                     DATABRICKS_CLI_UPSTREAM_VERSION: "1.0.0",
-                    DATABRICKS_HOST: "000000000000.00.azuredatabricks.net",
+                    DATABRICKS_HOST:
+                        "https://000000000000.00.azuredatabricks.net/",
                     DATABRICKS_AUTH_TYPE: "metadata-service",
                     DATABRICKS_METADATA_SERVICE_URL: "http://localhost:1234",
                     HOME: process.env.HOME,
@@ -123,7 +124,8 @@ describe(__filename, () => {
                     /* eslint-disable @typescript-eslint/naming-convention */
                     DATABRICKS_CLI_UPSTREAM: "databricks-vscode",
                     DATABRICKS_CLI_UPSTREAM_VERSION: "1.0.0",
-                    DATABRICKS_HOST: "000000000000.00.azuredatabricks.net",
+                    DATABRICKS_HOST:
+                        "https://000000000000.00.azuredatabricks.net/",
                     DATABRICKS_AUTH_TYPE: "metadata-service",
                     DATABRICKS_METADATA_SERVICE_URL: "http://localhost:1234",
                     HOME: process.env.HOME,

packages/databricks-vscode/src/utils/envVarGenerators.test.ts

@@ -0,0 +1,134 @@
+/* eslint-disable @typescript-eslint/naming-convention */
+import {anything, instance, mock, when} from "ts-mockito";
+import {ConnectionManager} from "../configuration/ConnectionManager";
+import {DatabricksWorkspace} from "../configuration/DatabricksWorkspace";
+import {ApiClient, Cluster, Config} from "@databricks/databricks-sdk";
+import {
+    getAuthEnvVars,
+    getDbConnectEnvVars,
+    getProxyEnvVars,
+} from "./envVarGenerators";
+import {Uri} from "vscode";
+import assert from "assert";
+import {AuthProvider} from "../configuration/auth/AuthProvider";
+
+describe(__filename, () => {
+    let mockConnectionManager: ConnectionManager;
+    let mockDatabricksWorkspace: DatabricksWorkspace;
+    let mockCluster: Cluster;
+    const mockClusterId = "clusterId";
+    const mockHost = "http://example.com";
+    let mockApiClient: ApiClient;
+
+    beforeEach(() => {
+        mockConnectionManager = mock(ConnectionManager);
+        mockDatabricksWorkspace = mock(DatabricksWorkspace);
+        mockCluster = mock(Cluster);
+        mockApiClient = mock(ApiClient);
+
+        when(mockConnectionManager.databricksWorkspace).thenReturn(
+            instance(mockDatabricksWorkspace)
+        );
+        when(mockDatabricksWorkspace.host).thenReturn(Uri.parse(mockHost));
+
+        when(mockCluster.id).thenReturn(mockClusterId);
+        when(mockConnectionManager.cluster).thenReturn(instance(mockCluster));
+
+        when(mockConnectionManager.apiClient).thenReturn(
+            instance(mockApiClient)
+        );
+        when(mockApiClient.host).thenResolve(new URL(mockHost));
+    });
+
+    it("should generate correct authEnvVars", () => {
+        when(mockConnectionManager.metadataServiceUrl).thenReturn(
+            "http://example.com/metadata-service"
+        );
+
+        const actual = getAuthEnvVars(instance(mockConnectionManager));
+        assert.deepEqual(actual, {
+            DATABRICKS_HOST: mockHost + "/",
+            DATABRICKS_AUTH_TYPE: "metadata-service",
+            DATABRICKS_METADATA_SERVICE_URL:
+                "http://example.com/metadata-service",
+            DATABRICKS_CLUSTER_ID: mockClusterId,
+        });
+    });
+
+    describe("getProxyEnvVars", () => {
+        it("should generate correct proxyEnvVars with lowerCase settings", () => {
+            process.env.http_proxy = "http://example.com";
+            process.env.https_proxy = "https://example.com";
+            const actual = getProxyEnvVars();
+            assert.deepEqual(actual, {
+                HTTP_PROXY: "http://example.com",
+                HTTPS_PROXY: "https://example.com",
+            });
+        });
+
+        it("should generate correct proxyEnvVars with upperCase settings", () => {
+            process.env.HTTP_PROXY = "http://example.com";
+            process.env.HTTPS_PROXY = "https://example.com";
+            const actual = getProxyEnvVars();
+            assert.deepEqual(actual, {
+                HTTP_PROXY: "http://example.com",
+                HTTPS_PROXY: "https://example.com",
+            });
+        });
+
+        after(() => {
+            delete process.env.http_proxy;
+            delete process.env.https_proxy;
+            delete process.env.HTTP_PROXY;
+            delete process.env.HTTPS_PROXY;
+        });
+    });
+
+    describe("getDbConnectEnvVars", () => {
+        const mockWorkspacePath = Uri.file("example");
+        let mockAuthProvider: AuthProvider;
+        beforeEach(() => {
+            when(mockApiClient.product).thenReturn("test");
+            when(mockApiClient.productVersion).thenReturn("0.0.1");
+            mockAuthProvider = mock(AuthProvider);
+            when(mockDatabricksWorkspace.authProvider).thenReturn(
+                instance(mockAuthProvider)
+            );
+        });
+
+        it("should generate correct dbconnect env vars when auth type is not profile", async () => {
+            when(mockAuthProvider.authType).thenReturn("azure-cli");
+
+            const actual = await getDbConnectEnvVars(
+                instance(mockConnectionManager),
+                mockWorkspacePath
+            );
+
+            assert.deepEqual(actual, {
+                SPARK_CONNECT_USER_AGENT: "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);
+            when(mockApiClient.config).thenReturn(instance(mockConfig));
+            when(mockConfig.authenticate(anything())).thenCall((headers) => {
+                headers["Authorization"] = "Bearer token";
+            });
+            const actual = await getDbConnectEnvVars(
+                instance(mockConnectionManager),
+                mockWorkspacePath
+            );
+
+            assert.deepEqual(actual, {
+                SPARK_CONNECT_USER_AGENT: "test/0.0.1",
+                DATABRICKS_PROJECT_ROOT: mockWorkspacePath.fsPath,
+                SPARK_REMOTE: `sc://${
+                    Uri.parse(mockHost).authority
+                }:443/;token=token;use_ssl=true;x-databricks-cluster-id=${mockClusterId}`,
+            });
+        });
+    });
+});

packages/databricks-vscode/src/utils/envVarGenerators.ts

@@ -55,7 +55,7 @@ export async function getNotebookEnvVars(
 }
 
 function getUserAgent(connectionManager: ConnectionManager) {
-    const client = connectionManager.workspaceClient?.apiClient;
+    const client = connectionManager.apiClient;
     if (!client) {
         return;
     }
@@ -64,7 +64,7 @@ function getUserAgent(connectionManager: ConnectionManager) {
 
 export function getAuthEnvVars(connectionManager: ConnectionManager) {
     const cluster = connectionManager.cluster;
-    const host = connectionManager.databricksWorkspace?.host.authority;
+    const host = connectionManager.databricksWorkspace?.host.toString();
     if (!host || !connectionManager.metadataServiceUrl) {
         return;
     }
@@ -92,9 +92,7 @@ export function getCommonDatabricksEnvVars(
 
 async function getPatToken(connectionManager: ConnectionManager) {
     const headers: Record<string, string> = {};
-    await connectionManager.workspaceClient?.apiClient.config.authenticate(
-        headers
-    );
+    await connectionManager.apiClient?.config.authenticate(headers);
     return headers["Authorization"]?.split(" ")[1];
 }