Run All Cells Notebook Implementation (#4713)

* runAllCells API

* add comment

* more run cells fixes

* Add integration test

* Add multiple cell SQL notebook test

* Comment out python tests as they fail in the lab

* remove unused imports

* PR comments

* Remove localize

* Return true instead of promise.resolve(true)

extensions/integration-tests/src/notebook.test.ts

@@ -10,7 +10,7 @@ import * as assert from 'assert';
 import * as azdata from 'azdata';
 import * as vscode from 'vscode';
 import { context } from './testContext';
-import { sqlNotebookContent, writeNotebookToFile, sqlKernelMetadata, getFileName, pySparkNotebookContent, pySpark3KernelMetadata, pythonKernelMetadata } from './notebook.util';
+import { sqlNotebookContent, writeNotebookToFile, sqlKernelMetadata, getFileName, pySparkNotebookContent, pySpark3KernelMetadata, pythonKernelMetadata, sqlNotebookMultipleCellsContent } from './notebook.util';
 import { getBdcServer } from './testConfig';
 import { connectToServer } from './utils';
 import * as fs from 'fs';
@@ -50,6 +50,22 @@ if (context.RunTest) {
 			assert(actualOutput2[0] === '1', `Expected result: 1, Actual: '${actualOutput2[0]}'`);
 		});
 
+		test('Sql NB multiple cells test', async function () {
+			let notebook = await openNotebook(sqlNotebookMultipleCellsContent, sqlKernelMetadata, this.test.title);
+			const expectedOutput0 = '(1 row affected)';
+			for (let i = 0; i < 3; i++) {
+				let cellOutputs = notebook.document.cells[i].contents.outputs;
+				console.log('Got cell outputs');
+				assert(cellOutputs.length === 3, `Expected length: 3, Actual: '${cellOutputs.length}'`);
+				let actualOutput0 = (<azdata.nb.IDisplayData>cellOutputs[0]).data['text/html'];
+				console.log('Got first output');
+				assert(actualOutput0 === expectedOutput0, `Expected row count: '${expectedOutput0}', Actual: '${actualOutput0}'`);
+				let actualOutput2 = (<azdata.nb.IExecuteResult>cellOutputs[2]).data['application/vnd.dataresource+json'].data[0];
+				assert(actualOutput2[0] === i.toString(), `Expected result: ${i.toString()}, Actual: '${actualOutput2[0]}'`);
+				console.log('Sql multiple cells NB done');
+			}
+		});
+
 		test('Clear all outputs - SQL notebook ', async function () {
 			let notebook = await openNotebook(sqlNotebookContent, sqlKernelMetadata, this.test.title);
 			await verifyClearAllOutputs(notebook);
@@ -81,7 +97,7 @@ if (context.RunTest) {
 	});
 }
 
-async function openNotebook(content: azdata.nb.INotebookContents, kernelMetadata: any, testName: string): Promise<azdata.nb.NotebookEditor> {
+async function openNotebook(content: azdata.nb.INotebookContents, kernelMetadata: any, testName: string, runAllCells?: boolean): Promise<azdata.nb.NotebookEditor> {
 	let notebookConfig = vscode.workspace.getConfiguration('notebook');
 	notebookConfig.update('pythonPath', process.env.PYTHON_TEST_PATH, 1);
 	let server = await getBdcServer();
@@ -91,12 +107,20 @@ async function openNotebook(content: azdata.nb.INotebookContents, kernelMetadata
 	console.log(uri);
 	let notebook = await azdata.nb.showNotebookDocument(uri);
 	console.log('Notebook is opened');
-	assert(notebook.document.cells.length === 1, 'Notebook should have 1 cell');
-	console.log('Before run notebook cell');
-	let ran = await notebook.runCell(notebook.document.cells[0]);
-	console.log('After run notebook cell');
-	assert(ran, 'Notebook runCell should succeed');
-	assert(notebook !== undefined && notebook !== null, 'Expected notebook object is defined');
+
+	if (!runAllCells) {
+		assert(notebook.document.cells.length === 1, 'Notebook should have 1 cell');
+		console.log('Before run notebook cell');
+		let ran = await notebook.runCell(notebook.document.cells[0]);
+		console.log('After run notebook cell');
+		assert(ran, 'Notebook runCell should succeed');
+	} else {
+		console.log('Before run all notebook cells');
+		let ran = await notebook.runAllCells();
+		assert(ran, 'Notebook runCell should succeed');
+		assert(notebook !== undefined && notebook !== null, 'Expected notebook object is defined');
+	}
+
 	return notebook;
 }
 async function verifyClearAllOutputs(notebook: azdata.nb.NotebookEditor) {

extensions/integration-tests/src/notebook.util.ts

@@ -35,6 +35,38 @@ export const pySparkNotebookContent: azdata.nb.INotebookContents = {
 	nbformat_minor: 2
 };
 
+export const pythonNotebookMultipleCellsContent: azdata.nb.INotebookContents = {
+	cells: [{
+		cell_type: CellTypes.Code,
+		source: '1+1',
+		metadata: { language: 'python' },
+		execution_count: 1
+	}, {
+		cell_type: CellTypes.Code,
+		source: '1+2',
+		metadata: { language: 'python' },
+		execution_count: 1
+	}, {
+		cell_type: CellTypes.Code,
+		source: '1+3',
+		metadata: { language: 'python' },
+		execution_count: 1
+	}, {
+		cell_type: CellTypes.Code,
+		source: '1+4',
+		metadata: { language: 'python' },
+		execution_count: 1
+	}],
+	metadata: {
+		'kernelspec': {
+			'name': 'python3',
+			'display_name': 'Python 3'
+		}
+	},
+	nbformat: 4,
+	nbformat_minor: 2
+};
+
 export const sqlNotebookContent: azdata.nb.INotebookContents = {
 	cells: [{
 		cell_type: CellTypes.Code,
@@ -52,6 +84,33 @@ export const sqlNotebookContent: azdata.nb.INotebookContents = {
 	nbformat_minor: 2
 };
 
+export const sqlNotebookMultipleCellsContent: azdata.nb.INotebookContents = {
+	cells: [{
+		cell_type: CellTypes.Code,
+		source: 'select 0',
+		metadata: { language: 'sql' },
+		execution_count: 1
+	}, {
+		cell_type: CellTypes.Code,
+		source: `WAITFOR DELAY '00:00:02'\nselect 1`,
+		metadata: { language: 'sql' },
+		execution_count: 1
+	}, {
+		cell_type: CellTypes.Code,
+		source: 'select 2',
+		metadata: { language: 'sql' },
+		execution_count: 1
+	}],
+	metadata: {
+		'kernelspec': {
+			'name': 'SQL',
+			'display_name': 'SQL'
+		}
+	},
+	nbformat: 4,
+	nbformat_minor: 2
+};
+
 export const pySpark3KernelMetadata = {
 	'kernelspec': {
 		'name': 'pyspark3kernel',

src/sql/azdata.proposed.d.ts

@@ -4035,7 +4035,11 @@ declare module 'azdata' {
 			runCell(cell?: NotebookCell): Thenable<boolean>;
 
 			/**
-			* Clears the outputs of all code cells in a Notebook
+			 * Kicks off execution of all code cells. Thenable will resolve only when full execution of all cells is completed.
+			 */
+			runAllCells(): Thenable<boolean>;
+
+			/** Clears the outputs of all code cells in a Notebook
 			* @return A promise that resolves with a value indicating if the outputs are cleared or not.
 			*/
 			clearAllOutputs(): Thenable<boolean>;

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

@@ -0,0 +1 @@
+<svg id="Layer_1" data-name="Layer 1" xmlns="http://www.w3.org/2000/svg" xmlns:xlink="http://www.w3.org/1999/xlink" viewBox="0 0 16 16"><defs><style>.cls-1{fill:none;}.cls-2{fill:#fff;}.cls-3{clip-path:url(#clip-path);}</style><clipPath id="clip-path"><rect class="cls-1" x="-552.55" y="-7.09" width="15.81" height="13.48"/></clipPath></defs><title>run_cells_inverse</title><path class="cls-2" d="M8,15.92a8,8,0,1,1,8-8A8,8,0,0,1,8,15.92ZM8,1a7,7,0,1,0,7,7A7,7,0,0,0,8,1Z"/><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/run_cells.svg

@@ -0,0 +1 @@
+<svg id="Layer_1" data-name="Layer 1" xmlns="http://www.w3.org/2000/svg" xmlns:xlink="http://www.w3.org/1999/xlink" viewBox="0 0 16 16"><defs><style>.cls-1{fill:none;}.cls-2{clip-path:url(#clip-path);}</style><clipPath id="clip-path"><rect class="cls-1" x="-579.3" y="-7.09" width="15.81" height="13.48"/></clipPath></defs><title>run_cells</title><path d="M8,15.92a8,8,0,1,1,8-8A8,8,0,0,1,8,15.92ZM8,1a7,7,0,1,0,7,7A7,7,0,0,0,8,1Z"/><polygon points="10.7 8 6.67 11 6.67 5 10.7 8 10.7 8"/></svg>

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

@@ -33,7 +33,7 @@ import * as notebookUtils from 'sql/parts/notebook/notebookUtils';
 import { Deferred } from 'sql/base/common/promise';
 import { IConnectionProfile } from 'sql/platform/connection/common/interfaces';
 import { Taskbar } from 'sql/base/browser/ui/taskbar/taskbar';
-import { KernelsDropdown, AttachToDropdown, AddCellAction, TrustedAction, ClearAllOutputsAction } from 'sql/parts/notebook/notebookActions';
+import { KernelsDropdown, AttachToDropdown, AddCellAction, TrustedAction, RunAllCellsAction, ClearAllOutputsAction } from 'sql/parts/notebook/notebookActions';
 import { IObjectExplorerService } from 'sql/workbench/services/objectExplorer/common/objectExplorerService';
 import * as TaskUtilities from 'sql/workbench/common/taskUtilities';
 import { ISingleNotebookEditOperation } from 'sql/workbench/api/common/sqlExtHostTypes';
@@ -63,6 +63,7 @@ export class NotebookComponent extends AngularDisposable implements OnInit, OnDe
 	private _modelRegisteredDeferred = new Deferred<NotebookModel>();
 	private profile: IConnectionProfile;
 	private _trustedAction: TrustedAction;
+	private _runAllCellsAction: RunAllCellsAction;
 	private _providerRelatedActions: IAction[] = [];
 	private _scrollTop: number;
 
@@ -372,6 +373,7 @@ export class NotebookComponent extends AngularDisposable implements OnInit, OnDe
 		let addTextCellButton = new AddCellAction('notebook.AddTextCell', localize('text', 'Text'), 'notebook-button icon-add');
 		addTextCellButton.cellType = CellTypes.Markdown;
 
+		this._runAllCellsAction = new RunAllCellsAction('notebook.runAllCells', localize('runAll', 'Run Cells'), 'notebook-button icon-run-cells');
 		let clearResultsButton = new ClearAllOutputsAction('notebook.ClearAllOutputs', localize('clearResults', 'Clear Results'), 'notebook-button icon-clear-results');
 
 		this._trustedAction = this.instantiationService.createInstance(TrustedAction, 'notebook.Trusted');
@@ -386,6 +388,7 @@ export class NotebookComponent extends AngularDisposable implements OnInit, OnDe
 			{ action: addCodeCellButton },
 			{ action: addTextCellButton },
 			{ action: this._trustedAction },
+			{ action: this._runAllCellsAction },
 			{ action: clearResultsButton }
 		]);
 
@@ -483,6 +486,20 @@ export class NotebookComponent extends AngularDisposable implements OnInit, OnDe
 		}
 	}
 
+	public async runAllCells(): Promise<boolean> {
+		await this.modelReady;
+		let codeCells = this._model.cells.filter(cell => cell.cellType === CellTypes.Code);
+		if (codeCells && codeCells.length) {
+			for (let i = 0; i < codeCells.length; i++) {
+				let cellStatus = await this.runCell(codeCells[i]);
+				if (!cellStatus) {
+					return Promise.reject(new Error(localize('cellRunFailed', 'running cell id {0} failed', codeCells[i].id)));
+				}
+			}
+		}
+		return true;
+	}
+
 	public async clearAllOutputs(): Promise<boolean> {
 		try {
 			await this.modelReady;

src/sql/parts/notebook/notebook.css

@@ -50,6 +50,15 @@
 	background-image: url("./media/dark/add_inverse.svg");
 }
 
+.notebookEditor .notebook-button.icon-run-cells{
+	background-image: url("./media/light/run_cells.svg");
+}
+
+.vs-dark .notebookEditor .notebook-button.icon-run-cells,
+.hc-black .notebookEditor .notebook-button.icon-run-cells{
+	background-image: url("./media/dark/run_cells_inverse.svg");
+}
+
 .notebookEditor .notebook-button.icon-trusted{
 	background-image: url("./media/light/trusted.svg");
 }

src/sql/parts/notebook/notebookActions.ts

@@ -227,6 +227,25 @@ export class TrustedAction extends ToggleableAction {
 	}
 }
 
+// Action to run all code cells in a notebook.
+export class RunAllCellsAction extends Action {
+	constructor(
+		id: string, label: string, cssClass: string
+	) {
+		super(id, label, cssClass);
+	}
+	public run(context: NotebookComponent): Promise<boolean> {
+		return new Promise<boolean>((resolve, reject) => {
+			try {
+				context.runAllCells();
+				resolve(true);
+			} catch (e) {
+				reject(e);
+			}
+		});
+	}
+}
+
 export class KernelsDropdown extends SelectBox {
 	private model: NotebookModel;
 	constructor(container: HTMLElement, contextViewProvider: IContextViewProvider, modelReady: Promise<INotebookModel>) {

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

@@ -156,6 +156,10 @@ export class ExtHostNotebookEditor implements azdata.nb.NotebookEditor, IDisposa
 		return this._proxy.$runCell(this._id, uri);
 	}
 
+	public runAllCells(): Thenable<boolean> {
+		return this._proxy.$runAllCells(this._id);
+	}
+
 	public clearAllOutputs(): Thenable<boolean> {
 		return this._proxy.$clearAllOutputs(this._id);
 	}

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

@@ -122,10 +122,16 @@ class MainThreadNotebookEditor extends Disposable {
 		if (!this.editor) {
 			return Promise.resolve(false);
 		}
-
 		return this.editor.runCell(cell);
 	}
 
+	public runAllCells(): Promise<boolean> {
+		if (!this.editor) {
+			return Promise.resolve(false);
+		}
+		return this.editor.runAllCells();
+	}
+
 	public clearAllOutputs(): Promise<boolean> {
 		if (!this.editor) {
 			return Promise.resolve(false);
@@ -366,6 +372,14 @@ export class MainThreadNotebookDocumentsAndEditors extends Disposable implements
 		return editor.runCell(cell);
 	}
 
+	$runAllCells(id: string): Promise<boolean> {
+		let editor = this.getEditor(id);
+		if (!editor) {
+			return Promise.reject(disposed(`TextEditor(${id})`));
+		}
+		return editor.runAllCells();
+	}
+
 	$clearAllOutputs(id: string): Promise<boolean> {
 		let editor = this.getEditor(id);
 		if (!editor) {

src/sql/workbench/api/node/sqlExtHost.protocol.ts

@@ -866,6 +866,7 @@ export interface MainThreadNotebookDocumentsAndEditorsShape extends IDisposable
 	$tryShowNotebookDocument(resource: UriComponents, options: INotebookShowOptions): Promise<string>;
 	$tryApplyEdits(id: string, modelVersionId: number, edits: ISingleNotebookEditOperation[], opts: IUndoStopOptions): Promise<boolean>;
 	$runCell(id: string, cellUri: UriComponents): Promise<boolean>;
+	$runAllCells(id: string): Promise<boolean>;
 	$clearAllOutputs(id: string): Promise<boolean>;
 }
 

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

@@ -109,5 +109,6 @@ export interface INotebookEditor {
 	isVisible(): boolean;
 	executeEdits(edits: ISingleNotebookEditOperation[]): boolean;
 	runCell(cell: ICellModel): Promise<boolean>;
+	runAllCells(): Promise<boolean>;
 	clearAllOutputs(): Promise<boolean>;
 }