From 736a8bfb5e160b1c6ce8f8dd5051be650f63f3e7 Mon Sep 17 00:00:00 2001
From: Kartik Gupta <88345179+kartikgupta-db@users.noreply.github.com>
Date: Thu, 3 Aug 2023 12:42:31 +0200
Subject: [PATCH] Export `SPARK_REMOTE` env var for `profile` authentication
 type (#825)

## Changes
<!-- Summary of your changes that are easy to understand -->

## Tests
<!-- How is this tested? -->
---
 .../file-managers/DatabricksEnvFileManager.ts |  8 ++---
 .../notebooks/NotebookInitScriptManager.ts    |  4 +--
 .../src/utils/envVarGenerators.ts             | 32 +++++++++++++++++--
 3 files changed, 36 insertions(+), 8 deletions(-)

diff --git a/packages/databricks-vscode/src/file-managers/DatabricksEnvFileManager.ts b/packages/databricks-vscode/src/file-managers/DatabricksEnvFileManager.ts
index da671eeef..c11c8f12d 100644
--- a/packages/databricks-vscode/src/file-managers/DatabricksEnvFileManager.ts
+++ b/packages/databricks-vscode/src/file-managers/DatabricksEnvFileManager.ts
@@ -179,10 +179,10 @@ export class DatabricksEnvFileManager implements Disposable {
 
         const data = Object.entries({
             ...(this.getDatabrickseEnvVars() || {}),
-            ...(EnvVarGenerators.getDbConnectEnvVars(
+            ...((await EnvVarGenerators.getDbConnectEnvVars(
                 this.connectionManager,
                 this.workspacePath
-            ) || {}),
+            )) || {}),
             ...this.getIdeEnvVars(),
             ...((await this.getUserEnvVars()) || {}),
             ...(await this.getNotebookEnvVars()),
@@ -219,10 +219,10 @@ export class DatabricksEnvFileManager implements Disposable {
         Object.entries({
             ...(this.getDatabrickseEnvVars() || {}),
             ...this.getIdeEnvVars(),
-            ...(EnvVarGenerators.getDbConnectEnvVars(
+            ...((await EnvVarGenerators.getDbConnectEnvVars(
                 this.connectionManager,
                 this.workspacePath
-            ) || {}),
+            )) || {}),
             ...(await this.getNotebookEnvVars()),
         }).forEach(([key, value]) => {
             if (value === undefined) {
diff --git a/packages/databricks-vscode/src/language/notebooks/NotebookInitScriptManager.ts b/packages/databricks-vscode/src/language/notebooks/NotebookInitScriptManager.ts
index bf8fa5d7b..9d9d4f05a 100644
--- a/packages/databricks-vscode/src/language/notebooks/NotebookInitScriptManager.ts
+++ b/packages/databricks-vscode/src/language/notebooks/NotebookInitScriptManager.ts
@@ -226,10 +226,10 @@ export class NotebookInitScriptManager implements Disposable {
                 ...(EnvVarGenerators.getCommonDatabricksEnvVars(
                     this.connectionManager
                 ) ?? {}),
-                ...(EnvVarGenerators.getDbConnectEnvVars(
+                ...((await EnvVarGenerators.getDbConnectEnvVars(
                     this.connectionManager,
                     this.workspacePath
-                ) ?? {}),
+                )) ?? {}),
                 ...(EnvVarGenerators.getIdeEnvVars() ?? {}),
                 ...((await this.getUserEnvVars()) ?? {}),
             };
diff --git a/packages/databricks-vscode/src/utils/envVarGenerators.ts b/packages/databricks-vscode/src/utils/envVarGenerators.ts
index 7bfb64efc..e07027ec5 100644
--- a/packages/databricks-vscode/src/utils/envVarGenerators.ts
+++ b/packages/databricks-vscode/src/utils/envVarGenerators.ts
@@ -90,16 +90,44 @@ export function getCommonDatabricksEnvVars(
     /* eslint-enable @typescript-eslint/naming-convention */
 }
 
-export function getDbConnectEnvVars(
+async function getPatToken(connectionManager: ConnectionManager) {
+    const headers: Record<string, string> = {};
+    await connectionManager.workspaceClient?.apiClient.config.authenticate(
+        headers
+    );
+    return headers["Authorization"]?.split(" ")[1];
+}
+
+async function getSparkRemoteEnvVar(connectionManager: ConnectionManager) {
+    const host = connectionManager.databricksWorkspace?.host.authority;
+    const authType =
+        connectionManager.databricksWorkspace?.authProvider.authType;
+
+    // We export spark remote only for profile auth type. This is to support
+    // SparkSession builder in oss spark connect (and also dbconnect).
+    // For all other auth types, we don't export spark remote and expect users
+    // to use DatabricksSession for full functionality.
+    if (host && connectionManager.cluster && authType === "profile") {
+        const pat = await getPatToken(connectionManager);
+        if (pat) {
+            return {
+                // eslint-disable-next-line @typescript-eslint/naming-convention
+                SPARK_REMOTE: `sc://${host}:443/;token=${pat};use_ssl=true;x-databricks-cluster-id=${connectionManager.cluster.id}`,
+            };
+        }
+    }
+}
+
+export async function getDbConnectEnvVars(
     connectionManager: ConnectionManager,
     workspacePath: Uri
 ) {
     const userAgent = getUserAgent(connectionManager);
-
     /* eslint-disable @typescript-eslint/naming-convention */
     return {
         SPARK_CONNECT_USER_AGENT: userAgent,
         DATABRICKS_PROJECT_ROOT: workspacePath.fsPath,
+        ...((await getSparkRemoteEnvVar(connectionManager)) || {}),
     };
     /* eslint-enable @typescript-eslint/naming-convention */
 }