Implement Session support through the extension host (#3228)

Full plumb through of Session support. Also fixed some test issues

- Load session and get necessary information in kernels list
- Run Cell button now works as expected
- Added a ToggleAction base class which can be used for anything that switches icons. I'd still prefer to have this be dynamic and as clean as the extension classes
- Fixed account test unhandled promise rejections (caused by incorrect / invalid tests) that made it hard to see all the test run output.

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

@@ -29,7 +29,7 @@ import { IModelService } from 'vs/editor/common/services/modelService';
 import { IContextMenuService, IContextViewService } from 'vs/platform/contextview/browser/contextView';
 import { ICellModel } from 'sql/parts/notebook/models/modelInterfaces';
 import { Taskbar } from 'sql/base/browser/ui/taskbar/taskbar';
-import { RunCellAction, DeleteCellAction, AddCellAction } from 'sql/parts/notebook/cellViews/codeActions';
+import { RunCellAction, DeleteCellAction, AddCellAction, CellContext } from 'sql/parts/notebook/cellViews/codeActions';
 import { NotebookModel } from 'sql/parts/notebook/models/notebookModel';
 import { ToggleMoreWidgetAction } from 'sql/parts/dashboard/common/actions';
 import { CellTypes } from 'sql/parts/notebook/models/contracts';
@@ -97,7 +97,7 @@ export class CodeComponent extends AngularDisposable implements OnInit {
 	get model(): NotebookModel {
 		return this._model;
 	}
-	
+
 	private createEditor(): void {
 		let instantiationService = this._instantiationService.createChild(new ServiceCollection([IProgressService, new SimpleProgressService()]));
 		this._editor = instantiationService.createInstance(QueryTextEditor);
@@ -130,12 +130,12 @@ export class CodeComponent extends AngularDisposable implements OnInit {
 	}
 
 	protected initActionBar() {
-
+		let context = new CellContext(this.model, this.cellModel);
 		let runCellAction = this._instantiationService.createInstance(RunCellAction);
 
 		let taskbar = <HTMLElement>this.toolbarElement.nativeElement;
 		this._actionBar = new Taskbar(taskbar, this.contextMenuService);
-		this._actionBar.context = this;
+		this._actionBar.context = context;
 		this._actionBar.setContent([
 			{ action: runCellAction }
 		]);
@@ -145,13 +145,13 @@ export class CodeComponent extends AngularDisposable implements OnInit {
 		this._moreActions.context = { target: moreActionsElement };
 
 		let actions: Action[] = [];
-		actions.push(this._instantiationService.createInstance(AddCellAction, 'codeBefore', localize('codeBefore', 'Insert Code before'), CellTypes.Code, false, this.notificationService));
-		actions.push(this._instantiationService.createInstance(AddCellAction, 'codeAfter', localize('codeAfter', 'Insert Code after'), CellTypes.Code, true, this.notificationService));
-		actions.push(this._instantiationService.createInstance(AddCellAction, 'markdownBefore', localize('markdownBefore', 'Insert Markdown before'), CellTypes.Markdown, false, this.notificationService));
-		actions.push(this._instantiationService.createInstance(AddCellAction, 'markdownAfter', localize('markdownAfter', 'Insert Markdown after'), CellTypes.Markdown, true, this.notificationService));
-		actions.push(this._instantiationService.createInstance(DeleteCellAction, 'delete', localize('delete', 'Delete'), CellTypes.Code, false, this.notificationService));
+		actions.push(this._instantiationService.createInstance(AddCellAction, 'codeBefore', localize('codeBefore', 'Insert Code before'), CellTypes.Code, false));
+		actions.push(this._instantiationService.createInstance(AddCellAction, 'codeAfter', localize('codeAfter', 'Insert Code after'), CellTypes.Code, true));
+		actions.push(this._instantiationService.createInstance(AddCellAction, 'markdownBefore', localize('markdownBefore', 'Insert Markdown before'), CellTypes.Markdown, false));
+		actions.push(this._instantiationService.createInstance(AddCellAction, 'markdownAfter', localize('markdownAfter', 'Insert Markdown after'), CellTypes.Markdown, true));
+		actions.push(this._instantiationService.createInstance(DeleteCellAction, 'delete', localize('delete', 'Delete')));
 
-		this._moreActions.push(this._instantiationService.createInstance(ToggleMoreWidgetAction, actions, this.model, this.contextMenuService), { icon: true, label: false });
+		this._moreActions.push(this._instantiationService.createInstance(ToggleMoreWidgetAction, actions, context), { icon: true, label: false });
 	}
 
 	private createUri(): URI {

src/sql/parts/notebook/cellViews/code.css

@@ -25,6 +25,17 @@ code-component .toolbarIconRun {
 	background-image: url('../media/dark/execute_cell_inverse.svg');
 }
 
+code-component .toolbarIconStop {
+	height: 20px;
+	background-image: url('../media/light/stop_cell.svg');
+	padding-bottom: 10px;
+}
+
+.vs-dark code-component .toolbarIconStop,
+.hc-black code-component .toolbarIconStop {
+	background-image: url('../media/dark/stop_cell_inverse.svg');
+}
+
 /* overview ruler */
 code-component .monaco-editor .decorationsOverviewRuler {
 	visibility: hidden;

src/sql/parts/notebook/cellViews/codeActions.ts

@@ -1,7 +1,9 @@
 /*---------------------------------------------------------------------------------------------
- *  Copyright (c) Microsoft Corporation. All rights reserved.
- *  Licensed under the Source EULA. See License.txt in the project root for license information.
- *--------------------------------------------------------------------------------------------*/
+*  Copyright (c) Microsoft Corporation. All rights reserved.
+*  Licensed under the Source EULA. See License.txt in the project root for license information.
+*--------------------------------------------------------------------------------------------*/
+
+import { nb } from 'sqlops';
 
 import { Action } from 'vs/base/common/actions';
 import { TPromise } from 'vs/base/common/winjs.base';
@@ -10,97 +12,176 @@ import { CellType } from 'sql/parts/notebook/models/contracts';
 import { NotebookModel } from 'sql/parts/notebook/models/notebookModel';
 import { getErrorMessage } from 'sql/parts/notebook/notebookUtils';
 import { INotificationService, Severity } from 'vs/platform/notification/common/notification';
-import { NOTFOUND } from 'dns';
-import { NsfwWatcherService } from 'vs/workbench/services/files/node/watcher/nsfw/nsfwWatcherService';
+import { ICellModel, FutureInternal } from 'sql/parts/notebook/models/modelInterfaces';
+import { ToggleableAction } from 'sql/parts/notebook/notebookActions';
 
 let notebookMoreActionMsg = localize('notebook.failed', "Please select active cell and try again");
-export class RunCellAction extends Action {
-	public static ID = 'jobaction.notebookRunCell';
+
+function hasModelAndCell(context: CellContext, notificationService: INotificationService): boolean  {
+	if (!context || !context.model) {
+		return false;
+	}
+	if (context.cell === undefined) {
+		notificationService.notify({
+			severity: Severity.Error,
+			message: notebookMoreActionMsg
+		});
+		return false;
+	}
+	return true;
+}
+
+export class CellContext {
+	constructor(public model: NotebookModel, private _cell?: ICellModel) {
+	}
+
+	public get cell(): ICellModel {
+		return this._cell ? this._cell : this.model.activeCell;
+	}
+}
+
+abstract class CellActionBase extends Action {
+
+	constructor(id: string, label: string, icon: string, protected notificationService: INotificationService) {
+		super(id, label, icon);
+	}
+
+	public run(context: CellContext): TPromise<boolean> {
+		if (hasModelAndCell(context, this.notificationService)) {
+			return TPromise.wrap(this.runCellAction(context).then(() => true));
+		}
+		return TPromise.as(true);
+	}
+
+	abstract runCellAction(context: CellContext): Promise<void>;
+}
+
+export class RunCellAction extends ToggleableAction {
+	public static ID = 'notebook.runCell';
 	public static LABEL = 'Run cell';
 
-	constructor(
-	) {
-		super(RunCellAction.ID, '', 'toolbarIconRun');
-		this.tooltip = localize('runCell', 'Run cell');
+	constructor(@INotificationService private notificationService: INotificationService) {
+		super(RunCellAction.ID, {
+			shouldToggleTooltip: true,
+			toggleOnLabel: localize('runCell', 'Run cell'),
+			toggleOnClass: 'toolbarIconRun',
+			toggleOffLabel: localize('stopCell', 'Cancel execution'),
+			toggleOffClass: 'toolbarIconStop',
+			isOn: true
+		});
 	}
 
-	public run(context: any): TPromise<boolean> {
-		return new TPromise<boolean>((resolve, reject) => {
-			try {
-				resolve(true);
-			} catch (e) {
-				reject(e);
+	public run(context: CellContext): TPromise<boolean> {
+		if (hasModelAndCell(context, this.notificationService)) {
+			return TPromise.wrap(this.runCellAction(context).then(() => true));
+		}
+		return TPromise.as(true);
+	}
+
+	public async runCellAction(context: CellContext): Promise<void> {
+        try {
+			let model = context.model;
+			let cell = context.cell;
+            let kernel = await this.getOrStartKernel(model);
+            if (!kernel) {
+                return;
 			}
-		});
+            // If cell is currently running and user clicks the stop/cancel button, call kernel.interrupt()
+            // This matches the same behavior as JupyterLab
+            if (cell.future && cell.future.inProgress) {
+                cell.future.inProgress = false;
+                await kernel.interrupt();
+            } else {
+                // TODO update source based on editor component contents
+                let content = cell.source;
+                if (content) {
+                    this.toggle(false);
+                    let future = await kernel.requestExecute({
+                        code: content,
+                        stop_on_error: true
+                    }, false);
+                    cell.setFuture(future as FutureInternal);
+                    // For now, await future completion. Later we should just track and handle cancellation based on model notifications
+                    let reply = await future.done;
+                }
+            }
+        } catch (error) {
+            let message = getErrorMessage(error);
+            this.notificationService.error(message);
+        } finally {
+            this.toggle(true);
+		}
+    }
+
+    private async getOrStartKernel(model: NotebookModel): Promise<nb.IKernel> {
+        let clientSession = model && model.clientSession;
+        if (!clientSession) {
+            this.notificationService.error(localize('notebookNotReady', 'The session for this notebook is not yet ready'));
+            return undefined;
+        } else if (!clientSession.isReady || clientSession.status === 'dead') {
+            this.notificationService.info(localize('sessionNotReady', 'The session for this notebook will start momentarily'));
+            await clientSession.kernelChangeCompleted;
+        }
+        if (!clientSession.kernel) {
+            let defaultKernel = model && model.defaultKernel && model.defaultKernel.name;
+            if (!defaultKernel) {
+                this.notificationService.error(localize('noDefaultKernel', 'No kernel is available for this notebook'));
+                return undefined;
+            }
+            await clientSession.changeKernel({
+                name: defaultKernel
+            });
+        }
+        return clientSession.kernel;
+    }
+}
+
+export class AddCellAction extends CellActionBase {
+	constructor(
+		id: string, label: string, private cellType: CellType, private isAfter: boolean,
+		@INotificationService notificationService: INotificationService
+	) {
+		super(id, label, undefined, notificationService);
+	}
+
+	runCellAction(context: CellContext): Promise<void> {
+		try {
+			let model = context.model;
+			let index = model.cells.findIndex((cell) => cell.id === context.cell.id);
+			if (index !== undefined && this.isAfter) {
+				index += 1;
+			}
+			model.addCell(this.cellType, index);
+		} catch (error) {
+			let message = getErrorMessage(error);
+
+			this.notificationService.notify({
+				severity: Severity.Error,
+				message: message
+			});
+		}
+		return Promise.resolve();
 	}
 }
 
-export class AddCellAction extends Action {
-	constructor(
-		id: string, label: string, private cellType: CellType, private isAfter: boolean, private notificationService: INotificationService
+export class DeleteCellAction extends CellActionBase {
+	constructor(id: string, label: string,
+		@INotificationService notificationService: INotificationService
 	) {
-		super(id, label);
+		super(id, label, undefined, notificationService);
 	}
-	public run(model: NotebookModel): TPromise<boolean> {
-		return new TPromise<boolean>((resolve, reject) => {
-			try {
-				if (!model) {
-					return;
-				}
-				if (model.activeCell === undefined) {
-					this.notificationService.notify({
-						severity: Severity.Error,
-						message: notebookMoreActionMsg
-					});
-				}
-				else {
-					let index = model.cells.findIndex((cell) => cell.id === model.activeCell.id);
-					if (index !== undefined && this.isAfter) {
-						index += 1;
-					}
-					model.addCell(this.cellType, index);
-				}
-			} catch (error) {
-				let message = getErrorMessage(error);
 
-				this.notificationService.notify({
-					severity: Severity.Error,
-					message: message
-				});
-			}
-		});
-	}
-}
+	runCellAction(context: CellContext): Promise<void> {
+		try {
+				context.model.deleteCell(context.cell);
+		} catch (error) {
+			let message = getErrorMessage(error);
 
-export class DeleteCellAction extends Action {
-	constructor(
-		id: string, label: string, private cellType: CellType, private isAfter: boolean, private notificationService: INotificationService
-	) {
-		super(id, label);
-	}
-	public run(model: NotebookModel): TPromise<boolean> {
-		return new TPromise<boolean>((resolve, reject) => {
-			try {
-				if (!model) {
-					return;
-				}
-				if (model.activeCell === undefined) {
-					this.notificationService.notify({
-						severity: Severity.Error,
-						message: notebookMoreActionMsg
-					});
-				}
-				else {
-					model.deleteCell(model.activeCell);
-				}
-			} catch (error) {
-				let message = getErrorMessage(error);
-
-				this.notificationService.notify({
-					severity: Severity.Error,
-					message: message
-				});
-			}
-		});
+			this.notificationService.notify({
+				severity: Severity.Error,
+				message: message
+			});
+		}
+		return Promise.resolve();
 	}
 }

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

@@ -39,6 +39,11 @@ export class CodeCellComponent extends CellView implements OnInit {
 	ngOnInit() {
 		this._register(this.themeService.onDidColorThemeChange(this.updateTheme, this));
 		this.updateTheme(this.themeService.getColorTheme());
+		if (this.cellModel) {
+			this.cellModel.onOutputsChanged(() => {
+				this._changeRef.detectChanges();
+			});
+		}
 	}
 
 	// Todo: implement layout

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

@@ -3,9 +3,8 @@
 *  Licensed under the Source EULA. See License.txt in the project root for license information.
 *--------------------------------------------------------------------------------------------*/
 import 'vs/css!./code';
-import { OnInit, Component, Input, Inject, forwardRef, ElementRef, ChangeDetectorRef, OnDestroy, ViewChild, Output, EventEmitter } from '@angular/core';
+import { OnInit, Component, Input, Inject, forwardRef, ChangeDetectorRef } from '@angular/core';
 import { AngularDisposable } from 'sql/base/common/lifecycle';
-import { IModeService } from 'vs/editor/common/services/modeService';
 import { ICellModel } from 'sql/parts/notebook/models/modelInterfaces';
 
 export const OUTPUT_AREA_SELECTOR: string = 'output-area-component';
@@ -20,11 +19,15 @@ export class OutputAreaComponent extends AngularDisposable implements OnInit {
 	private readonly _minimumHeight = 30;
 
 	constructor(
-		@Inject(IModeService) private _modeService: IModeService
+		@Inject(forwardRef(() => ChangeDetectorRef)) private _changeRef: ChangeDetectorRef,
 	) {
 		super();
 	}
 	ngOnInit(): void {
-
+		if (this.cellModel) {
+			this.cellModel.onOutputsChanged(() => {
+				this._changeRef.detectChanges();
+			});
+		}
 	}
 }

src/sql/parts/notebook/media/dark/stop_cell_inverse.svg

@@ -0,0 +1 @@
+<svg id="Layer_1" data-name="Layer 1" xmlns="http://www.w3.org/2000/svg" viewBox="0 0 16 16"><defs><style>.cls-1{fill:#fff;}</style></defs><title>stop_cell_inverse</title><circle class="cls-1" cx="8" cy="7.92" r="7.76"/><rect x="5" y="4.92" width="6" height="6"/></svg>

src/sql/parts/notebook/media/light/stop_cell.svg

@@ -0,0 +1 @@
+<svg id="Layer_1" data-name="Layer 1" xmlns="http://www.w3.org/2000/svg" viewBox="0 0 16 16"><defs><style>.cls-1{fill:#fff;}</style></defs><title>stop_cell</title><circle cx="8" cy="7.92" r="7.76"/><rect class="cls-1" x="5" y="4.92" width="6" height="6"/></svg>

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

@@ -10,7 +10,7 @@ import { Event, Emitter } from 'vs/base/common/event';
 import URI from 'vs/base/common/uri';
 
 import { nb } from 'sqlops';
-import { ICellModelOptions, IModelFactory } from './modelInterfaces';
+import { ICellModelOptions, IModelFactory, FutureInternal } from './modelInterfaces';
 import * as notebookUtils from '../notebookUtils';
 import { CellTypes, CellType, NotebookChangeType } from 'sql/parts/notebook/models/contracts';
 import { ICellModel } from 'sql/parts/notebook/models/modelInterfaces';
@@ -24,7 +24,7 @@ export class CellModel implements ICellModel {
 	private _cellType: nb.CellType;
 	private _source: string;
 	private _language: string;
-	private _future: nb.IFuture;
+	private _future: FutureInternal;
 	private _outputs: nb.ICellOutput[] = [];
 	private _isEditMode: boolean;
 	private _onOutputsChanged = new Emitter<ReadonlyArray<nb.ICellOutput>>();
@@ -69,7 +69,7 @@ export class CellModel implements ICellModel {
 		return this._isEditMode;
 	}
 
-	public get future(): nb.IFuture {
+	public get future(): FutureInternal {
 		return this._future;
 	}
 
@@ -137,7 +137,7 @@ export class CellModel implements ICellModel {
 	 * Sets the future which will be used to update the output
 	 * area for this cell
 	 */
-	setFuture(future: nb.IFuture): void {
+	setFuture(future: FutureInternal): void {
 		if (this._future === future) {
 			// Nothing to do
 			return;

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

@@ -299,7 +299,7 @@ export class ClientSession implements IClientSession {
 	public async shutdown(): Promise<void> {
 		// Always try to shut down session
 		if (this._session && this._session.id) {
-			this.notebookManager.sessionManager.shutdown(this._session.id);
+			await this.notebookManager.sessionManager.shutdown(this._session.id);
 		}
 		let serverManager = this.notebookManager.serverManager;
 		if (serverManager) {

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

@@ -338,10 +338,16 @@ export interface ICellModel {
 	cellType: CellType;
 	trustedMode: boolean;
 	active: boolean;
+	readonly future: FutureInternal;
 	readonly outputs: ReadonlyArray<nb.ICellOutput>;
+	readonly onOutputsChanged: Event<ReadonlyArray<nb.ICellOutput>>;
+	setFuture(future: FutureInternal): void;
 	equals(cellModel: ICellModel): boolean;
 	toJSON(): nb.ICell;
-	onOutputsChanged:  Event<ReadonlyArray<nb.ICellOutput>>;
+}
+
+export interface FutureInternal extends nb.IFuture {
+	inProgress: boolean;
 }
 
 export interface IModelFactory {

src/sql/parts/notebook/notebookActions.ts

@@ -43,41 +43,82 @@ export class AddCellAction extends Action {
 	}
 }
 
-export class TrustedAction extends Action {
+export interface IToggleableState {
+	baseClass?: string;
+	shouldToggleTooltip?: boolean;
+	toggleOnClass: string;
+	toggleOnLabel: string;
+	toggleOffLabel: string;
+	toggleOffClass: string;
+	isOn: boolean;
+}
+
+export abstract class ToggleableAction extends Action {
+
+	constructor(id: string, protected state: IToggleableState) {
+		super(id, '');
+		this.updateLabelAndIcon();
+	}
+
+	private updateLabelAndIcon() {
+		if (this.state.shouldToggleTooltip) {
+			this.tooltip = this.state.isOn ? this.state.toggleOnLabel : this.state.toggleOffLabel;
+		} else {
+			this.label = this.state.isOn ? this.state.toggleOnLabel : this.state.toggleOffLabel;
+		}
+		let classes = this.state.baseClass ? `${this.state.baseClass} ` : '';
+		classes += this.state.isOn ? this.state.toggleOnClass : this.state.toggleOffClass;
+		this.class = classes;
+	}
+
+	protected toggle(isOn: boolean): void {
+		this.state.isOn = isOn;
+		this.updateLabelAndIcon();
+	}
+}
+
+export class TrustedAction extends ToggleableAction {
 	// Constants
-	private static readonly trustLabel = localize('trustLabel', 'Trusted');
-	private static readonly notTrustLabel = localize('untrustLabel', 'Not Trusted');
+	private static readonly trustedLabel = localize('trustLabel', 'Trusted');
+	private static readonly notTrustedLabel = localize('untrustLabel', 'Not Trusted');
 	private static readonly alreadyTrustedMsg = localize('alreadyTrustedMsg', 'Notebook is already trusted.');
-	private static readonly trustedCssClass = 'notebook-button icon-trusted';
-	private static readonly notTrustedCssClass = 'notebook-button icon-notTrusted';
+	private static readonly baseClass = 'notebook-button';
+	private static readonly trustedCssClass = 'icon-trusted';
+	private static readonly notTrustedCssClass = 'icon-notTrusted';
+
 	// Properties
-	private _isTrusted: boolean = false;
-	public get trusted(): boolean {
-		return this._isTrusted;
-	}
-	public set trusted(value: boolean) {
-		this._isTrusted = value;
-		this._setClass(value ? TrustedAction.trustedCssClass : TrustedAction.notTrustedCssClass);
-		this._setLabel(value ? TrustedAction.trustLabel : TrustedAction.notTrustLabel);
-	}
 
 	constructor(
 		id: string,
 		@INotificationService private _notificationService: INotificationService
 	) {
-		super(id, TrustedAction.notTrustLabel, TrustedAction.notTrustedCssClass);
+		super(id, {
+			baseClass: TrustedAction.baseClass,
+			toggleOnLabel: TrustedAction.trustedLabel,
+			toggleOnClass: TrustedAction.trustedCssClass,
+			toggleOffLabel: TrustedAction.notTrustedLabel,
+			toggleOffClass: TrustedAction.notTrustedCssClass,
+			isOn: false
+		});
+	}
+
+	public get trusted(): boolean {
+		return this.state.isOn;
+	}
+	public set trusted(value: boolean) {
+		this.toggle(value);
 	}
 
 	public run(context: NotebookComponent): TPromise<boolean> {
 		let self = this;
 		return new TPromise<boolean>((resolve, reject) => {
 			try {
-				if (self._isTrusted) {
+				if (self.trusted) {
 					const actions: INotificationActions = { primary: [] };
 					self._notificationService.notify({ severity: Severity.Info, message: TrustedAction.alreadyTrustedMsg, actions });
 				}
 				else {
-					self.trusted = !self._isTrusted;
+					self.trusted = !self.trusted;
 					context.updateModelTrustDetails(self.trusted);
 				}
 				resolve(true);

src/sql/parts/notebook/outputs/common/outputProcessor.ts

@@ -40,7 +40,7 @@ export function getData(output: nb.ICellOutput): JSONObject {
             bundle['application/vnd.jupyter.stdout'] = output.text;
         }
     } else if (nbformat.isError(output)) {
-        let traceback = output.traceback.join('\n');
+        let traceback = output.traceback ? output.traceback.join('\n') : undefined;
         bundle['application/vnd.jupyter.stderr'] =
             traceback || `${output.ename}: ${output.evalue}`;
     }