diff --git a/packages/databricks-vscode/package.json b/packages/databricks-vscode/package.json index 3a9abfe25..20ffcebfe 100644 --- a/packages/databricks-vscode/package.json +++ b/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." } } } diff --git a/packages/databricks-vscode/resources/python/00-databricks-init.py b/packages/databricks-vscode/resources/python/00-databricks-init.py index 28b5ec854..c31e1845b 100644 --- a/packages/databricks-vscode/resources/python/00-databricks-init.py +++ b/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) diff --git a/packages/databricks-vscode/src/configuration/ConnectionCommands.ts b/packages/databricks-vscode/src/configuration/ConnectionCommands.ts index 9c65f8f48..c2c2ef100 100644 --- a/packages/databricks-vscode/src/configuration/ConnectionCommands.ts +++ b/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 ?? diff --git a/packages/databricks-vscode/src/extension.ts b/packages/databricks-vscode/src/extension.ts index 490e15277..df09dbbaf 100644 --- a/packages/databricks-vscode/src/extension.ts +++ b/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"); diff --git a/packages/databricks-vscode/src/file-managers/DatabricksEnvFileManager.ts b/packages/databricks-vscode/src/file-managers/DatabricksEnvFileManager.ts index bb7e36345..8fe337b73 100644 --- a/packages/databricks-vscode/src/file-managers/DatabricksEnvFileManager.ts +++ b/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; diff --git a/packages/databricks-vscode/src/language/notebooks/NotebookInitScriptManager.ts b/packages/databricks-vscode/src/language/notebooks/NotebookInitScriptManager.ts index 096f8561a..122022fed 100644 --- a/packages/databricks-vscode/src/language/notebooks/NotebookInitScriptManager.ts +++ b/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,13 +86,13 @@ 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(); @@ -96,12 +100,18 @@ export class NotebookInitScriptManager implements Disposable { 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(); } }) @@ -109,7 +119,11 @@ export class NotebookInitScriptManager implements Disposable { } 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 ) { diff --git a/packages/databricks-vscode/src/language/notebooks/restartNotebookDialogue.ts b/packages/databricks-vscode/src/language/notebooks/restartNotebookDialogue.ts new file mode 100644 index 000000000..5d08a7c85 --- /dev/null +++ b/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(); + } + } + ); +} diff --git a/packages/databricks-vscode/src/utils/envVarGenerators.ts b/packages/databricks-vscode/src/utils/envVarGenerators.ts index d08c0d6c7..535228813 100644 --- a/packages/databricks-vscode/src/utils/envVarGenerators.ts +++ b/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) { diff --git a/packages/databricks-vscode/src/utils/fileUtils.ts b/packages/databricks-vscode/src/utils/fileUtils.ts index 64d046211..825d54c97 100644 --- a/packages/databricks-vscode/src/utils/fileUtils.ts +++ b/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(); +} diff --git a/packages/databricks-vscode/src/vscode-objs/WorkspaceConfigs.ts b/packages/databricks-vscode/src/vscode-objs/WorkspaceConfigs.ts index dc511782b..09ba36f21 100644 --- a/packages/databricks-vscode/src/vscode-objs/WorkspaceConfigs.ts +++ b/packages/databricks-vscode/src/vscode-objs/WorkspaceConfigs.ts @@ -152,4 +152,14 @@ export const workspaceConfigs = { .get("wsfs.rearrangeCells") ?? true ); }, + + get ipythonDir(): string | undefined { + const dir = workspace + .getConfiguration("databricks") + .get("ipythonDir"); + if (dir === "") { + return undefined; + } + return dir; + }, }; diff --git a/yarn.lock b/yarn.lock index d4ce768e5..3075277ca 100644 --- a/yarn.lock +++ b/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