Add additional error handling to Python installation for Notebooks (#4891)

* Also enabled integration tests for python installation.

extensions/notebook/src/dialog/configurePythonDialog.ts

@@ -129,7 +129,7 @@ export class ConfigurePythonDialog {
 		}
 
 		// Don't wait on installation, since there's currently no Cancel functionality
-		this.jupyterInstallation.startInstallProcess(pythonLocation)
+		this.jupyterInstallation.startInstallProcess(false, pythonLocation)
 			.then(() => {
 				this._setupComplete.resolve();
 			})

extensions/notebook/src/extension.ts

@@ -13,6 +13,7 @@ import * as nls from 'vscode-nls';
 import { JupyterController } from './jupyter/jupyterController';
 import { AppContext } from './common/appContext';
 import { ApiWrapper } from './common/apiWrapper';
+import { IExtensionApi } from './types';
 
 const localize = nls.loadMessageBundle();
 
@@ -20,10 +21,9 @@ const JUPYTER_NOTEBOOK_PROVIDER = 'jupyter';
 const msgSampleCodeDataFrame = localize('msgSampleCodeDataFrame', 'This sample code loads the file into a data frame and shows the first 10 results.');
 const noNotebookVisible = localize('noNotebookVisible', 'No notebook editor is active');
 
+let controller: JupyterController;
 
-export let controller: JupyterController;
+export async function activate(extensionContext: vscode.ExtensionContext): Promise<IExtensionApi> {
-
-export function activate(extensionContext: vscode.ExtensionContext) {
 	extensionContext.subscriptions.push(vscode.commands.registerCommand('notebook.command.new', (context?: azdata.ConnectedContext) => {
 		let connectionProfile: azdata.IConnectionProfile = undefined;
 		if (context && context.connectionProfile) {
@@ -49,7 +49,16 @@ export function activate(extensionContext: vscode.ExtensionContext) {
 
 	let appContext = new AppContext(extensionContext, new ApiWrapper());
 	controller = new JupyterController(appContext);
-	controller.activate();
+	let result = await controller.activate();
+	if (!result) {
+		return undefined;
+	}
+
+	return {
+		getJupyterController() {
+			return controller;
+		}
+	};
 }
 
 function newNotebook(connectionProfile: azdata.IConnectionProfile) {

extensions/notebook/src/integrationTest/index.ts

@@ -0,0 +1,30 @@
+/*---------------------------------------------------------------------------------------------
+ *  Copyright (c) Microsoft Corporation. All rights reserved.
+ *  Licensed under the Source EULA. See License.txt in the project root for license information.
+ *--------------------------------------------------------------------------------------------*/
+
+const path = require('path');
+const testRunner = require('vscode/lib/testrunner');
+
+const suite = 'Notebook Extension Integration Tests';
+
+const options: any = {
+	ui: 'bdd',
+	useColors: true,
+	timeout: 600000
+};
+
+if (process.env.BUILD_ARTIFACTSTAGINGDIRECTORY) {
+	options.reporter = 'mocha-multi-reporters';
+	options.reporterOptions = {
+		reporterEnabled: 'spec, mocha-junit-reporter',
+		mochaJunitReporterReporterOptions: {
+			testsuitesTitle: `${suite} ${process.platform}`,
+			mochaFile: path.join(process.env.BUILD_ARTIFACTSTAGINGDIRECTORY, `test-results/${process.platform}-${suite.toLowerCase().replace(/[^\w]/g, '-')}-results.xml`)
+		}
+	};
+}
+
+testRunner.configure(options);
+
+export = testRunner;

extensions/notebook/src/integrationTest/notebookIntegration.test.ts

@@ -6,16 +6,17 @@
 'use strict';
 
 import * as should from 'should';
-import * as assert from 'assert';
 import * as vscode from 'vscode';
 import * as azdata from 'azdata';
 import * as tempWrite from 'temp-write';
+import * as assert from 'assert';
 import 'mocha';
 
 import { JupyterController } from '../jupyter/jupyterController';
 import { INotebook, CellTypes } from '../contracts/content';
+import JupyterServerInstallation from '../jupyter/jupyterServerInstallation';
 
-describe('Notebook Integration Test', function (): void {
+describe('Notebook Extension Integration Tests', function () {
 	this.timeout(600000);
 
 	let expectedNotebookContent: INotebook = {
@@ -35,12 +36,34 @@ describe('Notebook Integration Test', function (): void {
 		nbformat_minor: 2
 	};
 
+	let installComplete = false;
+	let pythonInstallDir = process.env.PYTHON_TEST_PATH;
+	let jupyterController: JupyterController;
+	before(async function () {
+		assert.ok(pythonInstallDir, 'Python install directory was not defined.');
+
+		let notebookExtension: vscode.Extension<any>;
+		while (true) {
+			notebookExtension = vscode.extensions.getExtension('Microsoft.notebook');
+			if (notebookExtension && notebookExtension.isActive) {
+				break;
+			} else {
+				await new Promise(resolve => { setTimeout(resolve, 1000); });
+			}
+		}
+
+		jupyterController = notebookExtension.exports.getJupyterController() as JupyterController;
+
+		await jupyterController.jupyterInstallation.startInstallProcess(false, pythonInstallDir);
+		installComplete = true;
+	});
 
 	it('Should connect to local notebook server with result 2', async function () {
-		this.timeout(60000);
+		should(installComplete).be.true('Python setup did not complete.');
+		should(JupyterServerInstallation.getPythonInstallPath(jupyterController.jupyterInstallation.apiWrapper)).be.equal(pythonInstallDir);
+
 		let pythonNotebook = Object.assign({}, expectedNotebookContent, { metadata: { kernelspec: { name: 'python3', display_name: 'Python 3' } } });
 		let uri = writeNotebookToFile(pythonNotebook);
-		await ensureJupyterInstalled();
 
 		let notebook = await azdata.nb.showNotebookDocument(uri);
 		should(notebook.document.cells).have.length(1);
@@ -51,82 +74,13 @@ describe('Notebook Integration Test', function (): void {
 		let result = (<azdata.nb.IExecuteResult>cellOutputs[0]).data['text/plain'];
 		should(result).equal('2');
 
-		try {
+		await vscode.commands.executeCommand('workbench.action.closeActiveEditor');
-			// TODO support closing the editor. Right now this prompts and there's no override for this. Need to fix in core
-			// Close the editor using the recommended vscode API
-			//await vscode.commands.executeCommand('workbench.action.closeActiveEditor');
-		}
-		catch (e) { }
-	});
-
-	it('Should connect to remote spark server with result 2', async function () {
-		this.timeout(240000);
-		let uri = writeNotebookToFile(expectedNotebookContent);
-		await ensureJupyterInstalled();
-
-		// Given a connection to a server exists
-		let connectionProfile = await connectToSparkIntegrationServer();
-
-		// When I open a Spark notebook and run the cell
-		let notebook = await azdata.nb.showNotebookDocument(uri, {
-			connectionProfile: connectionProfile
-		});
-		should(notebook.document.cells).have.length(1);
-		let ran = await notebook.runCell(notebook.document.cells[0]);
-		should(ran).be.true('Notebook runCell failed');
-
-		// Then I expect to get the output result of 1+1, executed remotely against the Spark endpoint
-		let cellOutputs = notebook.document.cells[0].contents.outputs;
-		should(cellOutputs).have.length(4);
-		let sparkResult = (<azdata.nb.IStreamResult>cellOutputs[3]).text;
-		should(sparkResult).equal('2');
-
-		try {
-			// TODO support closing the editor. Right now this prompts and there's no override for this. Need to fix in core
-			// Close the editor using the recommended vscode API
-			//await vscode.commands.executeCommand('workbench.action.closeActiveEditor');
-		}
-		catch (e) { }
 	});
 });
 
-async function connectToSparkIntegrationServer(): Promise<azdata.IConnectionProfile> {
-	assert.ok(process.env.BACKEND_HOSTNAME, 'BACKEND_HOSTNAME, BACKEND_USERNAME, BACKEND_PWD must be set using ./tasks/setbackenvariables.sh or .\\tasks\\setbackendvaraibles.bat');
-	let connInfo: azdata.connection.Connection = {
-		options: {
-			'host': process.env.BACKEND_HOSTNAME,
-			'groupId': 'C777F06B-202E-4480-B475-FA416154D458',
-			'knoxport': '',
-			'user': process.env.BACKEND_USERNAME,
-			'password': process.env.BACKEND_PWD
-		},
-		providerName: 'HADOOP_KNOX',
-		connectionId: 'abcd1234',
-	};
-	connInfo['savePassword'] = true;
-	let result = await azdata.connection.connect(<any>connInfo as azdata.IConnectionProfile);
-
-	should(result.connected).be.true();
-	should(result.connectionId).not.be.undefined();
-	should(result.connectionId).not.be.empty();
-	should(result.errorMessage).be.undefined();
-
-	let activeConnections = await azdata.connection.getActiveConnections();
-	should(activeConnections).have.length(1);
-
-	return <azdata.IConnectionProfile><any>connInfo;
-}
-
 function writeNotebookToFile(pythonNotebook: INotebook): vscode.Uri {
 	let notebookContentString = JSON.stringify(pythonNotebook);
 	let localFile = tempWrite.sync(notebookContentString, 'notebook.ipynb');
 	let uri = vscode.Uri.file(localFile);
 	return uri;
 }
-
-async function ensureJupyterInstalled(): Promise<void> {
-	let jupterControllerExports = vscode.extensions.getExtension('Microsoft.sql-vnext').exports;
-	let jupyterController = jupterControllerExports.getJupterController() as JupyterController;
-	await jupyterController.jupyterInstallation.installReady;
-}
-

extensions/notebook/src/jupyter/jupyterController.ts

@@ -175,13 +175,13 @@ export class JupyterController implements vscode.Disposable {
 	}
 
 	private async handleDependenciesReinstallation(): Promise<void> {
-		if (await this.confirmReinstall()) {
+		try {
-			this._jupyterInstallation = await JupyterServerInstallation.getInstallation(
+			let doReinstall = await this.confirmReinstall();
-				this.extensionContext.extensionPath,
+			if (doReinstall) {
-				this.outputChannel,
+				await this._jupyterInstallation.startInstallProcess(true);
-				this.apiWrapper,
+			}
-				undefined,
+		} catch (err) {
-				true);
+			this.apiWrapper.showErrorMessage(utils.getErrorMessage(err));
 		}
 	}
 

extensions/notebook/src/jupyter/jupyterServerInstallation.ts

@@ -31,6 +31,8 @@ const msgPythonUnpackError = localize('msgPythonUnpackError', 'Error while unpac
 const msgTaskName = localize('msgTaskName', 'Installing Notebook dependencies');
 const msgInstallPkgStart = localize('msgInstallPkgStart', 'Installing Notebook dependencies, see Tasks view for more information');
 const msgInstallPkgFinish = localize('msgInstallPkgFinish', 'Notebook dependencies installation is complete');
+const msgPythonRunningError = localize('msgPythonRunningError', 'Cannot overwrite existing Python installation while python is running.');
+const msgPendingInstallError = localize('msgPendingInstallError', 'Another Python installation is currently in progress.');
 function msgDependenciesInstallationFailed(errorMessage: string): string { return localize('msgDependenciesInstallationFailed', 'Installing Notebook dependencies failed with error: {0}', errorMessage); }
 function msgDownloadPython(platform: string, pythonDownloadUrl: string): string { return localize('msgDownloadPython', 'Downloading local python for platform: {0} to {1}', platform, pythonDownloadUrl); }
 
@@ -52,41 +54,19 @@ export default class JupyterServerInstallation {
 
 	// Allows dependencies to be installed even if an existing installation is already present
 	private _forceInstall: boolean;
+	private _installInProgress: boolean;
 
 	private static readonly DefaultPythonLocation = path.join(utils.getUserHome(), 'azuredatastudio-python');
 
-	private _installReady: Deferred<void>;
+	constructor(extensionPath: string, outputChannel: OutputChannel, apiWrapper: ApiWrapper, pythonInstallationPath?: string) {
-
-	constructor(extensionPath: string, outputChannel: OutputChannel, apiWrapper: ApiWrapper, pythonInstallationPath?: string, forceInstall?: boolean) {
 		this.extensionPath = extensionPath;
 		this.outputChannel = outputChannel;
 		this.apiWrapper = apiWrapper;
 		this._pythonInstallationPath = pythonInstallationPath || JupyterServerInstallation.getPythonInstallPath(this.apiWrapper);
-		this._forceInstall = !!forceInstall;
+		this._forceInstall = false;
+		this._installInProgress = false;
 
 		this.configurePackagePaths();
-
-		this._installReady = new Deferred<void>();
-		if (JupyterServerInstallation.isPythonInstalled(this.apiWrapper)) {
-			this._installReady.resolve();
-		}
-	}
-
-	public get installReady(): Promise<void> {
-		return this._installReady.promise;
-	}
-
-	public static async getInstallation(
-		extensionPath: string,
-		outputChannel: OutputChannel,
-		apiWrapper: ApiWrapper,
-		pythonInstallationPath?: string,
-		forceInstall?: boolean): Promise<JupyterServerInstallation> {
-
-		let installation = new JupyterServerInstallation(extensionPath, outputChannel, apiWrapper, pythonInstallationPath, forceInstall);
-		await installation.startInstallProcess();
-
-		return installation;
 	}
 
 	private async installDependencies(backgroundOperation: azdata.BackgroundOperation): Promise<void> {
@@ -217,7 +197,7 @@ export default class JupyterServerInstallation {
 		// Update python paths and properties to reference user's local python.
 		let pythonBinPathSuffix = process.platform === constants.winPlatform ? '' : 'bin';
 
-		this._pythonExecutable = path.join(pythonSourcePath, process.platform === constants.winPlatform ? 'python.exe' : 'bin/python3');
+		this._pythonExecutable = JupyterServerInstallation.getPythonExePath(this._pythonInstallationPath);
 		this.pythonBinPath = path.join(pythonSourcePath, pythonBinPathSuffix);
 
 		// Store paths to python libraries required to run jupyter.
@@ -230,6 +210,11 @@ export default class JupyterServerInstallation {
 		}
 		this.pythonEnvVarPath = this.pythonBinPath + delimiter + this.pythonEnvVarPath;
 
+		// Delete existing Python variables in ADS to prevent conflict with other installs
+		delete process.env['PYTHONPATH'];
+		delete process.env['PYTHONSTARTUP'];
+		delete process.env['PYTHONHOME'];
+
 		// Store the executable options to run child processes with env var without interfering parent env var.
 		let env = Object.assign({}, process.env);
 		delete env['Path']; // Delete extra 'Path' variable for Windows, just in case.
@@ -239,15 +224,51 @@ export default class JupyterServerInstallation {
 		};
 	}
 
-	public startInstallProcess(pythonInstallationPath?: string): Promise<void> {
+	private isPythonRunning(pythonInstallPath: string): Promise<boolean> {
-		if (pythonInstallationPath) {
+		let pythonExePath = JupyterServerInstallation.getPythonExePath(pythonInstallPath);
-			this._pythonInstallationPath = pythonInstallationPath;
+		return new Promise<boolean>(resolve => {
-			this.configurePackagePaths();
+			fs.open(pythonExePath, 'r+', (err, fd) => {
+				if (!err) {
+					fs.close(fd, err => {
+						this.apiWrapper.showErrorMessage(utils.getErrorMessage(err));
+					});
+					resolve(false);
+				} else {
+					resolve(err.code === 'EBUSY' || err.code === 'EPERM');
+				}
+			});
+		});
+	}
+
+	/**
+	 * Installs Python and associated dependencies to the specified directory.
+	 * @param forceInstall Indicates whether an existing installation should be overwritten, if it exists.
+	 * @param installationPath Optional parameter that specifies where to install python.
+	 * The previous path (or the default) is used if a new path is not specified.
+	 */
+	public async startInstallProcess(forceInstall: boolean, installationPath?: string): Promise<void> {
+		let isPythonRunning = await this.isPythonRunning(installationPath ? installationPath : this._pythonInstallationPath);
+		if (isPythonRunning) {
+			return Promise.reject(msgPythonRunningError);
 		}
-		let updateConfig = () => {
+
+		if (this._installInProgress) {
+			return Promise.reject(msgPendingInstallError);
+		}
+		this._installInProgress = true;
+
+		this._forceInstall = forceInstall;
+		if (installationPath) {
+			this._pythonInstallationPath = installationPath;
+		}
+		this.configurePackagePaths();
+
+		let updateConfig = async () => {
 			let notebookConfig = this.apiWrapper.getConfiguration(constants.notebookConfigKey);
-			notebookConfig.update(constants.pythonPathConfigKey, this._pythonInstallationPath, ConfigurationTarget.Global);
+			await notebookConfig.update(constants.pythonPathConfigKey, this._pythonInstallationPath, ConfigurationTarget.Global);
 		};
+
+		let installReady = new Deferred<void>();
 		if (!fs.existsSync(this._pythonExecutable) || this._forceInstall) {
 			this.apiWrapper.startBackgroundOperation({
 				displayName: msgTaskName,
@@ -255,27 +276,32 @@ export default class JupyterServerInstallation {
 				isCancelable: false,
 				operation: op => {
 					this.installDependencies(op)
-						.then(() => {
+						.then(async () => {
-							this._installReady.resolve();
+							await updateConfig();
-							updateConfig();
+							installReady.resolve();
+							this._installInProgress = false;
 						})
 						.catch(err => {
 							let errorMsg = msgDependenciesInstallationFailed(utils.getErrorMessage(err));
 							op.updateStatus(azdata.TaskStatus.Failed, errorMsg);
 							this.apiWrapper.showErrorMessage(errorMsg);
-							this._installReady.reject(errorMsg);
+							installReady.reject(errorMsg);
+							this._installInProgress = false;
 						});
 				}
 			});
 		} else {
 			// Python executable already exists, but the path setting wasn't defined,
 			// so update it here
-			this._installReady.resolve();
+			await updateConfig();
-			updateConfig();
+			installReady.resolve();
 		}
-		return this._installReady.promise;
+		return installReady.promise;
 	}
 
+	/**
+	 * Opens a dialog for configuring the installation path for the Notebook Python dependencies.
+	 */
 	public async promptForPythonInstall(): Promise<void> {
 		if (!JupyterServerInstallation.isPythonInstalled(this.apiWrapper)) {
 			let pythonDialog = new ConfigurePythonDialog(this.apiWrapper, this.outputChannel, this);
@@ -312,6 +338,10 @@ export default class JupyterServerInstallation {
 		return this._pythonExecutable;
 	}
 
+	/**
+	 * Checks if a python executable exists at the "notebook.pythonPath" defined in the user's settings.
+	 * @param apiWrapper An ApiWrapper to use when retrieving user settings info.
+	 */
 	public static isPythonInstalled(apiWrapper: ApiWrapper): boolean {
 		// Don't use _pythonExecutable here, since it could be populated with a default value
 		let pathSetting = JupyterServerInstallation.getPythonPathSetting(apiWrapper);
@@ -319,13 +349,15 @@ export default class JupyterServerInstallation {
 			return false;
 		}
 
-		let pythonExe = path.join(
+		let pythonExe = JupyterServerInstallation.getPythonExePath(pathSetting);
-			pathSetting,
-			constants.pythonBundleVersion,
-			process.platform === constants.winPlatform ? 'python.exe' : 'bin/python3');
 		return fs.existsSync(pythonExe);
 	}
 
+	/**
+	 * Returns the Python installation path defined in "notebook.pythonPath" in the user's settings.
+	 * Returns a default path if the setting is not defined.
+	 * @param apiWrapper An ApiWrapper to use when retrieving user settings info.
+	 */
 	public static getPythonInstallPath(apiWrapper: ApiWrapper): string {
 		let userPath = JupyterServerInstallation.getPythonPathSetting(apiWrapper);
 		return userPath ? userPath : JupyterServerInstallation.DefaultPythonLocation;
@@ -345,6 +377,11 @@ export default class JupyterServerInstallation {
 		return path;
 	}
 
+	/**
+	 * Returns the folder containing the python executable under the path defined in
+	 * "notebook.pythonPath" in the user's settings.
+	 * @param apiWrapper An ApiWrapper to use when retrieving user settings info.
+	 */
 	public static getPythonBinPath(apiWrapper: ApiWrapper): string {
 		let pythonBinPathSuffix = process.platform === constants.winPlatform ? '' : 'bin';
 
@@ -353,4 +390,11 @@ export default class JupyterServerInstallation {
 			constants.pythonBundleVersion,
 			pythonBinPathSuffix);
 	}
+
+	private static getPythonExePath(pythonInstallPath: string): string {
+		return path.join(
+			pythonInstallPath,
+			constants.pythonBundleVersion,
+			process.platform === constants.winPlatform ? 'python.exe' : 'bin/python3');
+	}
 }

extensions/notebook/src/test/index.ts

@@ -2,8 +2,6 @@
  *  Copyright (c) Microsoft Corporation. All rights reserved.
  *  Licensed under the Source EULA. See License.txt in the project root for license information.
  *--------------------------------------------------------------------------------------------*/
-// import * as vscode from 'vscode';
-// import { context } from './testContext';
 
 const path = require('path');
 const testRunner = require('vscode/lib/testrunner');

extensions/notebook/src/types.d.ts

@@ -0,0 +1,16 @@
+/*---------------------------------------------------------------------------------------------
+ *  Copyright (c) Microsoft Corporation. All rights reserved.
+ *  Licensed under the Source EULA. See License.txt in the project root for license information.
+ *--------------------------------------------------------------------------------------------*/
+
+import { JupyterController } from './jupyter/jupyterController';
+
+/**
+ * The API provided by this extension.
+ *
+ * @export
+ * @interface IExtensionApi
+ */
+export interface IExtensionApi {
+	getJupyterController(): JupyterController;
+}

scripts/sql-test-integration.bat

@@ -4,10 +4,14 @@ pushd %~dp0\..
 
 set VSCODEUSERDATADIR=%TMP%\adsuser-%RANDOM%-%TIME:~6,5%
 set VSCODEEXTENSIONSDIR=%TMP%\adsext-%RANDOM%-%TIME:~6,5%
+set PYTHON_TEST_PATH=%VSCODEUSERDATADIR%\TestPythonInstallation
 echo %VSCODEUSERDATADIR%
 echo %VSCODEEXTENSIONSDIR%
+echo %PYTHON_TEST_PATH%
 @echo OFF
 
+:: This first notebook test will do the python install that the later tests depend on
+call .\scripts\code.bat --extensionDevelopmentPath=%~dp0\..\extensions\notebook --extensionTestsPath=%~dp0\..\extensions\notebook\out\integrationTest --user-data-dir=%VSCODEUSERDATADIR% --extensions-dir=%VSCODEEXTENSIONSDIR% --remote-debugging-port=9222
 call .\scripts\code.bat --extensionDevelopmentPath=%~dp0\..\extensions\integration-tests --extensionTestsPath=%~dp0\..\extensions\integration-tests\out --user-data-dir=%VSCODEUSERDATADIR% --extensions-dir=%VSCODEEXTENSIONSDIR% --remote-debugging-port=9222
 
 if %errorlevel% neq 0 exit /b %errorlevel%

scripts/sql-test-integration.sh

@@ -12,12 +12,15 @@ else
 	VSCODEEXTDIR=`mktemp -d 2>/dev/null`
 fi
 
+export PYTHON_TEST_PATH=$VSCODEUSERDATADIR/TestPythonInstallation
+
 cd $ROOT
 echo $VSCODEUSERDATADIR
 echo $VSCODEEXTDIR
+echo $PYTHON_TEST_PATH
 
+./scripts/code.sh --extensionDevelopmentPath=$ROOT/extensions/notebook --extensionTestsPath=$ROOT/extensions/notebook/out/integrationTest --user-data-dir=$VSCODEUSERDATADIR --extensions-dir=$VSCODEEXTDIR --remote-debugging-port=9222
 ./scripts/code.sh --extensionDevelopmentPath=$ROOT/extensions/integration-tests --extensionTestsPath=$ROOT/extensions/integration-tests/out --user-data-dir=$VSCODEUSERDATADIR --extensions-dir=$VSCODEEXTDIR --remote-debugging-port=9222
 
-
+rm -r -f $VSCODEUSERDATADIR
-rm -r $VSCODEUSERDATADIR
 rm -r $VSCODEEXTDIR