Notebooks: Run all after/before (#6239)

* Run all above/below

* PR comments pre tests

* Added integration test

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

@@ -15,6 +15,7 @@ import { getBdcServer, getConfigValue, EnvironmentVariable_PYTHON_PATH } from '.
 import { connectToServer, sleep } from './utils';
 import * as fs from 'fs';
 import { stressify } from 'adstest';
+import { isNullOrUndefined } from 'util';
 
 if (context.RunTest) {
 	suite('Notebook integration test suite', function () {
@@ -36,6 +37,10 @@ if (context.RunTest) {
 			await (new NotebookTester()).sqlNbMultipleCellsTest(this.test.title);
 		});
 
+		test('Sql NB run cells above and below test', async function () {
+			await (new NotebookTester()).sqlNbRunCellsAboveBelowTest(this.test.title);
+		});
+
 		test('Clear cell output - SQL notebook', async function () {
 			await (new NotebookTester()).sqlNbClearOutputs(this.test.title);
 		});
@@ -108,6 +113,7 @@ class NotebookTester {
 	@stressify({ dop: NotebookTester.ParallelCount })
 	async pySpark3NbTest(title: string): Promise<void> {
 		let notebook = await this.openNotebook(pySparkNotebookContent, pySpark3KernelMetadata, title + this.invocationCount++);
+		await this.runCell(notebook);
 		let cellOutputs = notebook.document.cells[0].contents.outputs;
 		let sparkResult = (<azdata.nb.IStreamResult>cellOutputs[3]).text;
 		assert(sparkResult === '2', `Expected spark result: 2, Actual: ${sparkResult}`);
@@ -116,12 +122,14 @@ class NotebookTester {
 	@stressify({ dop: NotebookTester.ParallelCount })
 	async python3ClearAllOutputs(title: string): Promise<void> {
 		let notebook = await this.openNotebook(pySparkNotebookContent, pythonKernelMetadata, title + this.invocationCount++);
+		await this.runCell(notebook);
 		await this.verifyClearAllOutputs(notebook);
 	}
 
 	@stressify({ dop: NotebookTester.ParallelCount })
 	async python3NbTest(title: string): Promise<void> {
 		let notebook = await this.openNotebook(pySparkNotebookContent, pythonKernelMetadata, title + this.invocationCount++);
+		await this.runCell(notebook);
 		let cellOutputs = notebook.document.cells[0].contents.outputs;
 		console.log('Got cell outputs ---');
 		if (cellOutputs) {
@@ -134,17 +142,20 @@ class NotebookTester {
 	@stressify({ dop: NotebookTester.ParallelCount })
 	async sqlNbClearAllOutputs(title: string): Promise<void> {
 		let notebook = await this.openNotebook(sqlNotebookContent, sqlKernelMetadata, title + this.invocationCount++);
+		await this.runCell(notebook);
 		await this.verifyClearAllOutputs(notebook);
 	}
 
 	async sqlNbClearOutputs(title: string): Promise<void> {
 		let notebook = await this.openNotebook(sqlNotebookContent, sqlKernelMetadata, title + this.invocationCount++);
+		await this.runCell(notebook);
 		await this.verifyClearOutputs(notebook);
 	}
 
 	@stressify({ dop: NotebookTester.ParallelCount })
 	async sqlNbMultipleCellsTest(title: string): Promise<void> {
-		let notebook = await this.openNotebook(sqlNotebookMultipleCellsContent, sqlKernelMetadata, title + this.invocationCount++, true);
+		let notebook = await this.openNotebook(sqlNotebookMultipleCellsContent, sqlKernelMetadata, title + this.invocationCount++);
+		await this.runCells(notebook);
 		const expectedOutput0 = '(1 row affected)';
 		for (let i = 0; i < 3; i++) {
 			let cellOutputs = notebook.document.cells[i].contents.outputs;
@@ -162,9 +173,27 @@ class NotebookTester {
 		}
 	}
 
+	async sqlNbRunCellsAboveBelowTest(title: string): Promise<void> {
+		let notebook = await this.openNotebook(sqlNotebookMultipleCellsContent, sqlKernelMetadata, title + this.invocationCount++);
+		// When running all cells above a cell, ensure that only cells preceding current cell have output
+		await this.runCells(notebook, true, undefined, notebook.document.cells[1]);
+		assert(notebook.document.cells[0].contents.outputs.length === 3, `Expected length: '3', Actual: '${notebook.document.cells[0].contents.outputs.length}'`);
+		assert(notebook.document.cells[1].contents.outputs.length === 0, `Expected length: '0', Actual: '${notebook.document.cells[1].contents.outputs.length}'`);
+		assert(notebook.document.cells[2].contents.outputs.length === 0, `Expected length: '0', Actual: '${notebook.document.cells[2].contents.outputs.length}'`);
+
+		await notebook.clearAllOutputs();
+
+		// When running all cells below a cell, ensure that current cell and cells after have output
+		await this.runCells(notebook, undefined, true, notebook.document.cells[1]);
+		assert(notebook.document.cells[0].contents.outputs.length === 0, `Expected length: '0', Actual: '${notebook.document.cells[0].contents.outputs.length}'`);
+		assert(notebook.document.cells[1].contents.outputs.length === 3, `Expected length: '3', Actual: '${notebook.document.cells[1].contents.outputs.length}'`);
+		assert(notebook.document.cells[2].contents.outputs.length === 3, `Expected length: '3', Actual: '${notebook.document.cells[2].contents.outputs.length}'`);
+	}
+
 	@stressify({ dop: NotebookTester.ParallelCount })
 	async sqlNbTest(title: string): Promise<void> {
-		let notebook = await this.openNotebook(sqlNotebookContent, sqlKernelMetadata, title + this.invocationCount++, false, true);
+		let notebook = await this.openNotebook(sqlNotebookContent, sqlKernelMetadata, title + this.invocationCount++, true);
+		await this.runCell(notebook);
 		const expectedOutput0 = '(1 row affected)';
 		let cellOutputs = notebook.document.cells[0].contents.outputs;
 		console.log('Got cell outputs ---');
@@ -181,6 +210,7 @@ class NotebookTester {
 
 	async sqlNbChangeKernelDifferentProviderTest(title: string): Promise<void> {
 		let notebook = await this.openNotebook(sqlNotebookContent, sqlKernelMetadata, title);
+		await this.runCell(notebook);
 		assert(notebook.document.providerId === 'sql', `Expected providerId to be sql, Actual: ${notebook.document.providerId}`);
 		assert(notebook.document.kernelSpec.name === 'SQL', `Expected first kernel name: SQL, Actual: ${notebook.document.kernelSpec.name}`);
 
@@ -196,6 +226,7 @@ class NotebookTester {
 	async shouldNotBeDirtyAfterSavingNotebookTest(title: string): Promise<void> {
 		// Given a notebook that's been edited (in this case, open notebook runs the 1st cell and adds an output)
 		let notebook = await this.openNotebook(sqlNotebookContent, sqlKernelMetadata, title);
+		await this.runCell(notebook);
 		assert(notebook.document.providerId === 'sql', `Expected providerId to be sql, Actual: ${notebook.document.providerId}`);
 		assert(notebook.document.kernelSpec.name === 'SQL', `Expected first kernel name: SQL, Actual: ${notebook.document.kernelSpec.name}`);
 		assert(notebook.document.isDirty === true, 'Notebook should be dirty after edit');
@@ -230,6 +261,7 @@ class NotebookTester {
 
 	async pythonChangeKernelDifferentProviderTest(title: string): Promise<void> {
 		let notebook = await this.openNotebook(pySparkNotebookContent, pythonKernelMetadata, title);
+		await this.runCell(notebook);
 		assert(notebook.document.providerId === 'jupyter', `Expected providerId to be jupyter, Actual: ${notebook.document.providerId}`);
 		assert(notebook.document.kernelSpec.name === 'python3', `Expected first kernel name: python3, Actual: ${notebook.document.kernelSpec.name}`);
 
@@ -244,6 +276,7 @@ class NotebookTester {
 
 	async pythonChangeKernelSameProviderTest(title: string): Promise<void> {
 		let notebook = await this.openNotebook(pySparkNotebookContent, pythonKernelMetadata, title);
+		await this.runCell(notebook);
 		assert(notebook.document.providerId === 'jupyter', `Expected providerId to be jupyter, Actual: ${notebook.document.providerId}`);
 		assert(notebook.document.kernelSpec.name === 'python3', `Expected first kernel name: python3, Actual: ${notebook.document.kernelSpec.name}`);
 
@@ -348,7 +381,7 @@ class NotebookTester {
 		}
 	}
 
-	async openNotebook(content: azdata.nb.INotebookContents, kernelMetadata: any, testName: string, runAllCells?: boolean, connectToDifferentServer?: boolean): Promise<azdata.nb.NotebookEditor> {
+	async openNotebook(content: azdata.nb.INotebookContents, kernelMetadata: any, testName: string, connectToDifferentServer?: boolean): Promise<azdata.nb.NotebookEditor> {
 		let notebookConfig = vscode.workspace.getConfiguration('notebook');
 		notebookConfig.update('pythonPath', getConfigValue(EnvironmentVariable_PYTHON_PATH), 1);
 		if (!connectToDifferentServer) {
@@ -360,23 +393,30 @@ class NotebookTester {
 		let uri = writeNotebookToFile(notebookJson, testName);
 		console.log('Notebook uri ' + uri);
 		let notebook = await azdata.nb.showNotebookDocument(uri);
-		console.log('Notebook is opened');
-
-		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 runCells(notebook: azdata.nb.NotebookEditor, runCellsAbove?: boolean, runCellsBelow?: boolean, currentCell?: azdata.nb.NotebookCell) {
+		assert(notebook !== undefined && notebook !== null, 'Expected notebook object is defined');
+		let ran;
+		if (runCellsAbove) {
+			ran = await notebook.runAllCells(undefined, currentCell);
+		} else if (runCellsBelow) {
+			ran = await notebook.runAllCells(currentCell, undefined);
+		} else {
+			ran = await notebook.runAllCells();
+		}
+		assert(ran, 'Notebook runCell should succeed');
+	}
+
+	async runCell(notebook: azdata.nb.NotebookEditor, cell?: azdata.nb.NotebookCell) {
+		if (isNullOrUndefined(cell)) {
+			cell = notebook.document.cells[0];
+		}
+		let ran = await notebook.runCell(cell);
+		assert(ran, 'Notebook runCell should succeed');
+	}
+
 	async verifyClearAllOutputs(notebook: azdata.nb.NotebookEditor): Promise<void> {
 		let cellWithOutputs = notebook.document.cells.find(cell => cell.contents && cell.contents.outputs && cell.contents.outputs.length > 0);
 		assert(cellWithOutputs !== undefined, 'Could not find notebook cells with outputs');

extensions/notebook/src/extension.ts

@@ -176,11 +176,11 @@ async function clearActiveCellOutput(): Promise<void> {
 	}
 }
 
-async function runAllCells(): Promise<void> {
+async function runAllCells(startCell?: azdata.nb.NotebookCell, endCell?: azdata.nb.NotebookCell): Promise<void> {
 	try {
 		let notebook = azdata.nb.activeNotebookEditor;
 		if (notebook) {
-			await notebook.runAllCells();
+			await notebook.runAllCells(startCell, endCell);
 		} else {
 			throw new Error(noNotebookVisible);
 		}

src/sql/azdata.proposed.d.ts

@@ -4416,7 +4416,7 @@ declare module 'azdata' {
 			/**
 			 * Kicks off execution of all code cells. Thenable will resolve only when full execution of all cells is completed.
 			 */
-			runAllCells(): Thenable<boolean>;
+			runAllCells(startCell?: NotebookCell, endCell?: NotebookCell): Thenable<boolean>;
 
 			/**
 			 * Clears the outputs of the active code cell in a notebook.

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

@@ -155,8 +155,10 @@ export class ExtHostNotebookEditor implements azdata.nb.NotebookEditor, IDisposa
 		return this._proxy.$runCell(this._id, uri);
 	}
 
-	public runAllCells(): Thenable<boolean> {
+	public runAllCells(startCell?: azdata.nb.NotebookCell, endCell?: azdata.nb.NotebookCell): Thenable<boolean> {
-		return this._proxy.$runAllCells(this._id);
+		let startCellUri = startCell ? startCell.uri : undefined;
+		let endCellUri = endCell ? endCell.uri : undefined;
+		return this._proxy.$runAllCells(this._id, startCellUri, endCellUri);
 	}
 
 	public clearOutput(cell: azdata.nb.NotebookCell): Thenable<boolean> {

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

@@ -132,11 +132,11 @@ class MainThreadNotebookEditor extends Disposable {
 		return this.editor.runCell(cell);
 	}
 
-	public runAllCells(): Promise<boolean> {
+	public runAllCells(startCell?: ICellModel, endCell?: ICellModel): Promise<boolean> {
 		if (!this.editor) {
 			return Promise.resolve(false);
 		}
-		return this.editor.runAllCells();
+		return this.editor.runAllCells(startCell, endCell);
 	}
 
 	public clearOutput(cell: ICellModel): Promise<boolean> {
@@ -383,12 +383,22 @@ export class MainThreadNotebookDocumentsAndEditors extends Disposable implements
 		return editor.runCell(cell);
 	}
 
-	$runAllCells(id: string): Promise<boolean> {
+	$runAllCells(id: string, startCellUri?: UriComponents, endCellUri?: UriComponents): Promise<boolean> {
 		let editor = this.getEditor(id);
 		if (!editor) {
 			return Promise.reject(disposed(`TextEditor(${id})`));
 		}
-		return editor.runAllCells();
+		let startCell: ICellModel;
+		let endCell: ICellModel;
+		if (startCellUri) {
+			let uriString = URI.revive(startCellUri).toString();
+			startCell = editor.cells.find(c => c.cellUri.toString() === uriString);
+		}
+		if (endCellUri) {
+			let uriString = URI.revive(endCellUri).toString();
+			endCell = editor.cells.find(c => c.cellUri.toString() === uriString);
+		}
+		return editor.runAllCells(startCell, endCell);
 	}
 
 	$clearOutput(id: string, cellUri: UriComponents): Promise<boolean> {

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

@@ -919,7 +919,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>;
+	$runAllCells(id: string, startCellUri?: UriComponents, endCellUri?: UriComponents): Promise<boolean>;
 	$clearOutput(id: string, cellUri: UriComponents): Promise<boolean>;
 	$clearAllOutputs(id: string): Promise<boolean>;
 	$changeKernel(id: string, kernel: azdata.nb.IKernelInfo): Promise<boolean>;

src/sql/workbench/parts/notebook/cellToggleMoreActions.ts

@@ -18,6 +18,7 @@ import { NotebookModel } from 'sql/workbench/parts/notebook/models/notebookModel
 import { ToggleMoreWidgetAction } from 'sql/workbench/parts/dashboard/common/actions';
 import { CellTypes, CellType } from 'sql/workbench/parts/notebook/models/contracts';
 import { CellModel } from 'sql/workbench/parts/notebook/models/cell';
+import { INotebookService } from 'sql/workbench/services/notebook/common/notebookService';
 
 export const HIDDEN_CLASS = 'actionhidden';
 
@@ -33,6 +34,8 @@ export class CellToggleMoreActions {
 			instantiationService.createInstance(AddCellFromContextAction, 'codeAfter', localize('codeAfter', 'Insert Code After'), CellTypes.Code, true),
 			instantiationService.createInstance(AddCellFromContextAction, 'markdownBefore', localize('markdownBefore', 'Insert Text Before'), CellTypes.Markdown, false),
 			instantiationService.createInstance(AddCellFromContextAction, 'markdownAfter', localize('markdownAfter', 'Insert Text After'), CellTypes.Markdown, true),
+			instantiationService.createInstance(RunCellsAction, 'runAllBefore', localize('runAllBefore', "Run Cells Before"), false),
+			instantiationService.createInstance(RunCellsAction, 'runAllAfter', localize('runAllAfter', "Run Cells After"), true),
 			instantiationService.createInstance(ClearCellOutputAction, 'clear', localize('clear', 'Clear Output'))
 		);
 	}
@@ -141,3 +144,41 @@ export class ClearCellOutputAction extends CellActionBase {
 	}
 
 }
+
+export class RunCellsAction extends CellActionBase {
+	constructor(id: string,
+		label: string,
+		private isAfter: boolean,
+		@INotificationService notificationService: INotificationService,
+		@INotebookService private notebookService: INotebookService,
+	) {
+		super(id, label, undefined, notificationService);
+	}
+
+	public canRun(context: CellContext): boolean {
+		return context.cell && context.cell.cellType === CellTypes.Code;
+	}
+
+	async doRun(context: CellContext): Promise<void> {
+		try {
+			let cell = context.cell || context.model.activeCell;
+			if (cell) {
+				let editor = this.notebookService.findNotebookEditor(cell.notebookModel.notebookUri);
+				if (editor) {
+					if (this.isAfter) {
+						await editor.runAllCells(cell, undefined);
+					} else {
+						await editor.runAllCells(undefined, cell);
+					}
+				}
+			}
+		} catch (error) {
+			let message = getErrorMessage(error);
+			this.notificationService.notify({
+				severity: Severity.Error,
+				message: message
+			});
+		}
+		return Promise.resolve();
+	}
+}

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

@@ -50,6 +50,7 @@ import { ILogService } from 'vs/platform/log/common/log';
 import { ITextFileService } from 'vs/workbench/services/textfile/common/textfiles';
 import { LabeledMenuItemActionItem, fillInActions } from 'vs/platform/actions/browser/menuEntryActionViewItem';
 import { ITelemetryService } from 'vs/platform/telemetry/common/telemetry';
+import { isUndefinedOrNull } from 'vs/base/common/types';
 
 
 export const NOTEBOOK_SELECTOR: string = 'notebook-component';
@@ -523,11 +524,20 @@ export class NotebookComponent extends AngularDisposable implements OnInit, OnDe
 		}
 	}
 
-	public async runAllCells(): Promise<boolean> {
+	public async runAllCells(startCell?: ICellModel, endCell?: ICellModel): 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++) {
+			// For the run all cells scenario where neither startId not endId are provided, set defaults
+			let startIndex = 0;
+			let endIndex = codeCells.length;
+			if (!isUndefinedOrNull(startCell)) {
+				startIndex = codeCells.findIndex(c => c.id === startCell.id);
+			}
+			if (!isUndefinedOrNull(endCell)) {
+				endIndex = codeCells.findIndex(c => c.id === endCell.id);
+			}
+			for (let i = startIndex; i < endIndex; i++) {
 				let cellStatus = await this.runCell(codeCells[i]);
 				if (!cellStatus) {
 					return Promise.reject(new Error(localize('cellRunFailed', "Run Cells failed - See error in output of the currently selected cell for more information.")));

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

@@ -128,7 +128,7 @@ export interface INotebookEditor {
 	isVisible(): boolean;
 	executeEdits(edits: ISingleNotebookEditOperation[]): boolean;
 	runCell(cell: ICellModel): Promise<boolean>;
-	runAllCells(): Promise<boolean>;
+	runAllCells(startCell?: ICellModel, endCell?: ICellModel): Promise<boolean>;
 	clearOutput(cell: ICellModel): Promise<boolean>;
 	clearAllOutputs(): Promise<boolean>;
 }