Merge branch 'main' into fjakobs-patch-1

databricks · Jul 25, 2023 · 3ac158a · 3ac158a
2 parents bff466c + 6a67a75
commit 3ac158a
Show file tree

Hide file tree

Showing 9 changed files with 94 additions and 65 deletions.
diff --git a/packages/databricks-vscode/package.json b/packages/databricks-vscode/package.json
@@ -6,7 +6,7 @@
     "license": "LicenseRef-LICENSE",
     "version": "1.1.0",
     "engines": {
-        "vscode": "^1.76.0"
+        "vscode": "^1.78.0"
     },
     "categories": [
         "Data Science",
@@ -638,7 +638,8 @@
         ]
     },
     "extensionDependencies": [
-        "ms-python.python"
+        "ms-python.python",
+        "ms-toolsai.jupyter"
     ],
     "vsce": {
         "dependencies": false,

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

diff --git a/packages/databricks-vscode/src/cli/SyncTasks.ts b/packages/databricks-vscode/src/cli/SyncTasks.ts
@@ -19,6 +19,7 @@ import {Loggers} from "../logger";
 import {Context, context} from "@databricks/databricks-sdk/dist/context";
 import {PackageMetaData} from "../utils/packageJsonUtils";
 import {RWLock} from "../locking";
+import {EnvVarGenerators} from "../utils";
 
 export const TASK_SYNC_TYPE = {
     syncFull: "sync-full",
@@ -285,21 +286,6 @@ export class LazyCustomSyncTerminal extends CustomSyncTerminal {
             );
         }
 
-        // Pass through proxy settings to child process.
-        const proxySettings: {[key: string]: string | undefined} = {
-            /* eslint-disable @typescript-eslint/naming-convention */
-            HTTP_PROXY: process.env.HTTP_PROXY || process.env.http_proxy,
-            HTTPS_PROXY: process.env.HTTPS_PROXY || process.env.https_proxy,
-            /* eslint-enable @typescript-eslint/naming-convention */
-        };
-
-        // Remove undefined keys.
-        Object.keys(proxySettings).forEach((key) => {
-            if (proxySettings[key] === undefined) {
-                delete proxySettings[key];
-            }
-        });
-
         return {
             cwd: workspacePath,
             env: {
@@ -308,11 +294,9 @@ export class LazyCustomSyncTerminal extends CustomSyncTerminal {
                 DATABRICKS_CLI_UPSTREAM_VERSION: this.packageMetadata.version,
                 HOME: process.env.HOME,
                 PATH: process.env.PATH,
-                DATABRICKS_HOST: dbWorkspace.host.toString(),
-                DATABRICKS_AUTH_TYPE: "metadata-service",
-                DATABRICKS_METADATA_SERVICE_URL:
-                    this.connection.metadataServiceUrl || "",
-                ...proxySettings,
+                ...EnvVarGenerators.removeUndefinedKeys(
+                    EnvVarGenerators.getCommonDatabricksEnvVars(this.connection)
+                ),
                 /* eslint-enable @typescript-eslint/naming-convention */
             },
         } as SpawnOptions;

diff --git a/packages/databricks-vscode/src/configuration/ConnectionCommands.ts b/packages/databricks-vscode/src/configuration/ConnectionCommands.ts
@@ -76,13 +76,14 @@ export class ConnectionCommands implements Disposable {
 
     openDatabricksConfigFileCommand() {
         return async () => {
+            const homeDir = process.env.HOME ?? homedir();
             let filePath =
                 workspaceConfigs.databrickscfgLocation ??
                 process.env.DATABRICKS_CONFIG_FILE ??
-                path.join(homedir(), ".databrickscfg");
+                path.join(homeDir, ".databrickscfg");
 
             if (filePath.startsWith("~/")) {
-                filePath = path.join(homedir(), filePath.slice(2));
+                filePath = path.join(homeDir, filePath.slice(2));
             }
             const uri = Uri.file(path.normalize(filePath));
 

diff --git a/packages/databricks-vscode/src/file-managers/DatabricksEnvFileManager.ts b/packages/databricks-vscode/src/file-managers/DatabricksEnvFileManager.ts
@@ -151,37 +151,41 @@ export class DatabricksEnvFileManager implements Disposable {
         );
     }
 
-    private async getDatabrickseEnvVars() {
-        return EnvVarGenerators.getDatabrickseEnvVars(
-            this.connectionManager,
-            this.workspacePath
+    private getDatabrickseEnvVars() {
+        return EnvVarGenerators.getCommonDatabricksEnvVars(
+            this.connectionManager
         );
     }
 
-    private async getNotebookEnvVars() {
+    private getNotebookEnvVars() {
         return EnvVarGenerators.getNotebookEnvVars(
             this.featureManager,
             this.notebookInitScriptManager
         );
     }
 
-    private async getIdeEnvVars() {
+    private getIdeEnvVars() {
         return EnvVarGenerators.getIdeEnvVars();
     }
 
     //Get env variables from user's .env file
     private async getUserEnvVars() {
-        return EnvVarGenerators.getUserEnvVars(this.userEnvPath);
+        return await EnvVarGenerators.getUserEnvVars(this.userEnvPath);
     }
 
     @logging.withLogContext(Loggers.Extension)
     async writeFile(@context ctx?: Context) {
+        await this.connectionManager.waitForConnect();
+
         const data = Object.entries({
-            ...((await this.getDatabrickseEnvVars()) || {}),
-            ...(await this.getIdeEnvVars()),
+            ...(this.getDatabrickseEnvVars() || {}),
+            ...(EnvVarGenerators.getDbConnectEnvVars(
+                this.connectionManager,
+                this.workspacePath
+            ) || {}),
+            ...this.getIdeEnvVars(),
             ...((await this.getUserEnvVars()) || {}),
             ...(await this.getNotebookEnvVars()),
-            ...(await EnvVarGenerators.getProxyEnvVars()),
         })
             .filter(([, value]) => value !== undefined)
             .map(([key, value]) => `${key}=${value}`);
@@ -210,11 +214,16 @@ export class DatabricksEnvFileManager implements Disposable {
     }
 
     async emitToTerminal() {
+        await this.connectionManager.waitForConnect();
+
         Object.entries({
-            ...((await this.getDatabrickseEnvVars()) || {}),
-            ...(await this.getIdeEnvVars()),
+            ...(this.getDatabrickseEnvVars() || {}),
+            ...this.getIdeEnvVars(),
+            ...(EnvVarGenerators.getDbConnectEnvVars(
+                this.connectionManager,
+                this.workspacePath
+            ) || {}),
             ...(await this.getNotebookEnvVars()),
-            ...(await EnvVarGenerators.getProxyEnvVars()),
         }).forEach(([key, value]) => {
             if (value === undefined) {
                 return;

diff --git a/packages/databricks-vscode/src/language/notebooks/NotebookInitScriptManager.ts b/packages/databricks-vscode/src/language/notebooks/NotebookInitScriptManager.ts
@@ -218,15 +218,19 @@ export class NotebookInitScriptManager implements Disposable {
         executable: string,
         @context ctx?: Context
     ) {
+        await this.connectionManager.waitForConnect();
         let someScriptFailed = false;
         for (const fileBaseName of await this.sourceFiles) {
             const file = path.join(this.startupDir, fileBaseName);
             const env = {
-                ...((await EnvVarGenerators.getDatabrickseEnvVars(
+                ...(EnvVarGenerators.getCommonDatabricksEnvVars(
+                    this.connectionManager
+                ) ?? {}),
+                ...(EnvVarGenerators.getDbConnectEnvVars(
                     this.connectionManager,
                     this.workspacePath
-                )) ?? {}),
-                ...((await EnvVarGenerators.getIdeEnvVars()) ?? {}),
+                ) ?? {}),
+                ...(EnvVarGenerators.getIdeEnvVars() ?? {}),
                 ...((await this.getUserEnvVars()) ?? {}),
             };
             const {stderr} = await execFile(

diff --git a/packages/databricks-vscode/src/test/e2e/wdio.conf.ts b/packages/databricks-vscode/src/test/e2e/wdio.conf.ts
@@ -339,13 +339,17 @@ export const config: Options.Testrunner = {
         for (let i = 0; i < 2; i++) {
             try {
                 await new Promise((resolve, reject) => {
+                    const extensionDependencies =
+                        packageJson.extensionDependencies.flatMap((item) => [
+                            "--install-extension",
+                            item,
+                        ]);
                     execFile(
                         cli,
                         [
                             "--extensions-dir",
                             EXTENSION_DIR,
-                            "--install-extension",
-                            "ms-python.python",
+                            ...extensionDependencies,
                             "--install-extension",
                             VSIX_PATH,
                             "--force",

diff --git a/packages/databricks-vscode/src/utils/envVarGenerators.ts b/packages/databricks-vscode/src/utils/envVarGenerators.ts
@@ -30,7 +30,7 @@ export async function getUserEnvVars(userEnvPath: Uri) {
     }
 }
 
-export async function getIdeEnvVars() {
+export function getIdeEnvVars() {
     /* eslint-disable @typescript-eslint/naming-convention */
     return {
         //https://github.com/fabioz/PyDev.Debugger/blob/main/_pydevd_bundle/pydevd_constants.py
@@ -54,29 +54,18 @@ export async function getNotebookEnvVars(
     /* eslint-enable @typescript-eslint/naming-convention */
 }
 
-async function getUserAgent(connectionManager: ConnectionManager) {
+function getUserAgent(connectionManager: ConnectionManager) {
     const client = connectionManager.workspaceClient?.apiClient;
     if (!client) {
         return;
     }
     return `${client.product}/${client.productVersion}`;
 }
 
-export async function getDatabrickseEnvVars(
-    connectionManager: ConnectionManager,
-    workspacePath: Uri
-) {
-    await connectionManager.waitForConnect();
+export function getAuthEnvVars(connectionManager: ConnectionManager) {
     const cluster = connectionManager.cluster;
-    const userAgent = await getUserAgent(connectionManager);
-    const authProvider = connectionManager.databricksWorkspace?.authProvider;
     const host = connectionManager.databricksWorkspace?.host.authority;
-    if (
-        !userAgent ||
-        !authProvider ||
-        !host ||
-        !connectionManager.metadataServiceUrl
-    ) {
+    if (!host || !connectionManager.metadataServiceUrl) {
         return;
     }
 
@@ -85,18 +74,55 @@ export async function getDatabrickseEnvVars(
         DATABRICKS_HOST: host,
         DATABRICKS_AUTH_TYPE: "metadata-service",
         DATABRICKS_METADATA_SERVICE_URL: connectionManager.metadataServiceUrl,
-        SPARK_CONNECT_USER_AGENT: userAgent,
         DATABRICKS_CLUSTER_ID: cluster?.id,
+    };
+    /* eslint-enable @typescript-eslint/naming-convention */
+}
+
+export function getCommonDatabricksEnvVars(
+    connectionManager: ConnectionManager
+) {
+    /* eslint-disable @typescript-eslint/naming-convention */
+    return {
+        ...(getAuthEnvVars(connectionManager) || {}),
+        ...(getProxyEnvVars() || {}),
+    };
+    /* eslint-enable @typescript-eslint/naming-convention */
+}
+
+export 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,
     };
     /* eslint-enable @typescript-eslint/naming-convention */
 }
 
-export async function getProxyEnvVars() {
+export function getProxyEnvVars() {
     return {
         /* eslint-disable @typescript-eslint/naming-convention */
         HTTP_PROXY: process.env.HTTP_PROXY || process.env.http_proxy,
         HTTPS_PROXY: process.env.HTTPS_PROXY || process.env.https_proxy,
         /* eslint-enable @typescript-eslint/naming-convention */
     };
 }
+
+export function removeUndefinedKeys<
+    T extends Record<string, string | undefined>,
+>(envVarMap?: T): T | undefined {
+    if (envVarMap === undefined) {
+        return;
+    }
+
+    const filteredEntries = Object.entries(envVarMap).filter(
+        (entry) => entry[1] !== undefined
+    ) as [string, string][];
+
+    return Object.fromEntries<string>(filteredEntries) as T;
+}
diff --git a/packages/databricks-vscode/src/utils/urlUtils.ts b/packages/databricks-vscode/src/utils/urlUtils.ts
@@ -23,7 +23,7 @@ export function normalizeHost(host: string): URL {
     }
     if (
         !url.hostname.match(
-            /(\.databricks\.azure\.us|\.databricks\.azure\.cn|\.azuredatabricks\.net|\.gcp\.databricks\.com|\.cloud\.databricks\.com)$/
+            /(\.databricks\.azure\.us|\.databricks\.azure\.cn|\.azuredatabricks\.net|\.gcp\.databricks\.com|\.cloud\.databricks\.com|\.dev\.databricks\.com)$/
         )
     ) {
         throw new Error("Not a Databricks host");
@@ -43,5 +43,7 @@ export function isGcpHost(url: URL): boolean {
 }
 
 export function isAwsHost(url: URL): boolean {
-    return !!url.hostname.match(/\.cloud\.databricks\.com$/);
+    return !!url.hostname.match(
+        /(\.cloud\.databricks\.com|\.dev\.databricks\.com)$/
+    );
 }