Notebook: Re-Enable Attach to Dropdown Functionality (#3250)

* first attach to working

* Transfer changes from sqlopsstudioextensions PR 448

* Transfer changes from sqlopsstudioextensions PR 447

* Transfer changes from sqlopsstudioextensions PR 456

* Transfer changes from sqlopsstudioextensions PR 465

* Transfer changes from sqlopsstudioextensions PR 463

* Transfer changes from sqlopsstudioextensions PR 482

* Transfer changes from sqlopsstudioextensions PR 485

* Session and Kernel implementation except executeRequest

* Attach to port compiling

* Further tweaks to attach to dropdown, re-enable opening connection dialog

* Revert "Merge remote-tracking branch 'origin/Notebook/sessionExtension' into feature/workingAttachTo"

This reverts commit 94703db87c85416c4ae36762afc1094d6e71166a, reversing
changes made to e4dc25331036d259e9c762cfe8741f957bb5c590.

* Fix code formatting

* Fix for new Add new connection issue

src/sql/base/browser/ui/selectBox/selectBox.ts

@@ -122,6 +122,10 @@ export class SelectBox extends vsSelectBox {
 		return this._selectedOption;
 	}
 
+	public get values(): string[] {
+		return this._dialogOptions;
+	}
+
 	public enable(): void {
 		this.selectElement.disabled = false;
 		this.selectBackground = this.enabledSelectBackground;

src/sql/parts/notebook/models/cell.ts

@@ -156,7 +156,7 @@ export class CellModel implements ICellModel {
 		future.setIOPubHandler({ handle: (msg) => this.handleIOPub(msg) });
 	}
 
-	private clearOutputs(): void {
+	public clearOutputs(): void {
 		this._outputs = [];
 		this.fireOutputsChanged();
 	}
@@ -172,7 +172,7 @@ export class CellModel implements ICellModel {
 		}
 	}
 
-	public get outputs(): ReadonlyArray<nb.ICellOutput> {
+	public get outputs(): Array<nb.ICellOutput> {
 		return this._outputs;
 	}
 
@@ -220,6 +220,8 @@ export class CellModel implements ICellModel {
 		//     this._displayIdMap.set(displayId, targets);
 		// }
 		if (output) {
+			// deletes transient node in the serialized JSON
+			delete output['transient'];
 			this._outputs.push(output);
 			this.fireOutputsChanged();
 		}
@@ -241,7 +243,6 @@ export class CellModel implements ICellModel {
 			cellJson.metadata.language = this._language,
 			cellJson.outputs = this._outputs;
 			cellJson.execution_count = 1; // TODO: keep track of actual execution count
-
 		}
 		return cellJson as nb.ICell;
 	}

src/sql/parts/notebook/models/clientSession.ts

@@ -276,7 +276,7 @@ export class ClientSession implements IClientSession {
 
 	public async updateConnection(connection: NotebookConnection): Promise<void> {
 		if (!this.kernel) {
-			// TODO is there any case where skipping causes errors? Do far it seems like it gets called twice
+			// TODO is there any case where skipping causes errors? So far it seems like it gets called twice
 			return;
 		}
 		this._connection = (connection.connectionProfile.id !== '-1') ? connection : this._connection;

src/sql/parts/notebook/models/modelInterfaces.ts

@@ -288,7 +288,7 @@ export interface INotebookModel {
 	readonly contexts: IDefaultConnection | undefined;
 
 	/**
-	 * The trusted mode of the NoteBook
+	 * The trusted mode of the Notebook
 	 */
 	trustedMode: boolean;
 
@@ -301,7 +301,12 @@ export interface INotebookModel {
 	/**
 	 * Change the current context (if applicable)
 	 */
-	changeContext(host: string): void;
+	changeContext(host: string, connection?: IConnectionProfile): void;
+
+	/**
+	 * Find a cell's index given its model
+	 */
+	findCellIndex(cellModel: ICellModel): number;
 
 	/**
 	 * Adds a cell to the index of the model
@@ -380,5 +385,5 @@ export namespace notebookConstants {
 	export const python3 = 'python3';
 	export const python3DisplayName = 'Python 3';
 	export const defaultSparkKernel = 'pyspark3kernel';
-
+	export const hostPropName = 'host';
 }

src/sql/parts/notebook/models/notebookModel.ts

@@ -5,7 +5,7 @@
 
 'use strict';
 
-import { nb } from 'sqlops';
+import { nb, connection } from 'sqlops';
 
 import { localize } from 'vs/nls';
 import { Event, Emitter } from 'vs/base/common/event';
@@ -167,6 +167,10 @@ export class NotebookModel extends Disposable implements INotebookModel {
 		}
 	}
 
+	public get hadoopConnection(): NotebookConnection {
+		return this._hadoopConnection;
+	}
+
 	/**
 	 * Indicates the server has finished loading. It may have failed to load in
 	 * which case the view will be in an error state.
@@ -186,7 +190,7 @@ export class NotebookModel extends Disposable implements INotebookModel {
 		try {
 			this._trustedMode = isTrusted;
 			let contents = null;
-			if(this.notebookOptions.notebookUri.scheme !== Schemas.untitled) {
+			if (this.notebookOptions.notebookUri.scheme !== Schemas.untitled) {
 				contents = await this.notebookManager.contentManager.getNotebookContents(this.notebookOptions.notebookUri);
 			}
 			let factory = this.notebookOptions.factory;
@@ -208,6 +212,10 @@ export class NotebookModel extends Disposable implements INotebookModel {
 		}
 	}
 
+	public findCellIndex(cellModel: CellModel): number {
+		return this._cells.findIndex((cell) => cell.equals(cellModel));
+	}
+
 	public addCell(cellType: CellType, index?: number): void {
 		if (this.inErrorState || !this._cells) {
 			return;
@@ -289,6 +297,10 @@ export class NotebookModel extends Disposable implements INotebookModel {
 		});
 	}
 
+	private isValidKnoxConnection(profile: IConnectionProfile | connection.Connection) {
+		return profile && profile.providerName === notebookConstants.hadoopKnoxProviderName && profile.options[notebookConstants.hostPropName] !== undefined;
+	}
+
 	public get languageInfo(): nb.ILanguageInfo {
 		return this._defaultLanguageInfo;
 	}
@@ -304,8 +316,8 @@ export class NotebookModel extends Disposable implements INotebookModel {
 		this.doChangeKernel(spec);
 	}
 
-	private doChangeKernel(kernelSpec: nb.IKernelSpec): void {
+	public doChangeKernel(kernelSpec: nb.IKernelSpec): Promise<void> {
-		this._clientSession.changeKernel(kernelSpec)
+		return this._clientSession.changeKernel(kernelSpec)
 			.then((kernel) => {
 				kernel.ready.then(() => {
 					if (kernel.info) {
@@ -317,26 +329,40 @@ export class NotebookModel extends Disposable implements INotebookModel {
 				this.notifyError(localize('changeKernelFailed', 'Failed to change kernel: {0}', notebookUtils.getErrorMessage(err)));
 				// TODO should revert kernels dropdown
 			});
-
 	}
 
-	public changeContext(host: string): void {
+	public changeContext(host: string, newConnection?: IConnectionProfile): void {
 		try {
-			let newConnection: IConnectionProfile = this._activeContexts.otherConnections.find((connection) => connection.options['host'] === host);
+			if (!newConnection) {
+				newConnection = this._activeContexts.otherConnections.find((connection) => connection.options['host'] === host);
+			}
 			if (!newConnection && this._activeContexts.defaultConnection.options['host'] === host) {
 				newConnection = this._activeContexts.defaultConnection;
 			}
-			if (newConnection) {
+			SparkMagicContexts.configureContext(this.notebookOptions);
-				SparkMagicContexts.configureContext(newConnection, this.notebookOptions);
 			this._hadoopConnection = new NotebookConnection(newConnection);
+			this.refreshConnections(newConnection);
 			this._clientSession.updateConnection(this._hadoopConnection);
-			}
 		} catch (err) {
 			let msg = notebookUtils.getErrorMessage(err);
 			this.notifyError(localize('changeContextFailed', 'Changing context failed: {0}', msg));
 		}
 	}
 
+	private refreshConnections(newConnection: IConnectionProfile) {
+		if (this.isValidKnoxConnection(newConnection) &&
+			this._hadoopConnection.connectionProfile.id !== '-1' &&
+			this._hadoopConnection.connectionProfile.id !== this._activeContexts.defaultConnection.id) {
+			// Put the defaultConnection to the head of otherConnections
+			if (this.isValidKnoxConnection(this._activeContexts.defaultConnection)) {
+				this._activeContexts.otherConnections = this._activeContexts.otherConnections.filter(conn => conn.id !== this._activeContexts.defaultConnection.id);
+				this._activeContexts.otherConnections.unshift(this._activeContexts.defaultConnection);
+			}
+			// Change the defaultConnection to newConnection
+			this._activeContexts.defaultConnection = newConnection;
+		}
+	}
+
 	private loadKernelInfo(): void {
 		this.clientSession.kernelChanged(async (e) => {
 			await this.loadActiveContexts(e);
@@ -413,9 +439,11 @@ export class NotebookModel extends Disposable implements INotebookModel {
 	private async loadActiveContexts(kernelChangedArgs: nb.IKernelChangedArgs): Promise<void> {
 		this._activeContexts = await SparkMagicContexts.getContextsForKernel(this.notebookOptions.connectionService, kernelChangedArgs, this.connectionProfile);
 		this._contextsChangedEmitter.fire();
+		if (this.contexts.defaultConnection !== undefined && this.contexts.defaultConnection.options !== undefined) {
 			let defaultHadoopConnection = new NotebookConnection(this.contexts.defaultConnection);
 			this.changeContext(defaultHadoopConnection.host);
 		}
+	}
 
 	/**
 	 * Sanitizes display name to remove IP address in order to fairly compare kernels

src/sql/parts/notebook/models/sparkMagicContexts.ts

@@ -26,7 +26,7 @@ export class SparkMagicContexts {
 			id: '-1',
 			options:
 			{
-				host: localize('selectConnection', 'Select Connection')
+				host: localize('selectConnection', 'Select connection')
 			}
 		};
 
@@ -76,10 +76,13 @@ export class SparkMagicContexts {
 		let defaultConnection: IConnectionProfile = SparkMagicContexts.DefaultContext.defaultConnection;
 		let activeConnections: IConnectionProfile[] = await connectionService.getActiveConnections();
 		// If no connections exist, only show 'n/a'
+        if (activeConnections && activeConnections.length > 0) {
+            // Remove all non-Spark connections
+            activeConnections = activeConnections.filter(conn => conn.providerName === notebookConstants.hadoopKnoxProviderName);
+        }
         if (activeConnections.length === 0) {
             return SparkMagicContexts.DefaultContext;
         }
-		activeConnections = activeConnections.filter(conn => conn.providerName === notebookConstants.hadoopKnoxProviderName);
 
 		// If launched from the right click or server dashboard, connection profile data exists, so use that as default
 		if (profile && profile.options) {
@@ -109,13 +112,11 @@ export class SparkMagicContexts {
 		};
 	}
 
-	public static async configureContext(connection: IConnectionProfile, options: INotebookModelOptions): Promise<object> {
+	public static async configureContext(options: INotebookModelOptions): Promise<object> {
 		let sparkmagicConfDir = path.join(notebookUtils.getUserHome(), '.sparkmagic');
 		// TODO NOTEBOOK REFACTOR re-enable this or move to extension. Requires config files to be available in order to work
 		// await notebookUtils.mkDir(sparkmagicConfDir);
 
-		// let hadoopConnection = new Connection({ options: connection.options }, undefined, connection.connectionId);
-		// await hadoopConnection.getCredential();
 		// // Default to localhost in config file.
 		// let creds: ICredentials = {
 		//     'url': 'http://localhost:8088'

src/sql/parts/notebook/notebook.component.ts

@@ -20,7 +20,7 @@ import { AngularDisposable } from 'sql/base/common/lifecycle';
 
 import { CellTypes, CellType, NotebookChangeType } from 'sql/parts/notebook/models/contracts';
 import { ICellModel, IModelFactory } from 'sql/parts/notebook/models/modelInterfaces';
-import { IConnectionManagementService } from 'sql/parts/connection/common/connectionManagement';
+import { IConnectionManagementService, IConnectionDialogService } from 'sql/parts/connection/common/connectionManagement';
 import { INotebookService, INotebookParams, INotebookManager } from 'sql/services/notebook/notebookService';
 import { IBootstrapParams } from 'sql/services/bootstrap/bootstrapService';
 import { NotebookModel, ErrorInfo, MessageLevel, NotebookContentChange } from 'sql/parts/notebook/models/notebookModel';
@@ -66,7 +66,8 @@ export class NotebookComponent extends AngularDisposable implements OnInit {
 		@Inject(IBootstrapParams) private notebookParams: INotebookParams,
 		@Inject(IInstantiationService) private instantiationService: IInstantiationService,
 		@Inject(IContextMenuService) private contextMenuService: IContextMenuService,
-		@Inject(IContextViewService) private contextViewService: IContextViewService
+		@Inject(IContextViewService) private contextViewService: IContextViewService,
+		@Inject(IConnectionDialogService) private connectionDialogService: IConnectionDialogService
 	) {
 		super();
 		this.profile = this.notebookParams!.profile;
@@ -227,7 +228,8 @@ export class NotebookComponent extends AngularDisposable implements OnInit {
 		attachSelectBoxStyler(kernelDropdown, this.themeService);
 
 		let attachToContainer = document.createElement('div');
-		let attachTodropdwon = new AttachToDropdown(attachToContainer, this.contextViewService);
+		let attachTodropdwon = new AttachToDropdown(attachToContainer, this.contextViewService, this.modelRegistered,
+			this.connectionManagementService, this.connectionDialogService, this.notificationService);
 		attachTodropdwon.render(attachToContainer);
 		attachSelectBoxStyler(attachTodropdwon, this.themeService);
 

src/sql/parts/notebook/notebookActions.ts

@@ -9,18 +9,23 @@ import { Action } from 'vs/base/common/actions';
 import { TPromise } from 'vs/base/common/winjs.base';
 import { localize } from 'vs/nls';
 import { IContextViewProvider } from 'vs/base/browser/ui/contextview/contextview';
+import { INotificationService, Severity, INotificationActions } from 'vs/platform/notification/common/notification';
 
 import { SelectBox, ISelectBoxOptionsWithLabel } from 'sql/base/browser/ui/selectBox/selectBox';
-import { INotebookModel } from 'sql/parts/notebook/models/modelInterfaces';
+import { INotebookModel, notebookConstants } from 'sql/parts/notebook/models/modelInterfaces';
-import { CellTypes, CellType } from 'sql/parts/notebook/models/contracts';
+import { CellType } from 'sql/parts/notebook/models/contracts';
 import { NotebookComponent } from 'sql/parts/notebook/notebook.component';
-import { INotificationService, Severity, INotificationActions } from 'vs/platform/notification/common/notification';
+import { IConnectionProfile } from 'sql/parts/connection/common/interfaces';
-import { NotificationService } from 'vs/workbench/services/notification/common/notificationService';
+import { IConnectionManagementService, IConnectionDialogService } from 'sql/parts/connection/common/connectionManagement';
+import { getErrorMessage } from 'sql/parts/notebook/notebookUtils';
 
 const msgLoading = localize('loading', 'Loading kernels...');
 const kernelLabel: string = localize('Kernel', 'Kernel: ');
 const attachToLabel: string = localize('AttachTo', 'Attach to: ');
-const msgLocalHost: string = localize('localhost', 'Localhost');
+const msgLoadingContexts = localize('loadingContexts', 'Loading contexts...');
+const msgAddNewConnection = localize('addNewConnection', 'Add new connection');
+const msgSelectConnection = localize('selectConnection', 'Select connection');
+const msgConnectionNotApplicable = localize('connectionNotSupported', 'n/a');
 
 // Action to add a cell to notebook based on cell type(code/markdown).
 export class AddCellAction extends Action {
@@ -178,11 +183,128 @@ export class KernelsDropdown extends SelectBox {
 }
 
 export class AttachToDropdown extends SelectBox {
-	constructor(container: HTMLElement, contextViewProvider: IContextViewProvider) {
+	private model: INotebookModel;
-		let selectBoxOptionsWithLabel: ISelectBoxOptionsWithLabel = {
+
-			labelText: attachToLabel,
+	constructor(container: HTMLElement, contextViewProvider: IContextViewProvider, modelRegistered: Promise<INotebookModel>,
-			labelOnTop: false
+		@IConnectionManagementService private _connectionManagementService: IConnectionManagementService,
-		};
+		@IConnectionDialogService private _connectionDialogService: IConnectionDialogService,
-		super([msgLocalHost], msgLocalHost, contextViewProvider, container, selectBoxOptionsWithLabel);
+		@INotificationService private _notificationService: INotificationService) {
+		super([msgLoadingContexts], msgLoadingContexts, contextViewProvider, container, { labelText: attachToLabel, labelOnTop: false } as ISelectBoxOptionsWithLabel);
+		if (modelRegistered) {
+			modelRegistered
+				.then((model) => this.updateModel(model))
+				.catch((err) => {
+					// No-op for now
+				});
+		}
+		this.onDidSelect(e => {
+			let connection = this.model.contexts.otherConnections.find((c) => c.options.host === e.selected);
+			this.doChangeContext(connection);
+		});
+	}
+
+	public updateModel(model: INotebookModel): void {
+		this.model = model;
+		model.contextsChanged(() => {
+			if (this.model.clientSession.kernel && this.model.clientSession.kernel.name) {
+				let currentKernel = this.model.clientSession.kernel.name;
+				this.loadAttachToDropdown(this.model, currentKernel);
+			}
+		});
+	}
+
+	// Load "Attach To" dropdown with the values corresponding to Kernel dropdown
+	public async loadAttachToDropdown(model: INotebookModel, currentKernel: string): Promise<void> {
+		if (currentKernel === notebookConstants.python3) {
+			this.setOptions([msgConnectionNotApplicable]);
+			this.disable();
+		}
+		else {
+			let hadoopConnections = this.getHadoopConnections(model);
+			this.enable();
+			if (hadoopConnections.length === 1 && hadoopConnections[0] === msgAddNewConnection) {
+				hadoopConnections.unshift(msgSelectConnection);
+				this.selectWithOptionName(msgSelectConnection);
+			}
+			else {
+				hadoopConnections.push(msgAddNewConnection);
+			}
+			this.setOptions(hadoopConnections);
+		}
+
+	}
+
+	//Get hadoop connections from context
+	public getHadoopConnections(model: INotebookModel): string[] {
+		let otherHadoopConnections: IConnectionProfile[] = [];
+		model.contexts.otherConnections.forEach((conn) => { otherHadoopConnections.push(conn); });
+		this.selectWithOptionName(model.contexts.defaultConnection.options.host);
+		otherHadoopConnections = this.setHadoopConnectionsList(model.contexts.defaultConnection, model.contexts.otherConnections);
+		let hadoopConnections = otherHadoopConnections.map((context) => context.options.host);
+		return hadoopConnections;
+	}
+
+	private setHadoopConnectionsList(defaultHadoopConnection: IConnectionProfile, otherHadoopConnections: IConnectionProfile[]) {
+		if (defaultHadoopConnection.options.host !== msgSelectConnection) {
+			otherHadoopConnections = otherHadoopConnections.filter(conn => conn.options.host !== defaultHadoopConnection.options.host);
+			otherHadoopConnections.unshift(defaultHadoopConnection);
+			if (otherHadoopConnections.length > 1) {
+				otherHadoopConnections = otherHadoopConnections.filter(val => val.options.host !== msgSelectConnection);
+			}
+		}
+		return otherHadoopConnections;
+	}
+
+	public doChangeContext(connection?: IConnectionProfile): void {
+		if (this.value === msgAddNewConnection) {
+			this.openConnectionDialog();
+		} else {
+			this.model.changeContext(this.value, connection);
+		}
+	}
+
+	/**
+     * Open connection dialog
+     * Enter server details and connect to a server from the dialog
+     * Bind the server value to 'Attach To' drop down
+     * Connected server is displayed at the top of drop down
+     **/
+	public async openConnectionDialog(): Promise<void> {
+		try {
+			//TODO: Figure out how to plumb through the correct provider here
+			await this._connectionDialogService.openDialogAndWait(this._connectionManagementService, { connectionType: 1, providers: [notebookConstants.hadoopKnoxProviderName] }).then(connection => {
+				let attachToConnections = this.values;
+				if (!connection) {
+					this.loadAttachToDropdown(this.model, this.model.clientSession.kernel.name);
+					return;
+				}
+				let connectedServer = connection.options[notebookConstants.hostPropName];
+				//Check to see if the same host is already there in dropdown. We only have host names in dropdown
+				if (attachToConnections.some(val => val === connectedServer)) {
+					this.loadAttachToDropdown(this.model, this.model.clientSession.kernel.name);
+					this.doChangeContext();
+					return;
+				}
+				else {
+					attachToConnections.unshift(connectedServer);
+				}
+				//To ignore n/a after we have at least one valid connection
+				attachToConnections = attachToConnections.filter(val => val !== msgSelectConnection);
+
+				let index = attachToConnections.findIndex((connection => connection === connectedServer));
+				this.setOptions(attachToConnections);
+				if (!index || index < 0 || index >= attachToConnections.length) {
+					index = 0;
+				}
+				this.select(index);
+
+				// Call doChangeContext to set the newly chosen connection in the model
+				this.doChangeContext(connection);
+			});
+		}
+		catch (error) {
+			const actions: INotificationActions = { primary: [] };
+			this._notificationService.notify({ severity: Severity.Error, message: getErrorMessage(error), actions });
+		}
 	}
 }