Add support for low-privileged user to run spark notebooks (#12883)

* Add support for low-privileged user to run spark notebooks

* error

* fix test

extensions/big-data-cluster/src/bdc.d.ts

@@ -31,7 +31,9 @@ declare module 'bdc' {
 
 	export interface IClusterController {
 		getClusterConfig(): Promise<any>;
-		getKnoxUsername(clusterUsername: string): Promise<string>;
+		getKnoxUsername(defaultUsername: string): Promise<string>;
 		getEndPoints(promptConnect?: boolean): Promise<IEndPointsResponse>
+		username: string;
+		password: string;
 	}
 }

extensions/big-data-cluster/src/bigDataCluster/controller/clusterControllerApi.ts

@@ -173,14 +173,14 @@ export class ClusterController implements IClusterController {
 		}
 	}
 
-	public async getKnoxUsername(sqlLogin: string): Promise<string> {
+	public async getKnoxUsername(defaultUsername: string): Promise<string> {
 		// This all is necessary because prior to CU5 BDC deployments all had the same default username for
 		// accessing the Knox gateway. But in the allowRunAsRoot setting was added and defaulted to false - so
 		// if that exists and is false then we use the username instead.
 		// Note that the SQL username may not necessarily be correct here either - but currently this is what
 		// we're requiring to run Notebooks in a BDC
 		const config = await this.getClusterConfig();
-		return config.spec?.spec?.security?.allowRunAsRoot === false ? sqlLogin : DEFAULT_KNOX_USERNAME;
+		return config.spec?.spec?.security?.allowRunAsRoot === false ? defaultUsername : DEFAULT_KNOX_USERNAME;
 	}
 
 	public async getClusterConfig(promptConnect: boolean = false): Promise<any> {

extensions/notebook/src/book/bookTreeView.ts

@@ -8,7 +8,7 @@ import * as vscode from 'vscode';
 import * as path from 'path';
 import * as fs from 'fs-extra';
 import * as constants from '../common/constants';
-import { IPrompter, IQuestion, confirm } from '../prompts/question';
+import { IPrompter, IQuestion, QuestionTypes } from '../prompts/question';
 import CodeAdapter from '../prompts/adapter';
 import { BookTreeItem, BookTreeItemType } from './bookTreeItem';
 import { BookModel, BookVersion } from './bookModel';
@@ -580,7 +580,7 @@ export class BookTreeViewProvider implements vscode.TreeDataProvider<BookTreeIte
 	//Confirmation message dialog
 	private async confirmReplace(): Promise<boolean> {
 		return await this.prompter.promptSingle<boolean>(<IQuestion>{
-			type: confirm,
+			type: QuestionTypes.confirm,
 			message: loc.confirmReplace,
 			default: false
 		});

extensions/notebook/src/common/constants.ts

@@ -81,3 +81,12 @@ export const pythonMacInstallUrl = 'https://go.microsoft.com/fwlink/?linkid=2128
 export const pythonLinuxInstallUrl = 'https://go.microsoft.com/fwlink/?linkid=2110524';
 
 export const notebookLanguages = ['notebook', 'ipynb'];
+
+export const KNOX_ENDPOINT_SERVER = 'host';
+export const KNOX_ENDPOINT_PORT = 'knoxport';
+export const KNOX_ENDPOINT_GATEWAY = 'gateway';
+export const CONTROLLER_ENDPOINT = 'controller';
+export const SQL_PROVIDER = 'MSSQL';
+export const USER = 'user';
+export const AUTHTYPE = 'authenticationType';
+export const INTEGRATED_AUTH = 'integrated';

extensions/notebook/src/common/extensionContextHelper.ts

@@ -0,0 +1,14 @@
+/*---------------------------------------------------------------------------------------------
+ *  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';
+
+export class ExtensionContextHelper {
+	public static extensionContext: vscode.ExtensionContext;
+
+	public static setExtensionContext(extensionContext: vscode.ExtensionContext) {
+		ExtensionContextHelper.extensionContext = extensionContext;
+	}
+}

extensions/notebook/src/common/utils.ts

@@ -3,13 +3,14 @@
  *  Licensed under the Source EULA. See License.txt in the project root for license information.
  *--------------------------------------------------------------------------------------------*/
 
+import * as bdc from 'bdc';
 import * as childProcess from 'child_process';
 import * as fs from 'fs-extra';
 import * as nls from 'vscode-nls';
 import * as vscode from 'vscode';
 import * as azdata from 'azdata';
 import * as crypto from 'crypto';
-import { notebookLanguages, notebookConfigKey, pinnedBooksConfigKey } from './constants';
+import { notebookLanguages, notebookConfigKey, pinnedBooksConfigKey, AUTHTYPE, INTEGRATED_AUTH, KNOX_ENDPOINT_PORT, KNOX_ENDPOINT_SERVER } from './constants';
 
 const localize = nls.loadMessageBundle();
 
@@ -114,13 +115,6 @@ interface RawEndpoint {
 	port?: number;
 }
 
-export interface IEndpoint {
-	serviceName: string;
-	description: string;
-	endpoint: string;
-	protocol: string;
-}
-
 export function getOSPlatformId(): string {
 	let platformId = undefined;
 	switch (process.platform) {
@@ -193,15 +187,15 @@ export function isEditorTitleFree(title: string): boolean {
 	return !hasTextDoc && !hasNotebookDoc;
 }
 
-export function getClusterEndpoints(serverInfo: azdata.ServerInfo): IEndpoint[] {
+export function getClusterEndpoints(serverInfo: azdata.ServerInfo): bdc.IEndpointModel[] {
 	let endpoints: RawEndpoint[] = serverInfo.options['clusterEndpoints'];
 	if (!endpoints || endpoints.length === 0) { return []; }
 
 	return endpoints.map(e => {
 		// If endpoint is missing, we're on CTP bits. All endpoints from the CTP serverInfo should be treated as HTTPS
 		let endpoint = e.endpoint ? e.endpoint : `https://${e.ipAddress}:${e.port}`;
-		let updatedEndpoint: IEndpoint = {
-			serviceName: e.serviceName,
+		let updatedEndpoint: bdc.IEndpointModel = {
+			name: e.serviceName,
 			description: e.description,
 			endpoint: endpoint,
 			protocol: e.protocol
@@ -210,7 +204,6 @@ export function getClusterEndpoints(serverInfo: azdata.ServerInfo): IEndpoint[]
 	});
 }
 
-
 export type HostAndIp = { host: string, port: string };
 
 export function getHostAndPortFromEndpoint(endpoint: string): HostAndIp {
@@ -229,6 +222,26 @@ export function getHostAndPortFromEndpoint(endpoint: string): HostAndIp {
 	};
 }
 
+export function isIntegratedAuth(connection: azdata.IConnectionProfile): boolean {
+	return connection.options[AUTHTYPE] && connection.options[AUTHTYPE].toLowerCase() === INTEGRATED_AUTH.toLowerCase();
+}
+
+export function isSparkKernel(kernelName: string): boolean {
+	return kernelName && kernelName.toLowerCase().indexOf('spark') > -1;
+}
+
+export function setHostAndPort(delimeter: string, connection: azdata.IConnectionProfile): void {
+	let originalHost = connection.options[KNOX_ENDPOINT_SERVER];
+	if (!originalHost) {
+		return;
+	}
+	let index = originalHost.indexOf(delimeter);
+	if (index > -1) {
+		connection.options[KNOX_ENDPOINT_SERVER] = originalHost.slice(0, index);
+		connection.options[KNOX_ENDPOINT_PORT] = originalHost.slice(index + 1);
+	}
+}
+
 export async function exists(path: string): Promise<boolean> {
 	try {
 		await fs.access(path);

extensions/notebook/src/dialog/managePackages/installedPackagesTab.ts

@@ -11,7 +11,7 @@ import { JupyterServerInstallation, PythonPkgDetails } from '../../jupyter/jupyt
 import * as utils from '../../common/utils';
 import { ManagePackagesDialog } from './managePackagesDialog';
 import CodeAdapter from '../../prompts/adapter';
-import { IQuestion, confirm } from '../../prompts/question';
+import { IQuestion, QuestionTypes } from '../../prompts/question';
 import { IconPathHelper } from '../../common/iconHelper';
 
 const localize = nls.loadMessageBundle();
@@ -255,7 +255,7 @@ export class InstalledPackagesTab {
 
 		this.uninstallPackageButton.updateProperties({ enabled: false });
 		let doUninstall = await this.prompter.promptSingle<boolean>(<IQuestion>{
-			type: confirm,
+			type: QuestionTypes.confirm,
 			message: localize('managePackages.confirmUninstall', "Are you sure you want to uninstall the specified packages?"),
 			default: false
 		});

extensions/notebook/src/extension.ts

@@ -17,6 +17,8 @@ import { BuiltInCommands, unsavedBooksContextKey } from './common/constants';
 import { RemoteBookController } from './book/remoteBookController';
 import { RemoteBookDialog } from './dialog/remoteBookDialog';
 import { RemoteBookDialogModel } from './dialog/remoteBookDialogModel';
+import { IconPathHelper } from './common/iconHelper';
+import { ExtensionContextHelper } from './common/extensionContextHelper';
 
 const localize = nls.loadMessageBundle();
 
@@ -24,6 +26,9 @@ let controller: JupyterController;
 type ChooseCellType = { label: string, id: CellType };
 
 export async function activate(extensionContext: vscode.ExtensionContext): Promise<IExtensionApi> {
+	ExtensionContextHelper.setExtensionContext(extensionContext);
+	IconPathHelper.setExtensionContext(extensionContext);
+
 	const appContext = new AppContext(extensionContext);
 	const createBookPath: string = path.posix.join(extensionContext.extensionPath, 'resources', 'notebooks', 'JupyterBooksCreate.ipynb');
 	/**

extensions/notebook/src/jupyter/jupyterController.ts

@@ -14,7 +14,7 @@ import * as constants from '../common/constants';
 import * as localizedConstants from '../common/localizedConstants';
 import { JupyterServerInstallation } from './jupyterServerInstallation';
 import * as utils from '../common/utils';
-import { IPrompter, IQuestion, confirm } from '../prompts/question';
+import { IPrompter, IQuestion, QuestionTypes } from '../prompts/question';
 
 import { AppContext } from '../common/appContext';
 import { LocalJupyterServerManager, ServerInstanceFactory } from './jupyterServerManager';
@@ -28,7 +28,6 @@ import { LocalPipPackageManageProvider } from './localPipPackageManageProvider';
 import { LocalCondaPackageManageProvider } from './localCondaPackageManageProvider';
 import { ManagePackagesDialogModel, ManagePackageDialogOptions } from '../dialog/managePackages/managePackagesDialogModel';
 import { PyPiClient } from './pypiClient';
-import { IconPathHelper } from '../common/iconHelper';
 
 let untitledCounter = 0;
 
@@ -58,7 +57,6 @@ export class JupyterController {
 			this.extensionContext.extensionPath,
 			this.appContext.outputChannel);
 		await this._jupyterInstallation.configurePackagePaths();
-		IconPathHelper.setExtensionContext(this.extensionContext);
 
 		// Add command/task handlers
 		azdata.tasks.registerTask(constants.jupyterOpenNotebookTask, (profile: azdata.IConnectionProfile) => {
@@ -183,7 +181,7 @@ export class JupyterController {
 	//Confirmation message dialog
 	private async confirmReinstall(): Promise<boolean> {
 		return await this.prompter.promptSingle<boolean>(<IQuestion>{
-			type: confirm,
+			type: QuestionTypes.confirm,
 			message: localize('confirmReinstall', "Are you sure you want to reinstall?"),
 			default: true
 		});

extensions/notebook/src/jupyter/jupyterSessionManager.ts

@@ -3,7 +3,7 @@
  *  Licensed under the Source EULA. See License.txt in the project root for license information.
  *--------------------------------------------------------------------------------------------*/
 
-import { nb, ServerInfo, connection, IConnectionProfile } from 'azdata';
+import { nb, ServerInfo, connection, IConnectionProfile, credentials } from 'azdata';
 import { Session, Kernel } from '@jupyterlab/services';
 import * as fs from 'fs-extra';
 import * as nls from 'vscode-nls';
@@ -18,6 +18,10 @@ import { Deferred } from '../common/promise';
 import { JupyterServerInstallation } from './jupyterServerInstallation';
 import * as bdc from 'bdc';
 import { noBDCConnectionError, providerNotValidError } from '../common/localizedConstants';
+import { SQL_PROVIDER, CONTROLLER_ENDPOINT, KNOX_ENDPOINT_GATEWAY, KNOX_ENDPOINT_SERVER, KNOX_ENDPOINT_PORT } from '../common/constants';
+import CodeAdapter from '../prompts/adapter';
+import { IQuestion, QuestionTypes } from '../prompts/question';
+import { ExtensionContextHelper } from '../common/extensionContextHelper';
 
 const configBase = {
 	'kernel_python_credentials': {
@@ -55,15 +59,6 @@ const configBase = {
 	}
 };
 
-const KNOX_ENDPOINT_SERVER = 'host';
-const KNOX_ENDPOINT_PORT = 'knoxport';
-const KNOX_ENDPOINT_GATEWAY = 'gateway';
-const CONTROLLER_ENDPOINT = 'controller';
-const SQL_PROVIDER = 'MSSQL';
-const USER = 'user';
-const AUTHTYPE = 'authenticationType';
-const INTEGRATED_AUTH = 'integrated';
-
 export class JupyterSessionManager implements nb.SessionManager {
 	private _ready: Deferred<void>;
 	private _isReady: boolean;
@@ -183,7 +178,11 @@ export class JupyterSession implements nb.ISession {
 	private _kernel: nb.IKernel;
 	private _messagesComplete: Deferred<void> = new Deferred<void>();
 
-	constructor(private sessionImpl: Session.ISession, private _installation: JupyterServerInstallation, skipSettingEnvironmentVars?: boolean, private _pythonEnvVarPath?: string) {
+	constructor(
+		private sessionImpl: Session.ISession,
+		private _installation: JupyterServerInstallation,
+		skipSettingEnvironmentVars?: boolean,
+		private _pythonEnvVarPath?: string) {
 		this.setEnvironmentVars(skipSettingEnvironmentVars).catch(error => {
 			console.error(`Unexpected exception setting Jupyter Session variables : ${error}`);
 			// We don't want callers to hang forever waiting - it's better to continue on even if we weren't
@@ -275,53 +274,74 @@ export class JupyterSession implements nb.ISession {
 	}
 
 	public async configureConnection(connectionProfile: IConnectionProfile): Promise<void> {
-		if (connectionProfile && connectionProfile.providerName && this.isSparkKernel(this.sessionImpl.kernel.name)) {
+		if (connectionProfile && connectionProfile.providerName && utils.isSparkKernel(this.sessionImpl.kernel.name)) {
 			// %_do_not_call_change_endpoint is a SparkMagic command that lets users change endpoint options,
 			// such as user/profile/host name/auth type
 
-			let credentials;
-			if (!this.isIntegratedAuth(connectionProfile)) {
-				credentials = await connection.getCredentials(connectionProfile.id);
-			}
+			let knoxUsername = connectionProfile.userName || 'root';
+			let knoxPassword: string = '';
+
 			//Update server info with bigdata endpoint - Unified Connection
 			if (connectionProfile.providerName === SQL_PROVIDER) {
-				const endpoints = await this.getClusterEndpoints(connectionProfile.id);
-				const gatewayEndpoint: utils.IEndpoint = endpoints?.find(ep => ep.serviceName.toLowerCase() === KNOX_ENDPOINT_GATEWAY);
-				if (!gatewayEndpoint) {
+				const serverInfo: ServerInfo = await connection.getServerInfo(connectionProfile.id);
+				if (!serverInfo?.options['isBigDataCluster']) {
 					throw new Error(noBDCConnectionError);
 				}
+				const endpoints = utils.getClusterEndpoints(serverInfo);
+				const controllerEndpoint = endpoints.find(ep => ep.name.toLowerCase() === CONTROLLER_ENDPOINT);
+
+				// root is the default username for pre-CU5 instances, so while we prefer to use the connection username
+				// as a default now we'll still fall back to root if it's empty for some reason. (but the calls below should
+				// get the actual correct value regardless)
+				let clusterController: bdc.IClusterController | undefined = undefined;
+				if (!utils.isIntegratedAuth(connectionProfile)) {
+					// See if the controller creds have been saved already, otherwise fall back to using
+					// SQL creds as a default
+					const credentialProvider = await credentials.getProvider('notebook.bdc.password');
+					const usernameKey = `notebook.bdc.username::${connectionProfile.id}`;
+					const savedUsername = ExtensionContextHelper.extensionContext.globalState.get<string>(usernameKey) || connectionProfile.userName;
+					const connectionCreds = await connection.getCredentials(connectionProfile.id);
+					const savedPassword = (await credentialProvider.readCredential(connectionProfile.id)).password || connectionCreds.password;
+					clusterController = await getClusterController(controllerEndpoint.endpoint, 'basic', savedUsername, savedPassword);
+					// Now that we know that the username/password are valid store them for use later on with the same connection
+					await credentialProvider.saveCredential(connectionProfile.id, clusterController.password);
+					await ExtensionContextHelper.extensionContext.globalState.update(usernameKey, clusterController.username);
+					knoxPassword = clusterController.password;
+					try {
+						knoxUsername = await clusterController.getKnoxUsername(clusterController.username);
+					} catch (err) {
+						knoxUsername = clusterController.username;
+						console.log(`Unexpected error getting Knox username for Spark kernel: ${err}`);
+					}
+				} else {
+					clusterController = await getClusterController(controllerEndpoint.endpoint, 'integrated');
+				}
+
+				let gatewayEndpoint: bdc.IEndpointModel = endpoints?.find(ep => ep.name.toLowerCase() === KNOX_ENDPOINT_GATEWAY);
+				if (!gatewayEndpoint) {
+					// User doesn't have permission to see the gateway endpoint from the DMV so we need to query the controller instead
+					const allEndpoints = (await clusterController.getEndPoints()).endPoints;
+					gatewayEndpoint = allEndpoints?.find(ep => ep.name.toLowerCase() === KNOX_ENDPOINT_GATEWAY);
+					if (!gatewayEndpoint) {
+						throw new Error(localize('notebook.couldNotFindKnoxGateway', "Could not find Knox gateway endpoint"));
+					}
+				}
 				let gatewayHostAndPort = utils.getHostAndPortFromEndpoint(gatewayEndpoint.endpoint);
 				connectionProfile.options[KNOX_ENDPOINT_SERVER] = gatewayHostAndPort.host;
 				connectionProfile.options[KNOX_ENDPOINT_PORT] = gatewayHostAndPort.port;
-				// root is the default username for pre-CU5 instances, so while we prefer to use the connection username
-				// as a default now we'll still fall back to root if it's empty for some reason. (but the calls below should
-				// get the actual correct value regardless)
-				connectionProfile.options[USER] = connectionProfile.userName || 'root';
-				if (!this.isIntegratedAuth(connectionProfile)) {
-					try {
-						const bdcApi = <bdc.IExtension>await vscode.extensions.getExtension(bdc.constants.extensionName).activate();
-						const controllerEndpoint = endpoints.find(ep => ep.serviceName.toLowerCase() === CONTROLLER_ENDPOINT);
-						const controller = bdcApi.getClusterController(controllerEndpoint.endpoint, 'basic', connectionProfile.userName, credentials.password);
-						connectionProfile.options[USER] = await controller.getKnoxUsername(connectionProfile.userName);
-					} catch (err) {
-						console.log(`Unexpected error getting Knox username for Spark kernel: ${err}`);
-						// Optimistically use the SQL login name - that's going to normally be the case after CU5
-						connectionProfile.options[USER] = connectionProfile.userName;
-					}
-				}
 			}
 			else {
 				throw new Error(providerNotValidError);
 			}
-			this.setHostAndPort(':', connectionProfile);
-			this.setHostAndPort(',', connectionProfile);
+			utils.setHostAndPort(':', connectionProfile);
+			utils.setHostAndPort(',', connectionProfile);
 
 			let server = vscode.Uri.parse(utils.getLivyUrl(connectionProfile.options[KNOX_ENDPOINT_SERVER], connectionProfile.options[KNOX_ENDPOINT_PORT])).toString();
 			let doNotCallChangeEndpointParams: string;
-			if (this.isIntegratedAuth(connectionProfile)) {
+			if (utils.isIntegratedAuth(connectionProfile)) {
 				doNotCallChangeEndpointParams = `%_do_not_call_change_endpoint --server=${server} --auth=Kerberos`;
 			} else {
-				doNotCallChangeEndpointParams = `%_do_not_call_change_endpoint --username=${connectionProfile.options[USER]} --password=${credentials.password} --server=${server} --auth=Basic_Access`;
+				doNotCallChangeEndpointParams = `%_do_not_call_change_endpoint --username=${knoxUsername} --password=${knoxPassword} --server=${server} --auth=Basic_Access`;
 			}
 			let future = this.sessionImpl.kernel.requestExecute({
 				code: doNotCallChangeEndpointParams
@@ -330,26 +350,6 @@ export class JupyterSession implements nb.ISession {
 		}
 	}
 
-	private isIntegratedAuth(connection: IConnectionProfile): boolean {
-		return connection.options[AUTHTYPE] && connection.options[AUTHTYPE].toLowerCase() === INTEGRATED_AUTH.toLowerCase();
-	}
-
-	private isSparkKernel(kernelName: string): boolean {
-		return kernelName && kernelName.toLowerCase().indexOf('spark') > -1;
-	}
-
-	private setHostAndPort(delimeter: string, connection: IConnectionProfile): void {
-		let originalHost = connection.options[KNOX_ENDPOINT_SERVER];
-		if (!originalHost) {
-			return;
-		}
-		let index = originalHost.indexOf(delimeter);
-		if (index > -1) {
-			connection.options[KNOX_ENDPOINT_SERVER] = originalHost.slice(0, index);
-			connection.options[KNOX_ENDPOINT_PORT] = originalHost.slice(index + 1);
-		}
-	}
-
 	private updateConfig(config: ISparkMagicConfig, creds: ICredentials, homePath: string): void {
 		config.kernel_python_credentials = creds;
 		config.kernel_scala_credentials = creds;
@@ -358,14 +358,6 @@ export class JupyterSession implements nb.ISession {
 		config.ignore_ssl_errors = utils.getIgnoreSslVerificationConfigSetting();
 	}
 
-	private async getClusterEndpoints(profileId: string): Promise<utils.IEndpoint[]> {
-		let serverInfo: ServerInfo = await connection.getServerInfo(profileId);
-		if (!serverInfo || !serverInfo.options) {
-			return [];
-		}
-		return utils.getClusterEndpoints(serverInfo);
-	}
-
 	private async setEnvironmentVars(skip: boolean = false): Promise<void> {
 		// The PowerShell kernel doesn't define the %cd and %set_env magics; no need to run those here then
 		if (!skip && this.sessionImpl?.kernel?.name !== 'powershell') {
@@ -395,6 +387,55 @@ export class JupyterSession implements nb.ISession {
 	}
 }
 
+async function getClusterController(controllerEndpoint: string, authType: bdc.AuthType, username?: string, password?: string): Promise<bdc.IClusterController | undefined> {
+	const bdcApi = <bdc.IExtension>await vscode.extensions.getExtension(bdc.constants.extensionName).activate();
+	const controller = bdcApi.getClusterController(
+		controllerEndpoint,
+		authType,
+		username,
+		password);
+	try {
+		await controller.getClusterConfig();
+		return controller;
+	} catch (err) {
+		// Initial username/password failed so prompt user for username password until either user
+		// cancels out or we successfully connect
+		console.log(`Error connecting to cluster controller: ${err}`);
+		let errorMessage = '';
+		const prompter = new CodeAdapter();
+		while (true) {
+			const newUsername = await prompter.promptSingle<string>(<IQuestion>{
+				type: QuestionTypes.input,
+				name: 'inputPrompt',
+				message: localize('promptBDCUsername', "{0}Please provide the username to connect to the BDC Controller:", errorMessage),
+				default: username
+			});
+			if (!username) {
+				console.log(`User cancelled out of username prompt for BDC Controller`);
+				break;
+			}
+			const newPassword = await prompter.promptSingle<string>(<IQuestion>{
+				type: QuestionTypes.password,
+				name: 'passwordPrompt',
+				message: localize('promptBDCPassword', "Please provide the password to connect to the BDC Controller"),
+				default: ''
+			});
+			if (!password) {
+				console.log(`User cancelled out of password prompt for BDC Controller`);
+				break;
+			}
+			const controller = bdcApi.getClusterController(controllerEndpoint, authType, newUsername, newPassword);
+			try {
+				await controller.getClusterConfig();
+				return controller;
+			} catch (err) {
+				errorMessage = localize('bdcConnectError', "Error: {0}. ", err.message ?? err);
+			}
+		}
+		throw new Error(localize('clusterControllerConnectionRequired', "A connection to the cluster controller is required to run Spark jobs"));
+	}
+}
+
 interface ICredentials {
 	'url': string;
 }

extensions/notebook/src/prompts/confirm.ts

@@ -8,8 +8,8 @@ import * as vscode from 'vscode';
 
 export default class ConfirmPrompt extends Prompt {
 
-	constructor(question: any) {
-		super(question);
+	constructor(question: any, ignoreFocusOut?: boolean) {
+		super(question, ignoreFocusOut);
 	}
 
 	public render(): any {

extensions/notebook/src/prompts/factory.ts

@@ -3,13 +3,19 @@
 
 import Prompt from './prompt';
 import ConfirmPrompt from './confirm';
+import InputPrompt from './input';
+import PasswordPrompt from './password';
 
 export default class PromptFactory {
 
-	public static createPrompt(question: any): Prompt {
+	public static createPrompt(question: any, ignoreFocusOut?: boolean): Prompt {
 		switch (question.type) {
+			case 'input':
+				return new InputPrompt(question, ignoreFocusOut);
+			case 'password':
+				return new PasswordPrompt(question, ignoreFocusOut);
 			case 'confirm':
-				return new ConfirmPrompt(question);
+				return new ConfirmPrompt(question, ignoreFocusOut);
 			default:
 				throw new Error(`Could not find a prompt for question type ${question.type}`);
 		}

extensions/notebook/src/prompts/input.ts

@@ -0,0 +1,57 @@
+// This code is originally from https://github.com/DonJayamanne/bowerVSCode
+// License: https://github.com/DonJayamanne/bowerVSCode/blob/master/LICENSE
+
+import { window, InputBoxOptions } from 'vscode';
+import Prompt from './prompt';
+import EscapeException from './escapeException';
+
+const figures = require('figures');
+
+export default class InputPrompt extends Prompt {
+
+	protected _options: InputBoxOptions;
+
+	constructor(question: any, ignoreFocusOut?: boolean) {
+		super(question, ignoreFocusOut);
+
+		this._options = this.defaultInputBoxOptions;
+		this._options.prompt = this._question.message;
+	}
+
+	// Helper for callers to know the right type to get from the type factory
+	public static get promptType(): string { return 'input'; }
+
+	public render(): any {
+		// Prefer default over the placeHolder, if specified
+		let placeHolder = this._question.default ? this._question.default : this._question.placeHolder;
+
+		if (this._question.default instanceof Error) {
+			placeHolder = this._question.default.message;
+			this._question.default = undefined;
+		}
+
+		this._options.placeHolder = placeHolder;
+
+		return window.showInputBox(this._options)
+			.then(result => {
+				if (result === undefined) {
+					throw new EscapeException();
+				}
+
+				if (result === '') {
+					// Use the default value, if defined
+					result = this._question.default || '';
+				}
+
+				const validationError = this._question.validate ? this._question.validate(result || '') : undefined;
+
+				if (validationError) {
+					this._question.default = new Error(`${figures.warning} ${validationError}`);
+
+					return this.render();
+				}
+
+				return result;
+			});
+	}
+}

extensions/notebook/src/prompts/password.ts

@@ -0,0 +1,13 @@
+// This code is originally from https://github.com/DonJayamanne/bowerVSCode
+// License: https://github.com/DonJayamanne/bowerVSCode/blob/master/LICENSE
+
+import InputPrompt from './input';
+
+export default class PasswordPrompt extends InputPrompt {
+
+	constructor(question: any, ignoreFocusOut?: boolean) {
+		super(question, ignoreFocusOut);
+
+		this._options.password = true;
+	}
+}

extensions/notebook/src/prompts/question.ts

@@ -5,7 +5,11 @@
 
 import * as vscode from 'vscode';
 
-export const confirm = 'confirm';
+export const enum QuestionTypes {
+	input = 'input',
+	password = 'password',
+	confirm = 'confirm'
+}
 
 // Question interface to clarify how to use the prompt feature
 // based on Bower Question format: https://github.com/bower/bower/blob/89069784bb46bfd6639b4a75e98a0d7399a8c2cb/packages/bower-logger/README.md

extensions/notebook/src/protocol/notebookUriHandler.ts

@@ -11,7 +11,7 @@ import * as path from 'path';
 import * as querystring from 'querystring';
 import * as nls from 'vscode-nls';
 const localize = nls.loadMessageBundle();
-import { IQuestion, confirm } from '../prompts/question';
+import { IQuestion, QuestionTypes } from '../prompts/question';
 import CodeAdapter from '../prompts/adapter';
 import { getErrorMessage, isEditorTitleFree } from '../common/utils';
 import * as constants from '../common/constants';
@@ -61,7 +61,7 @@ export class NotebookUriHandler implements vscode.UriHandler {
 			}
 
 			let doOpen = await this.prompter.promptSingle<boolean>(<IQuestion>{
-				type: confirm,
+				type: QuestionTypes.confirm,
 				message: localize('notebook.confirmOpen', "Download and open '{0}'?", url),
 				default: true
 			});

extensions/notebook/src/test/common/prompt.test.ts

@@ -7,7 +7,7 @@ import * as should from 'should';
 import * as sinon from 'sinon';
 import * as vscode from 'vscode';
 
-import { IPrompter, confirm, IQuestion } from '../../prompts/question';
+import { IPrompter, IQuestion, QuestionTypes } from '../../prompts/question';
 import CodeAdapter from '../../prompts/adapter';
 
 
@@ -42,7 +42,7 @@ describe('Prompt', () => {
 	it('Should find prompt for confirm type', async function (): Promise<void> {
 		const quickPickSpy = sinon.stub(vscode.window, 'showQuickPick').returns(Promise.resolve('Yes') as any);
 		await prompter.promptSingle<boolean>(<IQuestion>{
-			type: confirm,
+			type: QuestionTypes.confirm,
 			message: 'sample message',
 			default: false
 		});

extensions/notebook/src/test/model/sessionManager.test.ts

@@ -19,10 +19,13 @@ import { JupyterSessionManager, JupyterSession } from '../../jupyter/jupyterSess
 import { Deferred } from '../../common/promise';
 import { SessionStub, KernelStub, FutureStub } from '../common';
 import { noBDCConnectionError, providerNotValidError } from '../../common/localizedConstants';
+import { ExtensionContextHelper } from '../../common/extensionContextHelper';
+import { AppContext } from '../../common/appContext';
+import uuid = require('uuid');
 
 export class TestClusterController implements bdc.IClusterController {
 	getClusterConfig(): Promise<any> {
-		throw new Error('Method not implemented.');
+		return Promise.resolve({});
 	}
 	getKnoxUsername(clusterUsername: string): Promise<string> {
 		return Promise.resolve('knoxUsername');
@@ -33,8 +36,17 @@ export class TestClusterController implements bdc.IClusterController {
 			endPoints: []
 		});
 	}
+	username: string;
+	password: string;
 }
 
+before(async function(): Promise<void> {
+	// We have to reset the extension context here since the test runner unloads the files before running the tests
+	// so the static state is lost
+	const api = await vscode.extensions.getExtension('Microsoft.notebook').activate();
+	ExtensionContextHelper.setExtensionContext((api.getAppContext() as AppContext).extensionContext);
+});
+
 describe('Jupyter Session Manager', function (): void {
 	let mockJupyterManager = TypeMoq.Mock.ofType<SessionManager>();
 	let sessionManager = new JupyterSessionManager();
@@ -276,25 +288,29 @@ describe('Jupyter Session', function (): void {
 			isCloud: false,
 			azureVersion: 0,
 			osVersion: '',
-			options: {}
+			options: {
+				isBigDataCluster: true
+			}
 		};
-		const mockGatewayEndpoint: utils.IEndpoint = {
-			serviceName: 'gateway',
+		const mockGatewayEndpoint: bdc.IEndpointModel = {
+			name: 'gateway',
 			description: '',
 			endpoint: '',
 			protocol: '',
 		};
-		const mockControllerEndpoint: utils.IEndpoint = {
-			serviceName: 'controller',
+		const mockControllerEndpoint: bdc.IEndpointModel = {
+			name: 'controller',
 			description: '',
 			endpoint: '',
 			protocol: '',
 		};
 		const mockHostAndIp: utils.HostAndIp = {
-			host: 'host',
-			port: 'port'
+			host: '127.0.0.1',
+			port: '1337'
 		};
 		const mockClustercontroller = new TestClusterController();
+		mockClustercontroller.username = 'admin';
+		mockClustercontroller.password = uuid.v4();
 		let mockBdcExtension: TypeMoq.IMock<bdc.IExtension> = TypeMoq.Mock.ofType<bdc.IExtension>();
 		let mockExtension: TypeMoq.IMock<vscode.Extension<any>> = TypeMoq.Mock.ofType<vscode.Extension<any>>();
 		mockBdcExtension.setup(m => m.getClusterController(TypeMoq.It.isAny(), TypeMoq.It.isAny(), TypeMoq.It.isAny(), TypeMoq.It.isAny())).returns(() => mockClustercontroller);
@@ -307,9 +323,8 @@ describe('Jupyter Session', function (): void {
 		sinon.stub(utils, 'getHostAndPortFromEndpoint').returns(mockHostAndIp);
 
 		await session.configureConnection(connectionProfile);
-		should(connectionProfile.options['host']).equal('host');
-		should(connectionProfile.options['knoxport']).equal('port');
-		should(connectionProfile.options['user']).equal('knoxUsername');
+		should(connectionProfile.options['host']).equal(mockHostAndIp.host);
+		should(connectionProfile.options['knoxport']).equal(mockHostAndIp.port);
 	});
 
 	it('configure connection should throw error if there is no connection to big data cluster', async function (): Promise<void> {

src/sql/workbench/api/common/sqlExtHostTypes.ts

@@ -418,17 +418,6 @@ export class TreeItem extends vsExtTypes.TreeItem {
 	providerHandle?: string;
 }
 
-export interface ServerInfoOption {
-	isBigDataCluster: boolean;
-	clusterEndpoints: ClusterEndpoint;
-}
-
-export interface ClusterEndpoint {
-	serviceName: string;
-	ipAddress: string;
-	port: number;
-}
-
 export class SqlThemeIcon {
 
 	static readonly Folder = new SqlThemeIcon('Folder');