Support execution count in run button and align correctly (#3979)

Fixes #3931 
- Align run button correctly so it's centered in new cell
- Refactor to support multi-state button.
  - Hidden state is set to show execution count
  - Stopped state shows run button
  - Running state shows stop button
  - Error state (will) show error button. This isn't fully handled right now
- Add execution count to model and to SqlKernel, verify serialization, loading, update matches other notebook viewers

**Notes on implementation**:
I think this is a decent solution for a) showing execution count as text, and b) perfectly centering the run button.
 
The below solution shows count correctly up to 999 runs (that’s clicking 999 times in a single session), the icon lines up juuust about right with [ ] but for other numbers it is pretty close but probably not exactly right. I wish I could solve this to work better but trying to change text float to center etc. really isn’t working.
 
**Screenshots**:
![image](https://user-images.githubusercontent.com/10819925/52466366-e8794200-2b36-11e9-9a50-86893e75d5af.png)

With running cell:
![image](https://user-images.githubusercontent.com/10819925/52466378-f333d700-2b36-11e9-9e6c-3cee098790fd.png)

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

@@ -5,7 +5,7 @@
  *--------------------------------------------------------------------------------------------*/
 -->
 <div style="width: 100%; height: 100%; display: flex; flex-flow: row" (mouseover)="hover=true" (mouseleave)="hover=false">
-	<div #toolbar class="toolbar" style="flex: 0 0 auto; display: flex; flex-flow:column; width: 40px; min-height: 40px; orientation: portrait">
+	<div #toolbar class="toolbar">
 	</div>
 	<div #editor class="editor" style="flex: 1 1 auto; overflow: hidden;">
 	</div>

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

@@ -55,10 +55,10 @@ export class CodeComponent extends AngularDisposable implements OnInit, OnChange
 	}
 
 	@Input() set hover(value: boolean) {
-		this._hover = value;
+		this.cellModel.hover = value;
 		if (!this.isActive()) {
 			// Only make a change if we're not active, since this has priority
-			this.toggleMoreActionsButton(this._hover);
+			this.toggleMoreActionsButton(this.cellModel.hover);
 		}
 	}
 
@@ -72,7 +72,6 @@ export class CodeComponent extends AngularDisposable implements OnInit, OnChange
 	private _model: NotebookModel;
 	private _activeCellId: string;
 	private _cellToggleMoreActions: CellToggleMoreActions;
-	private _hover: boolean;
 
 	constructor(
 		@Inject(forwardRef(() => CommonServiceInterface)) private _bootstrapService: CommonServiceInterface,

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

@@ -11,12 +11,18 @@ code-component {
 
 code-component .toolbar {
 	border-right-width: 1px;
+	flex: 0 0 auto;
+	display: flex;
+	flex-flow:column;
+	width: 40px;
+	min-height: 40px;
+	orientation: portrait
 }
 
 code-component .toolbar .carbon-taskbar {
 	position: sticky;
-	top: 10px;
-	bottom: 30px;
+	top: 0px;
+	margin-top: 5px;
 }
 
 code-component .toolbarIconRun {
@@ -30,15 +36,21 @@ code-component .toolbarIconRun {
 	background-image: url('../media/dark/execute_cell_inverse.svg');
 }
 
+code-component .toolbarIconRunError {
+	height: 20px;
+	background-image: url('../media/light/execute_cell_error.svg');
+	padding-bottom: 10px;
+}
+
 code-component .toolbarIconStop {
 	height: 20px;
-	background-image: url('../media/light/stop_cell.svg');
+	background-image: url('../media/light/stop_cell_solidanimation.svg');
 	padding-bottom: 10px;
 }
 
 .vs-dark code-component .toolbarIconStop,
 .hc-black code-component .toolbarIconStop {
-	background-image: url('../media/dark/stop_cell_inverse.svg');
+	background-image: url('../media/dark/stop_cell_solidanimation_inverse.svg');
 }
 
 code-component .editor {
@@ -54,6 +66,21 @@ code-component .carbon-taskbar .icon {
 	width: 40px;
 }
 
+code-component .carbon-taskbar .icon.hideIcon {
+	width: 0px;
+	padding-left: 0px;
+	padding-top: 6px;
+	font-family: monospace;
+	font-size: 12px;
+}
+
+code-component .carbon-taskbar .icon.hideIcon.execCountTen {
+	margin-left: -2px;
+}
+
+code-component .carbon-taskbar .icon.hideIcon.execCountHundred {
+	margin-left: -6px;
+}
 
 code-component .carbon-taskbar.monaco-toolbar .monaco-action-bar.animated .actions-container
 {

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

@@ -8,14 +8,17 @@ import { nb } from 'sqlops';
 import { Action } from 'vs/base/common/actions';
 import { TPromise } from 'vs/base/common/winjs.base';
 import { localize } from 'vs/nls';
+import { IDisposable } from 'vs/base/common/lifecycle';
+import * as types from 'vs/base/common/types';
+import { INotificationService, Severity } from 'vs/platform/notification/common/notification';
+
 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 { ICellModel } from 'sql/parts/notebook/models/modelInterfaces';
-import { ToggleableAction } from 'sql/parts/notebook/notebookActions';
-import { IDisposable } from 'vs/base/common/lifecycle';
+import { ICellModel, CellExecutionState } from 'sql/parts/notebook/models/modelInterfaces';
+import { MultiStateAction, IMultiStateData, IActionStateData } from 'sql/parts/notebook/notebookActions';
 
 let notebookMoreActionMsg = localize('notebook.failed', "Please select active cell and try again");
+const emptyExecutionCountLabel = '[ ]';
 
 function hasModelAndCell(context: CellContext, notificationService: INotificationService): boolean  {
 	if (!context || !context.model) {
@@ -60,28 +63,21 @@ export abstract class CellActionBase extends Action {
 	abstract doRun(context: CellContext): Promise<void>;
 }
 
-export class RunCellAction extends ToggleableAction {
+export class RunCellAction extends MultiStateAction<CellExecutionState> {
 	public static ID = 'notebook.runCell';
 	public static LABEL = 'Run cell';
 	private _executionChangedDisposable: IDisposable;
 	private _context: CellContext;
 	constructor(context: CellContext, @INotificationService private notificationService: INotificationService) {
-		super(RunCellAction.ID, {
-			shouldToggleTooltip: true,
-			toggleOffLabel: localize('runCell', 'Run cell'),
-			toggleOffClass: 'toolbarIconRun',
-			toggleOnLabel: localize('stopCell', 'Cancel execution'),
-			toggleOnClass: 'toolbarIconStop',
-			// On == running
-			isOn: false
-		});
+		super(RunCellAction.ID, new IMultiStateData<CellExecutionState>([
+			{ key: CellExecutionState.Hidden, value: { label: emptyExecutionCountLabel, className: '', tooltip: '', hideIcon: true }},
+			{ key: CellExecutionState.Stopped, value: { label: '', className: 'toolbarIconRun', tooltip: localize('runCell', 'Run cell') }},
+			{ key: CellExecutionState.Running, value: { label: '', className: 'toolbarIconStop', tooltip: localize('stopCell', 'Cancel execution') }},
+			{ key: CellExecutionState.Error, value: { label: '', className: 'toolbarIconRunError', tooltip: localize('errorRunCell', 'Error on last run. Click to run again') }},
+		], CellExecutionState.Hidden ));
 		this.ensureContextIsUpdated(context);
 	}
 
-	private _handleExecutionStateChange(running: boolean): void {
-		this.toggle(running);
-	}
-
 	public run(context?: CellContext): TPromise<boolean> {
 		return TPromise.wrap(this.doRun(context).then(() => true));
 	}
@@ -106,8 +102,30 @@ export class RunCellAction extends ToggleableAction {
 				this._executionChangedDisposable.dispose();
 			}
 			this._context = context;
-			this.toggle(context.cell.isRunning);
-			this._executionChangedDisposable = this._context.cell.onExecutionStateChange(this._handleExecutionStateChange, this);
+			this.updateStateAndExecutionCount(context.cell.executionState);
+			this._executionChangedDisposable = this._context.cell.onExecutionStateChange((state) => {
+				this.updateStateAndExecutionCount(state);
+			});
 		}
 	}
+
+	private updateStateAndExecutionCount(state: CellExecutionState) {
+		let label = emptyExecutionCountLabel;
+		let className = '';
+		if (!types.isUndefinedOrNull(this._context.cell.executionCount)) {
+			label = `[${this._context.cell.executionCount}]`;
+			// Heuristic to try and align correctly independent of execution count length. Moving left margin
+			// back by a few px seems to make things "work" OK, but isn't a super clean solution
+			if (label.length === 4) {
+				className = 'execCountTen';
+			} else if (label.length > 4) {
+				className = 'execCountHundred';
+			}
+		}
+		this.states.updateStateData(CellExecutionState.Hidden, (data) => {
+			data.label = label;
+			data.className = className;
+		});
+		this.updateState(state);
+	}
 }

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

@@ -0,0 +1,16 @@
+<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:#c1d6e6;}.cls-2{fill:#0078d4;}.cls-3{fill:#fff;}</style></defs>
+<title>stop_cell_solidanimation_inverse</title>
+<path class="cls-1" d="M8,16a8,8,0,1,1,8-8A8,8,0,0,1,8,16ZM8,1a7,7,0,1,0,7,7A7,7,0,0,0,8,1Z"/>
+<path class="cls-2" d="M8.51,0v1A7,7,0,0,1,15,8a6.87,6.87,0,0,1-1.07,3.7l.81.64A7.92,7.92,0,0,0,16,8,8,8,0,0,0,8.51,0Z">
+<animateTransform attributeName="transform"
+            type="rotate"
+            from="0 8 8"
+            to="360 8 8"
+            begin="0s"
+            dur="1.5s"
+            repeatCount="indefinite"
+        />
+</path>
+<circle class="cls-3" cx="8" cy="8" r="6.32"/>
+<rect x="4.91" y="4.91" width="6.18" height="6.18"/></svg>

src/sql/parts/notebook/media/light/execute_cell_error.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:#d02e00;}.cls-2{fill:#fff;}</style></defs><title>execute_cell_error</title><circle class="cls-1" cx="8" cy="7.92" r="7.76"/><polygon class="cls-2" points="10.7 8 6.67 11 6.67 5 10.7 8 10.7 8"/></svg>

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

@@ -0,0 +1,16 @@
+<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:#c1d6e6;}.cls-2{fill:#0078d4;}.cls-3{fill:#fff;}</style></defs>
+<title>stop_cell_solidanimation</title>
+<path class="cls-1" d="M8,16a8,8,0,1,1,8-8A8,8,0,0,1,8,16ZM8,1a7,7,0,1,0,7,7A7,7,0,0,0,8,1Z"/>
+<path class="cls-2" d="M8.51,0v1A7,7,0,0,1,15,8a6.87,6.87,0,0,1-1.07,3.7l.81.64A7.92,7.92,0,0,0,16,8,8,8,0,0,0,8.51,0Z">
+<animateTransform attributeName="transform"
+            type="rotate"
+            from="0 8 8"
+            to="360 8 8"
+            begin="0s"
+            dur="1.5s"
+            repeatCount="indefinite"
+        />
+</path>
+<circle cx="8" cy="8" r="6.32"/>
+<rect class="cls-3" x="4.91" y="4.91" width="6.18" height="6.18"/></svg>

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

@@ -12,7 +12,7 @@ import { Event, Emitter } from 'vs/base/common/event';
 import URI from 'vs/base/common/uri';
 import { localize } from 'vs/nls';
 
-import { ICellModelOptions, IModelFactory, FutureInternal } from './modelInterfaces';
+import { ICellModelOptions, IModelFactory, FutureInternal, CellExecutionState } 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';
@@ -30,11 +30,13 @@ export class CellModel implements ICellModel {
 	private _isEditMode: boolean;
 	private _onOutputsChanged = new Emitter<ReadonlyArray<nb.ICellOutput>>();
 	private _onCellModeChanged = new Emitter<boolean>();
-	private _onExecutionStateChanged = new Emitter<boolean>();
-	public id: string;
+	private _onExecutionStateChanged = new Emitter<CellExecutionState>();
 	private _isTrusted: boolean;
 	private _active: boolean;
+	private _hover: boolean;
+	private _executionCount: number | undefined;
 	private _cellUri: URI;
+	public id: string;
 
 	constructor(private factory: IModelFactory, cellData?: nb.ICellContents, private _options?: ICellModelOptions) {
 		this.id = `${modelId++}`;
@@ -97,6 +99,25 @@ export class CellModel implements ICellModel {
 
 	public set active(value: boolean) {
 		this._active = value;
+		this.fireExecutionStateChanged();
+	}
+
+	public get hover(): boolean {
+		return this._hover;
+	}
+
+	public set hover(value: boolean) {
+		this._hover = value;
+		this.fireExecutionStateChanged();
+	}
+
+	public get executionCount(): number | undefined {
+		return this._executionCount;
+	}
+
+	public set executionCount(value: number | undefined) {
+		this._executionCount = value;
+		this.fireExecutionStateChanged();
 	}
 
 	public get cellUri(): URI {
@@ -134,12 +155,23 @@ export class CellModel implements ICellModel {
 		this._language = newLanguage;
 	}
 
-	public get onExecutionStateChange(): Event<boolean> {
+	public get onExecutionStateChange(): Event<CellExecutionState> {
 		return this._onExecutionStateChanged.event;
 	}
 
-	public get isRunning(): boolean {
-		return !!(this._future && this._future.inProgress);
+	private fireExecutionStateChanged(): void {
+		this._onExecutionStateChanged.fire(this.executionState);
+	}
+
+	public get executionState(): CellExecutionState {
+		let isRunning = !!(this._future && this._future.inProgress);
+		if (isRunning) {
+			return CellExecutionState.Running;
+		} else if (this.active || this.hover) {
+			return CellExecutionState.Stopped;
+		}
+		// TODO save error state and show the error
+		return CellExecutionState.Hidden;
 	}
 
 	public async runCell(notificationService?: INotificationService): Promise<boolean> {
@@ -162,7 +194,7 @@ export class CellModel implements ICellModel {
 				// TODO update source based on editor component contents
 				let content = this.source;
 				if (content) {
-					this._onExecutionStateChanged.fire(true);
+					this.fireExecutionStateChanged();
 					let future = await kernel.requestExecute({
 						code: content,
 						stop_on_error: true
@@ -170,7 +202,13 @@ export class CellModel implements ICellModel {
 					this.setFuture(future as FutureInternal);
 					// For now, await future completion. Later we should just track and handle cancellation based on model notifications
 					let result: nb.IExecuteReplyMsg = <nb.IExecuteReplyMsg><any> await future.done;
-					return result && result.content.status === 'ok' ? true : false;
+					if (result && result.content) {
+						this.executionCount = result.content.execution_count;
+						if (result.content.status !== 'ok') {
+							// TODO track error state
+							return false;
+						}
+					}
 				}
 			}
 		} catch (error) {
@@ -179,9 +217,10 @@ export class CellModel implements ICellModel {
 			}
 			let message = notebookUtils.getErrorMessage(error);
 			this.sendNotification(notificationService, Severity.Error, message);
+			// TODO track error state for the cell
 			throw error;
 		} finally {
-			this._onExecutionStateChanged.fire(false);
+			this.fireExecutionStateChanged();
 		}
 
 		return true;
@@ -361,7 +400,7 @@ export class CellModel implements ICellModel {
 		if (this._cellType === CellTypes.Code) {
 			cellJson.metadata.language = this._language,
 			cellJson.outputs = this._outputs;
-			cellJson.execution_count = 1; // TODO: keep track of actual execution count
+			cellJson.execution_count = this.executionCount;
 		}
 		return cellJson as nb.ICellContents;
 	}
@@ -371,6 +410,7 @@ export class CellModel implements ICellModel {
 			return;
 		}
 		this._cellType = cell.cell_type;
+		this.executionCount = cell.execution_count;
 		this._source = Array.isArray(cell.source) ? cell.source.join('') : cell.source;
 		this.setLanguageFromContents(cell);
 		if (cell.outputs) {

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

@@ -411,6 +411,13 @@ export interface ICellModelOptions {
 	isTrusted: boolean;
 }
 
+export enum CellExecutionState {
+	Hidden = 0,
+	Stopped = 1,
+	Running = 2,
+	Error = 3
+}
+
 export interface ICellModel {
 	cellUri: URI;
 	id: string;
@@ -419,12 +426,14 @@ export interface ICellModel {
 	cellType: CellType;
 	trustedMode: boolean;
 	active: boolean;
+	hover: boolean;
+	executionCount: number | undefined;
 	readonly future: FutureInternal;
 	readonly outputs: ReadonlyArray<nb.ICellOutput>;
 	readonly onOutputsChanged: Event<ReadonlyArray<nb.ICellOutput>>;
-	readonly onExecutionStateChange: Event<boolean>;
+	readonly onExecutionStateChange: Event<CellExecutionState>;
 	setFuture(future: FutureInternal): void;
-	readonly isRunning: boolean;
+	readonly executionState: CellExecutionState;
 	runCell(notificationService?: INotificationService): Promise<boolean>;
 	equals(cellModel: ICellModel): boolean;
 	toJSON(): nb.ICellContents;

src/sql/parts/notebook/notebookActions.ts

@@ -28,6 +28,7 @@ const msgLoadingContexts = localize('loadingContexts', 'Loading contexts...');
 const msgAddNewConnection = localize('addNewConnection', 'Add new connection');
 const msgSelectConnection = localize('selectConnection', 'Select connection');
 const msgLocalHost = localize('localhost', 'localhost');
+const HIDE_ICON_CLASS = ' hideIcon';
 
 // Action to add a cell to notebook based on cell type(code/markdown).
 export class AddCellAction extends Action {
@@ -104,6 +105,80 @@ export abstract class ToggleableAction extends Action {
 	}
 }
 
+
+export interface IActionStateData {
+	className?: string;
+	label?: string;
+	tooltip?: string;
+	hideIcon?: boolean;
+}
+
+export class IMultiStateData<T> {
+	private _stateMap = new Map<T, IActionStateData>();
+	constructor(mappings: { key: T, value: IActionStateData}[], private _state: T, private _baseClass?: string) {
+		if (mappings) {
+			mappings.forEach(s => this._stateMap.set(s.key, s.value));
+		}
+	}
+
+	public set state(value: T) {
+		if (!this._stateMap.has(value)) {
+			throw new Error('State value must be in stateMap');
+		}
+		this._state = value;
+	}
+
+	public updateStateData(state: T, updater: (data: IActionStateData) => void): void {
+		let data = this._stateMap.get(state);
+		if (data) {
+			updater(data);
+		}
+	}
+
+	public get classes(): string {
+		let classVal = this.getStateValueOrDefault<string>((data) => data.className, '');
+		let classes = this._baseClass ? `${this._baseClass} ` : '';
+		classes += classVal;
+		if (this.getStateValueOrDefault<boolean>((data) => data.hideIcon, false)) {
+			classes += HIDE_ICON_CLASS;
+		}
+		return classes;
+	}
+
+	public get label(): string {
+		return this.getStateValueOrDefault<string>((data) => data.label, '');
+	}
+
+	public get tooltip(): string {
+		return this.getStateValueOrDefault<string>((data) => data.tooltip, '');
+	}
+
+	private getStateValueOrDefault<U>(getter: (data: IActionStateData) => U, defaultVal?: U): U {
+		let data = this._stateMap.get(this._state);
+		return data ? getter(data) : defaultVal;
+	}
+}
+
+
+export abstract class MultiStateAction<T> extends Action {
+
+	constructor(id: string, protected states: IMultiStateData<T>) {
+		super(id, '');
+		this.updateLabelAndIcon();
+	}
+
+	private updateLabelAndIcon() {
+		this.label = this.states.label;
+		this.tooltip = this.states.tooltip;
+		this.class = this.states.classes;
+	}
+
+	protected updateState(state: T): void {
+		this.states.state = state;
+		this.updateLabelAndIcon();
+	}
+}
+
 export class TrustedAction extends ToggleableAction {
 	// Constants
 	private static readonly trustedLabel = localize('trustLabel', 'Trusted');

src/sql/workbench/api/node/mainThreadNotebookDocumentsAndEditors.ts

@@ -542,7 +542,7 @@ export class MainThreadNotebookDocumentsAndEditors extends Disposable implements
 					uri: cell.cellUri,
 					contents: {
 						cell_type: cell.cellType,
-						execution_count: undefined,
+						execution_count: cell.executionCount,
 						metadata: {
 							language: cell.language
 						},

src/sql/workbench/services/notebook/common/sqlSessionManager.ts

@@ -131,6 +131,7 @@ class SqlKernel extends Disposable implements nb.IKernel {
 
 	private _id: string;
 	private _future: SQLFuture;
+	private _executionCount: number = 0;
 
 	constructor( @IConnectionManagementService private _connectionManagementService: IConnectionManagementService,
 		@IInstantiationService private _instantiationService: IInstantiationService,
@@ -212,7 +213,11 @@ class SqlKernel extends Disposable implements nb.IKernel {
 			canRun = false;
 		}
 
-		this._future = new SQLFuture(this._queryRunner);
+		// Only update execution count if this will run. if not, set as undefined in future so cell isn't shown as having run?
+		// TODO verify this is "canonical" behavior
+		let count = canRun ? ++this._executionCount : undefined;
+
+		this._future = new SQLFuture(this._queryRunner, count);
 		if (!canRun) {
 			// Complete early
 			this._future.handleDone(new Error(localize('connectionRequired', 'A connection must be chosen to run notebook cells')));
@@ -284,7 +289,7 @@ export class SQLFuture extends Disposable implements FutureInternal {
 	private doneHandler: nb.MessageHandler<nb.IShellMessage>;
 	private doneDeferred = new Deferred<nb.IShellMessage>();
 
-	constructor(private _queryRunner: QueryRunner) {
+	constructor(private _queryRunner: QueryRunner, private _executionCount: number | undefined) {
 		super();
 	}
 
@@ -305,10 +310,13 @@ export class SQLFuture extends Disposable implements FutureInternal {
 	}
 
 	public handleDone(err?: Error): void {
-		let msg: nb.IShellMessage = {
+		let msg: nb.IExecuteReplyMsg = {
 			channel: 'shell',
 			type: 'execute_reply',
-			content: { status: 'ok' },
+			content: {
+				status: 'ok',
+				execution_count: this._executionCount
+			},
 			header: undefined,
 			metadata: {},
 			parent_header: undefined