From caefbf5085e96fcd90bd02774cebacad815d7d9f Mon Sep 17 00:00:00 2001 From: kartikgupta-db Date: Thu, 3 Aug 2023 11:44:34 +0200 Subject: [PATCH 1/2] Export SPARK_REMOTE env var for profile authentication type --- .../file-managers/DatabricksEnvFileManager.ts | 8 ++--- .../notebooks/NotebookInitScriptManager.ts | 4 +-- .../src/utils/envVarGenerators.ts | 29 +++++++++++++++++-- 3 files changed, 32 insertions(+), 9 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 b32d20cd8..bf477070b 100644 --- a/packages/databricks-vscode/src/utils/envVarGenerators.ts +++ b/packages/databricks-vscode/src/utils/envVarGenerators.ts @@ -90,16 +90,39 @@ export function getCommonDatabricksEnvVars( /* eslint-enable @typescript-eslint/naming-convention */ } -export function getDbConnectEnvVars( +async function getPatToken(connectionManager: ConnectionManager) { + const headers: Record = {}; + 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; + 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 */ } @@ -114,7 +137,7 @@ export function getProxyEnvVars() { } export function removeUndefinedKeys< - T extends Record + T extends Record, >(envVarMap?: T): T | undefined { if (envVarMap === undefined) { return; From c945df682112b271311ec660dc900e21776e09e2 Mon Sep 17 00:00:00 2001 From: kartikgupta-db Date: Thu, 3 Aug 2023 11:46:08 +0200 Subject: [PATCH 2/2] Comment --- packages/databricks-vscode/src/utils/envVarGenerators.ts | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/packages/databricks-vscode/src/utils/envVarGenerators.ts b/packages/databricks-vscode/src/utils/envVarGenerators.ts index bf477070b..e07027ec5 100644 --- a/packages/databricks-vscode/src/utils/envVarGenerators.ts +++ b/packages/databricks-vscode/src/utils/envVarGenerators.ts @@ -102,6 +102,11 @@ 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) {