Traverse parent directories to find and load .databricks/.databricks.…

…env from notebook init script (#845)

## Changes
* also show restart-python-kernels dialogue box more prominently,

## Tests
* manual

packages/databricks-vscode/package.json

@@ -637,6 +637,10 @@
                         "type": "boolean",
                         "default": true,
                         "description": "Enable/disable rearranging cells in wrapper files created when using `workspace` as the sync destination. **Note:** It is recommended to NOT disable this setting. If you do disable it, you will need to manually handle sys.path for local imports in your notebooks."
+                    },
+                    "databricks.ipythonDir": {
+                        "type": "string",
+                        "description": "Absolute path to a directory for storing IPython files. Defaults to IPYTHONDIR environment variable (if set) or ~/.ipython."
                     }
                 }
             }

packages/databricks-vscode/resources/python/00-databricks-init.py

@@ -47,6 +47,22 @@ def wrapper(*args, **kwargs):
 
     return wrapper
 
+@logErrorAndContinue
+@disposable
+def load_env_from_leaf(path: str) -> bool:
+    curdir = path if os.path.isdir(path) else os.path.dirname(path)
+    env_file_path = os.path.join(curdir, ".databricks", ".databricks.env")
+    if os.path.exists(os.path.dirname(env_file_path)):
+        with open(env_file_path, "r") as f:
+            for line in f.readlines():
+                key, value = line.strip().split("=", 1)
+                os.environ[key] = value
+        return True
+
+    parent = os.path.dirname(curdir)
+    if parent == curdir:
+        return False
+    return load_env_from_leaf(parent)
 
 @logErrorAndContinue
 @disposable
@@ -347,6 +363,8 @@ def make_matplotlib_inline():
     import sys
 
     print(sys.modules[__name__])
+    if not load_env_from_leaf(os.getcwd()):
+        sys.exit(1)
     cfg = LocalDatabricksNotebookConfig()
     create_and_register_databricks_globals()
     register_magics(cfg)

packages/databricks-vscode/src/configuration/ConnectionCommands.ts

@@ -1,5 +1,4 @@
 import {Cluster, WorkspaceFsEntity, WorkspaceFsUtils} from "../sdk-extensions";
-import {homedir} from "node:os";
 import {
     Disposable,
     FileSystemError,
@@ -13,7 +12,7 @@ import {
 import {ClusterListDataProvider} from "../cluster/ClusterListDataProvider";
 import {ClusterModel} from "../cluster/ClusterModel";
 import {ConnectionManager} from "./ConnectionManager";
-import {UrlUtils} from "../utils";
+import {FileUtils, UrlUtils} from "../utils";
 import {workspaceConfigs} from "../vscode-objs/WorkspaceConfigs";
 import {WorkspaceFsCommands} from "../workspace-fs";
 import path from "node:path";
@@ -72,7 +71,7 @@ export class ConnectionCommands implements Disposable {
 
     openDatabricksConfigFileCommand() {
         return async () => {
-            const homeDir = process.env.HOME ?? homedir();
+            const homeDir = FileUtils.getHomedir();
             let filePath =
                 workspaceConfigs.databrickscfgLocation ??
                 process.env.DATABRICKS_CONFIG_FILE ??

packages/databricks-vscode/src/extension.ts

@@ -49,6 +49,7 @@ import {setDbnbCellLimits} from "./language/notebooks/DatabricksNbCellLimits";
 import {DbConnectStatusBarButton} from "./language/DbConnectStatusBarButton";
 import {NotebookAccessVerifier} from "./language/notebooks/NotebookAccessVerifier";
 import {NotebookInitScriptManager} from "./language/notebooks/NotebookInitScriptManager";
+import {showRestartNotebookDialogue} from "./language/notebooks/restartNotebookDialogue";
 
 export async function activate(
     context: ExtensionContext
@@ -314,13 +315,7 @@ export async function activate(
     databricksEnvFileManager.init();
     context.subscriptions.push(
         databricksEnvFileManager,
-        databricksEnvFileManager.onDidChangeEnvironmentVariables(() => {
-            if (workspace.notebookDocuments.length) {
-                window.showInformationMessage(
-                    "Environment variables have changed. Restart all jupyter kernels to pickup the latest environment variables."
-                );
-            }
-        })
+        showRestartNotebookDialogue(databricksEnvFileManager)
     );
     featureManager.isEnabled("debugging.dbconnect");
 

packages/databricks-vscode/src/file-managers/DatabricksEnvFileManager.ts

@@ -159,13 +159,6 @@ export class DatabricksEnvFileManager implements Disposable {
         );
     }
 
-    private getNotebookEnvVars() {
-        return EnvVarGenerators.getNotebookEnvVars(
-            this.featureManager,
-            this.notebookInitScriptManager
-        );
-    }
-
     private getIdeEnvVars() {
         return EnvVarGenerators.getIdeEnvVars();
     }
@@ -187,7 +180,6 @@ export class DatabricksEnvFileManager implements Disposable {
             )) || {}),
             ...this.getIdeEnvVars(),
             ...((await this.getUserEnvVars()) || {}),
-            ...(await this.getNotebookEnvVars()),
         })
             .filter(([, value]) => value !== undefined)
             .map(([key, value]) => `${key}=${value}`);
@@ -225,7 +217,6 @@ export class DatabricksEnvFileManager implements Disposable {
                 this.connectionManager,
                 this.workspacePath
             )) || {}),
-            ...(await this.getNotebookEnvVars()),
         }).forEach(([key, value]) => {
             if (value === undefined) {
                 return;

packages/databricks-vscode/src/language/notebooks/NotebookInitScriptManager.ts

@@ -29,11 +29,15 @@ const execFile = promisify(ef);
 const withLogContext = logging.withLogContext;
 
 async function isDbnbTextEditor(editor?: TextEditor) {
-    return (
-        editor?.document.languageId === "python" &&
-        (await FileUtils.isNotebook(new LocalUri(editor.document.uri))) ===
-            "PY_DBNB"
-    );
+    try {
+        return (
+            editor?.document.languageId === "python" &&
+            (await FileUtils.isNotebook(new LocalUri(editor.document.uri))) ===
+                "PY_DBNB"
+        );
+    } catch (e) {
+        return false;
+    }
 }
 
 export class NotebookInitScriptManager implements Disposable {
@@ -64,7 +68,7 @@ export class NotebookInitScriptManager implements Disposable {
         });
         this.disposables.push(
             this.connectionManager.onDidChangeState(async (e) => {
-                if (e !== "CONNECTED" || (await this.isKnowEnvironment())) {
+                if (e !== "CONNECTED" || (await this.isKnownEnvironment())) {
                     return;
                 }
                 this.initScriptSuccessfullyVerified = false;
@@ -82,34 +86,44 @@ export class NotebookInitScriptManager implements Disposable {
                 this.verifyInitScript();
             }),
             workspace.onDidOpenNotebookDocument(async () => {
-                if (await this.isKnowEnvironment()) {
+                if (await this.isKnownEnvironment()) {
                     return;
                 }
                 this.verifyInitScript();
             }),
             window.onDidChangeActiveNotebookEditor(async (activeNotebook) => {
-                if ((await this.isKnowEnvironment()) || !activeNotebook) {
+                if ((await this.isKnownEnvironment()) || !activeNotebook) {
                     return;
                 }
                 this.verifyInitScript();
             }),
             window.onDidChangeActiveTextEditor(async (activeTextEditor) => {
                 if (
                     activeTextEditor?.document.languageId !== "python" ||
-                    (await this.isKnowEnvironment())
+                    (await this.isKnownEnvironment())
                 ) {
                     return;
                 }
-                const localUri = new LocalUri(activeTextEditor?.document.uri);
-                if (await FileUtils.isNotebook(localUri)) {
+
+                if (
+                    activeTextEditor?.document.uri.scheme ===
+                        "vscode-notebook-cell" ||
+                    (await FileUtils.isNotebook(
+                        new LocalUri(activeTextEditor?.document.uri)
+                    ))
+                ) {
                     this.verifyInitScript();
                 }
             })
         );
     }
 
     get ipythonDir(): string {
-        return path.join(this.workspacePath.fsPath, ".databricks", "ipython");
+        return (
+            workspaceConfigs.ipythonDir ??
+            process.env.IPYTHONDIR ??
+            path.join(FileUtils.getHomedir(), ".ipython")
+        );
     }
 
     get startupDir(): string {
@@ -131,7 +145,7 @@ export class NotebookInitScriptManager implements Disposable {
         return readdir(this.generatedDir);
     }
 
-    async isKnowEnvironment() {
+    async isKnownEnvironment() {
         return (
             this.currentEnvPath ===
             (await this.pythonExtension.getPythonExecutable())
@@ -287,6 +301,8 @@ export class NotebookInitScriptManager implements Disposable {
         // If we are not in a jupyter notebook or a databricks notebook,
         // then we don't need to verify the init script
         if (
+            window.activeTextEditor?.document.uri.scheme !==
+                "vscode-notebook-cell" &&
             !isDbnbTextEditor(window.activeTextEditor) &&
             window.activeNotebookEditor === undefined
         ) {

packages/databricks-vscode/src/language/notebooks/restartNotebookDialogue.ts

@@ -0,0 +1,42 @@
+import {workspace, window, commands} from "vscode";
+import {DatabricksEnvFileManager} from "../../file-managers/DatabricksEnvFileManager";
+import {Mutex} from "../../locking";
+
+export function showRestartNotebookDialogue(
+    databricksEnvFileManager: DatabricksEnvFileManager
+) {
+    const mutex = new Mutex();
+    return databricksEnvFileManager.onDidChangeEnvironmentVariables(
+        async () => {
+            if (!workspace.notebookDocuments.length || mutex.locked) {
+                return;
+            }
+            await mutex.wait();
+            try {
+                const choice = await window.showInformationMessage(
+                    "Environment variables have changed. Restart all jupyter kernels to pickup the latest environment variables. ",
+                    {
+                        modal: true,
+                    },
+                    "Restart All Jupyter Kernels",
+                    "Cancel"
+                );
+
+                if (choice === "Restart All Jupyter Kernels") {
+                    for (const doc of workspace.notebookDocuments) {
+                        if (doc.isClosed) {
+                            return;
+                        }
+                        await doc.save();
+                        await window.showNotebookDocument(doc);
+                        await commands.executeCommand(
+                            "workbench.action.closeActiveEditor"
+                        );
+                    }
+                }
+            } finally {
+                mutex.signal();
+            }
+        }
+    );
+}

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

@@ -1,9 +1,7 @@
 import {Loggers} from "../logger";
 import {readFile} from "fs/promises";
 import {Uri} from "vscode";
-import {FeatureManager} from "../feature-manager/FeatureManager";
 import {logging} from "@databricks/databricks-sdk";
-import {NotebookInitScriptManager} from "../language/notebooks/NotebookInitScriptManager";
 import {ConnectionManager} from "../configuration/ConnectionManager";
 
 //Get env variables from user's .env file
@@ -39,21 +37,6 @@ export function getIdeEnvVars() {
     /* eslint-enable @typescript-eslint/naming-convention */
 }
 
-export async function getNotebookEnvVars(
-    featureManager: FeatureManager,
-    notebookInitScriptManager: NotebookInitScriptManager
-) {
-    if (!(await featureManager.isEnabled("notebooks.dbconnect")).avaliable) {
-        return;
-    }
-
-    /* eslint-disable @typescript-eslint/naming-convention */
-    return {
-        IPYTHONDIR: notebookInitScriptManager.ipythonDir,
-    };
-    /* eslint-enable @typescript-eslint/naming-convention */
-}
-
 function getUserAgent(connectionManager: ConnectionManager) {
     const client = connectionManager.apiClient;
     if (!client) {

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

@@ -4,6 +4,7 @@ import {LocalUri} from "../sync/SyncDestination";
 import {ConnectionManager} from "../configuration/ConnectionManager";
 import {exists} from "fs-extra";
 import path from "path";
+import {homedir} from "os";
 
 export type NotebookType = "IPYNB" | "PY_DBNB" | "OTHER_DBNB";
 export async function isNotebook(
@@ -43,3 +44,7 @@ export async function waitForDatabricksProject(
         await connectionManager.waitForConnect();
     }
 }
+
+export function getHomedir() {
+    return process.env.HOME ?? homedir();
+}

packages/databricks-vscode/src/vscode-objs/WorkspaceConfigs.ts

@@ -152,4 +152,14 @@ export const workspaceConfigs = {
                 .get<boolean>("wsfs.rearrangeCells") ?? true
         );
     },
+
+    get ipythonDir(): string | undefined {
+        const dir = workspace
+            .getConfiguration("databricks")
+            .get<string>("ipythonDir");
+        if (dir === "") {
+            return undefined;
+        }
+        return dir;
+    },
 };

yarn.lock

@@ -376,7 +376,7 @@ __metadata:
     inversify: ^6.0.1
     reflect-metadata: ^0.1.13
     semver: ^7.5.4
-  checksum: 2fbf7327009b6c0dd7e04ca2980c1a2245e115fcd60701d756824be128697b78307f4f857182810772df4a70e7fdab73313ff75a6034f153371f9f9434a73fa3
+  checksum: 2e28fb154da0743e5b26cd3bdb1dba6f286e1f34f41d64291fca872ae25678f44a7780bd965f8166632f5a2303a00ecc6def5a9570d09e8d37bcd6546d489cd3
   languageName: node
   linkType: hard
 
@@ -391,7 +391,7 @@ __metadata:
     inversify: ^6.0.1
     reflect-metadata: ^0.1.13
     semver: ^7.5.4
-  checksum: 2fbf7327009b6c0dd7e04ca2980c1a2245e115fcd60701d756824be128697b78307f4f857182810772df4a70e7fdab73313ff75a6034f153371f9f9434a73fa3
+  checksum: 2e28fb154da0743e5b26cd3bdb1dba6f286e1f34f41d64291fca872ae25678f44a7780bd965f8166632f5a2303a00ecc6def5a9570d09e8d37bcd6546d489cd3
   languageName: node
   linkType: hard